def test_wavelength_correction(fname, clight):
bc = 100
rv = 200
header = {"barycorr": bc, "radvel": rv}
wave = np.linspace(5000, 6000, 100)
wave = np.atleast_2d(wave)
echelle.save(fname, header, wave=wave)
# Read raw frame
ech = echelle.Echelle.read(fname, raw=True)
assert np.allclose(ech["wave"], wave)
assert ech.header["barycorr"] == bc
assert ech.header["radvel"] == rv
# Apply only barycentric correction
ech = echelle.Echelle.read(fname, barycentric_correction=True, radial_velociy_correction=False)
assert np.allclose(ech["wave"], wave * (1 + bc / clight))
assert ech.header["barycorr"] == 0
assert ech.header["radvel"] == rv
# Apply only radial velocity correction
ech = echelle.Echelle.read(fname, barycentric_correction=False, radial_velociy_correction=True)
assert np.allclose(ech["wave"], wave * (1 - rv / clight))
assert ech.header["barycorr"] == bc
assert ech.header["radvel"] == 0
# Apply both corrections
ech = echelle.Echelle.read(fname, barycentric_correction=True, radial_velociy_correction=True)
assert np.allclose(ech["wave"], wave * (1 + (bc - rv) / clight))
assert ech.header["barycorr"] == 0
assert ech.header["radvel"] == 0
def test_column_range_mask(fname):
spec = np.linspace(1, 1000, 100)
spec = np.stack([spec, spec])
sig = np.copy(spec)
cont = np.copy(spec)
wave = np.copy(spec)
columns = np.array([[10, 99], [0, 50]])
echelle.save(fname, {},
spec=spec,
sig=sig,
cont=cont,
wave=wave,
columns=columns)
ech = echelle.read(fname)
assert "mask" in ech
assert isinstance(ech["spec"], np.ma.masked_array)
for iord in range(2):
assert all(ech["mask"][iord, :columns[iord, 0]])
assert all(ech["mask"][iord, columns[iord, 1]:])
assert not any(ech["mask"][iord, columns[iord, 0]:columns[iord, 0]])
def test_continuum_normalization(fname):
spec = np.full((2, 100), 10, dtype=float)
sig = np.full((2, 100), 1, dtype=float)
cont = np.full((2, 100), 2, dtype=float)
echelle.save(fname, {}, spec=spec, sig=sig, cont=cont)
ech = echelle.read(fname)
assert np.allclose(ech["spec"], spec/cont)
assert np.allclose(ech["sig"], sig/cont)
assert np.allclose(ech["cont"], cont)
def test_read_write_echelle(fname):
data = np.linspace(1, 200, num=100, dtype=float)
header = {"BLA": "Blub"}
echelle.save(fname, header, test=data)
ech = echelle.Echelle.read(fname)
assert isinstance(ech, echelle.Echelle)
assert "test" in ech
assert np.allclose(ech["test"], data)
assert ech.header["BLA"] == "Blub"
def spec(
files,
instrument,
mode,
mask,
extension,
bias,
normflat,
orders,
settings,
output_dir,
order_range,
):
"""Load or create science spectrum
Parameters
----------
files : dict(str:str)
raw input files
instrument : str
instrument name
mode : str
observing mode
mask : array(bool)
Bad pixel mask
extension : int
fits data extension
bias : array(float)
bias calibration data
normflat : array(float)
normalized flat field
orders : array(float)
order tracing polynomials and column_ranges
settings : dict(str:obj)
run settings
output_dir : str
output data directory
Returns
-------
spec : array(float) of size (norders, ncol)
extracted science spectra
sigma : array(float) of size (norders, ncol)
uncertainty on the extracted science spectra
"""
orders, column_range = orders
specfile = os.path.join(output_dir, "test_spec.ech")
settings = settings["science"]
try:
science = echelle.read(specfile, raw=True)
head = science.header
spec = science["spec"]
sigma = science["sig"]
mask = np.full(spec.shape, True)
for iord in range(spec.shape[0]):
cr = column_range[iord]
mask[iord, cr[0]:cr[1]] = False
spec = np.ma.array(spec, mask=mask)
sigma = np.ma.array(sigma, mask=mask)
except FileNotFoundError:
flat, blaze = normflat
bias, _ = bias
# Fix column ranges
for i in range(blaze.shape[0]):
column_range[i] = np.where(blaze[i] != 0)[0][[0, -1]]
f = files["science"][0]
im, head = util.load_fits(f,
instrument,
mode,
extension,
mask=mask,
dtype=np.float32)
# Correct for bias and flat field
im -= bias
im /= flat
# Optimally extract science spectrum
spec, sigma, _, columns = extract(
im,
orders,
gain=head["e_gain"],
readnoise=head["e_readn"],
dark=head["e_drk"],
column_range=column_range,
order_range=order_range,
extraction_type=settings["extraction_method"],
extraction_width=settings["extraction_width"],
lambda_sf=settings["smooth_slitfunction"],
lambda_sp=settings["smooth_spectrum"],
osample=settings["oversampling"],
swath_width=settings["swath_width"],
plot=False,
)
echelle.save(specfile, head, spec=spec, sig=sigma, columns=columns)
return spec, sigma