Enable patchwise training and prediction

deepsensor/data/task.py

@@ -31,7 +31,9 @@ def __init__(self, task_dict: dict) -> None:
     @classmethod
     def summarise_str(cls, k, v):
         """Return string summaries for the _str__ method."""
-        if plum.isinstance(v, B.Numeric):
+        if isinstance(v, float):
+            return v
+        elif plum.isinstance(v, B.Numeric):
             return v.shape
         elif plum.isinstance(v, tuple):
             return tuple(vi.shape for vi in v)
@@ -57,6 +59,8 @@ def summarise_repr(cls, k, v) -> str:
         """
         if v is None:
             return "None"
+        elif isinstance(v, float):
+            return f"{type(v).__name__}"
         elif plum.isinstance(v, B.Numeric):
             return f"{type(v).__name__}/{v.dtype}/{v.shape}"
         if plum.isinstance(v, deepsensor.backend.nps.mask.Masked):

deepsensor/model/model.py

@@ -9,6 +9,7 @@
     Prediction,
     increase_spatial_resolution,
     infer_prediction_modality_from_X_t,
+    stitch_clipped_predictions,
 )
 from deepsensor.data.task import Task
 
@@ -648,6 +649,273 @@ def unnormalise_pred_array(arr, **kwargs):
 
         return pred
 
+    def predict_patchwise(
+        self,
+        tasks: Union[List[Task], Task],
+        X_t: Union[
+            xr.Dataset,
+            xr.DataArray,
+            pd.DataFrame,
+            pd.Series,
+            pd.Index,
+            np.ndarray,
+        ],
+        X_t_mask: Optional[Union[xr.Dataset, xr.DataArray]] = None,
+        **kwargs,
+    ) -> Prediction:
+        """Predict using tasks loaded using a sliding window patching strategy. Uses the `predict` method.
+
+        .. versionadded:: 0.4.3
+            :py:func:`predict_patchwise()` method.
+
+        Args:
+            tasks (List[Task] | Task):
+                List of tasks containing context data. Tasks for patchwise prediction must be generated by a PatchwiseTaskLoader using the "sliding" patching strategy.
+            X_t (:class:`xarray.Dataset` | :class:`xarray.DataArray` | :class:`pandas.DataFrame` | :class:`pandas.Series` | :class:`pandas.Index` | :class:`numpy:numpy.ndarray`):
+                Target locations to predict at. Can be an xarray object
+                containing on-grid locations or a pandas object containing off-grid locations.
+            X_t_mask: :class:`xarray.Dataset` | :class:`xarray.DataArray`, optional
+                2D mask to apply to gridded ``X_t`` (zero/False will be NaNs). Will be interpolated
+                to the same grid as ``X_t`` and patched in the same way. Default None (no mask).
+            **kwargs:
+                Keyword arguments as per ``predict``.
+
+        Returns:
+            :class:`~.model.pred.Prediction`):
+                A `dict`-like object mapping from target variable IDs to xarray or pandas objects
+                containing model predictions.
+                - If ``X_t`` is a pandas object, returns pandas objects
+                containing off-grid predictions.
+                - If ``X_t`` is an xarray object, returns xarray object
+                containing on-grid predictions.
+                - If ``n_samples`` == 0, returns only mean and std predictions.
+                - If ``n_samples`` > 0, returns mean, std and samples
+                predictions.
+
+        Raises:
+            AttributeError
+                If ``tasks`` are not generated using the "sliding" patching strategy of PatchwiseTaskLoader.
+            Errors
+                See `~.model.model.DeepSensorModel.predict`
+        """
+        # Get coordinate names of original unnormalised dataset.
+        orig_x1_name = self.data_processor.x1_name
+        orig_x2_name = self.data_processor.x2_name
+
+        def get_patches_per_row(preds) -> int:
+            """Calculate number of patches per row.
+            Required to stitch patches back together.
+
+            Args:
+                preds (List[class:`~.model.pred.Prediction`]):
+                        A list of `dict`-like objects containing patchwise predictions.
+
+            Returns:
+                patches_per_row: int
+                    Number of patches per row.
+            """
+            patches_per_row = 0
+            vars = list(preds[0][0].data_vars)
+            var = vars[0]
+            x1_val = preds[0][0][var].coords[orig_x1_name].min()
+
+            for pred in preds:
+                if pred[0][var].coords[orig_x1_name].min() == x1_val:
+                    patches_per_row = patches_per_row + 1
+
+            return patches_per_row
+
+        def get_patch_overlap(
+            overlap_norm, data_processor, X_t_ds, x1_ascend, x2_ascend
+        ) -> int:
+            """Calculate overlap between adjacent patches in pixels.
+
+            Parameters
+            ----------
+            overlap_norm : tuple[float].
+                Normalised size of overlap in x1/x2.
+
+            data_processor (:class:`~.data.processor.DataProcessor`):
+                Used for unnormalising the coordinates of the bounding boxes of patches.
+
+            X_t_ds (:class:`xarray.Dataset` | :class:`xarray.DataArray` | :class:`pandas.DataFrame` | :class:`pandas.Series` | :class:`pandas.Index` | :class:`numpy:numpy.ndarray`):
+                Data array containing target locations to predict at.
+
+            x1_ascend : str:
+                Boolean defining whether the x1 coords ascend (increase) from top to bottom, default = True.
+
+            x2_ascend : str:
+                Boolean defining whether the x2 coords ascend (increase) from left to right, default = True.
+
+            Returns:
+            -------
+            patch_overlap : tuple (int)
+                Unnormalised size of overlap between adjacent patches.
+            """
+            # Todo- check if there is simplier and more robust way to convert overlap into pixels.
+            # Place x1/x2 overlap values in Xarray to pass into unnormalise()
+            overlap_list = [0, overlap_norm[0], 0, overlap_norm[1]]
+            x1 = xr.DataArray([overlap_list[0], overlap_list[1]], dims="x1", name="x1")
+            x2 = xr.DataArray([overlap_list[2], overlap_list[3]], dims="x2", name="x2")
+            overlap_norm_xr = xr.Dataset(coords={"x1": x1, "x2": x2})
+
+            # Unnormalise coordinates of bounding boxes
+            overlap_unnorm_xr = data_processor.unnormalise(overlap_norm_xr)
+
+            unnorm_overlap_x1 = overlap_unnorm_xr.coords[orig_x1_name].values[1]
+            unnorm_overlap_x2 = overlap_unnorm_xr.coords[orig_x2_name].values[1]
+
+            def overlap_index(
+                coords: np.ndarray, ascend: bool, unnorm_overlap: float
+            ) -> int:
+                """Find size of overlap in a single coordinate direction, in units of pixels.
+
+                Parameters
+                ----------
+                coords : np.ndarray
+
+                ascend : bool
+                    Boolean defining whether coords ascend (increase) from top to bottom or left to right.
+
+                unnorm_overlap : float
+                    The patch overlap in unnormalised coordinates.
+
+                Returns:
+                -------
+                int : The number of pixels in the overlap.
+                """
+                pixel_coords_overlap_diffs = np.abs(coords - unnorm_overlap)
+                if ascend:
+                    trim_size = np.argmin(pixel_coords_overlap_diffs) / 2
+                    trim_size_rounded = int(
+                        np.floor(trim_size)
+                    )  # Always round down trim slide as stitching method can handle slight overlaps
+                    return trim_size_rounded
+
+                else:
+                    overlap_pixel_size = np.argmin(pixel_coords_overlap_diffs)
+                    overlap_pixel_size_rounded = np.ceil(overlap_pixel_size)
+                    trim_size = (
+                        (coords.size - int(overlap_pixel_size_rounded)) / 2
+                    )  # this extra step is so we get the overlap with respect to the largest value (i.e. is the number of pixels = 360, coords.size = 360)
+                    trim_size_rounded = int(np.floor(trim_size))
+                    return trim_size_rounded
+
+            return (
+                overlap_index(
+                    X_t_ds.coords[orig_x1_name].values, x1_ascend, unnorm_overlap_x1
+                ),
+                overlap_index(
+                    X_t_ds.coords[orig_x2_name].values, x2_ascend, unnorm_overlap_x2
+                ),
+            )
+
+        task_patch_strategy = tasks[0]["patch_strategy"]
+        if task_patch_strategy is not "sliding":
+            raise AttributeError(
+                f"For patchwise prediction, only tasks generated using a patch_strategy of 'sliding' are valid. \
+                    This task appears to have been created using the {task_patch_strategy} strategy."
+            )
+
+        # load patch_size and stride from task
+        patch_size = tasks[0]["patch_size"]
+        stride = tasks[0]["stride"]
+
+        # sanitise patch_size and stride arguments
+        if isinstance(patch_size, float) and patch_size is not None:
+            patch_size = (patch_size, patch_size)
+
+        if isinstance(stride, float) and stride is not None:
+            stride = (stride, stride)
+
+        if stride[0] > patch_size[0] or stride[1] > patch_size[1]:
+            raise ValueError(
+                f"stride must be smaller than patch_size in the corresponding dimensions for patchwise prediction. Got: patch_size: {patch_size}, stride: {stride}"
+            )
+
+        # patchwise prediction does not yet support more than a single date
+        num_task_dates = len(set([t["time"] for t in tasks]))
+        if num_task_dates > 1:
+            raise NotImplementedError(
+                f"Patchwise prediction does not yet support more than a single date at a time, got {num_task_dates}."
+            )
+
+        # tasks should be iterable, if only one is provided, make it a list
+        if type(tasks) is Task:
+            tasks = [tasks]
+
+        # Perform patchwise predictions
+        preds = []
+        for task in tasks:
+            bbox = task["bbox"]
+
+            # Unnormalise coordinates of bounding box of patch
+            x1 = xr.DataArray([bbox[0], bbox[1]], dims="x1", name="x1")
+            x2 = xr.DataArray([bbox[2], bbox[3]], dims="x2", name="x2")
+            bbox_norm = xr.Dataset(coords={"x1": x1, "x2": x2})
+            bbox_unnorm = self.data_processor.unnormalise(bbox_norm)
+            unnorm_bbox_x1 = (
+                bbox_unnorm[orig_x1_name].values.min(),
+                bbox_unnorm[orig_x1_name].values.max(),
+            )
+            unnorm_bbox_x2 = (
+                bbox_unnorm[orig_x2_name].values.min(),
+                bbox_unnorm[orig_x2_name].values.max(),
+            )
+
+            # Determine X_t for patch, however, cannot assume min/max ordering of slice coordinates
+            # Check the order of coordinates in X_t, sometimes they are increasing or decreasing in order.
+            x1_coords = X_t.coords[orig_x1_name].values
+            x2_coords = X_t.coords[orig_x2_name].values
+
+            if x1_coords[0] < x1_coords[-1]:
+                x1_slice = slice(unnorm_bbox_x1[0], unnorm_bbox_x1[1])
+                x1_ascending = True
+            else:
+                x1_slice = slice(unnorm_bbox_x1[1], unnorm_bbox_x1[0])
+                x1_ascending = False
+
+            if x2_coords[0] < x2_coords[-1]:
+                x2_slice = slice(unnorm_bbox_x2[0], unnorm_bbox_x2[1])
+                x2_ascending = True
+            else:
+                x2_slice = slice(unnorm_bbox_x2[1], unnorm_bbox_x2[0])
+                x2_ascending = False
+
+            # Determine X_t for patch with correct slice direction
+            task_X_t = X_t.sel(**{orig_x1_name: x1_slice, orig_x2_name: x2_slice})
+            task_X_t_mask = (
+                X_t_mask.sel(**{orig_x1_name: x1_slice, orig_x2_name: x2_slice})
+                if X_t_mask
+                else None
+            )
+
+            # Patchwise prediction
+            pred = self.predict(task, task_X_t, task_X_t_mask, **kwargs)
+            # Append patchwise DeepSensor prediction object to list
+            preds.append(pred)
+
+        overlap_norm = tuple(
+            patch - stride for patch, stride in zip(patch_size, stride)
+        )
+        patch_overlap_unnorm = get_patch_overlap(
+            overlap_norm, self.data_processor, X_t, x1_ascending, x2_ascending
+        )
+
+        patches_per_row = get_patches_per_row(preds)
+        prediction = stitch_clipped_predictions(
+            preds,
+            patch_overlap_unnorm,
+            patches_per_row,
+            X_t,
+            orig_x1_name,
+            orig_x2_name,
+            x1_ascending,
+            x2_ascending,
+        )
+
+        return prediction
+
 
 def add_valid_time_coord_to_pred_and_move_time_dims(pred: Prediction) -> Prediction:
     """Add a valid time coordinate "time" to a Prediction object based on the