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