# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
Defines some utility routines used to map a provided set of quantities.
----
.. include license and copyright
.. include:: ../include/copy.rst
----
.. include common links, assuming primary doc root is up one directory
.. include:: ../include/links.rst
"""
import warnings
from IPython import embed
import numpy
from scipy import ndimage
from astropy.wcs import WCS
from matplotlib import pyplot, patches
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.ticker import MaxNLocator
##############################################################################
[docs]
def masked_pixelized_image(x, y, z, pixelscale=1.0, zmin=None, zmax=None, imshow_prep=False,
fill_value=0.0):
r"""
Provided a set of pixelized data, return a masked image with a type
`numpy.ma.MaskedArray`_. The :math:`x` coordinates are organized
along rows and the :math:`y` coordinates are organized along
columns. I.e., i.e., img[0,1] = :math:`z(x_0,y_1)`, unless
*imshow_prep* is requested (see below).
Args:
x (numpy.ndarray): X coordinates of the pixels
y (numpy.ndarray): Y coordinates of the pixels
z (numpy.ndarray): Image values at :math:`x,y`.
pixelscale (float): (Optional) Pixelscale of the image in
arcsec/pixel.
zmin (float): (Optional) Minimum z value to include in the
output image. Default is to allow all pixel values.
zmax (float): (Optional) Maximum z value to include in the
output image. Default is to allow all pixel values.
imshow_prep (bool): (Optional) Prepare the matrix for use with
`pyplot.imshow`_. If *imshow_prep* is True, before output,
the matrix is reordered such that increasing :math:`x`
values are along columns and increasing :math:`y` values are
along rows; i.e., the output is the transpose of the default
behavior. The appropriate call to imshow that will then put
increasing x values along the abcissa and increasing y
values along the ordinate is, e.g.::
ext, img = masked_pixelized_image(x, y, z, imshow_prep=True)
pyplot.imshow(img, interpolation='nearest', extent=ext, origin='lower')
Note that origin **must** be "lower" when calling
`pyplot.imshow`_ on the array produced by this
routine for the :math:`x,y` ordering to be as expected!
fill_value (float): (Optional) The default value to use for
pixels without any data.
Returns:
numpy.ndarray, `numpy.ma.MaskedArray`_: The first object
returned is a four-element array with the extent of the image
from the bottom edge to the top edge of the x and y pixels,
respectively. The second object returned is the image data and
associated (boolean) mask.
"""
# Image extent
ext = numpy.empty(4, dtype=float)
ext[0], ext[1] = numpy.amin(x)-pixelscale/2., numpy.amax(x)+pixelscale/2.
ext[2], ext[3] = numpy.amin(y)-pixelscale/2., numpy.amax(y)+pixelscale/2.
# Image size
nx, ny = numpy.floor((ext[1]-ext[0])/pixelscale), numpy.floor((ext[3]-ext[2])/pixelscale)
# print('MAP: nx,ny: {0},{1}'.format(nx, ny))
# print('MAP: extent: {0}'.format(ext))
# Image and mask values
img = numpy.full((nx,ny), fill_value, dtype=float)
i = numpy.floor((x.ravel()-ext[0])/pixelscale).astype(int)
j = numpy.floor((y.ravel()-ext[2])/pixelscale).astype(int)
i[i == nx] = nx-1
j[j == ny] = ny-1
# print(i)
# print(j)
img[i,j] = z
mask = numpy.full((nx,ny), True, dtype=bool)
mask[i,j] = False
# Mask values
if zmin is not None:
mask[(img < zmin)] = True
if zmax is not None:
mask[(img > zmax)] = True
# # Flip the x axis (sky-right coordinates?)
# if xflip:
# img = numpy.fliplr(img)
# mask = numpy.fliplr(mask)
# ext[0], ext[1] = ext[1], ext[0]
# print('EXTENT: {0}'.format(ext))
# Prep the image for direct display using pyplot.imshow
if imshow_prep:
img = numpy.transpose(img)
mask = numpy.transpose(mask)
return ext, numpy.ma.MaskedArray(img, mask=mask, fill_value=fill_value)
[docs]
def map_quantity(x, y, z, zmin=None, zmax=None, ncolors=64, dots=False, cmap='coolwarm',
colorbar=False, nticks=7, clabel=None, flux=None, fixpdf=False, xlim=None,
xlabel=None, ylim=None, ylabel=None, pixelscale=None, fill_value=0.0,
contour_levels=None, **kwargs):
"""
Copyright (C) 2013-2014, Michele Cappellari
E-mail: cappellari_at_astro.ox.ac.uk
http://purl.org/cappellari/software
*Revision history*:
| V1.0: Michele Cappellari, Paranal, 11 November 2013
| V1.0.1: Clip values before contouring. MC, Oxford, 26 February
2014
| V1.0.2: Include SAURON colormap. MC, Oxford, 29 January 2014
| V1.0.3: Call set_aspect(1). MC, Oxford, 22 February 2014
| V1.0.4: Call autoscale_view(tight=True). Overplot small dots
by default. MC, Oxford, 25 February 2014
| V1.0.5: Use axis('image'). MC, Oxford, 29 March 2014
| V1.0.6: Allow changing colormap. MC, Oxford, 29 July 2014
| V1.0.7: Added optional fixpdf keyword to remove PDF artifacts
like `stackoverflow_15822159`_ . Make nice tick levels for
colorbar. Added nticks keyword for colorbar. MC, Oxford, 16
October 2014
.. _stackoverflow_15822159: http://stackoverflow.com/questions/15822159/aliasing-when-saving-matplotlib-filled-contour-plot-to-pdf-or-eps
"""
if zmin is None:
zmin = numpy.min(z)
if zmax is None:
zmax = numpy.max(z)
if cmap is None:
cmap = 'coolwarm'
# print('raveling')
x, y, z = map(numpy.ravel, [x, y, z])
# print('getting axes')
ax = pyplot.gca()
# Provide a pixelized map
if pixelscale is not None:
# print('pixelizing')
ext, img = masked_pixelized_image(x, y, z, pixelscale, zmin=zmin, zmax=zmax,
fill_value=fill_value, imshow_prep=True)
# print('made pixelized image')
cs = ax.imshow(img, interpolation='nearest', cmap=cmap, vmin=zmin,
vmax=zmax, extent=ext, origin='lower')
# print('output imshow')
if contour_levels is not None:
nx = img.shape[0]
dx = (ext[1]-ext[0])/nx
ny = img.shape[1]
dy = (ext[3]-ext[2])/ny
ximg, yimg = numpy.meshgrid(ext[0]+dx*(numpy.arange(nx)+0.5),
ext[2]+dy*(numpy.arange(ny)+0.5), indexing='xy')
cs1 = ax.contour(ximg, yimg, img, levels=contour_levels, colors='k', linewidths=1.5)
# Provide an interpolated map
else:
levels = numpy.linspace(zmin, zmax, ncolors)
cs = ax.tricontourf(x, y, z.clip(zmin, zmax), levels=levels, cmap=cmap)
ax.axis('image') # Equal axes and no rescaling
ax.minorticks_on()
ax.tick_params(length=10, which='major')
ax.tick_params(length=5, which='minor')
# print('axes limits: {0}'.format(ax.get_xlim()))
if xlim is not None:
ax.set_xlim(xlim)
if ylim is not None:
ax.set_ylim(ylim)
if xlabel is not None:
ax.set_xlabel(xlabel)
if ylabel is not None:
ax.set_ylabel(ylabel)
if flux is not None:
ax.tricontour(x, y, -2.5*numpy.log10(flux/numpy.max(flux).ravel()),
levels=numpy.arange(20), colors='k') # 1 mag contours
if fixpdf: # remove white contour lines in PDF at expense of larger file size
ax.tricontour(x, y, z.clip(zmin, zmax), levels=levels, zorder=0,
cmap=cmap)
if dots:
ax.plot(x, y, '.k', markersize=kwargs.get("markersize", 3))
if colorbar:
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.1)
ticks = MaxNLocator(nticks).tick_values(zmin, zmax)
cbar = pyplot.colorbar(cs, cax=cax, ticks=ticks)
if clabel:
cbar.set_label(clabel)
# print('done map')
return cs
# TODO: Move this into datacube?
[docs]
def map_extent(hdu, ext, offset=True):
"""
Get the on-sky extent of a map using the provided WCS coordinates.
Args:
hdu (`astropy.io.fits.HDUList`_): Fits HDU list.
ext (str): Extension of the file with the map.
offset (bool): (**Optional**) Return the extent as the offset
from the coordinates of the object in arcseconds. The WCS
is assumed to provide RA and declination in units of
degrees, and the header must contain the coordinates of the
object with keywords `OBJRA` and `OBJDEC`. If offset is
True and these keywords do not exist, the function will
throw a warning and proceed *without* applying any offset.
Default is just to return the extent of the WCS coordinates.
Returns:
list: List of four floats: minimum and maximum x and minimum and
maximum y.
"""
if offset:
try:
objra = hdu['PRIMARY'].header['OBJRA']
objdec = hdu['PRIMARY'].header['OBJDEC']
except Exception as e:
print(e)
warnings.warn('OBJRA and/or OBJDEC in the header of the PRIMARY extension. '
'Offseting to the center pixel.')
objra, objdec = None, None
# Does the extension contain a cube or a single map image
cube = len(hdu[ext].data.shape) > 2
# Get the shape of the image
if cube:
nc,ny,nx = hdu[ext].data.shape
else:
ny,nx = hdu[ext].data.shape
# Create a grid with the pixel coordinates
x,y = numpy.meshgrid( numpy.arange(nx)+1, numpy.arange(ny)+1,
indexing='ij')
# Reshape them into an array with X,Y coordinates
xy = numpy.array([x.reshape(nx*ny), y.reshape(nx*ny)]).T
if cube:
xy = numpy.append(xy, numpy.ones(nx*ny).reshape(-1,1), axis=1)
# Read the World Coordinate System information from the fits header
wcs = WCS(header=hdu[ext].header, fix=False)
if wcs.naxis == 3 and not cube:
warnings.warn('WCS is for 3D, but provided data is 2D. Attempting to drop the last axis.')
wcs = wcs.dropaxis(-1)
# Convert the pixel coordinates into world coordinates
XY = wcs.all_pix2world(xy, 1)
x = XY[:,0].reshape(nx,ny)
y = XY[:,1].reshape(nx,ny)
# So now the x coordinate of:
# hdu[ext].data[:,j,i]
# is
# x[j,i]
if not offset:
return [ numpy.amax(x), numpy.amin(x), numpy.amin(y), numpy.amax(y) ]
if objra is None or objdec is None:
center = tuple([_n // 2 for _n in x.shape])
objra = x[center]
objdec = y[center]
# Convert the world coordinates to offset from center in arcsec
x = (x-objra)*numpy.cos(numpy.radians(objdec))*3600.
y = (y-objdec)*3600.
# Return the extent
return [numpy.amax(x), numpy.amin(x), numpy.amin(y), numpy.amax(y)]
[docs]
def map_center_pixel_offset(hdu_1, hdu_2, ext, quiet=False):
"""
Determine the offset in pixels between the object center in two the
maps of two fits files.
Args:
hdu_1 (`astropy.io.fits.HDUList`_): Fits HDU list
for the reference map.
hdu_2 (`astropy.io.fits.HDUList`_): Fits HDU list
for the comparison map.
ext (str): Extension *in both files* with the maps.
quiet (bool): (**Optional**) Suppress terminal output.
Returns:
float: Two floats with the offset in x and y between the two maps.
"""
cube = len(hdu_1[ext].data.shape) > 2
objra = hdu_1['PRIMARY'].header['OBJRA']
objdec = hdu_1['PRIMARY'].header['OBJDEC']
if not quiet:
print('Coordinates of 1st object: {0:10.6f} {1:9.5f}'.format(objra, objdec))
print('Coordinates of 2nd object: {0:10.6f} {1:9.5f}'.format(
hdu_2['PRIMARY'].header['OBJRA'],
hdu_2['PRIMARY'].header['OBJDEC']))
wcs_1 = WCS(header=hdu_1[ext].header, fix=False)
wcs_2 = WCS(header=hdu_2[ext].header, fix=False)
XY = numpy.zeros(3 if cube else 2).reshape(1,-1)
XY[0,0] = objra
XY[0,1] = objdec
xy_1 = wcs_1.all_world2pix(XY, 1)
xy_2 = wcs_2.all_world2pix(XY, 1)
if not quiet:
print('Object pixel coordinates in 1st HDU: {0:6.2f} {1:6.2f}'.format(xy_1[0,0], xy_1[0,1]))
print('Object pixel coordinates in 2nd HDU: {0:6.2f} {1:6.2f}'.format(xy_2[0,0], xy_2[0,1]))
print('Pixel shifts: {0} {1}'.format(xy_2[0,0]-xy_1[0,0], xy_2[0,1]-xy_1[0,1]))
return (xy_2[0,0]-xy_1[0,0]), (xy_2[0,1]-xy_1[0,1])
[docs]
def match_map_arrays(arr1, ext1, arr2, ext2, dx, dy, tol=1e-6, truncate=True, pixelscale=0.5):
"""
Match the centers and extent of two map arrays by applying the
previously calculated pixel offsets. See
:func:`map_center_pixel_offset`.
This is just a wrapper for the two supporting functions below.
Both maps *must* be two dimensional; i.e., a single map, not a set
of maps arranged in different channels.
Args:
arr1 (`numpy.ma.MaskedArray`_): Reference map data array.
ext1 (list): Reference map extent in arcseconds.
arr2 (`numpy.ma.MaskedArray`_): Comparison map data array.
ext1 (list): Comparison map extent in arcseconds.
dx (float): **Pixel** offset in the *first* dimension (x).
dy (float): **Pixel** offset in the *second* dimension (y).
tol (float): (**Optional**) Tolerance for non-integer pixel
shifts.
truncate (float): (**Optional**) Truncate the shape of the
shifted array to match the unshifted array. If False, the
unshifted array is padded to match the shifted array.
pixelscale (float): (**Optional**) The pixel scale (extent units
per array pixel) that *must be* common to both input arrays.
Returns:
numpy.ma.MaskedArray, list: The two matched arrays and the
extent that is common to both.
"""
# Must apply a sub-pixel shift
if abs(dx-numpy.around(dx))>tol or abs(dy-numpy.around(dy))>tol:
return _match_map_arrays_sub_pixel_shift(arr1, ext1, arr2, ext2, -dx, -dy,
truncate=truncate, pixelscale=pixelscale) \
if arr1.size > arr2.size else \
_match_map_arrays_sub_pixel_shift(arr2, ext2, arr1, ext1, dx, dy, swap=True,
truncate=truncate, pixelscale=pixelscale)
# Can just apply a whole pixel shift
return _match_map_arrays_integer_pixel_shift(arr1, ext1, arr2, ext2, -int(numpy.around(dx)),
-int(numpy.around(dy))) \
if arr1.size > arr2.size else \
_match_map_arrays_integer_pixel_shift(arr2, ext2, arr1, ext1,
int(numpy.around(dx)),
int(numpy.around(dy)), swap=True)
[docs]
def _match_map_arrays_sub_pixel_shift(arr1, ext1, arr2, ext2, dx, dy, swap=False, truncate=False,
pixelscale=0.5):
# Pad the second array with an even number of pixels
add = int(abs(dx))+1 if abs(dx) > abs(dy) else int(abs(dy))+1
add *= 2
if arr2.shape[0]+add - arr1.shape[0] > 0 and (arr2.shape[0]+add - arr1.shape[0]) % 2 > 1:
add += 2
newshape = tuple(numpy.array(arr2.shape) + add)
# Pad and copy the array and its mask
_arr2 = numpy.ma.masked_all(newshape, dtype=float)
_arr2[add//2:add//2+arr2.shape[0],add//2:add//2+arr2.shape[1]] = arr2[:,:]
_gpm = numpy.invert(numpy.ma.getmaskarray(_arr2)).astype(float)
# Shifts are in pixels, so adjust given the new coordinates
dx -= add//2
dy -= add//2
newext2 = list(numpy.array(ext2) + numpy.array([add/2, -add/2, -add/2, add/2])*pixelscale)
# Get the input and output coordinates
newx, newy = numpy.meshgrid(numpy.arange(newshape[0]), numpy.arange(newshape[1]), indexing='ij')
outcoo = numpy.array([newx.ravel(), newy.ravel()])
incoo = numpy.array([newx.ravel()-dx, newy.ravel()-dy])
indx = ~(incoo[0,:] < 0) & ~(incoo[1,:] < 0)
# Shift the second array and set the mask
shifted_arr2 = numpy.ma.masked_all(newshape, dtype=float)
shifted_arr2[outcoo[0,indx], outcoo[1,indx]] = \
ndimage.map_coordinates(_arr2.filled(fill_value=0.0), incoo[:,indx], order=1)
shifted_gpm = numpy.zeros(newshape, dtype=float)
shifted_gpm[outcoo[0,indx], outcoo[1,indx]] = \
ndimage.map_coordinates(_gpm, incoo[:,indx], order=1)
shifted_arr2[ shifted_gpm < 0.8 ] = numpy.ma.masked
# Return the two arrays
sub = arr1.shape[0]-shifted_arr2.shape[0]
if sub == 0:
return (shifted_arr2, arr1, ext1) if swap else (arr1, shifted_arr2, ext1)
# Pad the shifted array up to the shape of the first array
if sub > 0:
_shifted_arr2 = numpy.ma.masked_all(arr1.shape, dtype=float)
_shifted_arr2[:shifted_arr2.shape[0],:shifted_arr2.shape[1]] = shifted_arr2[:,:]
return (_shifted_arr2, arr1, ext1) if swap else (arr1, _shifted_arr2, ext1)
sub *= -1
# If truncation is requested, set the shifted 2nd array to the same
# size as the first input array
if truncate:
_shifted_arr2 = numpy.ma.masked_all(arr1.shape, dtype=float)
_shifted_arr2[:,:] = shifted_arr2[:arr1.shape[0],:arr1.shape[1]]
return (_shifted_arr2, arr1, ext1) if swap else (arr1, _shifted_arr2, ext1)
# Pad the first array to the size of the shifted array
_arr1 = numpy.ma.masked_all(shifted_arr2.shape, dtype=float)
_arr1[:arr1.shape[0],:arr1.shape[1]] = arr1[:,:]
return (shifted_arr2, _arr1, newext2) if swap else (_arr1, shifted_arr2, newext2)
# Arr1 must be bigger or the same size as Arr2
[docs]
def _match_map_arrays_integer_pixel_shift(arr1, ext1, arr2, ext2, dx, dy, swap=False):
if arr1.size < arr2.size:
raise ValueError('First array must be larger than the second.')
if abs(dx) > (arr1.shape[0] - arr2.shape[0])/2 or abs(dy) > (arr1.shape[1] - arr2.shape[1])/2:
raise ValueError('First array is not large enough to accommodate second array and shift.')
_arr = numpy.ma.MaskedArray(numpy.zeros(arr1.shape, dtype=float),
mask=numpy.ones(arr1.shape, dtype=bool))
_arr[dx:dx+arr2.shape[0],dy:dy+arr2.shape[1]] = arr2[:,:]
return (_arr, arr1, ext1) if swap else (arr1, _arr, ext1)
[docs]
def map_beam_patch(extent, ax, fwhm=2.5, pos=(0.1,0.1), **kwargs):
width = extent[0]-extent[1]
return patches.Circle(pos, fwhm/width/2, transform=ax.transAxes, **kwargs)
[docs]
def permute_wcs_axes(wcs, axes):
r"""
Permute the axes in an `astropy.wcs.WCS`_ object.
The number of axes must match the number axes defined by the
`astropy.wcs.WCS`_ object. Axes are iteratively swapped until
they're in the requested order. For example, to transpose the
axes of a 3D WCS, set ``axes=[2,1,0]``.
.. note::
Method uses the :func:`deepcopy` method of the
`astropy.wcs.WCS`_ object when altering and returning a new
object. If the axes are already sorted, this function returns
a deepcopy of the input ``wcs``.
Args:
wcs (`astropy.wcs.WCS`_):
Object with the world-coordinate system.
axes (array-like):
New locations for the current axes; i.e., the axis
currently at index ``i`` (0-indexed) is moved to be at
index ``axis[i]``. For example, to swap the axes for a 2D
WCS, set ``axes=[1,0]``.
Returns:
`astropy.wcs.WCS`_: The world-coordinate system with
re-ordered axes.
Raises:
ValueError:
Raised if the length of the ``axes`` vector is not the
same as the number of axes defined by the
`astropy.wcs.WCS`_ object, or if ``axes`` includes
undefined axes (i.e., a value less than 0 or >= the
number of axes in the WCS.)
"""
if len(axes) != wcs.wcs.naxis:
raise ValueError('Number of specified axes must match the WCS object.')
if not numpy.array_equal(numpy.arange(len(axes)), numpy.sort(axes)):
raise ValueError('Must select defined axes; 0 ... wcs.wcs.naxis-1.')
_wcs = wcs.deepcopy()
if numpy.array_equal(axes, numpy.sort(axes)):
return _wcs
_axes = numpy.atleast_1d(axes).copy()
for i,a in enumerate(_axes):
if a == i:
continue
_wcs = _wcs.swapaxes(a,i)
_axes[i], _axes[a] = i, a
return _wcs