Commit 7fe30704 authored by Roland Bacon's avatar Roland Bacon

Add wcs.coord()

parent d5b69237
Pipeline #3854 passed with stages
in 11 minutes
......@@ -725,6 +725,70 @@ class WCS:
np.clip(res, (0, 0), (self.naxis2 - 1, self.naxis1 - 1),
out=res)
return res
def coord(self, spaxel=False, relative=False, center=None,
polar=False, unit=None, reshape=False, mask=None):
"""Return the full coordinate array of the image
Parameters
----------
unit : `astropy.units.Unit`
Unit of the spatial coordinates
if None return pixels
relative : bool
If True, the coordinate are given relative to the center
center : tuple
center coordinate use for relative and radial mode
if none, the image center is used
must be given in unit
polar : bool
if True, return polar coordinate (r, theta)
unit : `astropy.units.Unit`
The units of the world coordinates
mask : 2D array of bool
If not None, return only non masked spaxels
reshape : bool
if True, return two 1D array,
if False, return two 2D array (meshgrid)
Returns
-------
out : array of array of float
array of [dec,ra], [p,q], [delta_dec, delta_ra], [delta_p, delta_q],
[r, theta]
"""
x,y = np.indices((self.naxis2,self.naxis1))
if not spaxel:
coord = np.mgrid[:self.naxis2, :self.naxis1].reshape(2, -1).T
coord = self.pix2sky(coord)
x,y = coord.T
x,y = np.meshgrid(x,y)
if mask is not None:
x = x[~mask]
y = y[~mask]
if relative or polar:
if center is None:
if spaxel:
center = 0.5*np.array([self.naxis2-1,self.naxis1-1])
else:
center = self.get_center()
x0,y0 = center
x = x - x0
y = y - y0
if (not spaxel) and (unit is not None):
x = UnitArray(x, self.unit, unit)
y = UnitArray(y, self.unit, unit)
if polar:
theta = np.arctan2(y, x)
rho = np.hypot(x, y)
x = rho
y = theta
if reshape:
x,y = np.vstack([x.ravel(), y.ravel()])
x = np.sort(np.unique(x))
y = np.sort(np.unique(y))
return x,y
def pix2sky(self, x, unit=None):
"""Convert image pixel indexes (y,x) to world coordinates (dec,ra).
......
......@@ -128,11 +128,13 @@ class TestWCS:
assert_array_equal(wcs.get_start(), [0.0, 0.0])
assert_array_equal(wcs.get_end(), [4.0, 5.0])
assert_array_equal(wcs.get_range(), [0.0, 0.0, 4.0, 5.0])
assert_array_equal(wcs.get_center(), [2.0, 2.5])
wcs2 = WCS(crval=(0, 0), shape=(5, 6))
assert_array_equal(wcs2.get_step(), [1.0, 1.0])
assert_array_equal(wcs2.get_start(), [-2.0, -2.5])
assert_array_equal(wcs2.get_end(), [2.0, 2.5])
assert_array_equal(wcs2.get_center(), [0, 0])
wcs2.set_step([0.5, 2.5])
assert_array_equal(wcs2.get_step(), [0.5, 2.5])
......@@ -144,7 +146,38 @@ class TestWCS:
assert wcs2.get_crval1(unit=2*u.pix) == -2.0
wcs2.set_crval2(-2, unit=2*u.pix)
assert wcs2.get_crval2(unit=2*u.pix) == -2.0
def test_coord(self):
"""WCS class: testing coord"""
wcs = WCS(shape=(4,5), crval=(-23.0,5.0), cdelt=(0.2/3600.,-0.2/3600.),
deg=True)
p,q = wcs.coord(spaxel=True, relative=True)
assert (p.mean() == 0) & (q.mean() == 0)
p,q = wcs.coord(spaxel=True, relative=True, reshape=True)
assert (p.shape == (4,)) & (q.shape == (5,))
p,q = wcs.coord(spaxel=True, relative=True, center=(5.2,6.3))
assert (p.min() == -5.2) & (q.min() == -6.3)
dec,ra = wcs.coord()
assert_allclose(ra.min(), 4.99987929, rtol=1e-6)
assert_allclose(dec.min(), -23.000083333, rtol=1e-6)
dec,ra = wcs.coord(relative=True, unit='arcsec')
assert_allclose(ra.min(), -0.4345444, rtol=1e-6)
assert_allclose(dec.min(),-0.3000001, rtol=1e-6)
r,theta = wcs.coord(polar=True, unit='arcsec')
assert_allclose(r.max(), 0.5280425, rtol=1e-6)
assert_allclose(theta.mean(),0.314159265, rtol=1e-6)
mask = np.empty(shape=(4,5), dtype=bool)
mask[:,:] = False
mask[0,0] = True
p,q = wcs.coord(spaxel=True, mask=mask)
assert (p.shape == (19,)) & (q.shape == (19,))
class TestWaveCoord:
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment