""" Module for telluric corrections.
.. include common links, assuming primary doc root is up one directory
.. include:: ../include/links.rst
"""
from IPython import embed
import matplotlib.pyplot as plt
import numpy as np
import scipy.interpolate
import scipy.optimize
import scipy.signal
import scipy.special
from astropy import table
from astropy.io import fits
from pypeit.core import flux_calib
from pypeit.core.wavecal import wvutils
from pypeit.core import coadd
from pypeit.core import fitting
from pypeit import data
from pypeit import specobjs
from pypeit import utils
from pypeit import msgs
from pypeit import onespec
from pypeit.spectrographs.util import load_spectrograph
from pypeit import datamodel
##############################
# Telluric model functions #
##############################
ZP_UNIT_CONST = flux_calib.zp_unit_const()
# TODO These codes should probably be in a separate qso_pca module. Also pickle functionality needs to be removed.
# The prior is not used (that was the reason for pickling), so the components could be stored in fits format.
# npca is not actually required here.
[docs]def qso_init_pca(filename,wave_grid,redshift,npca):
"""
This routine reads in the pickle file created by coarse_pca.create_coarse_pca.
The relevant pieces are the wavelengths (wave_pca_c), the PCA components (pca_comp_c), and the Gaussian mixture
model prior (mix_fit).
FILL THESE IN
Args:
filename (str):
Pickle filename
wave_grid (`numpy.ndarray`_):
redshift: (float):
Approximate redshift of the quasar.
npca:
Numer of pca components.
Returns:
dict:
"""
# Read in the pickle file from coarse_pca.create_coarse_pca
# The relevant pieces are the wavelengths (wave_pca_c), the PCA components (pca_comp_c),
# and the Gaussian mixture model prior (mix_fit)
# The PCA file location is provided by data.Paths.tel_model
file_with_path = data.Paths.tel_model / filename
loglam = np.log10(wave_grid)
dloglam = np.median(loglam[1:] - loglam[:-1])
pca_table = table.Table.read(file_with_path)
wave_pca_c = pca_table['WAVE_PCA'][0].flatten()
pca_comp_c = pca_table['PCA_COMP'][0][0,:,:]
coeffs_c = pca_table['PCA_COEFFS'][0][0,:,:]
num_comp = pca_comp_c.shape[0] # number of PCA components
# Interpolate PCA components onto wave_grid
pca_interp = scipy.interpolate.interp1d(wave_pca_c*(1.0+redshift), pca_comp_c, bounds_error=False, fill_value=0.0, axis=1)
pca_comp_new = pca_interp(wave_grid)
# Generate a mixture model for the coefficients prior, what should ngauss be?
#prior = mixture.GaussianMixture(n_components = npca-1).fit(coeffs_c[:, 1:npca])
# Construct the PCA dict
qso_pca_dict = {'npca': npca, 'components': pca_comp_new, 'coeffs': coeffs_c, 'z_fid': redshift, 'dloglam': dloglam}
return qso_pca_dict
[docs]def qso_pca_eval(theta,qso_pca_dict):
"""
Function for evaluating the quasar PCA
Args:
theta (`numpy.ndarray`_):
Parameter vector, where theta_pca[0] is redshift, theta_pca[1] is the normalization, and
theta_pca[2:npca+1] are the PCA coefficients, where npca is the PCA dimensionality
qso_pca_dict (dict):
Dictionary continaing the PCA information generated by qso_init_pca
Returns:
`numpy.ndarray`_: Evaluated PCA for the QSO
"""
C = qso_pca_dict['components']
z_fid = qso_pca_dict['z_fid']
dloglam = qso_pca_dict['dloglam']
npca = qso_pca_dict['npca'] # Size of the PCA currently being used, original PCA in the dict could be larger
z_qso = theta[0]
norm = theta[1]
A = theta[2:]
dshift = int(np.round(np.log10((1.0 + z_qso)/(1.0 + z_fid))/dloglam))
C_now = np.roll(C[:npca,:], dshift, axis=1)
return norm*np.exp(np.dot(np.append(1.0,A),C_now))
# TODO The prior is not currently used, but this is left in here anyway.
#def qso_pca_lnprior(theta,qso_pca_dict):
# """
# Routine to evaluate the the ln of the prior probability lnPrior,
# from the Gaussian mixture model fit to the distriution of PCA
# coefficients.
#
# Args:
# theta (`numpy.ndarray`_):
# Parameter vector, where ``theta_pca[0]`` is redshift,
# ``theta_pca[1]`` is the normalization, and
# ``theta_pca[2:npca+1]`` are the PCA coefficients, where npca
# is the PCA dimensionality
# qso_pca_dict (dict):
# Dictionary continaing the PCA information generated by
# ``qso_init_pca``
#
#
# Returns:
# float: The log of the prior probability evaluated at the current
# set of coefficients in the parameter vector, i.e. ``A =
# theta_pca[2:npca+1]``
#
# """
# gaussian_mixture_model = qso_pca_dict['prior']
# A = theta[2:]
# return gaussian_mixture_model.score_samples(A.reshape(1,-1))
[docs]def read_telluric_pca(filename, wave_min=None, wave_max=None, pad_frac=0.10):
"""
Reads in the telluric PCA components from a file.
Optionally, this method also trims wavelength to be in within ``wave_min``
and ``wave_max`` and pads the data (see ``pad_frac``).
Args:
filename (:obj:`str`):
Telluric PCA filename
wave_min (:obj:`float`, optional):
Minimum wavelength at which the grid is desired
wave_max (:obj:`float`, optional):
Maximum wavelength at which the grid is desired.
pad_frac (:obj:`float`, optional):
Percentage padding to be added to the grid boundaries if
``wave_min`` or ``wave_max`` are input; ignored otherwise. The
resulting grid will extend from ``(1.0 - pad_frac)*wave_min`` to
``(1.0 + pad_frac)*wave_max``.
Returns:
:obj:`dict`: Dictionary containing the telluric PCA components.
- wave_grid: Telluric model wavelength grid
- dloglam: Wavelength sampling of telluric model
- tell_pad_pix: Number of pixels to pad at edges for convolution
- ncomp_tell_pca: Number of PCA components
- tell_pca: PCA component vectors
- bounds_tell_pca: Maximum/minimum coefficient
- coefs_tell_pca: Set of model coefficient values (for prior in future)
- teltype: Type of telluric model, i.e. 'pca'
"""
# load_telluric_grid() takes care of path and existance check
hdul = data.load_telluric_grid(filename)
wave_grid_full = hdul[1].data
pca_comp_full = hdul[0].data
nspec_full = wave_grid_full.size
ncomp = hdul[0].header['NCOMP']
bounds = hdul[2].data
model_coefs = hdul[3].data
ind_lower = np.argmin(np.abs(wave_grid_full - (1.0 - pad_frac)*wave_min)) \
if wave_min is not None else 0
ind_upper = np.argmin(np.abs(wave_grid_full - (1.0 + pad_frac)*wave_max)) \
if wave_max is not None else nspec_full
wave_grid = wave_grid_full[ind_lower:ind_upper]
pca_comp_grid = pca_comp_full[:,ind_lower:ind_upper]
dwave, dloglam, resln_guess, pix_per_sigma = wvutils.get_sampling(wave_grid)
tell_pad_pix = int(np.ceil(10.0 * pix_per_sigma))
return dict(wave_grid=wave_grid, dloglam=dloglam,
tell_pad_pix=tell_pad_pix, ncomp_tell_pca=ncomp,
tell_pca=pca_comp_grid, bounds_tell_pca=bounds,
coefs_tell_pca=model_coefs, teltype='pca')
[docs]def read_telluric_grid(filename, wave_min=None, wave_max=None, pad_frac=0.10):
"""
Reads in the telluric grid from a file. This method is no longer the
preferred approach; see "read_telluric_pca" for the PCA mode.
Optionally, this method also trims the grid to be in within ``wave_min``
and ``wave_max`` and pads the data (see ``pad_frac``).
Args:
filename (:obj:`str`):
Telluric grid filename
wave_min (:obj:`float`, optional):
Minimum wavelength at which the grid is desired
wave_max (:obj:`float`, optional):
Maximum wavelength at which the grid is desired.
pad_frac (:obj:`float`, optional):
Percentage padding to be added to the grid boundaries if
``wave_min`` or ``wave_max`` are input; ignored otherwise. The
resulting grid will extend from ``(1.0 - pad_frac)*wave_min`` to
``(1.0 + pad_frac)*wave_max``.
Returns:
:obj:`dict`: Dictionary containing the telluric grid
- wave_grid: Telluric model wavelength grid
- dloglam: Wavelength sampling of telluric model
- tell_pad_pix: Number of pixels to pad at edges for convolution
- pressure_grid: Atmospheric pressure values in telluric grid [mbar]
- temp_grid: Temperature values in telluric grid [degrees C]
- h2o_grid: Humidity values in telluric grid [%]
- airmass_grid: Airmass values in telluric grid
- tell_grid: Grid of telluric models
- teltype: Type of telluric model, i.e. 'grid'
"""
# load_telluric_grid() takes care of path and existance check
hdul = data.load_telluric_grid(filename)
wave_grid_full = 10.0*hdul[1].data
model_grid_full = hdul[0].data
nspec_full = wave_grid_full.size
ind_lower = np.argmin(np.abs(wave_grid_full - (1.0 - pad_frac)*wave_min)) \
if wave_min is not None else 0
ind_upper = np.argmin(np.abs(wave_grid_full - (1.0 + pad_frac)*wave_max)) \
if wave_max is not None else nspec_full
wave_grid = wave_grid_full[ind_lower:ind_upper]
model_grid = model_grid_full[...,ind_lower:ind_upper]
pg = hdul[0].header['PRES0']+hdul[0].header['DPRES']*np.arange(0,hdul[0].header['NPRES'])
tg = hdul[0].header['TEMP0']+hdul[0].header['DTEMP']*np.arange(0,hdul[0].header['NTEMP'])
hg = hdul[0].header['HUM0']+hdul[0].header['DHUM']*np.arange(0,hdul[0].header['NHUM'])
if hdul[0].header['NAM'] > 1:
ag = hdul[0].header['AM0']+hdul[0].header['DAM']*np.arange(0,hdul[0].header['NAM'])
else:
ag = hdul[0].header['AM0']+1*np.arange(0,1)
dwave, dloglam, resln_guess, pix_per_sigma = wvutils.get_sampling(wave_grid)
tell_pad_pix = int(np.ceil(10.0 * pix_per_sigma))
return dict(wave_grid=wave_grid, dloglam=dloglam, tell_pad_pix=tell_pad_pix,
pressure_grid=pg, temp_grid=tg, h2o_grid=hg, airmass_grid=ag,
tell_grid=model_grid, teltype='grid')
[docs]def interp_telluric_grid(theta, tell_dict):
"""
Interpolate the telluric model grid to the specified location in
parameter space.
The interpolation is only performed over the 4D parameter space specified
by pressure, temperature, humidity, and airmass. This routine performs
nearest-gridpoint interpolation to evaluate the telluric model at an
arbitrary location in this 4-d space. The telluric grid is assumed to be
uniformly sampled in this parameter space.
Args:
theta (`numpy.ndarray`_):
A 4-element vector with the telluric model parameters: pressure,
temperature, humidity, and airmass.
tell_dict (dict):
Dictionary containing the telluric grid. See
:func:`read_telluric_grid`.
Returns:
`numpy.ndarray`_: Telluric model evaluated at the provided 4D
position in parameter space. The telluric model is provided over all
available wavelengths in ``tell_dict``.
"""
if len(theta) != 4:
msgs.error('Input parameter vector must have 4 and only 4 values.')
pg = tell_dict['pressure_grid']
tg = tell_dict['temp_grid']
hg = tell_dict['h2o_grid']
ag = tell_dict['airmass_grid']
p, t, h, a = theta
pi = int(np.round((p-pg[0])/(pg[1]-pg[0]))) if len(pg) > 1 else 0
ti = int(np.round((t-tg[0])/(tg[1]-tg[0]))) if len(tg) > 1 else 0
hi = int(np.round((h-hg[0])/(hg[1]-hg[0]))) if len(hg) > 1 else 0
ai = int(np.round((a-ag[0])/(ag[1]-ag[0]))) if len(ag) > 1 else 0
return tell_dict['tell_grid'][pi,ti,hi,ai]
[docs]def conv_telluric(tell_model, dloglam, res):
"""
Routine to convolve the telluric model to desired resolution.
Args:
tell_model (`numpy.ndarray`_):
Input telluric model at the native resolution of the telluric model grid. The shape of this input is in
general different from the size of the raw telluric model (read in by read_telluric_* above) because it is
trimmed to relevant wavelengths using ind_lower, ind_upper. See eval_telluric below.
dloglam (float):
Wavelength spacing of the telluric grid expressed as a dlog10(lambda), i.e. stored in the
tell_dict as tell_dict['dloglam']
res (float):
Desired resolution expressed as lambda/dlambda. Note that here dlambda is linear, whereas dloglam is
the delta of the log10.
Returns:
`numpy.ndarray`_:
Resolution convolved telluric model. Shape = same size as input tell_model.
"""
pix_per_sigma = 1.0/res/(dloglam*np.log(10.0))/(2.0 * np.sqrt(2.0 * np.log(2))) # number of dloglam pixels per 1 sigma dispersion
sig2pix = 1.0/pix_per_sigma # number of sigma per 1 pix
if sig2pix > 2.0:
msgs.warn('The telluric model grid is not sampled finely enough to properly convolve to the desired resolution. '
'Skipping resolution convolution for now. Create a higher resolution telluric model grid')
return tell_model
# x = loglam/sigma on the wavelength grid from -4 to 4, symmetric, centered about zero.
x = np.hstack([-1*np.flip(np.arange(sig2pix,4,sig2pix)),np.arange(0,4,sig2pix)])
# g = Gaussian evaluated at x, sig2pix multiplied in to properly normalize the convolution
#g = (1.0/(np.sqrt(2*np.pi)))*np.exp(-0.5*np.square(x))*sig2pix
g=np.exp(-0.5*np.square(x))
g /= g.sum()
conv_model = scipy.signal.convolve(tell_model,g,mode='same')
return conv_model
[docs]def shift_telluric(tell_model, loglam, dloglam, shift, stretch):
"""
Routine to apply a shift to the telluric model. Note that the shift can be sub-pixel, i.e this routine interpolates.
Note that the shift units are pixels *of the telluric model*.
Args:
tell_model (`numpy.ndarray`_):
Input telluric model. The shape of this input is in general different from the size of the telluric models
(read in by read_telluric_* above) because it is trimmed to relevant wavelengths using ind_lower, ind_upper.
See eval_telluric_* below.
loglam (`numpy.ndarray`_):
The log10 of the wavelength grid on which the tell_model is evaluated.
dloglam (float):
Wavelength spacing of the telluric grid expressed as a a dlog10(lambda), i.e. stored in the
tell_dict as tell_dict['dloglam']
shift (float):
Desired shift. Note that this shift can be sub-pixel.
stretch (float):
Desired stretch.
Returns:
`numpy.ndarray`_:
Shifted telluric model. Shape = same size as input tell_model.
"""
loglam_shift = loglam[0] + np.arange(len(loglam)) * dloglam * stretch + shift * dloglam
#loglam_shift = loglam + shift*dloglam
tell_model_shift = np.interp(loglam_shift, loglam, tell_model)
return tell_model_shift
[docs]def eval_telluric(theta_tell, tell_dict, ind_lower=None, ind_upper=None):
"""
Evaluate the telluric model.
Parameters from ``theta_tell`` are: (if teltype == 'pca') the PCA coefficients
or (if teltype == 'grid') the telluric grid parameters pressure, temperature,
humidity, and airmass, in both cases followed by spectral resolution, shift,
and stretch.
This routine performs the following steps:
1. summation of telluric PCA components multiplied by coefficients
2. transformation of telluric PCA model from arsinh(tau) to transmission
3. convolution of the atmosphere model to the resolution set by
the spectral resolution.
4. shift and stretch the telluric model.
Args:
theta_tell (`numpy.ndarray`_):
Vector with tell_npca PCA coefficients (if teltype='pca')
or pressure, temperature, humidity, and airmass (if teltype='grid'),
followed by spectral resolution, shift, and stretch.
Final length is then tell_npca+3 or 7.
tell_dict (:obj:`dict`):
Dictionary containing the telluric data. See
:func:`read_telluric_pca` if teltype=='pca'.
or
:func:`read_telluric_grid` if teltype=='grid'.
ind_lower (:obj:`int`, optional):
The index of the first pixel to include in the model. Selecting a
wavelength region for the modeling makes things faster because we
only need to convolve the portion that is needed for the current
model fit.
ind_upper (:obj:`int`, optional):
The index (inclusive) of the last pixel to include in the model.
Selecting a wavelength region for the modeling makes things
faster because we only need to convolve the portion that is
needed for the current model fit.
Returns:
`numpy.ndarray`_: Telluric model evaluated at the desired location
theta_tell in model atmosphere parameter space. Shape is given by
the size of ``wave_grid`` plus ``tell_pad_pix`` padding from the input
tell_dict.
"""
ntheta = len(theta_tell)
teltype = tell_dict['teltype']
# FD: Currently assumes that shift and stretch are on.
# TODO: Make this work without shift and stretch.
if teltype == 'pca':
ntell = ntheta-3
elif teltype == 'grid':
ntell = 4
# Set the wavelength range if not provided
ind_lower = 0 if ind_lower is None else ind_lower
ind_upper = tell_dict['wave_grid'].size - 1 if ind_upper is None else ind_upper
# Deal with padding for the convolutions
ind_lower_pad = np.fmax(ind_lower - tell_dict['tell_pad_pix'], 0)
ind_upper_pad = np.fmin(ind_upper + tell_dict['tell_pad_pix'], tell_dict['wave_grid'].size - 1)
## FW: There is an extreme case with ind_upper == ind_upper_pad, the previous -0 won't work
ind_lower_final = ind_lower_pad if ind_lower_pad == ind_lower else ind_lower - ind_lower_pad
ind_upper_final = ind_upper_pad if ind_upper_pad == ind_upper else ind_upper - ind_upper_pad
if teltype == 'pca':
# Evaluate PCA model after truncating the wavelength range
tellmodel_hires = np.zeros_like(tell_dict['tell_pca'][0])
tellmodel_hires[ind_lower_pad:ind_upper_pad+1] = np.dot(np.append(1,theta_tell[:ntell]),
tell_dict['tell_pca'][:ntell+1][:,ind_lower_pad:ind_upper_pad+1])
# PCA model is inverse sinh of the optical depth, convert to transmission here
tellmodel_hires[ind_lower_pad:ind_upper_pad+1] = np.sinh(tellmodel_hires[ind_lower_pad:ind_upper_pad+1])
# It should generally be very rare, but trim negative optical depths here just in case.
clip = tellmodel_hires < 0
tellmodel_hires[clip] = 0
tellmodel_hires[ind_lower_pad:ind_upper_pad+1] = np.exp(-tellmodel_hires[ind_lower_pad:ind_upper_pad+1])
elif teltype == 'grid':
# Interpolate within the telluric grid
tellmodel_hires = interp_telluric_grid(theta_tell[:ntell], tell_dict)
tellmodel_conv = conv_telluric(tellmodel_hires[ind_lower_pad:ind_upper_pad+1],
tell_dict['dloglam'], theta_tell[-3])
# Stretch and shift telluric wavelength grid
tellmodel_out = shift_telluric(tellmodel_conv,
np.log10(tell_dict['wave_grid'][ind_lower_pad:ind_upper_pad+1]),
tell_dict['dloglam'], theta_tell[-2], theta_tell[-1])
return tellmodel_out[ind_lower_final:ind_upper_final]
############################
# Fitting routines #
############################
[docs]def tellfit_chi2(theta, flux, thismask, arg_dict):
"""
Loss function which is optimized by differential evolution to perform the object + telluric model fitting for
telluric corrections. This is a general abstracted routine that provides the loss function for any object model
that the user provides.
Args:
theta (`numpy.ndarray`_):
Parameter vector for the object + telluric model.
This is actually two concatenated parameter vectors, one for
the object and one for the telluric, i.e.:
(in PCA mode)
theta_obj = theta[:-(tell_npca+3)]
theta_tell = theta[-(tell_npca+3):]
(in grid mode)
theta_obj = theta[:-7]
theta_tell = theta[-7:]
The telluric model theta_tell includes a either user-specified
number of PCA coefficients (in PCA mode) or ambient pressure,
temperature, humidity, and airmass (in grid mode) followed by
spectral resolution, shift, and stretch.
That is, in PCA mode,
pca_coeffs = theta_tell[:tell_npca]
while in grid mode,
pressure = theta_tell[0]
temperature = theta_tell[1]
humidity = theta_tell[2]
airmass = theta_tell[3]
with the last three indices of the array corresponding to
resolution = theta_tell[-3]
shift = theta_tell[-2]
stretch = theta_tell[-1]
The object model theta_obj can have an arbitrary size and is
provided as an argument to obj_model_func
flux (`numpy.ndarray`_):
The flux of the object being fit
thismask (`numpy.ndarray`_, boolean):
A mask indicating which values are to be fit. This is a good pixel mask, i.e. True=Good
arg_dict (dict):
A dictionary containing the parameters needed to evaluate the telluric model and the object model. See
documentation of tellfit for a detailed description.
Returns:
float:
The value of the loss function at the location in parameter space theta. This is loss function is the thing
that is minimized to perform the fit.
"""
obj_model_func = arg_dict['obj_model_func']
flux_ivar = arg_dict['ivar']
teltype = arg_dict['tell_dict']['teltype']
# TODO: make this work without shift and stretch?
# Number of telluric model parameters, plus shift, stretch, and resolution
if teltype == 'pca':
nfit = arg_dict['tell_npca']+3
elif teltype == 'grid':
nfit = 4+3
theta_obj = theta[:-nfit]
theta_tell = theta[-nfit:]
tell_model = eval_telluric(theta_tell, arg_dict['tell_dict'],
ind_lower=arg_dict['ind_lower'], ind_upper=arg_dict['ind_upper'])
obj_model, model_gpm = obj_model_func(theta_obj, arg_dict['obj_dict'])
totalmask = thismask & model_gpm
if not np.any(totalmask):
return np.inf # If everyting is masked retrun infinity
else:
chi_vec = totalmask * (flux - tell_model*obj_model) * np.sqrt(flux_ivar)
robust_scale = 2.0
huber_vec = scipy.special.huber(robust_scale, chi_vec)
loss_function = np.sum(huber_vec * totalmask)
return loss_function
[docs]def tellfit(flux, thismask, arg_dict, init_from_last=None):
"""
Routine to perform the object + telluric model fitting for telluric
corrections. This is a general abstracted routine that performs the
fits for any object model that the user provides.
Args:
flux (`numpy.ndarray`_):
The flux of the object being fit
thismask (`numpy.ndarray`_, boolean):
A mask indicating which values are to be fit. This is a good
pixel mask, i.e. True=Good
arg_dict (dict):
A dictionary containing the parameters needed to evaluate
the telluric model and the object model. The required keys
are:
- ``arg_dict['flux_ivar']``: Inverse variance for the
flux array
- ``arg_dict['tell_dict']``: Dictionary containing the
telluric model and its parameters read in by
read_telluric_pca or read_telluric_grid
- ``arg_dict['ind_lower']``: Lower index into the
telluric model wave_grid to trim down the telluric
model.
- ``arg_dict['ind_upper']``: Upper index into the
telluric model wave_grid to trim down the telluric
model.
- ``arg_dict['obj_model_func']``: User provided function
for evaluating the object model
- ``arg_dict['obj_dict']``: Dictionary containing the
object model arguments which is passed to the
obj_model_func
init_from_last (object, optional):
Optional. Result object returned by the differential
evolution optimizer for the last iteration. If this is passed the code
will initialize from the previous best-fit for faster convergence.
Returns:
tuple: Returns three objects:
- result (obj): Result object returned by the differential
evolution optimizer
- modelfit (`numpy.ndarray`_): Modelfit to the input flux.
This has the same size as the flux
- ivartot (`numpy.ndarray`_): Corrected inverse variances
for the flux. This has the same size as the flux. The
errors are renormalized using the renormalize_errors
function by a correction factor, i.e. ivartot =
flux_ivar/sigma_corr**2
"""
# Unpack arguments
obj_model_func = arg_dict['obj_model_func'] # Evaluation function
flux_ivar = arg_dict['ivar'] # Inverse variance of flux or counts
bounds = arg_dict['bounds'] # bounds for differential evolution optimization
rng = arg_dict['rng'] # Seed for differential evolution optimizaton
maxiter = arg_dict['diff_evol_maxiter'] # Maximum number of iterations
ballsize = arg_dict['ballsize'] # Ballsize for initialization from a previous optimum
nparams = len(bounds) # Number of parameters in the model
popsize = arg_dict['popsize'] # Note this does nothing if the init is done from a previous iteration or optimum
nsamples = arg_dict['popsize']*nparams
teltype = arg_dict['tell_dict']['teltype']
# FD: Assumes shift and stretch are turned on.
if teltype == 'pca':
ntheta_tell = arg_dict['tell_npca']+3 # Total number of telluric model parameters in PCA mode
elif teltype == 'grid':
ntheta_tell = 4+3 # Total number of telluric model parameters in grid mode
# Decide how to initialize
if init_from_last is not None:
# Use a Gaussian ball about the optimum from a previous iteration
init = np.array([[np.clip(param + ballsize*(bounds[i][1] - bounds[i][0]) * rng.standard_normal(1)[0],
bounds[i][0], bounds[i][1])
for i, param in enumerate(init_from_last.x)] for jsamp in range(nsamples)])
elif 'init_obj_opt_theta' in arg_dict['obj_dict']:
# Initialize from the object parameters. Use a Gaussian ball about the best object model, and latin hypercube
# for the telluric parameters
bounds_obj = arg_dict['obj_dict']['bounds_obj']
init_obj = np.array([[np.clip(param + ballsize*(bounds_obj[i][1] - bounds_obj[i][0]) * rng.standard_normal(1)[0],
bounds_obj[i][0], bounds_obj[i][1]) for i, param in enumerate(arg_dict['obj_dict']['init_obj_opt_theta'])]
for jsamp in range(nsamples)])
tell_lhs = utils.lhs(ntheta_tell, samples=nsamples)
init_tell = np.array([[bounds[-idim][0] + tell_lhs[isamp, idim] * (bounds[-idim][1] - bounds[-idim][0])
for idim in range(ntheta_tell)] for isamp in range(nsamples)])
init = np.hstack((init_obj, init_tell))
else:
# If this is the first iteration and no object model optimum is presented, use a latin hypercube which is the default
init = 'latinhypercube'
result = scipy.optimize.differential_evolution(tellfit_chi2, bounds, args=(flux, thismask, arg_dict,), seed=rng,
init = init, updating='immediate', popsize=popsize,
recombination=arg_dict['recombination'], maxiter=arg_dict['diff_evol_maxiter'],
polish=arg_dict['polish'], disp=arg_dict['disp'])
theta_obj = result.x[:-ntheta_tell]
theta_tell = result.x[-ntheta_tell:]
tell_model = eval_telluric(theta_tell, arg_dict['tell_dict'],
ind_lower=arg_dict['ind_lower'], ind_upper=arg_dict['ind_upper'])
obj_model, modelmask = obj_model_func(theta_obj, arg_dict['obj_dict'])
totalmask = thismask & modelmask
chi_vec = totalmask*(flux - tell_model*obj_model)*np.sqrt(flux_ivar)
debug = arg_dict['debug'] if 'debug' in arg_dict else False
# Name of function for title in case QA requested
obj_model_func_name = getattr(obj_model_func, '__name__', repr(obj_model_func))
sigma_corr, maskchi = coadd.renormalize_errors(chi_vec, mask=totalmask, title = obj_model_func_name,
debug=debug)
ivartot = flux_ivar/sigma_corr**2
return result, tell_model*obj_model, ivartot
# TODO This should be a general reader once we get our act together with the data model.
# For echelle: read in all the orders into a (nspec, nporders) array
# For longslit: read in the standard into a (nspec, 1) array
[docs]def unpack_orders(sobjs, ret_flam=False):
"""
Utility function to unpack the sobjs object and return the
arrays necessary for telluric fitting.
Args:
sobjs (:class:`~pypeit.specobjs.SpecObjs`):
SpecObjs object
ret_flam (bool):
If true return the FLAM, otherwise return COUNTS
Returns:
tuple: wave, flam, flam_ivar, flam_mask
wave (`numpy.ndarray`_) Wavelength grids;
flam (`numpy.ndarray`_) Flambda or counts;
flam_ivar (`numpy.ndarray`_) Inverse variance (of Flambda or counts);
flam_mask (`numpy.ndarray`_) Good pixel mask. True=Good.
All return values have shape (nspec, norders)
"""
# Read in the spec1d file
norders = len(sobjs) # ToDO: This is incorrect if you have more than one object in the sobjs
if ret_flam:
nspec = sobjs[0].OPT_FLAM.size
else:
nspec = sobjs[0].OPT_COUNTS.size
# Allocate arrays and unpack spectrum
wave = np.zeros((nspec, norders))
flam = np.zeros((nspec, norders))
flam_ivar = np.zeros((nspec, norders))
flam_mask = np.zeros((nspec, norders),dtype=bool)
for iord in range(norders):
wave[:, iord] = sobjs[iord].OPT_WAVE
#wave_mask[:,iord] = sobjs[iord].optimal['WAVE'] > 0.0
flam_mask[:, iord] = sobjs[iord].OPT_MASK
if ret_flam:
flam[:, iord] = sobjs[iord].OPT_FLAM
flam_ivar[:, iord] = sobjs[iord].OPT_FLAM_IVAR
else:
flam[:, iord] = sobjs[iord].OPT_COUNTS
flam_ivar[:,iord] = sobjs[iord].OPT_COUNTS_IVAR
return wave.squeeze(), flam.squeeze(), flam_ivar.squeeze(), flam_mask.squeeze()
# TODO: This function needs to be revisited. Better yet, it would useful to
# brainstorm about whether or not it's worth revisiting the spec1d datamodel.
[docs]def general_spec_reader(specfile, ret_flam=False, chk_version=True):
"""
Read a spec1d file or a coadd spectrum file.
Args:
specfile (:obj:`str`):
File with the data
ret_flam (:obj:`bool`, optional):
Return FLAM instead of COUNTS.
chk_version (:obj:`bool`, optional):
When reading in existing files written by PypeIt, perform strict
version checking to ensure a valid file. If False, the code will
try to keep going, but this may lead to faults and quiet failures.
User beware!
Returns:
:obj:`tuple`: Seven objects are returned. The first five are
`numpy.ndarray`_ objects with the wavelength, bin centers of wavelength grid,
flux, inverse variance, and good-pixel mask for each spectral pixel.
The 6th is a :obj:`dict` of metadata. And the 7th is an
`astropy.io.fits.Header`_ object with the primary header from the input file.
"""
# Place holder routine that provides a generic spectrum reader
bonus = {}
# Figure out which flavor input file
hdul = fits.open(specfile)
if 'DMODCLS' in hdul[1].header and hdul[1].header['DMODCLS'] == 'OneSpec':
# Load
spec = onespec.OneSpec.from_file(specfile, chk_version=chk_version)
# Unpack
wave = spec.wave
# wavelength grid evaluated at the bin centers, uniformly-spaced in lambda or log10-lambda/velocity.
# see core.wavecal.wvutils.py for more info.
# variable defaults to None if datamodel for this is also None (which is the case for spec1d file).
wave_grid_mid = spec.wave_grid_mid
counts = spec.flux
counts_ivar = spec.ivar
counts_gpm = spec.mask.astype(bool)
spect_dict = spec.spect_meta
head = spec.head0
else:
sobjs = specobjs.SpecObjs.from_fitsfile(specfile, chk_version=chk_version)
# TODO: What bug? Is it fixed now? How can we test if it's fixed?
if np.sum(sobjs.OPT_WAVE) is None:
raise ValueError("This is an ugly hack until the DataContainer bug is fixed")
head = sobjs.header
wave, counts, counts_ivar, counts_gpm = unpack_orders(sobjs, ret_flam=ret_flam)
# Made a change to the if statement to account for unpack_orders now squeezing returned arrays
#if (head['PYPELINE'] !='Echelle') and (wave.shape[1]>1)
if (head['PYPELINE'] !='Echelle') and (wave.ndim>1):
idx = flux_calib.find_standard(sobjs)
npix = head['NPIX']
wave, counts = np.reshape(wave[:,idx],(npix,1)), np.reshape(counts[:,idx],(npix,1))
counts_ivar = np.reshape(counts_ivar[:,idx],(npix,1))
counts_gpm = np.reshape(counts_gpm[:,idx],(npix,1))
bonus['ECH_ORDER'] = sobjs.ECH_ORDER if sobjs.ECH_ORDER[0] is None else (sobjs.ECH_ORDER).astype(int)
bonus['ECH_ORDERINDX'] = sobjs.ECH_ORDERINDX if sobjs.ECH_ORDERINDX[0] is None else (sobjs.ECH_ORDERINDX).astype(int)
bonus['ECH_SNR'] = sobjs.ech_snr if sobjs.ech_snr[0] is None else (sobjs.ech_snr).astype(int)
bonus['NORDERS'] = 1 if wave.ndim == 1 else wave.shape[1] # Again accounting for unpack_orders squeeze
try:
spectrograph = load_spectrograph(head['INSTRUME'])
except:
spectrograph = load_spectrograph('shane_kast_blue')
spect_dict = spectrograph.parse_spec_header(head)
head['PYP_SPEC'] = spectrograph.name
# Build this
meta_spec = dict(bonus=bonus)
meta_spec['core'] = spect_dict
return wave, wave_grid_mid, counts, counts_ivar, counts_gpm, meta_spec, head
[docs]def save_coadd1d_tofits(outfile, wave, flux, ivar, gpm, wave_grid_mid=None, spectrograph=None, telluric=None,
obj_model=None, header=None, ex_value='OPT', overwrite=True):
"""
Write final spectrum to disk using the OneSpec write command
Args:
outfile (:obj:`str`):
File to output telluric corrected spectrum to.
wave (`numpy.ndarray`_):
wavelength in units of Angstrom
flux (`numpy.ndarray`_):
flux array
ivar (`numpy.ndarray`_):
inverse variance array.
gpm (`numpy.ndarray`_):
good pixel mask for your spectrum.
wave_grid_mid (`numpy.ndarray`_):
wavelength grid evaluated at the bin centers, uniformly-spaced in lambda or log10-lambda/velocity.
See core.wavecal.wvutils.py for more info. Default is None because not all files that uses telluric.py
has this keyword parameter (spec1dfile does not but coadd1d files do).
spectrograph (:obj:`str`, optional):
spectrograph name
telluric (`numpy.ndarray`_, optional):
telluric model
obj_model (`numpy.ndarray`_, optional):
object model used for telluric fitting
header (`astropy.io.fits.Header`_, optional):
Primary header
ex_value (:obj:`str`, optional):
extraction mode, OPT or BOX
overwrite (:obj:`bool`, optional):
Overwrite existing file?
"""
wave_gpm = wave > 1.0
spec = onespec.OneSpec(wave=wave[wave_gpm], wave_grid_mid=None if wave_grid_mid is None else wave_grid_mid[wave_gpm],
flux=flux[wave_gpm], PYP_SPEC=spectrograph, ivar=ivar[wave_gpm], sigma=np.sqrt(utils.inverse(ivar[wave_gpm])),
mask=gpm[wave_gpm].astype(int),
telluric=None if telluric is None else telluric[wave_gpm],
obj_model=None if obj_model is None else obj_model[wave_gpm],
ext_mode=ex_value, fluxed=True)
# TODO: We need a better way of assigning the header...
spec.head0 = header
spec.to_file(outfile, overwrite=overwrite)
##############
# Sensfunc Model #
##############
[docs]def init_sensfunc_model(obj_params, iord, wave, counts_per_ang, ivar, gpm, tellmodel):
"""
Initializes a sensitivity function model fit for joint sensitivity function
and telluric fitting by setting up the obj_dict and bounds.
Parameters
----------
obj_params : :obj:`dict`
Dictionary of object parameters
iord : :obj:`int`
The slit/order in question
wave : `numpy.ndarray`_, float, shape is (nspec,)
Wavelength array for the standard star
counts_per_ang : `numpy.ndarray`_, float, shape is (nspec,)
Counts per Angstrom array for the standard star
ivar : `numpy.ndarray`_, float, shape is (nspec,)
Inverse variance of the counts per Angstrom array for the standard star
gpm : `numpy.ndarray`_, bool, shape is (nspec,)
Good pixel mask for the counts per Angstrom array for the standard star
tellmodel : `numpy.ndarray`_, float, shape is (nspec,)
Telluric absorption model guess
Returns
-------
obj_dict : :obj:`dict`
Dictionary of object paramaters for joint sensitivity function telluric
model fitting
bounds_obj : :obj:`list`
List of bounds for each parameter in the joint sensitivity function and
telluric model fit.
"""
# Model parameter guess for starting the optimizations
flam_true = scipy.interpolate.interp1d(obj_params['std_dict']['wave'].value,
obj_params['std_dict']['flux'].value, kind='linear',
bounds_error=False, fill_value=-1e20)(wave)
flam_true_gpm = (wave >= obj_params['std_dict']['wave'].value.min()) \
& (wave <= obj_params['std_dict']['wave'].value.max())
if np.any(np.logical_not(flam_true_gpm)):
msgs.warn('Your data extends beyond the range covered by the standard star spectrum. '
'Proceeding by masking these regions, but consider using another standard star')
N_lam = counts_per_ang/obj_params['exptime']
zeropoint_data, zeropoint_data_gpm \
= flux_calib.compute_zeropoint(wave, N_lam, (gpm & flam_true_gpm), flam_true,
tellmodel=tellmodel)
zeropoint_poly = zeropoint_data + 5.0*np.log10(wave) - ZP_UNIT_CONST
if obj_params['log10_blaze_func_per_ang'] is not None:
zeropoint_poly -= 2.5*obj_params['log10_blaze_func_per_ang']
# Perform an initial fit to the sensitivity function to set the starting
# point for optimization
wave_min, wave_max = wave.min(), wave.max()
pypeitFit = fitting.robust_fit(wave, zeropoint_poly, obj_params['polyorder_vec'][iord],
function=obj_params['func'], minx=wave_min, maxx=wave_max,
in_gpm=zeropoint_data_gpm, lower=obj_params['sigrej'],
upper=obj_params['sigrej'], use_mad=True)
zeropoint_fit = flux_calib.eval_zeropoint(pypeitFit.fitc, obj_params['func'], wave, wave_min, wave_max,
log10_blaze_func_per_ang=obj_params['log10_blaze_func_per_ang'])
zeropoint_fit_gpm = pypeitFit.bool_gpm
# Polynomial coefficient bounds
bounds_obj = [(np.fmin(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][0],
obj_params['minmax_coeff_bounds'][0]),
np.fmax(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][1],
obj_params['minmax_coeff_bounds'][1])) for this_coeff in pypeitFit.fitc]
# Create the obj_dict
obj_dict = dict(wave=wave, wave_min=wave.min(), wave_max=wave.max(),
log10_blaze_func_per_ang=obj_params['log10_blaze_func_per_ang'],
exptime=obj_params['exptime'], flam_true=flam_true,
flam_true_gpm=flam_true_gpm, func=obj_params['func'],
polyorder=obj_params['polyorder_vec'][iord], bounds_obj=bounds_obj,
init_obj_opt_theta = pypeitFit.fitc)
if obj_params['debug']:
title = 'Zeropoint Initialization Guess for order/det={:d}'.format(iord + 1) # +1 to account 0-index starting
flux_calib.zeropoint_qa_plot(wave, zeropoint_data, zeropoint_data_gpm, zeropoint_fit,
zeropoint_fit_gpm, title=title, show=True)
return obj_dict, bounds_obj
# Sensitivity function evaluation function. Model for counts is flam_true/sensfunc
[docs]def eval_sensfunc_model(theta, obj_dict):
"""
Utility routine to evaluate a sensitivity function model
Parameters
----------
theta : `numpy.ndarray`_
Array containing the parameters that describe the model.
shape is (ntheta,)
obj_dict : :obj:`dict`
Dictionary containing additional arguments needed to evaluate the model
Returns
-------
counts_per_angstrom_model : `numpy.ndarray`_, shape is the same as obj_dict['wave_star']
Model of the quantity being fit for the sensitivity function which is
``exptime*f_lam_true/sensfunc``. Note that this routine is used to fit the
counts per angstrom. The zeropoint is defined to be zeropoint
2.5log10(S_lam) where S_lam = F_lam/(counts/s/A) where F_lam is in units
of 1e-17 erg/s/cm^2/A, and so S_lam has units of
(erg/s/cm^2/A)/(counts/s/A) = erg/cm^2/counts, and the zeropoint is then
the magnitude of the astronomical source that produces one count per
second in a one-angstrom wide filter on the detector. And zerpoint +
2.5log10(delta_lambda/1 A) is the magnitude that produces one count per
spectral pixel of width delta_lambda
gpm : `numpy.ndarray`_, bool, shape is the same as obj_dict['wave_star']
Good pixel mask indicating where the model is valid
"""
zeropoint = flux_calib.eval_zeropoint(theta, obj_dict['func'], obj_dict['wave'],
obj_dict['wave_min'], obj_dict['wave_max'],
log10_blaze_func_per_ang=obj_dict['log10_blaze_func_per_ang'])
counts_per_angstrom_model = obj_dict['exptime'] \
* flux_calib.Flam_to_Nlam(obj_dict['wave'], zeropoint) \
* obj_dict['flam_true'] * obj_dict['flam_true_gpm']
return counts_per_angstrom_model, obj_dict['flam_true_gpm']
##############
# QSO Model #
##############
[docs]def init_qso_model(obj_params, iord, wave, flux, ivar, mask, tellmodel):
"""
Routine used to initialize the quasar spectrum telluric fits
Parameters
----------
obj_params : dict
Dictionary containing parameters necessary for initializing the quasar model
iord : int
Order in question. This is not used for initializing the qso model, but is kept here for compatibility
with the standard function argument list
wave : array shape (nspec,)
Wavelength array for the object in question
flux : array shape (nspec,)
Flux array for the object in question
ivar : array shape (nspec,)
Inverse variance array for the oejct in question
mask : array shape (nspec,)
Good pixel mask for the object in question
tellmodel : array shape (nspec,)
This is a telluric model computed on the wave wavelength grid. Initialization usually requires some initial
best guess for the telluric absorption, which is computed from the mean of the telluric model grid using
the resolution of the spectrograph.
Returns
-------
obj_dict : dict
Dictionary containing the meta-information and variables that are used for the object model evaluations.
For example for the quasar model which based on a PCA decomposition, this dictionary holds the number of PCA
basis vectors and those basis vectors.
bounds_obj : tuple
Tuple of bounds for each parameter that will be fit for the object model
"""
qso_pca_dict = qso_init_pca(obj_params['pca_file'], wave, obj_params['z_qso'], obj_params['npca'])
qso_pca_mean = np.exp(qso_pca_dict['components'][0, :])
tell_mask = tellmodel > obj_params['tell_norm_thresh']
# Create a reference model and bogus noise
flux_ref = qso_pca_mean * tellmodel
ivar_ref = utils.inverse((qso_pca_mean/100.0) ** 2)
flam_norm_inv = coadd.robust_median_ratio(flux, ivar, flux_ref, ivar_ref, mask=mask, mask_ref=tell_mask)
flam_norm = 1.0/flam_norm_inv
# Set the bounds for the PCA and truncate to the right dimension
coeffs = qso_pca_dict['coeffs'][:,1:obj_params['npca']]
# Compute the min and max arrays of the coefficients which are not the norm, i.e. grab the coeffs that aren't the first one
coeff_min = np.amin(coeffs, axis=0) # only
coeff_max = np.amax(coeffs, axis=0)
# QSO redshift: can vary within delta_zqso
bounds_z = [(obj_params['z_qso'] - obj_params['delta_zqso'], obj_params['z_qso'] + obj_params['delta_zqso'])]
bounds_flam = [(flam_norm*obj_params['lbound_norm'], flam_norm*obj_params['ubound_norm'])] # Norm: bounds determined from estimate above
bounds_pca = [(i, j) for i, j in zip(coeff_min, coeff_max)] # Coefficients: determined from PCA model
bounds_obj = bounds_z + bounds_flam + bounds_pca
# Create the obj_dict
obj_dict = dict(npca=obj_params['npca'], pca_dict=qso_pca_dict)
return obj_dict, bounds_obj
# QSO evaluation function. Model for QSO is a PCA spectrum
[docs]def eval_qso_model(theta, obj_dict):
"""
Routine to evaluate a sensitivity function model
for a QSO spectrum.
Parameters
----------
theta : `numpy.ndarray`_
Array containing the PCA coefficients
shape=(ntheta,)
obj_dict : dict
Dictionary containing additional arguments needed to evaluate the PCA model
Returns
-------
qso_pca_model : array with same shape as the PCA vectors (stored in the obj_dict['pca_dict'])
PCA vectors were already interpolated onto the telluric model grid by init_qso_model
gpm : `numpy.ndarray`_ : array with same shape as the qso_pca_model
Good pixel mask indicating where the model is valid
"""
qso_pca_model = qso_pca_eval(theta, obj_dict['pca_dict'])
# TODO Is the prior evaluation slowing things down??
# TODO Disablingthe prior for now as I think it slows things down for no big gain
#ln_pca_pri = qso_pca.pca_lnprior(theta_PCA, arg_dict['pca_dict'])
#ln_pca_pri = 0.0
#flux_model, tell_model, spec_model, modelmask
return qso_pca_model, (qso_pca_model > 0.0)
##############
# Star Model #
##############
[docs]def init_star_model(obj_params, iord, wave, flux, ivar, mask, tellmodel):
"""
Routine used to initialize the star spectrum model for telluric fits. The star model is the true spectrum
of the standard star times a polynomial to accomodate situations where the star-model is not perfect.
Parameters
----------
obj_params : dict
Dictionary containing parameters necessary for initializing the quasar model
iord : int
Order in question. This is used here because each echelle order can have a different polynomial order
wave : array shape (nspec,)
Wavelength array for the object in question
flux : array shape (nspec,)
Flux array for the object in question
ivar : array shape (nspec,)
Inverse variance array for the oejct in question
mask : array shape (nspec,)
Good pixel mask for the object in question
tellmodel : array shape (nspec,)
This is a telluric model computed on the wave wavelength grid. Initialization usually requires some initial
best guess for the telluric absorption, which is computed from the midpoint of the telluric model grid parameter
space using the resolution of the spectrograph and the airmass of the observations.
Returns
-------
obj_dict : dict
Dictionary containing the meta-information and variables that are used for the object model evaluations.
bounds_obj : tuple
Tuple of bounds for each parameter that will be fit for the object model, which are here the polynomial
coefficients.
"""
# Model parameter guess for starting the optimizations
flam_true = scipy.interpolate.interp1d(obj_params['std_dict']['wave'].value,
obj_params['std_dict']['flux'].value, kind='linear',
bounds_error=False, fill_value=np.nan)(wave)
flam_model = flam_true*tellmodel
flam_model_ivar = (100.0*utils.inverse(flam_model))**2 # This is just a bogus noise to give S/N of 100
flam_model_mask = np.isfinite(flam_model)
# As solve_poly_ratio is designed to multiply a scale factor into the flux, and not the flux_ref, we
# set the flux_ref to be the data here, i.e. flux
scale, fit_tuple, flux_scale, ivar_scale, outmask = coadd.solve_poly_ratio(
wave, flam_model, flam_model_ivar, flux, ivar, obj_params['polyorder_vec'][iord],
mask=flam_model_mask, mask_ref=mask, func=obj_params['func'], model=obj_params['model'])
coeff, wave_min, wave_max = fit_tuple
if(wave_min != wave.min()) or (wave_max != wave.max()):
msgs.error('Problem with the wave_min or wave_max')
# Polynomial coefficient bounds
bounds_obj = [(np.fmin(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][0], obj_params['minmax_coeff_bounds'][0]),
np.fmax(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][1], obj_params['minmax_coeff_bounds'][1]))
for this_coeff in coeff]
# Create the obj_dict
obj_dict = dict(wave=wave, wave_min=wave_min, wave_max=wave_max, flam_true=flam_true, func=obj_params['func'],
model=obj_params['model'], polyorder=obj_params['polyorder_vec'][iord])
if obj_params['debug']:
plt.plot(wave, flux, drawstyle='steps-mid', alpha=0.7, zorder=5, label='star spectrum')
plt.plot(wave, flux_scale, drawstyle='steps-mid', alpha=0.7, zorder=4, label='poly_model*star_model*telluric')
plt.plot(wave, flam_model, label='star_model*telluric')
plt.plot(wave, flam_true, label='star_model')
plt.ylim(-0.1 * flam_model.min(), 1.3 * flam_model.max())
plt.legend()
plt.title('Sensitivity Function Guess for iord={:d}'.format(iord+1)) # +1 to account 0-index starting
plt.show()
return obj_dict, bounds_obj
# Star evaluation function.
[docs]def eval_star_model(theta, obj_dict):
"""
Routine to evaluate a star spectrum model as a true model spectrum times a polynomial.
Parameters
----------
theta : `numpy.ndarray`_
Array containing the polynomial coefficients.
shape (ntheta,)
obj_dict : dict
Dictionary containing additional arguments needed to evaluate the star model.
Returns
-------
star_model : `numpy.ndarray`_
PCA vectors were already interpolated onto the telluric model grid by init_qso_model.
array with same shape obj_dict['wave']
gpm : `numpy.ndarray`_
Good pixel mask indicating where the model is valid.
array with same shape as the star_model
"""
wave_star = obj_dict['wave']
wave_min = obj_dict['wave_min']
wave_max = obj_dict['wave_max']
flam_true = obj_dict['flam_true']
func = obj_dict['func']
model = obj_dict['model']
ymult = coadd.poly_model_eval(theta, func, model, wave_star, wave_min, wave_max)
star_model = ymult*flam_true
return star_model, (star_model > 0.0)
####################
# Polynomial Model #
####################
[docs]def init_poly_model(obj_params, iord, wave, flux, ivar, mask, tellmodel):
"""
Routine used to initialize a polynomial object model for telluric fits.
Parameters
----------
obj_params : dict
Dictionary containing parameters necessary for initializing the quasar model
iord : int
Order in question. This is used here because each echelle order can have a different polynomial order
wave : array shape (nspec,)
Wavelength array for the object in question
flux : array shape (nspec,)
Flux array for the object in question
ivar : array shape (nspec,)
Inverse variance array for the oejct in question
mask : array shape (nspec,)
Good pixel mask for the object in question
tellmodel : array shape (nspec,)
This is a telluric model computed on the wave wavelength grid. Initialization usually requires some initial
best guess for the telluric absorption, which is computed from the midpoint of the telluric model grid parameter
space using the resolution of the spectrograph and the airmass of the observations.
Returns
-------
obj_dict : dict
Dictionary containing the meta-information and variables that are used for the object model evaluations.
bounds_obj : tuple
Tuple of bounds for each parameter that will be fit for the object model, which are here the polynomial
coefficients.
"""
tellmodel_ivar = (100.0*utils.inverse(tellmodel))**2 # This is just a bogus noise to give S/N of 100
tellmodel_mask = np.isfinite(tellmodel) & mask
if obj_params['mask_lyman_a']:
mask = mask & (wave>1216.15*(1+obj_params['z_obj']))
# As solve_poly_ratio is designed to multiply a scale factor into the flux, and not the flux_ref, we
# set the flux_ref to be the data here, i.e. flux
scale, fit_tuple, flux_scale, ivar_scale, outmask = coadd.solve_poly_ratio(
wave, tellmodel, tellmodel_ivar, flux, ivar, obj_params['polyorder_vec'][iord],
mask=tellmodel_mask, mask_ref=mask, func=obj_params['func'], model=obj_params['model'], scale_max=1e5)
# TODO JFH Sticky = False seems to recover better from bad initial fits. Maybe we should change this since poly ratio
# uses a different optimizer.
coeff, wave_min, wave_max = fit_tuple
if(wave_min != wave.min()) or (wave_max != wave.max()):
msgs.error('Problem with the wave_min or wave_max')
# Polynomial model
polymodel = coadd.poly_model_eval(coeff, obj_params['func'], obj_params['model'], wave, wave_min, wave_max)
# Polynomial coefficient bounds
bounds_obj = [(np.fmin(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][0], obj_params['minmax_coeff_bounds'][0]),
np.fmax(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][1], obj_params['minmax_coeff_bounds'][1]))
for this_coeff in coeff]
# Create the obj_dict
obj_dict = dict(wave=wave, wave_min=wave_min, wave_max=wave_max, polymodel=polymodel, func=obj_params['func'],
model=obj_params['model'], polyorder=obj_params['polyorder_vec'][iord])
if obj_params['debug']:
plt.plot(wave, flux, drawstyle='steps-mid', alpha=0.7, zorder=5, label='observed spectrum')
plt.plot(wave, flux_scale, drawstyle='steps-mid', alpha=0.7, zorder=4, label='poly_model*telluric')
plt.plot(wave, tellmodel, label='telluric')
plt.plot(wave, polymodel, label='poly_model')
plt.xlim(wave[mask].min(), wave[mask].max())
plt.ylim(-0.3 * flux[mask].min(), 1.3 * flux[mask].max())
plt.legend()
plt.title('Sensitivity Function Guess for iord={:d}'.format(iord + 1)) # +1 to account 0-index starting
plt.show()
return obj_dict, bounds_obj
# Polynomial evaluation function.
[docs]def eval_poly_model(theta, obj_dict):
"""
Routine to evaluate a star spectrum model as a true
model spectrum times a polynomial.
Parameters
----------
theta : `numpy.ndarray`_
Array containing the polynomial coefficients.
shape=(ntheta,)
obj_dict : dict
Dictionary containing additional arguments needed to evaluate the star model.
Returns
-------
star_model : `numpy.ndarray`_
array with same shape obj_dict['polymodel']
gpm : `numpy.ndarray`_
Good pixel mask indicating where the model is valid.
array with same shape as the star_model.
"""
polymodel = coadd.poly_model_eval(theta, obj_dict['func'], obj_dict['model'],
obj_dict['wave'], obj_dict['wave_min'], obj_dict['wave_max'])
return polymodel, (polymodel > 0.0)
[docs]def sensfunc_telluric(wave, counts, counts_ivar, counts_mask, exptime, airmass, std_dict,
telgridfile, log10_blaze_function=None, ech_orders=None, polyorder=8,
tell_npca=4, teltype='pca',
mask_hydrogen_lines=True, mask_helium_lines=False, hydrogen_mask_wid=10.,
resln_guess=None, resln_frac_bounds=(0.6, 1.4), pix_shift_bounds=(-5.0, 5.0),
delta_coeff_bounds=(-20.0, 20.0), minmax_coeff_bounds=(-5.0, 5.0),
sn_clip=30.0, ballsize=5e-4, only_orders=None, maxiter=3, lower=3.0,
upper=3.0, tol=1e-3, popsize=30, recombination=0.7, polish=True, disp=False,
debug_init=False, debug=False):
r"""
Compute a sensitivity function from a standard star spectrum by
simultaneously fitting a polynomial sensitivity function and a telluric
model for atmospheric absorption.
This method is primarily used with :class:`~pypeit.sensfunc.SensFunc` to
compute sensitivity functions.
The sensitivity function is defined to be :math:`S_\lambda = F_\lambda/({\rm
counts/s/A})` where :math:`F_\lambda` is in units of 1e-17 :math:`{\rm
erg/s/cm^2/A}`, and so the sensitivity function has units of
:math:`{\rm (erg/s/cm^2/A)/(counts/s/A) = erg/cm^2/counts}`
Parameters
----------
wave : `numpy.ndarray`_, shape is (nspec,) or (nspec, norders)
Wavelength array. If array is 2D, the data is likely from an echelle
spectrograph.
counts : `numpy.ndarray`_, shape must match ``wave``
Counts for the object in question.
counts_ivar : `numpy.ndarray`_, shape must match ``wave``
Inverse variance for the object in question.
counts_mask : `numpy.ndarray`_, shape must match ``wave``
Good pixel mask for the object in question.
exptime : :obj:`float`
Exposure time
airmass : :obj:`float`
Airmass of the observation
std_dict : :obj:`dict`
Dictionary containing the information for the true flux of the standard
star.
log10_blaze_function : `numpy.ndarray`_ , optional
The log10 blaze function determined from a flat field image. If this is
passed in the sensitivity function model will be a (parametric)
polynomial fit multiplied into the (non-parametric)
log10_blaze_function. Shape must match ``wave``, i.e. (nspec,) or
(nspec, norddet).
telgridfile : :obj:`str`
File containing grid of HITRAN atmosphere models; see
:class:`~pypeit.par.pypeitpar.TelluricPar`.
ech_orders : `numpy.ndarray`_, shape is (norders,), optional
If passed, provides the true order numbers for the spectra provided.
polyorder : :obj:`int`, optional, default = 8
Polynomial order for the sensitivity function fit.
teltype : :obj:`str`, optional, default = 'pca'
Method for evaluating telluric models, either `pca` or `grid`.
tell_npca : :obj:`int`, optional, default = 4
Number of telluric PCA components used, must be <= 10
mask_hydrogen_lines : :obj:`bool`, optional
If True, mask stellar hydrogen absorption lines before fitting sensitivity function. Default = True
mask_helium_lines : :obj:`bool`, optional
If True, mask stellar helium absorption lines before fitting sensitivity function. Default = False
hydrogen_mask_wid : :obj:`float`, optional
Parameter describing the width of the mask for or stellar absorption lines (`i.e.`, ``mask_hydrogen_lines=True``)
in Angstroms. A region equal to ``hydrogen_mask_wid`` on either side of the line center is masked.
Default = 10A
resln_guess : :obj:`float`, optional
A guess for the resolution of your spectrum expressed as
lambda/dlambda. The resolution is fit explicitly as part of the
telluric model fitting, but this guess helps determine the bounds
for the optimization (see ``resln_frac_bounds``). If not
provided, the wavelength sampling of your spectrum will be used
and the resolution calculated using a typical sampling of 3
spectral pixels per resolution element.
resln_frac_bounds : :obj:`tuple`, optional
A two-tuple with the bounds for the resolution fit optimization
which is part of the telluric model. This range is in units of
``resln_guess``. For example, ``(0.5, 1.5)`` would bound the
spectral resolution fit to be within the range
``(0.5*resln_guess, 1.5*resln_guess)``.
delta_coeff_bounds : :obj:`tuple`, optional, default = (-20.0, 20.0)
Parameters setting the polynomial coefficient bounds for sensfunc
optimization.
minmax_coeff_bounds : :obj:`tuple`, optional, default = (-5.0, 5.0)
Parameters setting the polynomial coefficient bounds for sensfunc
optimization. Bounds are currently determined as follows. We compute an
initial fit to the sensitivity function using
:func:`init_sensfunc_model`, which determines a set of coefficients. The
bounds are then determined according to::
[(np.fmin(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][0],
obj_params['minmax_coeff_bounds'][0]),
np.fmax(np.abs(this_coeff)*obj_params['delta_coeff_bounds'][1],
obj_params['minmax_coeff_bounds'][1]))]
sn_clip : :obj:`float`, optional, default=30.0
Errors are capped during rejection so that the S/N is never greater than
``sn_clip``. This prevents overly aggressive rejection in high S/N ratio
spectrum, which neverthless differ at a level greater than the implied
S/N due to systematics. These systematics are most likely due to the
inadequecy of our telluric model to fit the data at the the precision
required at very high S/N ratio.
only_orders : :obj:`list` of int, optional, default = None
Only fit these specific orders
tol : :obj:`float`, optional, default = 1e-3
Relative tolerance for convergence of the differential evolution
optimization. See `scipy.optimize.differential_evolution`_ for details.
popsize : :obj:`int`, optional, default = 30
A multiplier for setting the total population size for the differential
evolution optimization. See `scipy.optimize.differential_evolution`_ for
details.
recombination : :obj:`float`, optional, default = 0.7
The recombination constant for the differential evolution optimization.
This should be in the range ``[0, 1]``. See
`scipy.optimize.differential_evolution`_ for details.
polish : :obj:`bool`, optional, default=True
If True then differential evolution will perform an additional
optimization at the end to polish the best fit, which can improve the
optimization slightly. See `scipy.optimize.differential_evolution`_ for
details.
disp : :obj:`bool`, optional, default=True
Argument for `scipy.optimize.differential_evolution`_, which will
display status messages to the screen indicating the status of the
optimization. See above for a description of the output and how to know
if things are working well.
debug_init : :obj:`bool`, optional, default=False
Show plots to the screen useful for debugging model initialization
debug : :obj:`bool`, optional, default=False
Show plots to the screen useful for debugging the telluric/object model
fits.
Returns
-------
TelObj : :class:`Telluric`
Best-fitting telluric model
"""
# Turn on disp for the differential_evolution if debug mode is turned on.
if debug:
disp = True
norders = counts.shape[1] if counts.ndim == 2 else 1
# Create the polyorder_vec
if np.size(polyorder) > 1:
if np.size(polyorder) != norders:
msgs.error('polyorder must have either have norder elements or be a scalar')
# TODO: Should this be np.asarray?
polyorder_vec = np.array(polyorder)
else:
polyorder_vec = np.full(norders, polyorder)
func ='legendre'
# Initalize the object parameters
obj_params = dict(std_dict=std_dict, airmass=airmass, delta_coeff_bounds=delta_coeff_bounds,
minmax_coeff_bounds=minmax_coeff_bounds, polyorder_vec=polyorder_vec,
exptime=exptime, func=func, sigrej=3.0, std_src=std_dict['std_source'],
std_ra=std_dict['std_ra'], std_dec=std_dict['std_dec'],
std_name=std_dict['name'], std_cal=std_dict['cal_file'],
output_meta_keys=('airmass', 'exptime', 'polyorder_vec', 'func', 'std_src',
'std_ra', 'std_dec', 'std_name', 'std_cal'),
debug=debug_init)
# Optionally, mask prominent stellar absorption features
mask_bad, mask_recomb, mask_tell = flux_calib.get_mask(wave, counts, counts_ivar, counts_mask,
mask_hydrogen_lines=mask_hydrogen_lines,
mask_helium_lines=mask_helium_lines,
mask_telluric=False, hydrogen_mask_wid=hydrogen_mask_wid)
mask_tot = mask_bad & mask_recomb & mask_tell
# Since we are fitting a sensitivity function, first compute counts per second per angstrom.
TelObj = Telluric(wave, counts, counts_ivar, mask_tot, telgridfile, obj_params,
init_sensfunc_model, eval_sensfunc_model, log10_blaze_function=log10_blaze_function,
teltype=teltype, tell_npca=tell_npca,
ech_orders=ech_orders, pix_shift_bounds=pix_shift_bounds,
resln_guess=resln_guess, resln_frac_bounds=resln_frac_bounds, sn_clip=sn_clip,
maxiter=maxiter, lower=lower, upper=upper, tol=tol,
popsize=popsize, recombination=recombination, polish=polish, disp=disp,
sensfunc=True, debug=debug)
TelObj.run(only_orders=only_orders)
return TelObj
[docs]def create_bal_mask(wave, bal_wv_min_max):
"""
Example of a utility function for creating a BAL mask for QSOs with BAL features. Can also be used to mask other
features that the user does not want to fit.
Parameters
----------
wave : `numpy.ndarray`_
Wavelength array for the quasar in question
shape=(nspec,)
Returns
-------
gpm : `numpy.ndarray`_
Good pixel (non-bal pixels) mask for the fits.
shape=(nspec,)
"""
if np.size(bal_wv_min_max) % 2 !=0:
msgs.error('bal_wv_min_max must be a list/array with even numbers.')
bal_bpm = np.zeros_like(wave, dtype=bool)
nbal = int(np.size(bal_wv_min_max) / 2)
if isinstance(bal_wv_min_max, list):
bal_wv_min_max = np.array(bal_wv_min_max)
wav_min_max = np.reshape(bal_wv_min_max,(nbal,2))
for ibal in range(nbal):
bal_bpm |= (wave > wav_min_max[ibal,0]) & (wave < wav_min_max[ibal,1])
return np.logical_not(bal_bpm)
[docs]def qso_telluric(spec1dfile, telgridfile, pca_file, z_qso, telloutfile, outfile, npca=8,
pca_lower=1220.0, pca_upper=3100.0, bal_wv_min_max=None, delta_zqso=0.1,
teltype='pca', tell_npca=4,
bounds_norm=(0.1, 3.0), tell_norm_thresh=0.9, sn_clip=30.0, only_orders=None,
maxiter=3, tol=1e-3, popsize=30, recombination=0.7, polish=True, disp=False,
pix_shift_bounds=(-5.0,5.0), debug_init=False, debug=False, show=False,
chk_version=True):
"""
Fit and correct a QSO spectrum for telluric absorption.
Parameters
----------
spec1dfile : :obj:`str`
File name of the input 1D spectrum. Can be a ``PypeIt`` spec1d file or
the output from a ``PypeIt`` 1D coadd.
telgridfile : :obj:`str`
File name with the grid of telluric spectra.
pca_file: :obj:`str`
File name of the QSO PCA model fits file.
z_qso : :obj:`float`
QSO redshift.
telloutfile : :obj:`str`
Output file name for the best-fit telluric model.
outfile : :obj:`str`
Output file name of the telluric corrected spectrum.
npca : :obj:`int`, optional
Numer of pca components.
pca_lower: :obj:`float`, optional
Wavelength lower bounds of the PCA model
pca_upper: :obj:`float`, optional
Wavelength upper bounds of the PCA model
bal_wv_min_max : :obj:`list`, optional
Broad absorption line feature mask. If there are several BAL features,
the format for this mask is ``[wave_min_bal1, wave_max_bal1,
wave_min_bal2, wave_max_bal2,...]``. These wavelength ranges will be
ignored during the fitting.
delta_zqso : :obj:`float`, optional
During the fit, the QSO redshift is allowed to vary within
``+/-delta_zqso``.
teltype : :obj:`str`, optional, default = 'pca'
Method for evaluating telluric models, either `pca` or `grid`.
tell_npca : :obj:`int`, optional, default = 4
Number of telluric PCA components used, must be <=10
bounds_norm : :obj:`tuple`, optional
A two-tuple with the lower and upper bounds on the fractional adjustment
of the flux in the QSO model spectrum. For example, a value of ``(0.1,
3.0)`` means the best-fit QSO model flux is allowed to vary between 10%
and a factor of 3 times the input QSO model.
tell_norm_thresh : :obj:`float`, optional
When initializing the model, pixels in the reference telluric spectrum
below this value are not included in an initial guess of the model flux
normalization.
sn_clip : :obj:`float`, optional, default=30.0
Errors are capped during rejection so that the S/N is never greater than
``sn_clip``. This prevents overly aggressive rejection in high S/N ratio
spectrum, which neverthless differ at a level greater than the implied
S/N due to systematics. These systematics are most likely due to the
inadequecy of our telluric model to fit the data at the the precision
required at very high S/N ratio.
only_orders : :obj:`list` of int, optional, default = None
Only fit these specific orders
maxiter : :obj:`int`, optional, default = 3
Maximum number of iterations for the telluric + object model fitting.
The code performs multiple iterations rejecting outliers at each step.
The fit is then performed anew to the remaining good pixels. For this
reason if you run with the disp=True option, you will see that the f(x)
loss function gets progressively better during the iterations.
tol : :obj:`float`, default = 1e-3
Relative tolerance for converage of the differential evolution
optimization. See `scipy.optimize.differential_evolution`_ for details.
popsize : :obj:`int`, default = 30
A multiplier for setting the total population size for the differential
evolution optimization. See `scipy.optimize.differential_evolution`_ for
details.
recombination : :obj:`float`, optional, default = 0.7
The recombination constant for the differential evolution optimization.
This should be in the range [0, 1]. See
`scipy.optimize.differential_evolution`_ for details.
polish : :obj:`bool`, optional, default=True
If True then differential evolution will perform an additional
optimizatino at the end to polish the best fit at the end, which can
improve the optimization slightly. See
`scipy.optimize.differential_evolution`_ for details.
disp : :obj:`bool`, optional, default=True
Argument for `scipy.optimize.differential_evolution`_ that will display
status messages to the screen indicating the status of the optimization.
See above for a description of the output and how to know if things are
working well. Note: this is automatically set to True if debug is True.
debug_init : :obj:`bool`, optional, default=False
Show plots to the screen useful for debugging model initialization
debug : :obj:`bool`, optional, default=False
Show plots to the screen useful for debugging the telluric/object model
fits.
show : :obj:`bool`, optional
Show a QA plot of the final fit.
chk_version : :obj:`bool`, optional
When reading in existing files written by PypeIt, perform strict version
checking to ensure a valid file. If False, the code will try to keep
going, but this may lead to faults and quiet failures. User beware!
Returns
-------
TelObj : :class:`Telluric`
Object with the telluric modeling results
"""
# Turn on disp for the differential_evolution if debug mode is turned on.
if debug:
disp = True
obj_params = dict(pca_file=pca_file, npca=npca, z_qso=z_qso, delta_zqso=delta_zqso,
lbound_norm=bounds_norm[0], ubound_norm=bounds_norm[1],
tell_norm_thresh=tell_norm_thresh,
output_meta_keys=('pca_file', 'npca', 'z_qso', 'delta_zqso',
'lbound_norm', 'ubound_norm', 'tell_norm_thresh'),
debug_init=debug_init)
wave, wave_grid_mid, flux, ivar, mask, meta_spec, header \
= general_spec_reader(spec1dfile, ret_flam=True, chk_version=chk_version)
header = fits.getheader(spec1dfile) # clean this up!
# Mask the IGM and mask wavelengths that extend redward of our PCA
qsomask = (wave > (1.0 + z_qso)*pca_lower) & (wave < pca_upper*(1.0 +z_qso))
# Mask BAL if there is any BAL absorptions.
mask_tot = mask & qsomask & create_bal_mask(wave, bal_wv_min_max) \
if bal_wv_min_max is not None else mask & qsomask
# parameters lowered for testing
TelObj = Telluric(wave, flux, ivar, mask_tot, telgridfile, obj_params, init_qso_model,
eval_qso_model, pix_shift_bounds=pix_shift_bounds,
sn_clip=sn_clip, maxiter=maxiter, tol=tol, popsize=popsize, teltype=teltype,
tell_npca=tell_npca, recombination=recombination, polish=polish,
disp=disp, debug=debug)
TelObj.run(only_orders=only_orders)
TelObj.to_file(telloutfile, overwrite=True)
# Apply the telluric correction
telluric = TelObj.model['TELLURIC'][0,:]
qso_pca_model = TelObj.model['OBJ_MODEL'][0,:]
# Plot the telluric corrected and rescaled orders
flux_corr = flux/(telluric + (telluric == 0.0))
ivar_corr = (telluric > 0.0) * ivar * telluric * telluric
mask_corr = (telluric > 0.0) * mask
sig_corr = np.sqrt(utils.inverse(ivar_corr))
if show:
# TODO: This should get moved into a Telluric.show() method.
# Median filter
med_width = int(flux.size*0.001)
flux_med = utils.fast_running_median(flux_corr, med_width)
fig = plt.figure(figsize=(12, 8))
plt.plot(wave, flux_corr, drawstyle='steps-mid', color='0.7', label='corrected data',
alpha=0.7, zorder=5)
plt.plot(wave, flux_med, drawstyle='steps-mid', color='k', label='corrected data',
alpha=0.7, zorder=5)
plt.plot(wave, sig_corr, drawstyle='steps-mid', color='r', label='noise', alpha=0.3,
zorder=1)
plt.plot(wave, qso_pca_model, color='cornflowerblue', linewidth=1.0, label='PCA model',
zorder=7, alpha=0.7)
plt.plot(wave, qso_pca_model.max()*0.9*telluric, color='magenta', drawstyle='steps-mid',
label='telluric', alpha=0.4)
plt.ylim(-0.1*qso_pca_model.max(), 1.5*qso_pca_model.max())
plt.legend()
plt.xlabel('Wavelength')
plt.ylabel('Flux')
plt.show()
# save the telluric corrected spectrum
save_coadd1d_tofits(outfile, wave, flux_corr, ivar_corr, mask_corr, wave_grid_mid=wave_grid_mid,
spectrograph=header['PYP_SPEC'], telluric=telluric, obj_model=qso_pca_model,
header=header, ex_value='OPT', overwrite=True)
return TelObj
[docs]def star_telluric(spec1dfile, telgridfile, telloutfile, outfile, star_type=None,
star_mag=None, star_ra=None, star_dec=None, func='legendre', model='exp',
polyorder=5, teltype='pca', tell_npca=4, mask_hydrogen_lines=True,
mask_helium_lines=False, hydrogen_mask_wid=10., delta_coeff_bounds=(-20.0, 20.0),
minmax_coeff_bounds=(-5.0, 5.0), only_orders=None, sn_clip=30.0, maxiter=3,
tol=1e-3, popsize=30, recombination=0.7, polish=True, disp=False,
pix_shift_bounds=(-5.0,5.0), debug_init=False, debug=False, show=False,
chk_version=True):
"""
This needs a doc string.
USE the one from qso_telluric() as a starting point
Returns
-------
TelObj : :class:`Telluric`
Object with the telluric modeling results
"""
# Turn on disp for the differential_evolution if debug mode is turned on.
if debug:
disp = True
# Read in the data
wave, wave_grid_mid, flux, ivar, mask, meta_spec, header \
= general_spec_reader(spec1dfile, ret_flam=False, chk_version=chk_version)
# Read in standard star dictionary and interpolate onto regular telluric wave_grid
star_ra = meta_spec['core']['RA'] if star_ra is None else star_ra
star_dec = meta_spec['core']['DEC'] if star_dec is None else star_dec
std_dict = flux_calib.get_standard_spectrum(star_type=star_type, star_mag=star_mag, ra=star_ra,
dec=star_dec)
if flux.ndim == 2:
norders = flux.shape[1]
else:
norders = 1
# Create the polyorder_vec
if np.size(polyorder) > 1:
if np.size(polyorder) != norders:
msgs.error('polyorder must have either have norder elements or be a scalar')
polyorder_vec = np.array(polyorder)
else:
polyorder_vec = np.full(norders, polyorder)
# Initalize the object parameters
obj_params = dict(std_dict=std_dict, airmass=meta_spec['core']['AIRMASS'],
delta_coeff_bounds=delta_coeff_bounds,
minmax_coeff_bounds=minmax_coeff_bounds, polyorder_vec=polyorder_vec,
exptime=meta_spec['core']['EXPTIME'], func=func, model=model, sigrej=3.0,
std_ra=std_dict['std_ra'], std_dec=std_dict['std_dec'],
std_name=std_dict['name'], std_cal=std_dict['cal_file'],
output_meta_keys=('airmass', 'polyorder_vec', 'exptime', 'func', 'std_ra',
'std_dec', 'std_cal'),
debug=debug_init)
# Optionally, mask prominent stellar absorption features
mask_bad, mask_recomb, mask_tell = flux_calib.get_mask(wave, flux, ivar, mask,
mask_hydrogen_lines=mask_hydrogen_lines,
mask_helium_lines=mask_helium_lines,
mask_telluric=False, hydrogen_mask_wid=hydrogen_mask_wid)
mask_tot = mask_bad & mask_recomb & mask_tell
# parameters lowered for testing
TelObj = Telluric(wave, flux, ivar, mask_tot, telgridfile, obj_params, init_star_model,
eval_star_model, pix_shift_bounds=pix_shift_bounds,
teltype=teltype, tell_npca=tell_npca,
sn_clip=sn_clip, tol=tol, popsize=popsize,
recombination=recombination, polish=polish, disp=disp, debug=debug)
TelObj.run(only_orders=only_orders)
TelObj.to_file(telloutfile, overwrite=True)
# Apply the telluric correction
telluric = TelObj.model['TELLURIC'][0,:]
star_model = TelObj.model['OBJ_MODEL'][0,:]
# Plot the telluric corrected and rescaled spectrum
flux_corr = flux*utils.inverse(telluric)
ivar_corr = (telluric > 0.0) * ivar * telluric * telluric
mask_corr = (telluric > 0.0) * mask
sig_corr = np.sqrt(utils.inverse(ivar_corr))
if show:
# TODO: This should get moved into a Telluric.show() method.
# Median filter
fig = plt.figure(figsize=(12, 8))
plt.plot(wave, flux_corr*mask_corr, drawstyle='steps-mid', color='k',
label='corrected data', alpha=0.7, zorder=5)
plt.plot(wave, flux*mask_corr, drawstyle='steps-mid', color='0.7',
label='uncorrected data', alpha=0.7, zorder=3)
plt.plot(wave, sig_corr*mask_corr, drawstyle='steps-mid', color='r', label='noise',
alpha=0.3, zorder=1)
plt.plot(wave, star_model, color='cornflowerblue', linewidth=1.0,
label='poly scaled star model', zorder=7, alpha=0.7)
plt.plot(std_dict['wave'].value, std_dict['flux'].value, color='green', linewidth=1.0,
label='original star model', zorder=8, alpha=0.7)
plt.plot(wave, star_model.max()*0.9*telluric, color='magenta', drawstyle='steps-mid',
label='telluric', alpha=0.4)
plt.ylim(-np.median(sig_corr[mask_corr]).max(), 1.5*star_model.max())
plt.xlim(wave[wave > 1.0].min(), wave[wave > 1.0].max())
plt.legend()
plt.xlabel('Wavelength')
plt.ylabel('Flux')
plt.show()
# save the telluric corrected spectrum
save_coadd1d_tofits(outfile, wave, flux_corr, ivar_corr, mask_corr, wave_grid_mid=wave_grid_mid,
spectrograph=header['PYP_SPEC'], telluric=telluric,
obj_model=star_model, header=header, ex_value='OPT', overwrite=True)
return TelObj
[docs]def poly_telluric(spec1dfile, telgridfile, telloutfile, outfile, z_obj=0.0, func='legendre',
model='exp', polyorder=3, fit_wv_min_max=None, mask_lyman_a=True, teltype='pca',
tell_npca=4, delta_coeff_bounds=(-20.0, 20.0), minmax_coeff_bounds=(-5.0, 5.0),
only_orders=None, sn_clip=30.0, maxiter=3, tol=1e-3, popsize=30,
recombination=0.7, polish=True, disp=False, pix_shift_bounds=(-5.0,5.0),
debug_init=False, debug=False, show=False, chk_version=True):
"""
This needs a doc string.
FILL IN
COULD USE THE qso_telluric() doc string as a starting point
"""
# Turn on disp for the differential_evolution if debug mode is turned on.
if debug:
disp = True
# Read in the data
wave, wave_grid_mid, flux, ivar, mask, meta_spec, header \
= general_spec_reader(spec1dfile, ret_flam=False, chk_version=chk_version)
if flux.ndim == 2:
norders = flux.shape[1]
else:
norders = 1
# Create the polyorder_vec
if np.size(polyorder) > 1:
if np.size(polyorder) != norders:
msgs.error('polyorder must have either have norder elements or be a scalar')
polyorder_vec = np.array(polyorder)
else:
polyorder_vec = np.full(norders, polyorder)
# Initalize the object parameters
obj_params = dict(z_obj=z_obj, mask_lyman_a=mask_lyman_a, airmass=meta_spec['core']['AIRMASS'],
delta_coeff_bounds=delta_coeff_bounds,
minmax_coeff_bounds=minmax_coeff_bounds, polyorder_vec=polyorder_vec,
exptime=meta_spec['core']['EXPTIME'], func=func, model=model, sigrej=3.0,
output_meta_keys=('airmass', 'polyorder_vec', 'exptime', 'func'),
debug=debug_init)
# Optionally, only using the redward of Lyman-alpha line to do the fitting
if mask_lyman_a:
inmask = wave > 1216.15 * (1+z_obj)
mask_tot = inmask & mask
else:
mask_tot = mask
if fit_wv_min_max is not None:
mask_tot &= np.logical_not(create_bal_mask(wave, fit_wv_min_max))
# parameters lowered for testing
TelObj = Telluric(wave, flux, ivar, mask_tot, telgridfile, obj_params, init_poly_model,
eval_poly_model, pix_shift_bounds=pix_shift_bounds,
sn_clip=sn_clip, maxiter=maxiter, tol=tol, popsize=popsize, teltype=teltype,
tell_npca=tell_npca, recombination=recombination, polish=polish, disp=disp, debug=debug)
TelObj.run(only_orders=only_orders)
TelObj.to_file(telloutfile, overwrite=True)
# Apply the telluric correction
telluric = TelObj.model['TELLURIC'][0,:]
poly_model = TelObj.model['OBJ_MODEL'][0,:]
# Plot the telluric corrected and rescaled spectrum
flux_corr = flux*utils.inverse(telluric)
ivar_corr = (telluric > 0.0) * ivar * telluric * telluric
mask_corr = (telluric > 0.0) * mask
sig_corr = np.sqrt(utils.inverse(ivar_corr))
if show:
# TODO: This should get moved into a Telluric.show() method.
# Median filter
fig = plt.figure(figsize=(12, 8))
plt.plot(wave, flux_corr*mask_corr, drawstyle='steps-mid', color='k',
label='corrected data', alpha=0.8, zorder=5)
plt.plot(wave, flux*mask_corr, drawstyle='steps-mid', color='0.7',
label='uncorrected data', alpha=0.5, zorder=3)
plt.plot(wave, sig_corr*mask_corr, drawstyle='steps-mid', color='b', label='noise',
alpha=0.3, zorder=1)
#plt.plot(wave, poly_model, color='cornflowerblue', linewidth=1.0,
# label='polynomial model', zorder=7, alpha=0.7)
plt.plot(wave, poly_model[mask_corr].max()*0.9*telluric, color='magenta',
drawstyle='steps-mid', label='telluric', alpha=0.3)
plt.ylim(-np.median(sig_corr[mask_corr]).max(), 1.5*flux_corr[mask_corr].max())
plt.xlim(wave[mask_corr].min(), wave[mask_corr].max())
plt.legend()
plt.xlabel('Wavelength')
plt.ylabel('Flux')
plt.show()
# save the telluric corrected spectrum
save_coadd1d_tofits(outfile, wave, flux_corr, ivar_corr, mask_corr, wave_grid_mid=wave_grid_mid,
spectrograph=header['PYP_SPEC'], telluric=telluric, obj_model=poly_model,
header=header, ex_value='OPT', overwrite=True)
return TelObj
[docs]class Telluric(datamodel.DataContainer):
r"""
Simultaneously fit model object and telluric spectra to an observed
spectrum.
The model telluric absorption spectra for the atmosphere are generated by
HITRAN atmosphere models with a baseline atmospheric profile model of the
observatory in question. One can interpolate over grids of these absorption
spectra directly, which have been created for each observatory, or instead
(by default) use a PCA decomposition of the telluric models across all
observatories, which is more flexible and much lighter in file size.
The modeling can be applied to both multi-slit
or echelle data. The code performs a differential evolution optimization
using `scipy.optimize.differential_evolution`_. Several iterations are
performed with rejection of outliers (governed by the maxiter parameter
below). In general it produces very good results, but is rather slow.
Progress can be monitored using the ``disp=True`` option, which will
print out the progress to the screen. This will look like::
differential_evolution step 101: f(x)= 105392
Here ``f(x)`` represents the value of the loss function, which is
essentially :math:`\chi^2` (but with the non-Gaussian tails remapped and
outliers rejected). Things are going well if the result of differential
evolution yields a value of ``f(x)`` comparable to the number of pixels in
your spectrum (or in the order in question for echelle data), which are the
number of degrees of freedom. If ``debug=True``, it will show the residuals
from the fit at each iteration. Note that for spectra with S/N > 30, the
residual distribution may look narrower because a floor is added to the
noise (see ``sn_clip``) to prevent excessive rejection and poorly behaved
fits. Basically these telluric models are only good to about 3%, and so we
cannot fit the data within the noise better than that, which is why
``sn_clip`` is implemented.
The object model is specified by the user through two user provided
functions ``init_obj_model`` and ``eval_obj_model``. The
``init_obj_model`` function is required because differential evolution
can only perform bounded optimization. The bounds for the telluric are
set by the grid and initialized without user involvement. But for the
object model, the bounds depend on the nature of the object model, which
is why ``init_obj_model`` must be provided.
The telluric model is governed by seven parameters --- for the `grid`
approach: ``pressure``, ``temperature``, ``humidity``, ``airmass``; for
the `pca` approach 4 PCA coefficients; plus ``resln``, ``shift``, and
``stretch`` --- where ``resln`` is the resolution of the spectrograph and
``shift`` is a shift in pixels. The ``stretch`` is a stretch of the pixel
scale. The airmass of the object will be used to initalize the fit (this
helps with initalizing the object model), but the models are sufficiently
flexible that often the best-fit airmass actually differs from the
airmass of the spectrum. In the `pca` approach, the number of PCA
components used can be adjusted between 1 and 10, with a tradeoff between
speed and accuracy.
This code can be run on stacked spectra covering a large range of
airmasses and will still provide good results. The resulting airmass
(in the `grid` approach) will be an effective value, and as per above may
not have much relation to the true airmass of the observation. The
exception to this rule is extremely high signal-to-noise ratio data
(S/N > 30), for which it can be difficult to obtain a good fit within the
noise of the data. In such cases, the user should split up the data into
chunks taken around the same airmass, fit those individually with this
class, and then combine the resulting telluric corrected spectra. This
will, in general, result in better fits, and will also average down the
residuals from the telluric model fit in the final averaged spectrum.
The datamodel attributes are:
.. include:: ../include/class_datamodel_telluric.rst
.. todo::
- List the elements of ``obj_params``.
Args:
wave (`numpy.ndarray`_):
Wavelength array. Must either have shape (nspec,) or (nspec,
norders), the latter being for echelle data.
flux (`numpy.ndarray`_):
Flux for the object in question. Same shape as ``wave``.
ivar (`numpy.ndarray`_):
Inverse variance for the object in question. Same shape as
``wave``.
gpm (`numpy.ndarray`_):
Good pixel gpm for the object in question. Same shape as
``wave``.
telgridfile (:obj:`str`):
File containing grid of HITRAN atmosphere models or PCA
decomposition of such models, see
:class:`~pypeit.par.pypeitpar.TelluricPar`.
teltype (:obj:`str`, optional):
Method for evaluating telluric models, either `pca` or `grid`.
The `grid` method directly interpolates a pre-computed grid of
HITRAN atmospheric models with nearest grid-point interpolation,
while the `pca` method instead fits for the coefficients of a
PCA decomposition of HITRAN atmospheric models, enabling a more
flexible and far more file-size efficient interpolation
of the telluric absorption model space.
tell_npca (:obj:`int`, optional):
Number of telluric PCA components used, must be <=10
obj_params (:obj:`dict`):
Dictionary of parameters for initializing the object model.
init_obj_model (callable):
User defined function for initializing the object model. This
function must follow the calling sequence::
obj_dict, bounds_obj = init_obj_model(obj_params, iord, wave, flux, ivar, gpm,
tellmodel)
See, e.g., :func:`init_star_model` for a detailed explanation of
these paramaters and return values.
eval_obj_model (callable):
User defined function for evaluating the object model at a set of
parameter values (``theta_obj``). This function must follow the
calling sequence::
obj_model, model_gpm = eval_obj_model(theta_obj, obj_dict)
Where ``obj_dict`` is one of the return values from the
``init_obj_model`` above. See, e.g., :func:`eval_star_model` for
a detailed explanation of these paramaters and return values.
log10_blaze_function (`numpy.ndarray`_, optional):
The log10 blaze function determined from a flat field image. If
this is passed in the sensitivity function model will be a
(parametric) polynomial fit multiplied into the (non-parametric)
log10_blaze_function. Shape = (nspec,) or (nspec, norddet), i.e.
the same as ``wave``.
ech_orders (`numpy.ndarray`_, optional):
If passed the echelle orders will be included in the output data.
Must be a numpy array of integers with the shape (norders,)
giving the order numbers.
sn_clip (:obj:`float`, optional):
This adds an error floor to the ``ivar``, preventing too much
rejection at high-S/N (i.e. standard stars, bright objects) using
the function :func:`~pypeit.utils.clip_ivar`. A small error is
added to the input ``ivar`` so that the output ``ivar_out`` will
never give S/N greater than ``sn_clip``. This prevents overly
aggressive rejection in high S/N ratio spectra that neverthless
differ at a level greater than the formal S/N due to the fact
that our telluric models are only good to about 3%.
airmass_guess (:obj:`float`, optional):
A guess for the airmass of your object. The code fits the airmass
as part of the telluric model, but this initial guess is useful
for initializing the object model to determine the bounds for the
object model parameter optimization via ``init_obj_model``, since
typically that is done by dividing out a guess for the telluric
absorption and then performing some kind of initial fit.
resln_guess (:obj:`float`, optional):
A guess for the resolution of your spectrum expressed as
lambda/dlambda. The resolution is fit explicitly as part of the
telluric model fitting, but this guess helps determine the bounds
for the optimization (see ``resln_frac_bounds``). If not
provided, the wavelength sampling of your spectrum will be used
and the resolution calculated using a typical sampling of 3
spectral pixels per resolution element.
resln_frac_bounds (:obj:`tuple`, optional):
A two-tuple with the bounds for the resolution fit optimization
which is part of the telluric model. This range is in units of
``resln_guess``. For example, ``(0.5, 1.5)`` would bound the
spectral resolution fit to be within the range
``(0.5*resln_guess, 1.5*resln_guess)``.
pix_shift_bounds (:obj:`tuple`, optional):
A two-tuple with the bounds for the pixel shift optimization in
telluric model fit in units of pixels. The atmosphere will be
allowed to shift within this range during the fit.
pix_stretch_bounds (:obj:`tuple`, optional):
A two-tuple with the bounds for the pixel scale stretch
optimization in telluric model fit in units of pixels. The
atmosphere spectrum will be allowed to stretch within this range
during the fit.
maxiter (:obj:`int`, optional):
Maximum number of iterations for the telluric + object model
fitting. The code performs multiple iterations rejecting outliers
at each step. The fit is then performed anew to the remaining
good pixels. For this reason if you run with the disp=True
option, you will see that the f(x) loss function gets
progressively better during the iterations.
sticky (:obj:`bool`, optional):
Sticky parameter for the :func:`~pypeit.core.pydl.djs_reject`
algorithm for iterative model fit rejection. If True then points
rejected from a previous iteration are kept rejected, in other
words the bad pixel mask is the OR of all previous iterations and
rejected pixels accumulate. If False, the bad pixel mask is the
mask from the previous iteration, and if the model fit changes
between iterations, points can alternate from being rejected to
not rejected. At present this code only performs optimizations
with differential evolution and experience shows that sticky
needs to be True in order for these to converge. This is because
the outliers can be so large that they dominate the loss
function, and one never iteratively converges to a good model
fit. In other words, the deformations in the model between
iterations with sticky=False are too small to approach a
reasonable fit.
lower (:obj:`float`, optional):
Lower rejection threshold in units of ``sigma_corr*sigma``, where
``sigma`` is the formal noise of the spectrum and ``sigma_corr``
is an empirically determined correction to the formal error. The
distribution of input ``chi`` (defined by ``chi = (data -
model)/sigma``) values is analyzed, and a correction factor to
the formal error ``sigma_corr`` is returned, which is multiplied
into the formal errors. In this way, a rejection threshold of
i.e. 3-sigma, will always correspond to roughly the same
percentile. This renormalization is performed with
:func:`~pypeit.core.coadd.renormalize_errors`, and guarantees that
rejection is not too agressive in cases where the empirical
errors determined from the chi-distribution differ significantly
from the formal noise, which is used to determine ``chi``.
upper (:obj:`float`, optional):
Upper rejection threshold in units of ``sigma_corr*sigma``, where
``sigma`` is the formal noise of the spectrum and ``sigma_corr``
is an empirically determined correction to the formal error. See
``lower`` for more detail.
seed (:obj:`int`, optional):
An initial seed for the differential evolution optimization,
which is a random process. The default is a seed = 777 which will
be used to seed a `numpy.random.Generator`_ object. A specific
seed is used because otherwise the random number generator will
use the time for the seed and the results will not be
reproducible. TODO: (JFH) Note that differential evolution seems
to have some other source of stochasticity that I have not yet
figured out.
ballsize (:obj:`float`, optional):
This parameter governs how the differential evolution random
population is initialized for the object model and for subsequent
iterations. The idea is that after the first iteration we can
generate a new random population in model space about the optimum
from the previous iteration. This significantly speeds up
convergence. The new population of parameters is created by
perturbing the previous optimum by ``theta = theta_opt +
ballsize*(bounds[1]-bounds[0])*standard_normal_deviate``, where
``bounds`` are the bounds for the optimization and
``standard_normal_deviate`` is a random draw from a unit variance
Gaussian. If the ballsize is too small, the optimization may get
stuck in a local maximum. If it is too large it will not speed up
convergence. This choice is highly dependent on how the bounds
are chosen using in the ``init_obj_model`` function and from the
bound parameters for the atmospheric, shift, and resolution
parameters. If the ``init_obj_model`` functions also have a way
of isolating a good guess for the object model, then this can be
put assigned to ``obj_dict['init_obj_opt_theta']`` and the object
model will similarly be initialized about this optimum using the
ballsize for the first iteration. See :func:`init_sensfunc_model`
for an example.
tol (:obj:`float`, optional):
Relative tolerance for converage of the differential evolution
optimization. See `scipy.optimize.differential_evolution`_ for
details.
popsize (:obj:`int`, optional):
A multiplier for setting the total population size for the
differential evolution optimization. See
`scipy.optimize.differential_evolution`_ for details.
recombination (:obj:`float`, optional):
The recombination constant for the differential evolution
optimization. This should be in the range [0, 1]. See
`scipy.optimize.differential_evolution`_ for details.
polish (:obj:`bool`, optional):
If True then differential evolution will perform an additional
optimization at the end to polish the best fit at the end, which
can improve the optimization slightly. See
`scipy.optimize.differential_evolution`_ for details.
disp (:obj:`bool`, optional):
Argument for `scipy.optimize.differential_evolution`_ that prints
status messages to the screen indicating the status of the
optimization. See description above.
sensfunc (:obj:`bool`, optional):
This option is used for usage of this class for joint telluric
fitting and sensitivity function computation. If True, the input
flux is in counts and then converted to counts per angstrom,
since the sensfunc is obtained by fitting counts per angstrom.
TODO: This explanation is unclear.
debug (:obj:`bool`, optional):
If True, QA plots will be shown to the screen indicating the
quality of the fits. Specifically, the residual distributions
will be shown at each iteration, and the fit will be shown at the
end (for each order). This is useful if you are running the code
for the first time, but since the algorithm is slow, particularly
for fitting multi-order echelle data, it will require lots of
clicking to close interactive matplotlib windows which halt code
flow.
"""
version = '1.0.0'
"""Datamodel version."""
datamodel = {'telgrid': dict(otype=str,
descr='File containing PCA components or grid of '
'HITRAN atmosphere models'),
'teltype': dict(otype=str,
descr='Type of telluric model, `pca` or `grid`'),
'tell_npca': dict(otype=int,
descr='Number of telluric PCA components used'),
'std_src': dict(otype=str, descr='Name of the standard source'),
'std_name': dict(otype=str, descr='Type of standard source'),
'std_cal': dict(otype=str,
descr='File name (or shorthand) with the standard flux data'),
'func': dict(otype=str, descr='Polynomial function used'),
'pca_file': dict(otype=str, descr='Name of the QSO PCA file'),
'tol': dict(otype=float,
descr='Relative tolerance for converage of the differential '
'evolution optimization.'),
'popsize': dict(otype=int,
descr='A multiplier for setting the total population size for '
'the differential evolution optimization.'),
'recombination': dict(otype=float,
descr='The recombination constant for the differential '
'evolution optimization. Should be in the range '
'[0, 1].'),
'polish': dict(otype=bool,
descr='Perform a final optimization to tweak the best solution; '
'see scipy.optimize.differential_evolution.'),
'airmass': dict(otype=float, descr='Airmass of the observation'),
'exptime': dict(otype=float, descr='Exposure time (s)'),
# TODO: Is it possible/useful to force the coordinates to always be floats
'std_ra': dict(otype=float, descr='RA of the standard source'),
'std_dec': dict(otype=float, descr='DEC of the standard source'),
'npca': dict(otype=int, descr='Number of PCA components'),
'z_qso': dict(otype=float, descr='Redshift of the QSO'),
'delta_zqso': dict(otype=float,
descr='Allowed range for the QSO redshift about z_qso'),
'lbound_norm': dict(otype=float, descr='Flux normalization lower bound'),
'ubound_norm': dict(otype=float, descr='Flux normalization upper bound'),
'tell_norm_thresh': dict(otype=float, descr='??'),
'model': dict(otype=table.Table, descr='Table with the best-fitting model data')}
"""DataContainer datamodel."""
internals = ['obj_params',
'init_obj_model',
'airmass_guess',
'eval_obj_model',
'ech_orders',
'sn_clip',
'resln_frac_bounds',
'pix_shift_bounds',
'pix_stretch_bounds',
'maxiter',
'sticky',
'lower',
'upper',
'seed',
'rng',
'ballsize',
'diff_evol_maxiter',
'disp',
'sensfunc',
'debug',
'wave_in_arr',
'flux_in_arr',
'ivar_in_arr',
'mask_in_arr',
'log10_blaze_func_in_arr',
'nspec_in',
'norders',
'tell_dict',
'wave_grid',
'ngrid',
'resln_guess',
'tell_guess',
'bounds_tell',
'flux_arr',
'ivar_arr',
'mask_arr',
'log10_blaze_func_arr',
'wave_mask_arr',
'ind_lower',
'ind_upper',
'srt_order_tell',
'obj_dict_list',
'bounds_obj_list',
'bounds_list',
'arg_dict_list',
'max_ntheta_obj',
'result_list',
'outmask_list',
'obj_model_list',
'tellmodel_list',
'theta_obj_list',
'theta_tell_list',
]
[docs] @staticmethod
def empty_model_table(norders, nspec, tell_npca=4, n_obj_par=0):
"""
Construct an empty `astropy.table.Table`_ for the telluric model
results.
Args:
norders (:obj:`int`):
The number of slits/orders on the detector.
nspec (:obj:`int`):
The number of spectral pixels on the detector.
tell_npca (:obj:`int`):
Number of telluric model parameters
n_obj_par (:obj:`int`, optional):
The number of parameters used to construct the object model
spectrum.
Returns:
`astropy.table.Table`_: Instance of the empty model table.
"""
return table.Table(data=[
table.Column(name='WAVE', dtype=float, length=norders, shape=(nspec,),
description='Wavelength vector'),
table.Column(name='TELLURIC', dtype=float, length=norders, shape=(nspec,),
description='Best-fitting telluric model spectrum'),
table.Column(name='OBJ_MODEL', dtype=float, length=norders, shape=(nspec,),
description='Best-fitting object model spectrum'),
# TODO: Why do we need both TELL_THETA and all the individual parameters...
table.Column(name='TELL_THETA', dtype=float, length=norders, shape=(tell_npca+3,),
description='Best-fitting telluric model parameters'),
table.Column(name='TELL_PARAM', dtype=float, length=norders, shape=(tell_npca,),
description='Best-fitting telluric atmospheric parameters or PCA coefficients'),
table.Column(name='TELL_RESLN', dtype=float, length=norders,
description='Best-fitting telluric model spectral resolution'),
table.Column(name='TELL_SHIFT', dtype=float, length=norders,
description='Best-fitting shift applied to the telluric model spectrum'),
table.Column(name='TELL_STRETCH', dtype=float, length=norders,
description='Best-fitting stretch applied to the telluric model spectrum'),
table.Column(name='OBJ_THETA', dtype=float, length=norders, shape=(n_obj_par,),
description='Best-fitting object model parameters'),
table.Column(name='CHI2', dtype=float, length=norders,
description='Chi-square of the best-fit model'),
table.Column(name='SUCCESS', dtype=bool, length=norders,
description='Flag that fit was successful'),
table.Column(name='NITER', dtype=int, length=norders,
description='Number of fit iterations'),
table.Column(name='ECH_ORDERS', dtype=int, length=norders,
description='Echelle order for this specrum (echelle data only)'),
table.Column(name='POLYORDER_VEC', dtype=int, length=norders,
description='Polynomial order for each slit/echelle (if applicable)'),
table.Column(name='IND_LOWER', dtype=int, length=norders,
description='Lowest index of a spectral pixel included in the fit'),
table.Column(name='IND_UPPER', dtype=int, length=norders,
description='Highest index of a spectral pixel included in the fit'),
table.Column(name='WAVE_MIN', dtype=float, length=norders,
description='Minimum wavelength included in the fit'),
table.Column(name='WAVE_MAX', dtype=float, length=norders,
description='Maximum wavelength included in the fit')])
def __init__(self, wave, flux, ivar, gpm, telgridfile, obj_params, init_obj_model, eval_obj_model,
log10_blaze_function=None, ech_orders=None, sn_clip=30.0, teltype='pca', tell_npca=4,
airmass_guess=1.5, resln_guess=None, resln_frac_bounds=(0.3, 1.5), pix_shift_bounds=(-5.0, 5.0),
pix_stretch_bounds=(0.9,1.1), maxiter=2, sticky=True, lower=3.0, upper=3.0,
seed=777, ballsize = 5e-4, tol=1e-3, diff_evol_maxiter=1000, popsize=30,
recombination=0.7, polish=True, disp=False, sensfunc=False, debug=False):
# Instantiate as an empty DataContainer
super().__init__()
# This init function performs the following steps:
# 1) assignement of relevant input arguments
# 2) reshape all spectra to be shape (nspec, norders) which the code operates on
# 3) read in and initalize the telluric grid
# 4) Interpolate spectra onto the fixed telluric wavelength grid, clip S/N
# 5) Loop over orders to initialize object models, and determine index range of fits
# 6) Initalize the output tables
# 1) Assign arguments
self.telgrid = telgridfile
self.teltype = teltype
self.obj_params = obj_params
self.init_obj_model = init_obj_model
self.tell_npca = tell_npca
self.airmass_guess = airmass_guess
self.eval_obj_model = eval_obj_model
self.ech_orders = ech_orders
self.sn_clip = sn_clip
self.resln_frac_bounds = resln_frac_bounds
self.pix_shift_bounds = pix_shift_bounds
self.pix_stretch_bounds = pix_stretch_bounds
self.maxiter = maxiter
self.sticky = sticky
self.lower = lower
self.upper = upper
# Optimizer requires a seed. This guarantees that the fit will be
# deterministic and hence reproducible
self.seed = seed
self.rng = np.random.default_rng(seed)
self.ballsize= ballsize
self.tol = tol
self.diff_evol_maxiter = diff_evol_maxiter
self.popsize = popsize
self.recombination = recombination
self.polish = polish
# Turn on disp for the differential_evolution if debug mode is turned on.
self.disp = disp or debug
self.sensfunc = sensfunc
self.debug = debug
self.log10_blaze_func_in_arr = None
# 2) Reshape all spectra to be (nspec, norders)
if log10_blaze_function is not None:
self.wave_in_arr, self.flux_in_arr, self.ivar_in_arr, self.mask_in_arr, self.log10_blaze_func_in_arr, \
self.nspec_in, self.norders = utils.spec_atleast_2d(
wave, flux, ivar, gpm, log10_blaze_function=log10_blaze_function)
else:
self.wave_in_arr, self.flux_in_arr, self.ivar_in_arr, self.mask_in_arr, _, \
self.nspec_in, self.norders = utils.spec_atleast_2d(
wave, flux, ivar, gpm)
# 3) Read the telluric grid and initalize associated parameters
wv_gpm = self.wave_in_arr > 1.0
if self.teltype == 'pca':
self.tell_dict = read_telluric_pca(self.telgrid, wave_min=self.wave_in_arr[wv_gpm].min(),
wave_max=self.wave_in_arr[wv_gpm].max())
elif self.teltype == 'grid':
self.tell_npca = 4
self.tell_dict = read_telluric_grid(self.telgrid, wave_min=self.wave_in_arr[wv_gpm].min(),
wave_max=self.wave_in_arr[wv_gpm].max())
self.wave_grid = self.tell_dict['wave_grid']
self.ngrid = self.wave_grid.size
self.resln_guess = wvutils.get_sampling(self.wave_in_arr)[2] \
if resln_guess is None else resln_guess
# Model parameter guess for determining the bounds with the init_obj_model function
self.tell_guess = self.get_tell_guess()
# Set the bounds for the telluric optimization
self.bounds_tell = self.get_bounds_tell()
# 4) Interpolate the input values onto the fixed telluric wavelength
# grid, clip S/N and process inmask
if log10_blaze_function is not None:
self.flux_arr, self.ivar_arr, self.mask_arr, self.log10_blaze_func_arr \
= coadd.interp_spec(self.wave_grid, self.wave_in_arr, self.flux_in_arr,
self.ivar_in_arr, self.mask_in_arr, log10_blaze_function=self.log10_blaze_func_in_arr,
sensfunc=self.sensfunc)
else:
self.flux_arr, self.ivar_arr, self.mask_arr, _ \
= coadd.interp_spec(self.wave_grid, self.wave_in_arr, self.flux_in_arr,
self.ivar_in_arr, self.mask_in_arr,
sensfunc=self.sensfunc)
# This is a hack to get an interpolate mask indicating where wavelengths
# are good on each order
_, _, self.wave_mask_arr, _ = coadd.interp_spec(self.wave_grid, self.wave_in_arr,
np.ones_like(self.flux_in_arr),
np.ones_like(self.ivar_in_arr),
(self.wave_in_arr > 1.0).astype(float))
# Clip the ivar if that is requested (sn_clip = None simply returns the ivar otherwise)
self.ivar_arr = utils.clip_ivar(self.flux_arr, self.ivar_arr, self.sn_clip,
gpm=self.mask_arr)
# 5) Loop over orders to initialize object models, and determine index range of fits
# sort the orders by the strength of their telluric absorption
self.ind_lower, self.ind_upper = self.get_ind_lower_upper()
self.srt_order_tell = self.sort_telluric()
# Loop over the data to:
# 1) determine the ind_lower, ind_upper for every order/spectrum
# 2) initialize the obj_dict, and bounds by running the init_obj_model callable
self.obj_dict_list = [None]*self.norders
self.bounds_obj_list = [None]*self.norders
self.bounds_list = [None]*self.norders
self.arg_dict_list = [None]*self.norders
self.max_ntheta_obj = 0
for counter, iord in enumerate(self.srt_order_tell):
msgs.info(f'Initializing object model for order: {iord}, {counter}/{self.norders}'
+ f' with user supplied function: {self.init_obj_model.__name__}')
tellmodel = eval_telluric(self.tell_guess, self.tell_dict,
ind_lower=self.ind_lower[iord],
ind_upper=self.ind_upper[iord])
# TODO This is a pretty ugly way to pass in the blaze function. Particularly since now all the other models
# (star, qso, poly) are going to have this parameter set in their obj_params dictionary.
# Is there something more elegant that can be done with e.g. functools.partial?
obj_params['log10_blaze_func_per_ang'] = \
self.log10_blaze_func_arr[self.ind_lower[iord]:self.ind_upper[iord] + 1, iord] \
if log10_blaze_function is not None else None
obj_dict, bounds_obj \
= init_obj_model(obj_params, iord,
self.wave_grid[self.ind_lower[iord]:self.ind_upper[iord]+1],
self.flux_arr[self.ind_lower[iord]:self.ind_upper[iord]+1,iord],
self.ivar_arr[self.ind_lower[iord]:self.ind_upper[iord]+1,iord],
self.mask_arr[self.ind_lower[iord]:self.ind_upper[iord]+1,iord],
tellmodel)
self.obj_dict_list[iord] = obj_dict
self.bounds_obj_list[iord] = bounds_obj
self.max_ntheta_obj = np.fmax(self.max_ntheta_obj, len(bounds_obj))
bounds_iord = bounds_obj + self.bounds_tell
self.bounds_list[iord] = bounds_iord
arg_dict_iord = dict(ivar=self.ivar_arr[self.ind_lower[iord]:self.ind_upper[iord]+1,iord],
tell_dict=self.tell_dict, ind_lower=self.ind_lower[iord],
ind_upper=self.ind_upper[iord],
tell_npca=self.tell_npca,
obj_model_func=self.eval_obj_model, obj_dict=obj_dict,
ballsize=self.ballsize, bounds=bounds_iord, rng=self.rng,
diff_evol_maxiter=self.diff_evol_maxiter, tol=self.tol,
popsize=self.popsize, recombination=self.recombination,
polish=self.polish, disp=self.disp, debug=debug)
self.arg_dict_list[iord] = arg_dict_iord
# 6) Initalize the output tables
self.init_output()
[docs] def run(self, only_orders=None):
"""
Loops over orders/slits, runs the telluric correction, and evaluates the
object and telluric models.
Parameters
----------
only_orders : :obj:`list`, :obj:`int`, optional
A single integer or list of integers setting the orders to correct.
If None, all orders are corrected.
"""
only_orders = [only_orders] if (only_orders is not None and
isinstance(only_orders, (int, np.integer))) \
else only_orders
good_orders = self.srt_order_tell if only_orders is None else only_orders
# Run the fits
self.result_list = [None]*self.norders
self.outmask_list = [None]*self.norders
self.obj_model_list = [None]*self.norders
self.tellmodel_list = [None]*self.norders
self.theta_obj_list = [None]*self.norders
self.theta_tell_list = [None]*self.norders
for counter, iord in enumerate(self.srt_order_tell):
if iord not in good_orders:
continue
msgs.info(f'Fitting object + telluric model for order: {iord}, {counter}/{self.norders}'
+ f' with user supplied function: {self.init_obj_model.__name__}')
self.result_list[iord], ymodel, ivartot, self.outmask_list[iord] \
= fitting.robust_optimize(self.flux_arr[self.ind_lower[iord]:self.ind_upper[iord]+1,iord],
tellfit, self.arg_dict_list[iord],
inmask=self.mask_arr[self.ind_lower[iord]:self.ind_upper[iord]+1,iord],
maxiter=self.maxiter, lower=self.lower,
upper=self.upper, sticky=self.sticky)
if self.teltype == 'pca':
self.theta_obj_list[iord] = self.result_list[iord].x[:-(self.tell_npca+3)]
self.theta_tell_list[iord] = self.result_list[iord].x[-(self.tell_npca+3):]
elif self.teltype == 'grid':
self.theta_obj_list[iord] = self.result_list[iord].x[:-(4+3)]
self.theta_tell_list[iord] = self.result_list[iord].x[-(4+3):]
self.obj_model_list[iord], modelmask \
= self.eval_obj_model(self.theta_obj_list[iord], self.obj_dict_list[iord])
self.tellmodel_list[iord] = eval_telluric(self.theta_tell_list[iord], self.tell_dict,
ind_lower=self.ind_lower[iord],
ind_upper=self.ind_upper[iord])
self.assign_output(iord)
if self.debug:
self.show_fit_qa(iord)
[docs] def show_fit_qa(self, iord):
"""
Generates QA plot for telluric fitting
Args:
iord: the order being currently fit
"""
wave_now = self.wave_grid[self.ind_lower[iord]:self.ind_upper[iord]+1]
flux_now = self.flux_arr[self.ind_lower[iord]:self.ind_upper[iord]+1, iord]
sig_now = np.sqrt(utils.inverse(self.ivar_arr[self.ind_lower[iord]:self.ind_upper[iord]+1,iord]))
mask_now = self.mask_arr[self.ind_lower[iord]:self.ind_upper[iord]+1, iord]
model_now = self.tellmodel_list[iord]*self.obj_model_list[iord]
rejmask = mask_now & np.logical_not(self.outmask_list[iord])
fig = plt.figure(figsize=(12, 8))
plt.plot(wave_now, flux_now, drawstyle='steps-mid',
color='k', label='data', alpha=0.7, zorder=5)
plt.plot(wave_now, sig_now, drawstyle='steps-mid', color='0.7', label='noise', alpha=0.7,
zorder=1)
plt.plot(wave_now, model_now, drawstyle='steps-mid', color='red', linewidth=1.0,
label='model', zorder=7, alpha=0.7)
plt.plot(wave_now[rejmask], flux_now[rejmask], 's', zorder=10, mfc='None', mec='blue',
label='rejected pixels')
plt.plot(wave_now[np.logical_not(mask_now)], flux_now[np.logical_not(mask_now)], 'v',
zorder=9, mfc='None', mec='orange', label='originally masked')
plt.ylim(-0.1 * model_now[mask_now].max(), 1.3 * model_now[mask_now].max())
plt.legend()
plt.xlabel('Wavelength')
plt.ylabel('Flux or Counts')
plt.title('QA plot for order/det: {:d}/{:d}'.format(iord + 1, self.norders)) # +1 to account 0-index starting
plt.show()
[docs] def init_output(self):
"""
Method to initialize the outputs
"""
self.model = self.empty_model_table(self.norders, self.nspec_in,
tell_npca=self.tell_npca, n_obj_par=self.max_ntheta_obj)
if 'output_meta_keys' in self.obj_params:
for key in self.obj_params['output_meta_keys']:
if key.lower() in self.datamodel.keys():
setattr(self, key.lower(), self.obj_params[key])
elif key.upper() in self.model.keys():
self.model[key.upper()] = self.obj_params[key]
if self.ech_orders is not None:
self.model['ECH_ORDERS'] = self.ech_orders
self.model['IND_LOWER'] = self.ind_lower
self.model['IND_UPPER'] = self.ind_upper
self.model['WAVE_MIN'] = self.wave_grid[self.ind_lower]
self.model['WAVE_MAX'] = self.wave_grid[self.ind_upper]
[docs] def assign_output(self, iord):
"""
Routine to assign outputs to self.model for the order in question.
Args:
iord (int):
The order for which the output table should bbe assigned.
"""
## TODO Store the outmask with rejected pixels??
gdwave = self.wave_in_arr[:,iord] > 1.0
wave_in_gd = self.wave_in_arr[gdwave,iord]
wave_grid_now = self.wave_grid[self.ind_lower[iord]:self.ind_upper[iord]+1]
if self.teltype == 'pca':
ntell = self.tell_npca
elif self.teltype == 'grid':
ntell = 4
self.model['WAVE'][iord] = self.wave_in_arr[:,iord]
self.model['TELLURIC'][iord][gdwave] \
= scipy.interpolate.interp1d(wave_grid_now, self.tellmodel_list[iord],
kind='linear', bounds_error=False,
fill_value=0.0)(wave_in_gd)
self.model['OBJ_MODEL'][iord][gdwave] \
= scipy.interpolate.interp1d(wave_grid_now, self.obj_model_list[iord],
kind='linear', bounds_error=False,
fill_value=0.0)(wave_in_gd)
self.model['TELL_THETA'][iord] = self.theta_tell_list[iord]
self.model['TELL_PARAM'][iord] = self.theta_tell_list[iord][:ntell]
self.model['TELL_RESLN'][iord] = self.theta_tell_list[iord][ntell]
self.model['TELL_SHIFT'][iord] = self.theta_tell_list[iord][ntell+1]
self.model['TELL_STRETCH'][iord] = self.theta_tell_list[iord][ntell+2]
ntheta_iord = len(self.theta_obj_list[iord])
self.model['OBJ_THETA'][iord][:ntheta_iord] = self.theta_obj_list[iord]
self.model['CHI2'][iord] = self.result_list[iord].fun
self.model['SUCCESS'][iord] = self.result_list[iord].success
self.model['NITER'][iord] = self.result_list[iord].nit
# TODO Purge? This does not appear to be used at the moment.
[docs] def interpolate_inmask(self, mask, wave_inmask, inmask):
"""
Utitlity routine to interpolate the input mask.
"""
if inmask is not None:
if wave_inmask is None:
msgs.error('If you are specifying a mask you need to pass in the corresponding '
'wavelength grid')
# TODO we shoudld consider refactoring the interpolator to take a
# list of images and masks to remove the the fake zero images in the
# call below
_, _, inmask_int = coadd.interp_spec(self.wave_grid, wave_inmask,
np.ones_like(wave_inmask),
np.ones_like(wave_inmask), inmask)
# If the data mask is 2d, and inmask is 1d, tile to create the
# inmask aligned with the data
if mask.ndim == 2 & inmask.ndim == 1:
inmask_out = np.tile(inmask_int, (self.norders, 1)).T
# If the data mask and inmask have the same dimensionlaity,
# interpolated mask has correct dimensions
elif mask.ndim == inmask.ndim:
inmask_out = inmask_int
else:
msgs.error('Unrecognized shape for data mask')
return (mask & inmask_out)
else:
return mask
[docs] def get_ind_lower_upper(self):
"""
Utiltity routine to determine the ind_lower and ind_upper for each
order. This trimming makes things faster because then we only need to
convolve the portion of the telluric model that is needed for the model
fit to each order, rather than convolving the entire telluric model
grid.
Returns:
ind_lower, ind_upper
ind_lower (int):
Lower index into the telluric model wave_grid to trim down the telluric model.
ind_upper (int):
Upper index into the telluric model wave_grid to trim down the telluric model.
"""
wave_ma = np.ma.MaskedArray(np.tile(self.wave_grid, (self.norders,1)).T,
mask=np.logical_not(self.wave_mask_arr))
# TODO: It would be useful if the upper index is actually the stopping
# slice value. That way the slices are [ind_lower:ind_upper] instead of
# always having to do [ind_lower:ind_upper+1]. It also means that the
# length of the fitted region is just ind_upper-ind_lower.
return np.ma.argmin(wave_ma, axis=0), np.ma.argmax(wave_ma, axis=0)
##########################
## telluric grid methods #
##########################
[docs] def get_tell_guess(self):
"""
Return guess parameters for the telluric model.
This is also used to determine the bounds with the ``init_obj_model``
function.
Returns:
:obj:`tuple`: The guess telluric model parameters,
resolution, shift, and stretch parameters.
"""
if self.teltype == 'grid':
guess = [np.median(self.tell_dict['pressure_grid'])]
guess.append(np.median(self.tell_dict['temp_grid']))
guess.append(np.median(self.tell_dict['h2o_grid']))
guess.append(self.airmass_guess)
else:
guess = list(np.zeros(self.tell_npca))
guess.append(self.resln_guess)
guess.append(0.0)
guess.append(1.0)
return tuple(guess)
[docs] def get_bounds_tell(self):
"""
Return the parameter boundaries for the telluric model.
Returns:
:obj:`list`: A list of two-tuples, where each two-tuple proives
the minimum and maximum allowed model values for the PCA coefficients,
(or for the `grid` approach: pressure, temperature, humidity, airmass)
as well as the resolution, shift, and stretch parameters.
"""
# Set the bounds for the optimization
bounds = []
if self.teltype == 'grid':
bounds.append((self.tell_dict['pressure_grid'].min(),
self.tell_dict['pressure_grid'].max()))
bounds.append((self.tell_dict['temp_grid'].min(),
self.tell_dict['temp_grid'].max()))
bounds.append((self.tell_dict['h2o_grid'].min(),
self.tell_dict['h2o_grid'].max()))
bounds.append((self.tell_dict['airmass_grid'].min(),
self.tell_dict['airmass_grid'].max()))
else: # i.e. for teltype == 'pca'
for ii in range(self.tell_npca):
bounds.append((self.tell_dict['bounds_tell_pca'][0][ii+1],
self.tell_dict['bounds_tell_pca'][1][ii+1]))
bounds.append((self.resln_guess * self.resln_frac_bounds[0],
self.resln_guess * self.resln_frac_bounds[1]))
bounds.append(self.pix_shift_bounds)
bounds.append(self.pix_stretch_bounds)
return bounds
[docs] def sort_telluric(self):
"""
Sort the model telluric spectra by the strength of their telluric
features.
This is used to set the order in which the telluric model is fit for
multiple spectra (either multiple orders or multiple slits). The
strength of the telluric features is determined by computing the
median telluric absorption at the midpoint of parameters governing
the telluric grid, specifically for the wavelengths included when
fitting each spectrum.
Returns:
`numpy.ndarray`_: Array of sorted indices from strongest telluric
absorption to weakest.
"""
tell_med = np.zeros(self.norders)
# Do a quick loop over all the orders to sort them in order of strongest
# to weakest telluric absorption
for iord in range(self.norders):
if self.teltype == 'grid':
tm_grid = self.tell_dict['tell_grid'][...,self.ind_lower[iord]:self.ind_upper[iord]+1]
tell_model_mid = tm_grid[tm_grid.shape[0]//2, tm_grid.shape[1]//2,
tm_grid.shape[2]//2, tm_grid.shape[3]//2, :]
tell_med[iord] = np.mean(tell_model_mid)
else:
tell_model_mean = self.tell_dict['tell_pca'][0,self.ind_lower[iord]:self.ind_upper[iord]+1]
tell_med[iord] = np.mean(np.exp(-np.sinh(tell_model_mean)))
# Perform fits in order of telluric strength
return tell_med.argsort()