WIP

Author: astrofrog

reproject/common.py

@@ -60,6 +60,7 @@ def _reproject_dispatcher(
     shape_out,
     wcs_out,
     block_size=None,
+    non_reprojected_dims=None,
     array_out=None,
     return_footprint=True,
     output_footprint=None,
@@ -92,6 +93,11 @@ def _reproject_dispatcher(
         the block size automatically determined. If ``block_size`` is not
         specified or set to `None`, the reprojection will not be carried out in
         blocks.
+    non_reprojected_dims : tuple
+        Dimensions that should not be reprojected but instead for which a
+        1-to-1 mapping between input and output pixel space should be assumed.
+        By default, this is any leading extra dimensions if the input WCS has
+        fewer dimensions than the input data.
     array_out : `~numpy.ndarray`, optional
         An array in which to store the reprojected data.  This can be any numpy
         array including a memory map, which may be helpful when dealing with
@@ -198,9 +204,32 @@ def _reproject_dispatcher(
         # shape_out will be the full size of the output array as this is updated
         # in parse_output_projection, even if shape_out was originally passed in as
         # the shape of a single image.
-        broadcasting = wcs_in.low_level_wcs.pixel_n_dim < len(shape_out)
+        if non_reprojected_dims is None:
+            non_reprojected_dims = list(range(len(shape_out) - wcs_in.low_level_wcs.pixel_n_dim))
+        else:
+            non_reprojected_dims = list(non_reprojected_dims)
+
+        broadcasting = len(non_reprojected_dims) > 0
+
+        reprojected_dims = [x for x in range(len(shape_out)) if x not in non_reprojected_dims]
 
         logger.info(f"Broadcasting is {'' if broadcasting else 'not '}being used")
+        logger.info(f"Dimensions being reprojected: {reprojected_dims}")
+        logger.info(f"Dimensions not being reprojected: {non_reprojected_dims}")
+
+        if len(block_size) < len(shape_out):
+            block_size = [-1] * (len(shape_out) - len(block_size)) + list(block_size)
+        elif len(block_size) > len(shape_out):
+            raise ValueError(
+                f"block_size {len(block_size)} cannot have more elements "
+                f"than the dimensionality of the output ({len(shape_out)})"
+            )
+
+        block_size = np.array(block_size)
+        shape_out = np.array(shape_out)
+
+        # TODO: replace block size of -1 by actual value for logic below to work
+        # TODO: re-implement block_size auto
 
         # Check block size and determine whether block size indicates we should
         # parallelize over broadcasted dimension. The logic is as follows: if
@@ -212,33 +241,23 @@ def _reproject_dispatcher(
         # don't make any assumptions for now and assume a single chunk in the
         # missing dimensions.
         broadcasted_parallelization = False
-        if broadcasting and block_size is not None and block_size != "auto":
-            if len(block_size) == len(shape_out):
-                if (
-                    block_size[-wcs_in.low_level_wcs.pixel_n_dim :]
-                    == shape_out[-wcs_in.low_level_wcs.pixel_n_dim :]
-                ):
-                    broadcasted_parallelization = True
-                    block_size = (
-                        block_size[: -wcs_in.low_level_wcs.pixel_n_dim]
-                        + (-1,) * wcs_in.low_level_wcs.pixel_n_dim
-                    )
-                else:
-                    for i in range(len(shape_out) - wcs_in.low_level_wcs.pixel_n_dim):
-                        if block_size[i] != -1 and block_size[i] != shape_out[i]:
-                            raise ValueError(
-                                "block shape should either match output data shape along broadcasted dimension or non-broadcasted dimensions"
-                            )
-            elif len(block_size) < len(shape_out):
-                block_size = [-1] * (len(shape_out) - len(block_size)) + list(block_size)
-            else:
-                raise ValueError(
-                    f"block_size {len(block_size)} cannot have more elements "
-                    f"than the dimensionality of the output ({len(shape_out)})"
+        if broadcasting and block_size is not None:
+            if np.all(block_size[reprojected_dims] == shape_out[reprojected_dims]):
+                broadcasted_parallelization = True
+                block_size = np.array(
+                    tuple(block_size[non_reprojected_dims].tolist())
+                    + (-1,) * len(reprojected_dims)
                 )
+            elif np.all(block_size[non_reprojected_dims] != shape_out[non_reprojected_dims]):
+                raise ValueError(
+                    "block shape should either match output data shape along broadcasted dimension or non-broadcasted dimensions"
+                        )
 
             # TODO: check for shape_out not matching shape_in along broadcasted dimensions
 
+        block_size = tuple(block_size.tolist())
+        shape_out = tuple(shape_out.tolist())
+
         logger.info(
             f"{'P' if broadcasted_parallelization else 'Not p'}arallelizing along "
             f"broadcasted dimension ({block_size=}, {shape_out=})"
@@ -270,17 +289,38 @@ def reproject_single_block(a, array_or_path, block_info=None):
             wcs_in_cp = wcs_in.deepcopy() if isinstance(wcs_in, WCS) else wcs_in
             wcs_out_cp = wcs_out.deepcopy() if isinstance(wcs_out, WCS) else wcs_out
 
-            slices = [
-                slice(*x) for x in block_info[None]["array-location"][-wcs_out_cp.pixel_n_dim :]
-            ]
+            print(block_info[None]["array-location"])
 
-            if isinstance(wcs_out, BaseHighLevelWCS):
+            slices = []
+            for i in reprojected_dims:
+                slices.append(slice(*block_info[None]["array-location"][i]))
+
+            print(slices)
+
+            if isinstance(wcs_out_cp, BaseHighLevelWCS):
                 low_level_wcs = SlicedLowLevelWCS(wcs_out_cp.low_level_wcs, slices=slices)
             else:
                 low_level_wcs = SlicedLowLevelWCS(wcs_out_cp, slices=slices)
 
+            print(low_level_wcs.pixel_n_dim, low_level_wcs.world_n_dim)
+
             wcs_out_sub = HighLevelWCSWrapper(low_level_wcs)
 
+            slices = []
+            for i in range(wcs_in_cp.pixel_n_dim):
+                if i in non_reprojected_dims:
+                    # slices.append(slice(*block_info[None]["array-location"][i]))
+                    slices.append(block_info[None]["array-location"][i][0])
+                else:
+                    slices.append(slice(None))
+
+            if isinstance(wcs_in_cp, BaseHighLevelWCS):
+                low_level_wcs_in = SlicedLowLevelWCS(wcs_in_cp.low_level_wcs, slices=slices)
+            else:
+                low_level_wcs_in = SlicedLowLevelWCS(wcs_in_cp, slices=slices)
+
+            wcs_in_sub = HighLevelWCSWrapper(low_level_wcs_in)
+
             if isinstance(array_or_path, tuple):
                 array_in = np.memmap(array_or_path[0], **array_or_path[1], mode="r")
             elif isinstance(array_or_path, str):
@@ -295,7 +335,7 @@ def reproject_single_block(a, array_or_path, block_info=None):
 
             array, footprint = reproject_func(
                 array_in,
-                wcs_in_cp,
+                wcs_in_sub,
                 wcs_out_sub,
                 shape_out=shape_out,
                 array_out=np.zeros(shape_out),
@@ -308,10 +348,11 @@ def reproject_single_block(a, array_or_path, block_info=None):
 
             array_out_dask = da.empty(shape_out, chunks=block_size)
             if isinstance(array_in, da.core.Array):
+                # FIXME: Should take into account -1s here
                 if array_in.chunksize != block_size:
                     logger.info(
                         f"Rechunking input dask array as chunks ({array_in.chunksize}) "
-                        "do not match block size ({block_size})"
+                        f"do not match block size ({block_size})"
                     )
                     array_in = array_in.rechunk(block_size)
             else:

reproject/interpolation/core.py

@@ -10,7 +10,7 @@
 
 def _validate_wcs(wcs_in, wcs_out, shape_in, shape_out):
     if wcs_in.low_level_wcs.pixel_n_dim != wcs_out.low_level_wcs.pixel_n_dim:
-        raise ValueError("Number of dimensions in input and output WCS should match")
+        raise ValueError(f"Number of dimensions in input and output WCS should match (got {wcs_in.low_level_wcs.pixel_n_dim} and {wcs_out.low_level_wcs.pixel_n_dim})")
     elif len(shape_out) < wcs_out.low_level_wcs.pixel_n_dim:
         raise ValueError("Too few dimensions in shape_out")
     elif len(shape_in) < wcs_in.low_level_wcs.pixel_n_dim:

reproject/interpolation/high_level.py

@@ -25,6 +25,7 @@ def reproject_interp(
     output_footprint=None,
     return_footprint=True,
     block_size=None,
+    non_reprojected_dims=None,
     parallel=False,
     return_type=None,
 ):
@@ -142,6 +143,7 @@ def reproject_interp(
         array_out=output_array,
         parallel=parallel,
         block_size=block_size,
+        non_reprojected_dims=non_reprojected_dims,
         return_footprint=return_footprint,
         output_footprint=output_footprint,
         reproject_func_kwargs=dict(