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