def test_split_matrix(): assert _split_matrix(np.array([[1]])) == [([0], [0])] assert _split_matrix(np.array([[1, 1], [1, 1]])) == [([0, 1], [0, 1])] assert _split_matrix(np.array([[1, 1, 0], [1, 1, 0], [0, 0, 1]])) == [([0, 1], [0, 1]), ([2], [2])] assert _split_matrix(np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]])) == [([0], [1]), ([1], [0]), ([2], [2])] assert _split_matrix(np.array([[0, 1, 1], [1, 0, 0], [1, 0, 1]])) == [([0, 1, 2], [0, 1, 2])]
def _generate_world_coords(self, pixel_corners, wcs): # TODO: We can improve this by not always generating all coordinates # To make our lives easier here we generate all the coordinates for all # pixels and then choose the ones we want to return to the user based # on the axes argument. We could be smarter by integrating this logic # into the main loop, this would potentially reduce the number of calls # to pixel_to_world_values # Create meshgrid of all pixel coordinates. # If user, wants pixel_corners, set pixel values to pixel pixel_corners. # Else make pixel centers. pixel_shape = self.data.shape[::-1] if pixel_corners: pixel_shape = tuple(np.array(pixel_shape) + 1) ranges = [np.arange(i) - 0.5 for i in pixel_shape] else: ranges = [np.arange(i) for i in pixel_shape] # Limit the pixel dimensions to the ones present in the ExtraCoords if isinstance(wcs, ExtraCoords): ranges = [ranges[i] for i in wcs.mapping] wcs = wcs.wcs if wcs is None: return [] world_coords = [None] * wcs.world_n_dim for (pixel_axes_indices, world_axes_indices) in _split_matrix(wcs.axis_correlation_matrix): # First construct a range of pixel indices for this set of coupled dimensions sub_range = [ranges[idx] for idx in pixel_axes_indices] # Then get a set of non correlated dimensions non_corr_axes = set(list(range( wcs.pixel_n_dim))) - set(pixel_axes_indices) # And inject 0s for those coordinates for idx in non_corr_axes: sub_range.insert(idx, 0) # Generate a grid of broadcastable pixel indices for all pixel dimensions grid = np.meshgrid(*sub_range, indexing='ij') # Convert to world coordinates world = wcs.pixel_to_world_values(*grid) # TODO: this isinstance check is to mitigate https://github.com/spacetelescope/gwcs/pull/332 if wcs.world_n_dim == 1 and not isinstance(world, tuple): world = [world] # Extract the world coordinates of interest and remove any non-correlated axes # Transpose the world coordinates so they match array ordering not pixel for idx in world_axes_indices: array_slice = np.zeros((wcs.pixel_n_dim, ), dtype=object) array_slice[wcs.axis_correlation_matrix[idx]] = slice(None) tmp_world = world[idx][tuple(array_slice)].T world_coords[idx] = tmp_world for i, (coord, unit) in enumerate(zip(world_coords, wcs.world_axis_units)): world_coords[i] = coord << u.Unit(unit) return world_coords