__author__ = "Matt Gomer"
import numpy as np
from scipy.interpolate import RectBivariateSpline
import matplotlib.pyplot as plt
[docs]
class KinNNImageAlign(object):
"""Class to rotate and interpolate SKiNN image (which is built along x-axis) to be
on the grid of the spectral data (e.g., MUSE), rotated to match light distribution
as defined by the lens model main function is interp_image() which will output a 2D
image interpolated on the spectra grid."""
[docs]
def __init__(self, spectra_inputs, imaging_inputs, kin_nn_inputs):
"""Initialize input data.
:param spectra_inputs: dictionary which encodes grid and transformation
information for kinematic data :'image': contains 2d image used to calculate
grid coordinates :'transform_pix2angle': transformation matrix to convert
from pixel xy to ra/dec :'ra_at_xy0': ra coordinate at pixel (0,0)
:'dec_at_xy0': dec coordinate at pixel (0,0)
:param imaging_inputs: dictionary which encodes grid and transformation
information for imaging data :'image': contains 2d image used to calculate
grid coordinates :'transform_pix2angle': transformation matrix to convert
from pixel xy to ra/dec :'ra_at_xy0': ra coordinate at pixel (0,0)
:'dec_at_xy0': dec coordinate at pixel (0,0) :'ellipse_PA': position angle
of ellipse axis relative to x direction :'offset_x': how many pixels to
offset the center of the grid to match the kinNN center (x-direction)
:'offset_y': how many pixels to offset the center of the grid to match the
kinNN center (y-direction)
:param kin_nn_inputs: dictionary which encodes grid information for NN output
data :'image': contains 2d image used to calculate grid coordinates
:'delta_pix': pixel size
"""
self.spectra_data = spectra_inputs
self.imaging_data = imaging_inputs
self.imaging_delta_pix = np.sqrt(
np.abs(np.linalg.det(self.imaging_data["transform_pix2angle"]))
)
self.kinNN_data = kin_nn_inputs
self.write_npix()
[docs]
def update(
self,
spectra_inputs=None,
imaging_inputs=None,
kin_nn_inputs=None,
update_npix=False,
):
"""Update with inputs."""
if spectra_inputs is not None:
self.spectra_data = spectra_inputs
if imaging_inputs is not None:
self.imaging_data = imaging_inputs
if kin_nn_inputs is not None:
self.kinNN_data = kin_nn_inputs
if update_npix:
self.write_npix()
[docs]
def write_npix(self):
"""Check that images are squared and write the keyword 'npix'."""
for input_set in [self.spectra_data, self.imaging_data, self.kinNN_data]:
# make sure each image is square and add npix to each dictionary
if "image" in input_set.keys():
if np.shape(input_set["image"])[0] != np.shape(input_set["image"])[1]:
raise ValueError("current version only works for square images")
npix = np.shape(input_set["image"])[0]
input_set["npix"] = npix
[docs]
def pix_coords(self, input_set, flatten=True):
"""Simple function to give pixel coordinates of grid.
:param input_set: dictionary from above (e.g. spectra_inputs) which completely
describes grid transformation :boolean flatten: default True; if True,
return 1D flattened output, if False, return 2D grid :return pixel
coordinates of grid
"""
x_grid = np.tile(np.arange(input_set["npix"]), input_set["npix"])
y_grid = np.repeat(np.arange(input_set["npix"]), input_set["npix"])
if flatten is True:
return x_grid, y_grid
else:
return x_grid.reshape(input_set["npix"], input_set["npix"]), y_grid.reshape(
input_set["npix"], input_set["npix"]
)
[docs]
def radec_to_xy(self, ra, dec, xy_to_radec_matrix, ra_atxy0, dec_atxy0):
"""Converts from radec to pixel coordinates.
:param ra: ra coordinate to transform
:param dec: dec coordinate to transform
:param xy_to_radec_matrix: transformation matrix to convert from pixel xy to
ra/dec
:param ra_atxy0: ra coordinate at pixel (0,0)
:param dec_atxy0: dec coordinate at pixel (0,0)
:return: x and y coordinates
"""
ra = ra - ra_atxy0
dec = dec - dec_atxy0
x, y = np.linalg.inv(xy_to_radec_matrix).dot(np.array([ra, dec]))
return x, y
[docs]
def xy_to_radec(self, x, y, xy_to_radec_matrix, ra_atxy0, dec_atxy0):
"""Converts from pixel coordinates to radec.
:param x: x coordinate to transform
:param y: y coordinate to transform
:param xy_to_radec_matrix: transformation matrix to convert from pixel xy to
ra/dec
:param ra_atxy0: ra coordinate at pixel (0,0)
:param dec_atxy0: dec coordinate at pixel (0,0)
:return: ra and dec coordinates
"""
ra, dec = xy_to_radec_matrix.dot(np.array([x, y]))
return ra + ra_atxy0, dec + dec_atxy0
[docs]
def rotate_imaging_into_kin_nn(
self,
imaging_x,
imaging_y,
ellipse_pa_to_imagingx_angle,
delta_pix_imaging,
delta_pix_kin_nn,
npix_imaging,
npix_kin_nn,
offsetx=0,
offsety=0,
):
"""Rotates and rescales from the x,y imaging coordinate system into the NN
coordinate system.
:param imaging_x: imaging x coordinate to transform
:param imaging_y: imaging y coordinate to transform
:param ellipse_pa_to_imagingx_angle: (radians) position angle of ellipse major
axis relative to x in the imaging coordinate system
:param delta_pix_imaging: pixel size of imaging image
:param delta_pix_kin_nn: pixel size of NN image
:param npix_imaging: number of pixels on a side of the imaging image
:param npix_kin_nn: number of pixels on a side of the NN image
:param offsetx: how many pixels to offset the center of the grid to match the
kinNN center (x-direction)
:param offsety: how many pixels to offset the center of the grid to match the
kinNN center (y-direction)
:return: x and y coordinates in NN coordinate system
"""
# define rotation matrix to rotate back into alignment
counterrotation = -ellipse_pa_to_imagingx_angle
cd1_1 = np.cos(counterrotation)
cd1_2 = -np.sin(counterrotation)
cd2_1 = np.sin(counterrotation)
cd2_2 = np.cos(counterrotation)
# rotation matrix, applied to matching centers ()
rotation_by_ellipse_angle = np.array([[cd1_1, cd1_2], [cd2_1, cd2_2]]) * (
delta_pix_imaging / delta_pix_kin_nn
)
(
kin_nn_x_at_imagingcenter,
kin_nn_y_at_imagingcenter,
) = rotation_by_ellipse_angle.dot(
np.array(
[
-npix_imaging / 2 - offsetx / delta_pix_imaging,
-npix_imaging / 2 - offsety / delta_pix_imaging,
]
)
) + [
npix_kin_nn / 2,
npix_kin_nn / 2,
]
kin_nn_x, kin_nn_y = rotation_by_ellipse_angle.dot(
np.array([imaging_x, imaging_y])
)
return (
kin_nn_x + kin_nn_x_at_imagingcenter,
kin_nn_y + kin_nn_y_at_imagingcenter,
)
[docs]
def plot_contour_and_grid(self, xcoords, ycoords, orig_image, color, alpha=0.4):
"""Plotting function for visualization of a grid with a single ellipse contour.
:param xcoords: grid x coordinates, flattened
:param ycoords: grid y coordinates, flattened
:param orig_image: original 2D image
:param color: plotting color
:param alpha: plotting alpha
"""
ell_cond_2d = np.isclose(orig_image, 0.1, rtol=5e-02)
ell_cond = ell_cond_2d.flatten()
plt.scatter(
xcoords[ell_cond], ycoords[ell_cond], color=color, alpha=alpha, s=10
)
plt.scatter(xcoords, ycoords, color=color, alpha=alpha, s=1)
[docs]
def spectragrid_in_radec(self):
"""Calculates ra and dec coordinates of spectra input coordinate grid.
:return: ra and dec coordinates
"""
spectra_x, spectra_y = self.pix_coords(self.spectra_data, flatten=True)
spectra_ra, spectra_dec = self.xy_to_radec(
spectra_x,
spectra_y,
self.spectra_data["transform_pix2angle"],
self.spectra_data["ra_at_xy0"],
self.spectra_data["dec_at_xy0"],
)
return spectra_ra, spectra_dec
[docs]
def spectragrid_in_imagingxy(self):
"""Calculates x and y coordinates in the imaging coordinate system of the
original spectra input grid.
:return: x and y coordinates
"""
spectra_ra, spectra_dec = self.spectragrid_in_radec()
spectra_coords_in_imaging_x, spectra_coords_in_imaging_y = self.radec_to_xy(
spectra_ra,
spectra_dec,
self.imaging_data["transform_pix2angle"],
self.imaging_data["ra_at_xy0"],
self.imaging_data["dec_at_xy0"],
)
return spectra_coords_in_imaging_x, spectra_coords_in_imaging_y
[docs]
def spectragrid_in_kin_nn_xy(self):
"""Calculates x and y coordinates in the NN coordinate system of the original
spectra input grid.
:return: x and y coordinates
"""
(
spectra_coords_in_imaging_x,
spectra_coords_in_imaging_y,
) = self.spectragrid_in_imagingxy()
kin_nn_x, kin_nn_y = self.rotate_imaging_into_kin_nn(
spectra_coords_in_imaging_x,
spectra_coords_in_imaging_y,
self.imaging_data["ellipse_PA"],
self.imaging_delta_pix,
self.kinNN_data["delta_pix"],
self.imaging_data["npix"],
self.kinNN_data["npix"],
offsetx=self.imaging_data["offset_x"],
offsety=self.imaging_data["offset_y"],
)
return kin_nn_x, kin_nn_y
[docs]
def interp_image(self):
"""Interpolates kinNN image at the coordinates of the transformed spectra grid.
:return: interpolated image which lines up with spectra coordinates
"""
spectra_kin_nn_x, spectra_kin_nn_y = self.spectragrid_in_kin_nn_xy()
x_axis = np.arange(self.kinNN_data["npix"])
y_axis = np.arange(self.kinNN_data["npix"])
interp_fcn = RectBivariateSpline(x_axis, y_axis, self.kinNN_data["image"])
# y and x are flipped in RectBivariateSpline call:
flux_interp = interp_fcn.ev(spectra_kin_nn_y, spectra_kin_nn_x).reshape(
self.spectra_data["npix"], self.spectra_data["npix"]
)
return flux_interp