class CCM89Extinction(StellarOperationModel):
operation_name = 'ccm89_extinction'
a_v = modeling.Parameter(default=0.0)
r_v = modeling.Parameter(default=3.1, fixed=True)
@property
def ebv(self):
return self.a_v / self.r_v
def __init__(self, a_v=0.0, r_v=3.1):
super(CCM89Extinction, self).__init__(a_v=a_v, r_v=r_v)
def evaluate(self, wavelength, flux, a_v, r_v):
from specutils import extinction
extinction_factor = np.ones_like(wavelength)
valid_wavelength = ((wavelength > 910) & (wavelength < 33333))
extinction_factor[valid_wavelength] = 10**(
-0.4 * extinction.extinction_ccm89(
wavelength[valid_wavelength] * u.angstrom,
a_v=np.abs(a_v),
r_v=np.abs(r_v)))
return wavelength, extinction_factor * flux
class RotationalBroadening(StellarOperationModel):
"""
The rotational broadening kernel was taken from
Observation and Analysis of Stellar Photospheres
by David Gray
"""
operation_name = 'rotation'
vrot = modeling.Parameter()
limb_darkening = modeling.Parameter(fixed=True, default=0.6)
@classmethod
def from_grid(cls, grid, vrot=0):
velocity_per_pix = getattr(grid, 'velocity_per_pix', None)
return cls(velocity_per_pix=velocity_per_pix, vrot=vrot)
def __init__(self, velocity_per_pix=None, vrot=0):
super(RotationalBroadening, self).__init__(vrot=vrot)
self.c_in_kms = const.c.to(u.km / u.s).value
if velocity_per_pix is not None:
self.log_sampling = True
self.velocity_per_pix = u.Quantity(velocity_per_pix,
u.km / u.s).value
else:
self.log_sampling = False
self.velocity_per_pix = None
def rotational_profile(self, vrot, limb_darkening):
vrot = float(vrot)
limb_darkening = float(limb_darkening)
vrot_by_c = np.maximum(0.0001, np.abs(vrot)) / self.c_in_kms
half_width_pix = np.round((vrot / self.velocity_per_pix)).astype(int)
profile_velocity = (np.linspace(-half_width_pix, half_width_pix,
2 * half_width_pix + 1) *
self.velocity_per_pix)
profile = np.maximum(0., 1. - (profile_velocity / vrot)**2)
profile = ((2 * (1 - limb_darkening) * np.sqrt(profile) +
0.5 * np.pi * limb_darkening * profile) /
(np.pi * vrot_by_c * (1. - limb_darkening / 3.)))
return profile / profile.sum()
def evaluate(self, wavelength, flux, v_rot, limb_darkening):
v_rot = np.asscalar(v_rot)
limb_darkening = np.asscalar(limb_darkening)
if self.velocity_per_pix is None:
raise NotImplementedError('Regridding not implemented yet')
if np.abs(v_rot) < 1e-5:
return wavelength, flux
profile = self.rotational_profile(v_rot, limb_darkening)
return wavelength, nd.convolve1d(flux, profile)
class single_schechter_model(modeling.Fittable1DModel):
lmstar = modeling.Parameter(default=10.8)
alpha = modeling.Parameter(default=-1.)
phistar = modeling.Parameter(default=10**-2.)
@staticmethod
def evaluate(x, lmstar, alpha, phistar):
return pylab.log(10) * phistar * 10**(
(x - lmstar) * (1 + alpha)) * pylab.exp(-10**(x - lmstar))
class double_schechter_model(modeling.Fittable1DModel):
lmstar = modeling.Parameter(default=10.8)
alpha1 = modeling.Parameter(default=-1.)
alpha2 = modeling.Parameter(default=-2.54)
phistar1 = modeling.Parameter(default=10**-2.)
phistar2 = modeling.Parameter(default=10**-4.)
@staticmethod
def evaluate(x, lmstar, alpha1, alpha2, phistar1, phistar2):
factor1 = pylab.log(10) * pylab.exp(-10**(x - lmstar)) * 10**(x -
lmstar)
factor2 = phistar1 * 10**(alpha1 *
(x - lmstar)) + phistar2 * 10**(alpha2 *
(x - lmstar))
return factor1 * factor2
class GenericPSFBackground(GenericPSFModel):
background_psf_amplitude = modeling.Parameter(bounds=(0, np.inf))
resolution = modeling.Parameter(default=10000)
matrix_parameter = ['background_psf_amplitude']
inputs = ()
outputs = ('row_id', 'column_id', 'value')
def __init__(self, pixel_to_wavelength, wavelength, sigma_impact_width=10.):
background_psf_amplitude = np.empty_like(wavelength) * np.nan
super(GenericPSFBackground, self).__init__(
background_psf_amplitude=background_psf_amplitude)
self._initialize_model(pixel_to_wavelength, wavelength)
impact_range = (
(wavelength.mean() / self.resolution) * FWHM_TO_SIGMA *
sigma_impact_width)
self.row_ids, self.column_ids = self._initialize_matrix_coordinates(
pixel_to_wavelength, wavelength, impact_range)
self.pixel_wavelength = pixel_to_wavelength[self.row_ids]
self.virt_pixel_wavelength = self.wavelength[self.column_ids]
def generate_design_matrix_coordinates(self, resolution):
row_ids = self.row_ids
column_ids = self.column_ids
sigma = (self.virt_pixel_wavelength / resolution) * FWHM_TO_SIGMA
norm_factor = 1 / (sigma * np.sqrt(2 * np.pi))
pixel_wavelength = self.pixel_wavelength
virt_pixel_wavelength = self.virt_pixel_wavelength
matrix_values = norm_factor * ne.evaluate(
'exp(-0.5 * '
'(pixel_wavelength - virt_pixel_wavelength)**2 '
'/ sigma**2)', )
return row_ids, column_ids, matrix_values
def generate_design_matrix(self, resolution):
row_ids, column_ids, matrix_values = (
self.generate_design_matrix_coordinates(resolution))
return sparse.coo_matrix((matrix_values, (row_ids, column_ids)))
def evaluate(self, background_level):
row_ids = self.pixel_table.pixel_id.values.astype(np.int64)
column_ids = self.pixel_table.column_ids.values.astype(np.int64)
matrix_values = self.pixel_table.sub_x.values
return row_ids, column_ids, matrix_values * 1.
class SpectralScaledChi2Likelihood(SpectralChi2Likelihood):
inputs = ('wavelength', 'flux')
outputs = ('loglikelihood', )
lnf = modeling.Parameter()
def __init__(self, observed, lnf=0.0):
super(SpectralScaledChi2Likelihood, self).__init__(observed, lnf=lnf)
def evaluate(self, wavelength, flux, lnf):
"""
Calculating likelihood and scaling the uncertainties (as shown in http://dfm.io/emcee/current/user/line/)
This likelihood function is simply a Gaussian where the variance is underestimated by some fractional amount: f.
One can fit for the natural logarithm of f.
Parameters
----------
wavelength : numpy.ndarray
wavelength
flux : numpy.ndarray
flux
Returns
-------
: float
log likelihood
"""
inv_sigma2 = 1.0 / (self.observed_uncertainty**2 +
self.observed_flux**2 * np.exp(2 * lnf))
loglikelihood = -0.5 * (np.sum(
(flux - self.observed_flux)**2 * inv_sigma2 - np.log(inv_sigma2)))
if np.isnan(loglikelihood):
return -1e300
else:
return loglikelihood
class SpectralChi2LikelihoodAddErr(StarKitModel):
## additive error model
inputs = ('wavelength', 'flux')
outputs = ('loglikelihood', )
add_err = modeling.Parameter(default=0.0)
def __init__(self, observed, add_err=0.0):
super(SpectralChi2LikelihoodAddErr, self).__init__(add_err=add_err)
self.observed_wavelength = observed.wavelength.to(u.angstrom).value
self.observed_flux = observed.flux.value
self.observed_uncertainty = getattr(observed, 'uncertainty', None)
if self.observed_uncertainty is not None:
self.observed_uncertainty = self.observed_uncertainty.value
else:
self.observed_uncertainty = np.ones_like(self.observed_wavelength)
def evaluate(self, wavelength, flux, add_err):
norm = 1.0 / np.sqrt(2.0 * np.pi *
(self.observed_uncertainty**2 + add_err**2))
loglikelihood = np.sum(
np.log(norm) + -0.5 *
(((self.observed_flux - flux)**2 /
(self.observed_uncertainty**2 + add_err**2))))
if np.isnan(loglikelihood):
return -1e300
return loglikelihood
class PolynomialBackground(LinearLeastSquaredModel):
background_level = modeling.Parameter(bounds=(0, np.inf))
matrix_parameter = ['background_level']
inputs = ()
outputs = ('row_id', 'column_id', 'value')
def __init__(self, pixel_table, wavelength_pixels):
background_level = np.empty_like(
pixel_table.wavelength_pixel_id.unique())
super(GenericBackground, self).__init__(background_level)
self._initialize_lls_model(pixel_table, wavelength_pixels)
def generate_design_matrix_coordinates(self):
row_ids = self.pixel_table.pixel_id.values
column_ids = self.pixel_table.wavelength_pixel_id.values
matrix_values = self.pixel_table.sub_x.values
return row_ids, column_ids, matrix_values * 1.
def generate_design_matrix(self):
row_ids, column_ids, matrix_values = (
self.generate_design_matrix_coordinates())
return sparse.coo_matrix((matrix_values, (row_ids, column_ids)))
def evaluate(self, background_level):
row_ids = self.pixel_table.pixel_id.values.astype(np.int64)
column_ids = self.pixel_table.column_ids.values.astype(np.int64)
matrix_values = self.pixel_table.sub_x.values
return row_ids, column_ids, matrix_values * 1.
def load_telluric_grid(hdf_fname,
stellar_grid=None,
wavelength_type=None,
base_class=BaseTelluricGrid):
"""
Load the grid from an HDF file
Parameters
----------
hdf_fname: ~str
filename and path to the HDF file
stellar_grid: BaseSpectralGrid
spectral_grid to adapt to
wavelength_type: str
use 'air' or 'vacuum' wavelength and convert if necessary (by inspecting
what the grid uses in meta)
Returns
-------
: SpectralGrid object
"""
wavelength, meta, index, fluxes = read_grid(hdf_fname)
class_dict = {}
class_dict['vrad_telluric'] = modeling.Parameter(default=0.0)
class_dict, initial_parameters = construct_grid_class_dict(
meta, index, class_dict=class_dict)
class_dict['__init__'] = base_class.__init__
if stellar_grid is not None:
class_dict['inputs'] = ('wavelength', 'flux')
class_dict['evaluate'] = base_class.evaluate_transmission
initial_parameters['target_wavelength'] = stellar_grid.wavelength
initial_parameters['target_R'] = stellar_grid.R
if wavelength_type is not None and stellar_grid is not None:
logger.warn(
'Ignoring requested wavelength_type in favour of spectral_grid wavelength_type'
)
initial_parameters['wavelength_type'] = meta['wavelength_type']
initial_parameters['target_wavelength_type'] = stellar_grid.meta_grid[
'wavelength_type']
else:
class_dict['inputs'] = tuple()
class_dict['evaluate'] = base_class.evaluate_raw
initial_parameters['wavelength_type'] = wavelength_type
TelluricGrid = type('TelluricGrid', (base_class, ), class_dict)
logger.info('Initializing spec grid')
telluric_grid = TelluricGrid(wavelength, index[meta['parameters']].values,
fluxes, meta, **initial_parameters)
return telluric_grid
class InstrumentConvolveGrism(SpectrographOperationModel):
"""
Convolve with a gaussian with given resolution to mimick an instrument
assuming delta_lambda being constant
Parameters
----------
R : float or astropy.units.Quantity (unitless)
resolution of the spectrum R = lambda/delta_lambda at wavelength
sampling: float
number of pixels per resolution element (default=2.)
"""
operation_name = 'resolution'
R = modeling.Parameter()
requires_observed_spectrum = False
@classmethod
def from_grid(cls, wavelength, grid, R=np.inf):
grid_R = getattr(grid, 'R', None)
return cls(wavelength, R=R, grid_R=grid_R)
def __init__(
self,
wavelength,
R,
grid_R,
sampling=4,
):
super(InstrumentConvolveGrism, self).__init__(R=R)
self.wavelength = wavelength
self.sampling = sampling
self.fwhm2sigma = 1 / (2 * np.sqrt(np.log(2) * 2))
self.grid_R = grid_R
def evaluate(self, wavelength, flux, R):
if np.isinf(R):
return wavelength, flux
rescaled_R = 1 / np.sqrt((1 / R)**2 - (1 / self.grid_R)**2)
delta_lambda = (self.wavelength / rescaled_R) * self.fwhm2sigma
new_wavelength = np.arange(wavelength[0], wavelength[-1],
delta_lambda / self.sampling)
new_flux = np.interp(new_wavelength, wavelength, flux)
return new_wavelength, nd.gaussian_filter1d(new_flux, self.sampling)
class DopplerShift(StellarOperationModel):
operation_name = 'doppler'
vrad = modeling.Parameter()
def __init__(self, vrad):
super(DopplerShift, self).__init__(vrad=vrad)
self.c_in_kms = const.c.to(u.km / u.s).value
def evaluate(self, wavelength, flux, vrad):
beta = vrad / self.c_in_kms
doppler_factor = np.sqrt((1 + beta) / (1 - beta))
return wavelength * doppler_factor, flux
class InstrumentDeltaLambdaConstant(SpectrographOperationModel):
"""
Convolve with a gaussian with given resolution to mimick an instrument
assuming delta_lambda being constant
Parameters
----------
R : float or astropy.units.Quantity (unitless)
resolution of the spectrum R = lambda/delta_lambda at wavelength
sampling: float
number of pixels per resolution element (default=2.)
"""
operation_name = 'delta_lambda_constant'
delta_lambda = modeling.Parameter()
requires_observed_spectrum = False
@classmethod
def from_grid(cls, grid, delta_lambda=0.0):
grid_R = getattr(grid, 'R', None)
return cls(delta_lambda, grid_R=grid_R)
def __init__(self, delta_lambda, grid_R, sampling=4):
super(InstrumentDeltaLambdaConstant,
self).__init__(delta_lambda=delta_lambda)
self.sampling = sampling
self.fwhm2sigma = 1 / (2 * np.sqrt(np.log(2) * 2))
self.grid_R = grid_R
def evaluate(self, wavelength, flux, delta_lambda):
if np.isclose(delta_lambda, 0.0):
return wavelength, flux
sigma_lambda = delta_lambda * self.fwhm2sigma
new_wavelength = np.arange(wavelength[0], wavelength[-1],
sigma_lambda / self.sampling)
new_flux = np.interp(new_wavelength, wavelength, flux)
return new_wavelength, nd.gaussian_filter1d(new_flux, self.sampling)
class InstrumentConvolveGrating(SpectrographOperationModel):
"""
Convolve with a gaussian with given resolution to mimick an instrument
assuming lambda / delta_lambda being constant
Parameters
----------
R : float or astropy.units.Quantity (unitless)
resolution of the spectrum R = lambda/delta_lambda
sampling: float
number of pixels per resolution element (default=2.)
"""
operation_name = 'resolution'
R = modeling.Parameter()
requires_observed_spectrum = False
@classmethod
def from_grid(cls, grid, R=np.inf):
grid_R = getattr(grid, 'R', None)
grid_sampling = getattr(grid, 'R_sampling', None)
return cls(R=R, grid_R=grid_R, grid_sampling=grid_sampling)
def __init__(self, R=np.inf, grid_R=None, grid_sampling=None):
super(InstrumentConvolveGrating, self).__init__(R=R)
self.grid_sampling = grid_sampling
self.grid_R = grid_R
def evaluate(self, wavelength, flux, R):
if np.isinf(R):
return wavelength, flux
if self.grid_R is None:
raise NotImplementedError('grid_R not given - this mode is not '
'implemented yet')
rescaled_R = 1 / np.sqrt((1 / R)**2 - (1 / self.grid_R)**2)
sigma = ((self.grid_R / rescaled_R) * self.grid_sampling /
(2 * np.sqrt(2 * np.log(2))))
return wavelength, nd.gaussian_filter1d(flux, sigma)
class InstrumentRConstant(SpectrographOperationModel):
"""
Convolve with a gaussian with given resolution to mimick an instrument
assuming lambda / delta_lambda being constant
Parameters
----------
R : float or astropy.units.Quantity (unitless)
resolution of the spectrum R = lambda/delta_lambda
sampling: float
number of pixels per resolution element (default=2.)
"""
operation_name = 'resolution_constant'
R = modeling.Parameter()
requires_observed_spectrum = False
@classmethod
def from_grid(cls, grid, R=np.inf):
grid_R = getattr(grid, 'R', None)
grid_sampling = getattr(grid, 'R_sampling', None)
return cls(R=R, grid_R=grid_R, grid_sampling=grid_sampling)
def __init__(self, R=np.inf, grid_R=None, grid_sampling=None):
super(InstrumentRConstant, self).__init__(R=R)
self.grid_sampling = grid_sampling
self.grid_R = grid_R
def evaluate(self, wavelength, flux, R):
R = float(np.squeeze(R))
if np.isinf(R):
return wavelength, flux
if self.grid_R is None:
raise NotImplementedError('grid_R not given - '
'this mode is not implemented yet')
convolved_flux = convolve_to_resolution(flux, self.grid_R,
self.grid_sampling, R)
return wavelength, convolved_flux
class Distance(StellarOperationModel):
operation_name = 'distance'
distance = modeling.Parameter(default=10.0) # in pc
@classmethod
def from_grid(cls, grid, distance=10.):
lum_density2cgs = grid.flux_unit.to('erg / (s * angstrom)')
return cls(distance=distance, lum_density2cgs=lum_density2cgs)
def __init__(self, distance, lum_density2cgs=1.):
super(Distance, self).__init__(distance=distance)
self.pc2cm = u.pc.to(u.cm)
self.lum_density2cgs = lum_density2cgs
def evaluate(self, wavelength, flux, distance):
conversion = self.lum_density2cgs / (4 * np.pi *
(distance * self.pc2cm)**2)
return wavelength, flux * conversion
class SlopedMoffatTrace(MoffatTrace):
amplitude = modeling.Parameter(bounds=(0, np.inf))
trace_pos = modeling.Parameter(default=0.0, bounds=(-6, 6))
trace_slope = modeling.Parameter(default=0.0, bounds=(-0.5, 0.5))
sigma = modeling.Parameter(default=1.0, bounds=(0, 99))
sigma_slope = modeling.Parameter(default=0.0, bounds=(-0.5, 0.5))
beta = modeling.Parameter(default=1.5, fixed=True, bounds=(1.1, 3.))
matrix_parameter = ['amplitude']
def __init__(self, pixel_table, wavelength_pixels):
amplitude = np.empty_like(pixel_table.wavelength_pixel_id.unique())
super(MoffatTrace, self).__init__(amplitude=amplitude * np.nan)
self._initialize_lls_model(pixel_table, wavelength_pixels)
def generate_design_matrix_coordinates(self, trace_pos, trace_slope, sigma,
sigma_slope, beta):
row_ids = self.pixel_table.pixel_id.values
column_ids = self.pixel_table.wavelength_pixel_id.values
matrix_values = self.pixel_table.sub_x.values
normed_wavelength = self.pixel_table.normed_wavelength.values.copy()
varying_trace_pos = (trace_pos + trace_slope * normed_wavelength)
varying_sigma = (sigma + sigma_slope * 0.5 * (normed_wavelength + 1))
moffat_profile = self._moffat(self.pixel_table.slit_pos.values,
varying_trace_pos, varying_sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
def generate_design_matrix(self, trace_pos, trace_slope, sigma,
sigma_slope, beta):
row_ids, column_ids, matrix_values = (
self.generate_design_matrix_coordinates(trace_pos, trace_slope,
sigma, sigma_slope, beta))
return sparse.coo_matrix((matrix_values, (row_ids, column_ids)))
def evaluate(self, amplitude, trace_pos, trace_slope, sigma, beta):
row_ids = self.pixel_table.pixel_id.values.astype(np.int64)
column_ids = self.pixel_table.wavelength_pixel_id.values.astype(
np.int64)
matrix_values = self.pixel_table.sub_x.values
varying_trace_pos = (
trace_pos + trace_slope * self.pixel_table.normaled_wavelength)
moffat_profile = self._moffat(self.pixel_table.slit_pos.values,
varying_trace_pos, sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
class PolynomialMoffatTrace(MoffatTrace):
amplitude = modeling.Parameter(bounds=(0, np.inf))
trace0 = modeling.Parameter(default=1.0, bounds=(-6, 6))
trace1 = modeling.Parameter(default=0.0, bounds=(-6, 6))
trace2 = modeling.Parameter(default=0.0, bounds=(-6, 6))
trace3 = modeling.Parameter(default=0.0, bounds=(-6, 6))
trace4 = modeling.Parameter(default=0.0, bounds=(-6, 6))
sigma0 = modeling.Parameter(default=1.0, bounds=(0, 99))
sigma1 = modeling.Parameter(default=0.0, bounds=(0, 99))
sigma2 = modeling.Parameter(default=1.0, bounds=(0, 99))
sigma3 = modeling.Parameter(default=1.0, bounds=(0, 99))
sigma4 = modeling.Parameter(default=1.0, bounds=(0, 99))
beta = modeling.Parameter(default=1.5, fixed=True, bounds=(1.1, 3.))
matrix_parameter = ['amplitude']
def __init__(self, pixel_table, wavelength_pixels, **kwargs):
amplitude = np.empty_like(pixel_table.wavelength_pixel_id.unique())
super(MoffatTrace, self).__init__(amplitude=amplitude * np.nan,
**kwargs)
self._initialize_lls_model(pixel_table, wavelength_pixels)
def generate_design_matrix_coordinates(self, trace0, trace1, trace2,
trace3, trace4, sigma0, sigma1,
sigma2, sigma3, sigma4, beta):
row_ids = self.pixel_table.pixel_id.values
column_ids = self.pixel_table.wavelength_pixel_id.values
matrix_values = self.pixel_table.sub_x.values
normed_wavelength = self.pixel_table.normed_wavelength.values.copy()
varying_trace_pos = np.polyval(
[trace4, trace3, trace2, trace1, trace0], normed_wavelength)
varying_sigma = np.abs(
np.polyval([sigma3, sigma2, sigma1, sigma0], normed_wavelength))
moffat_profile = self._moffat(self.pixel_table.slit_pos.values,
varying_trace_pos, varying_sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
def generate_design_matrix(self, trace0, trace1, trace2, trace3, trace4,
sigma0, sigma1, sigma2, sigma3, sigma4, beta):
row_ids, column_ids, matrix_values = (
self.generate_design_matrix_coordinates(trace0, trace1, trace2,
trace3, trace4, sigma0,
sigma1, sigma2, sigma3,
sigma4, beta))
return sparse.coo_matrix((matrix_values, (row_ids, column_ids)))
class MoffatTrace(LinearLeastSquaredModel):
amplitude = modeling.Parameter(bounds=(0, np.inf))
trace_pos = modeling.Parameter(default=0.0, bounds=(-6, 6))
sigma = modeling.Parameter(default=1.0, bounds=(0, 99))
beta = modeling.Parameter(default=1.5, fixed=True, bounds=(1.1, 3))
matrix_parameter = ['amplitude']
def __init__(self, pixel_table, wavelength_pixels):
amplitude = np.empty_like(pixel_table.wavelength_pixel_id.unique())
super(MoffatTrace, self).__init__(amplitude=amplitude * np.nan)
self._initialize_lls_model(pixel_table, wavelength_pixels)
def generate_design_matrix_coordinates(self, trace_pos, sigma, beta):
row_ids = self.pixel_table.pixel_id.values
column_ids = self.pixel_table.wavelength_pixel_id.values
matrix_values = self.pixel_table.sub_x.values.copy()
moffat_profile = self._moffat(self.pixel_table.slit_pos.values,
trace_pos, sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
def generate_design_matrix(self, trace_pos, sigma, beta):
row_ids, column_ids, matrix_values = (
self.generate_design_matrix_coordinates(trace_pos, sigma, beta))
return sparse.coo_matrix((matrix_values, (row_ids, column_ids)))
def evaluate(self, amplitude, trace_pos, sigma, beta):
row_ids = self.pixel_table.pixel_id.values.astype(np.int64)
column_ids = self.pixel_table.wavelength_pixel_id.values.astype(
np.int64)
matrix_values = self.pixel_table.sub_x.values.copy()
moffat_profile = self._moffat(self.pixel_table.slit_pos.values,
trace_pos, sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
def to_spectrum(self):
try:
from specutils import Spectrum1D
except ImportError:
raise ImportError('specutils needed for this functionality')
from xtool.fix_spectrum1d import Spectrum1D
if getattr(self, 'amplitude_uncertainty', None) is None:
uncertainty = None
else:
uncertainty = self.amplitude_uncertainty
spec = Spectrum1D.from_array(self.wavelength_pixels * u.nm,
self.amplitude.value,
uncertainty=uncertainty)
return spec
@staticmethod
def _moffat(s, s0, sigma, beta=1.5):
"""
Calculate the moffat profile
Parameters
----------
s : ndarray
slit position
s0 : float
center of the Moffat profile
sigma : float
sigma of the Moffat profile
beta : float, optional
beta parameter of the Moffat profile (default = 1.5)
Returns
-------
"""
beta = getattr(beta, 'value', beta)
fwhm = sigma * SIGMA_TO_FWHM
alpha = fwhm / (2 * np.sqrt(2**(1.0 / float(beta)) - 1.0))
norm_factor = (beta - 1.0) / (np.pi * alpha**2)
return norm_factor * (1.0 + ((s - s0) / alpha)**2)**(-beta)
class PSFPolynomialMoffatTrace(PSFMoffatTrace):
psf_amplitude = modeling.Parameter(bounds=(0, np.inf))
resolution_trace = modeling.Parameter(default=10000, bounds=(1000, 20000))
trace0 = modeling.Parameter(default=1.0, bounds=(-6, 6))
trace1 = modeling.Parameter(default=0.0, bounds=(-6, 6))
trace2 = modeling.Parameter(default=0.0, bounds=(-6, 6))
trace3 = modeling.Parameter(default=0.0, bounds=(-6, 6))
trace4 = modeling.Parameter(default=0.0, bounds=(-6, 6))
sigma0 = modeling.Parameter(default=1.0, bounds=(0, 99))
sigma1 = modeling.Parameter(default=0.0, bounds=(0, 99))
sigma2 = modeling.Parameter(default=0.0, bounds=(0, 99))
sigma3 = modeling.Parameter(default=0.0, bounds=(0, 99))
sigma4 = modeling.Parameter(default=0.0, bounds=(0, 99))
beta = modeling.Parameter(default=1.5, fixed=True, bounds=(1.1, 3.))
matrix_parameter = ['psf_amplitude']
def __init__(self, pixel_to_wavelength, pixel_to_slit_pos, wavelength,
sigma_impact_width=10.):
psf_amplitude = np.empty_like(wavelength) * np.nan
super(PSFMoffatTrace, self).__init__(
psf_amplitude=psf_amplitude)
self._initialize_model(pixel_to_wavelength, wavelength)
impact_range = (
(wavelength.mean() / self.resolution_trace) * FWHM_TO_SIGMA *
sigma_impact_width)
self.row_ids, self.column_ids = self._initialize_matrix_coordinates(
pixel_to_wavelength, wavelength, impact_range)
self.pixel_wavelength = pixel_to_wavelength[self.row_ids]
self.normed_pixel_wavelength = (
(self.pixel_to_wavelength - self.pixel_wavelength.min()) /
(self.pixel_wavelength.max() - self.pixel_wavelength.min()))
self.normed_pixel_wavelength = (
(self.normed_pixel_wavelength - 0.5) * 2)[self.row_ids]
self.virt_pixel_wavelength = self.wavelength[self.column_ids]
self.slit_pos = pixel_to_slit_pos[self.row_ids]
def generate_design_matrix_coordinates(
self, resolution_trace, trace0, trace1, trace2, trace3, trace4,
sigma0, sigma1, sigma2, sigma3, sigma4, beta):
row_ids = self.row_ids
column_ids = self.column_ids
psf_sigma = (
(self.virt_pixel_wavelength / resolution_trace) * FWHM_TO_SIGMA)
norm_factor = 1 / (psf_sigma * np.sqrt(2 * np.pi))
pixel_wavelength = self.pixel_wavelength
virt_pixel_wavelength = self.virt_pixel_wavelength
matrix_values = norm_factor * ne.evaluate(
'exp(-0.5 * '
'(pixel_wavelength - virt_pixel_wavelength)**2 '
'/ psf_sigma**2)', )
normed_wavelength = self.normed_pixel_wavelength
varying_trace_pos = np.polyval([trace4, trace3, trace2, trace1, trace0],
normed_wavelength)
varying_sigma = np.abs(np.polyval([sigma4, sigma3, sigma2, sigma1, sigma0],normed_wavelength))
moffat_profile = self._moffat(self.slit_pos,
varying_trace_pos, varying_sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
def generate_design_matrix(
self, resolution_trace, trace0, trace1, trace2, trace3, trace4,
sigma0, sigma1, sigma2, sigma3, sigma4, beta):
row_ids, column_ids, matrix_values = (
self.generate_design_matrix_coordinates(
resolution_trace, trace0, trace1, trace2, trace3, trace4,
sigma0, sigma1, sigma2, sigma3, sigma4, beta))
return sparse.coo_matrix((matrix_values, (row_ids, column_ids)))
class PSFMoffatTrace(GenericPSFModel):
psf_amplitude = modeling.Parameter(bounds=(0, np.inf))
resolution_trace = modeling.Parameter(default=10000)
trace_pos = modeling.Parameter(default=0.0, bounds=(-6, 6))
sigma = modeling.Parameter(default=1.0, bounds=(0, 99))
beta = modeling.Parameter(default=1.5, fixed=True, bounds=(1.1, 3))
matrix_parameter = ['psf_amplitude']
def __init__(self, pixel_to_wavelength, pixel_to_slit_pos, wavelength, sigma_impact_width=10.):
psf_amplitude = np.empty_like(wavelength) * np.nan
super(PSFMoffatTrace, self).__init__(
psf_amplitude=psf_amplitude)
self._initialize_model(pixel_to_wavelength, wavelength)
impact_range = (
(wavelength.mean() / self.resolution_trace) * FWHM_TO_SIGMA *
sigma_impact_width)
self.row_ids, self.column_ids = self._initialize_matrix_coordinates(
pixel_to_wavelength, wavelength, impact_range)
self.pixel_wavelength = pixel_to_wavelength[self.row_ids]
self.virt_pixel_wavelength = self.wavelength[self.column_ids]
self.slit_pos = pixel_to_slit_pos[self.row_ids]
def generate_design_matrix_coordinates(self, resolution_trace, trace_pos,
sigma, beta):
row_ids = self.row_ids
column_ids = self.column_ids
psf_sigma = (
(self.virt_pixel_wavelength / resolution_trace) * FWHM_TO_SIGMA)
norm_factor = 1 / (psf_sigma * np.sqrt(2 * np.pi))
pixel_wavelength = self.pixel_wavelength
virt_pixel_wavelength = self.virt_pixel_wavelength
matrix_values = norm_factor * ne.evaluate(
'exp(-0.5 * '
'(pixel_wavelength - virt_pixel_wavelength)**2 '
'/ sigma**2)', )
moffat_profile = self._moffat(self.slit_pos, trace_pos, sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
def generate_design_matrix(self, resolution_trace, trace_pos, sigma, beta):
row_ids, column_ids, matrix_values = (
self.generate_design_matrix_coordinates(
resolution_trace, trace_pos, sigma, beta))
return sparse.coo_matrix((matrix_values, (row_ids, column_ids)))
def evaluate(self, amplitude, trace_pos, sigma, beta):
row_ids = self.pixel_table.pixel_id.values.astype(np.int64)
column_ids = self.pixel_table.wavelength_pixel_id.values.astype(np.int64)
matrix_values = self.pixel_table.sub_x.values.copy()
moffat_profile = self._moffat(self.pixel_table.slit_pos.values, trace_pos, sigma, beta)
return row_ids, column_ids, matrix_values * moffat_profile
def to_spectrum(self):
try:
from specutils import Spectrum1D
except ImportError:
raise ImportError('specutils needed for this functionality')
from xtool.fix_spectrum1d import Spectrum1D
if getattr(self, 'amplitude_uncertainty', None) is None:
uncertainty = None
else:
uncertainty = self.amplitude_uncertainty
spec = Spectrum1D.from_array(
self.wavelength * u.nm, self.psf_amplitude.value,
uncertainty=uncertainty)
return spec
@staticmethod
def _moffat(s, s0, sigma, beta=1.5):
"""
Calculate the moffat profile
Parameters
----------
s : ndarray
slit position
s0 : float
center of the Moffat profile
sigma : float
sigma of the Moffat profile
beta : float, optional
beta parameter of the Moffat profile (default = 1.5)
Returns
-------
"""
fwhm = sigma * SIGMA_TO_FWHM
alpha = fwhm / (2 * np.sqrt(2**(1.0 / float(beta)) - 1.0))
norm_factor = (beta - 1.0) / (np.pi * alpha**2)
return norm_factor * (1.0 + ((s - s0) / alpha)**2)**(-beta)
class LUTOrderWCS(modeling.Model):
inputs = ('x', 'y')
outputs = ('wave', 'slit')
pix_to_wave_grid = modeling.Parameter()
pix_to_slit_grid = modeling.Parameter()
def __init__(self, transform_pixel_to_wavelength, transform_pixel_to_slit,
mask):
"""
A Look Up Table (LUT) World Coordinate System (WCS) that can transform
from pixel coordinate input to wavelength and slit coordinates
Parameters
----------
transform_pixel_to_wavelength : astropy.units.Quantity
transform_pixel_to_slit : astropy.units.Quantity
mask : np.ndarray
"""
self.mask = mask
self.x_grid, self.y_grid = self._generate_coordinate_grid(mask)
self.pix_to_wave_ma = np.ma.MaskedArray(
transform_pixel_to_wavelength.value, ~mask, fill_value=np.nan)
self.pix_to_slit_ma = np.ma.MaskedArray(transform_pixel_to_slit.value,
~mask,
fill_value=np.nan)
self.pix_to_wave_unit = transform_pixel_to_wavelength.unit
self.pix_to_slit_unit = transform_pixel_to_slit.unit
super(LUTOrderWCS, self).__init__(self.pix_to_wave_ma.filled(),
self.pix_to_slit_ma.filled())
self.standard_broadcasting = False
@property
def x(self):
return self.x_grid.compressed()
@property
def y(self):
return self.y_grid.compressed()
@property
def pix_to_wave(self):
return self.pix_to_wave_ma.compressed()
@property
def pix_to_slit(self):
return self.pix_to_slit_ma.compressed()
def evaluate(self, x, y, pix_to_wave, pix_to_slit):
x = np.int64(x)
y = np.int64(y)
return np.squeeze(pix_to_wave)[y, x], np.squeeze(pix_to_slit)[y, x]
@staticmethod
def _generate_coordinate_grid(mask):
"""
generate a coordinate grid for the order slice
Parameters
----------
mask : numpy.ndarray
boolean masked with True for valid data and false for invalid data
Returns
-------
x_grid : numpy.ma.MaskedArray
y_grid : numpy.ma.MaskedArray
"""
y_grid, x_grid = np.mgrid[:mask.shape[0], :mask.shape[1]]
return (np.ma.MaskedArray(x_grid, ~mask, fill_value=-1),
np.ma.MaskedArray(y_grid, ~mask, fill_value=-1))
def _update_mask(self, mask):
self.pix_to_wave_ma.mask = mask
self.pix_to_slit_ma.mask = mask
self.x_grid.mask = mask
self.y_grid.mask = mask
self.pix_to_wave_grid = self.pix_to_wave_ma.filled()
self.pix_to_slit_grid = self.pix_to_slit_ma.filled()
class PolynomialOrderWCS(modeling.Model):
inputs = ('x', 'y')
outputs = ('wave', 'slit')
wave_transform_coef = modeling.Parameter()
slit_transform_coef = modeling.Parameter()
@classmethod
def from_lut_order_wcs(cls, lut_order_wcs, poly_order=(2, 3)):
"""
Fits a polynomial on n-th degree to the WCS transform
Parameters
----------
lut_order_wcs : LUTOrderWCS
Lookup Table WCS to be fit
poly_order : int or tuple, optional
tuple consisting of two integers describing the polynomial in
wave and slit default=(2, 3)
Returns
-------
polynomial_order_wcs: PolynomialOrderWCS
"""
if not hasattr(poly_order, '__iter__'):
poly_order = (poly_order, poly_order)
wave_model_coef = cls.polyfit2d(lut_order_wcs.x, lut_order_wcs.y,
lut_order_wcs.pix_to_wave, poly_order)
slit_model_coef = cls.polyfit2d(lut_order_wcs.x, lut_order_wcs.y,
lut_order_wcs.pix_to_slit, poly_order)
return cls(wave_model_coef, slit_model_coef)
def __init__(self, wave_transform_coef, slit_transform_coef):
super(PolynomialOrderWCS, self).__init__(wave_transform_coef,
slit_transform_coef)
self.standard_broadcasting = False
@staticmethod
def polyfit2d(x, y, f, deg):
"""
Fits a 2D polynomial and returns the coefficients
Parameters
----------
x : numpy.ndarray
y : numpy.ndarray
f : numpy.ndarray
deg : tuple
Returns
-------
poly_2d_coef : numpy.ndarray
"""
deg = np.asarray(deg)
vander = polynomial.polyvander2d(x, y, deg)
vander = vander.reshape((-1, vander.shape[-1]))
f = f.reshape((vander.shape[0], ))
c = np.linalg.lstsq(vander, f)[0]
return c.reshape(deg + 1)
def evaluate(self, x, y, wave_transform_coef, slit_transform_coef):
x = np.squeeze(x)
y = np.squeeze(x)
wave_transform_coef = np.squeeze(wave_transform_coef)
slit_transform_coef = np.squeeze(slit_transform_coef)
return (polynomial.polyval2d(x, y, wave_transform_coef),
polynomial.polyval2d(x, y, slit_transform_coef))
class VirtualPixelWavelength(modeling.Model):
inputs = ()
outputs = ('pixel_table', )
wave_transform_coef = modeling.Parameter()
wavelength_sampling_defaults = {'UVB': 0.04, 'VIS': 0.04, 'NIR': 0.03}
@classmethod
def from_order(cls,
order,
poly_order=(2, 3),
wavelength_sampling=None,
sub_sampling=5):
"""
Instantiate a Virtualpixel table from the order raw Lookup Table WCS
and then fitting a Polynomial WCS with given orde
Parameters
----------
order : xtool.data.Order
order data or object
poly_order : tuple or int
wavelength_sampling: float, optional
float for wavelength spacing. If `None` will use for different arms
* UVB - 0.04 nm
* VIS - 0.04 nm
* NIR - 0.03 nm
Returns
-------
VirtualPixelWavelength
"""
polynomial_wcs = PolynomialOrderWCS.from_lut_order_wcs(
order.wcs, poly_order)
if wavelength_sampling is None:
wavelength_sampling = cls.wavelength_sampling_defaults[
order.instrument_arm]
wavelength_bins = np.arange(
order.wcs.pix_to_wave.min(),
order.wcs.pix_to_wave.max() + wavelength_sampling,
wavelength_sampling)
return VirtualPixelWavelength(polynomial_wcs,
order.wcs,
wavelength_bins,
sub_sampling=sub_sampling)
def __init__(self,
polynomial_wcs,
lut_wcs,
raw_wavelength_pixels,
sub_sampling=5):
"""
A model that represents the virtual pixels of the model
Parameters
----------
polynomial_wcs : xtool.wcs.PolynomialOrderWCS
lut_wcs : xtool.wcs.LUTOrderWCS
raw_wavelength_pixels : numpy.ndarray
sub_sampling : int
how many subsamples to create in each dimension
"""
super(VirtualPixelWavelength,
self).__init__(polynomial_wcs.wave_transform_coef)
self.lut_wcs = lut_wcs
self.polynomial_wcs = polynomial_wcs
self.raw_wavelength_pixels = raw_wavelength_pixels
self.standard_broadcasting = False
self.sub_sampling = sub_sampling
@staticmethod
def _initialize_sub_pixel_table(x, y, sub_sampling):
"""
Generate a subgrid for x and y with sampling `sub_sampling`.
For example, for a subsampling of 5 each real x y coordinate will have
25 additional entries.
Parameters
----------
x : numpy.ndarray
y : numpy.ndarray
sub_sampling : int
Returns
-------
pandas.DataFrame
subpixel table with pixel index, sub_x, sub_y
"""
sub_edges = sub_sampling + 1
sub_y_delta, sub_x_delta = np.mgrid[-0.5:0.5:sub_edges * 1j,
-0.5:0.5:sub_edges * 1j]
sub_pixel_size = 1 / np.float(sub_sampling)
sub_x_delta = sub_x_delta[:-1, :-1] + 0.5 * sub_pixel_size
sub_y_delta = sub_y_delta[:-1, :-1] + 0.5 * sub_pixel_size
sub_x = np.add.outer(x.flatten(), sub_x_delta.flatten()).flatten()
sub_y = np.add.outer(y.flatten(), sub_y_delta.flatten()).flatten()
pixel_id = np.multiply.outer(
np.arange(len(x.flatten())),
np.ones_like(sub_x_delta.flatten(), dtype=np.int64)).flatten()
sub_pixel_table = pd.DataFrame(columns=[
'pixel_id',
'sub_x',
'sub_y',
],
data={
'pixel_id': pixel_id,
'sub_x': sub_x,
'sub_y': sub_y
})
return sub_pixel_table
@staticmethod
def _add_wavelength_sub_pixel_table(sub_pixel_table, wavelength_bins,
wave_transform_coef):
"""
Add a 'sub_wavelength' column that is calculated for each of the
sub-pixel locations using the PolynomialWCS and add a
'wavelength_bin_id' column that identifies the wavelength
bin that this sub pixel corresponds to.
Parameters
----------
sub_pixel_table : pandas.DataFrame
wavelength_bins : numpy.ndarray
wave_transform_coef : numpy.ndarray
Returns
-------
"""
kdt = cKDTree(wavelength_bins[np.newaxis].T)
sub_pixel_table['sub_wavelength'] = np.polynomial.polynomial.polyval2d(
sub_pixel_table.sub_x.values, sub_pixel_table.sub_y.values,
wave_transform_coef)
_, sub_pixel_table['wavelength_pixel_id'] = kdt.query(
sub_pixel_table.sub_wavelength.values[np.newaxis].T)
return sub_pixel_table
@staticmethod
def _generate_pixel_table(sub_pixel_table, sub_sampling):
pixel_table = sub_pixel_table.groupby(
['pixel_id', 'wavelength_pixel_id']).sub_x.count()
pixel_table /= sub_sampling**2
pixel_table = pixel_table.reset_index()
return pixel_table
def evaluate(self, wave_transform_coef):
sub_pixel_table = self._initialize_sub_pixel_table(
self.lut_wcs.x, self.lut_wcs.y, self.sub_sampling)
sub_pixel_table = self._add_wavelength_sub_pixel_table(
sub_pixel_table, self.raw_wavelength_pixels,
np.squeeze(wave_transform_coef))
pixel_table = self._generate_pixel_table(sub_pixel_table,
self.sub_sampling)
wavelength_pixel_contain_real_pixel_id = np.sort(
pixel_table.wavelength_pixel_id.unique())
self.wavelength_pixels = self.raw_wavelength_pixels[
wavelength_pixel_contain_real_pixel_id]
pixel_table['wavelength_pixel_id'] = (
wavelength_pixel_contain_real_pixel_id.searchsorted(
pixel_table.wavelength_pixel_id))
pixel_table['wavelength'] = self.wavelength_pixels[
pixel_table.wavelength_pixel_id].astype(np.float64)
pixel_table['slit_pos'] = self.lut_wcs.pix_to_slit[
pixel_table.pixel_id].astype(np.float64)
pixel_table['normed_wavelength'] = (
(pixel_table.wavelength - pixel_table.wavelength.min()) /
(pixel_table.wavelength.max() -
pixel_table.wavelength.min())) * 2 - 1
return pixel_table
