def test_trace_polynomial():
"""Test trace_polynomial function"""
# No order specified
assert len(lt.trace_polynomial(order=None, evaluate=False)) == 2
# Single order
assert len(lt.trace_polynomial(order=1, evaluate=False)) == 5
# Test evaluate
assert len(lt.trace_polynomial(order=1, evaluate=True)) == 2048
def plot(self, scale='linear', coeffs=None, **kwargs):
"""
Plot the frames of data
Parameters
----------
scale: str
The scale to plot, ['linear', 'log']
coeffs: sequence
The polynomial coefficients of the traces
"""
# Reshape the data
dim = self.data.shape
if self.data.ndim == 4:
data = self.data.reshape(dim[0] * dim[1], dim[2], dim[3])
elif self.data.ndim == 2:
data = self.data.reshape(1, dim[0], dim[1])
else:
data = self.data
# Make the figure
title = '{} Frames'.format(self.ext)
coeffs = lt.trace_polynomial()
fig = plt.plot_frames(data,
scale=scale,
trace_coeffs=coeffs,
wavecal=self.wavecal,
title=title,
**kwargs)
return fig
def subarray(self, subarr):
"""Setter for the subarray
Properties
----------
subarr: str
The name of the subarray to use,
['SUBSTRIP256', 'SUBSTRIP96', 'FULL']
"""
subs = ['SUBSTRIP256', 'SUBSTRIP96', 'FULL']
# Check the value
if subarr not in subs:
raise ValueError("'{}' not a supported subarray. Try {}".format(
subarr, subs))
# Set the subarray
self._subarray = subarr
self.subarray_specs = utils.subarray_specs(subarr)
# Set the dependent quantities
self._ncols = 2048
self._nrows = self.subarray_specs.get('y')
self.wave = utils.wave_solutions(subarr)
self.avg_wave = np.mean(self.wave, axis=1)
self.coeffs = locate_trace.trace_polynomial(subarray=subarr)
# Reset the data and time arrays
self._reset_data()
self._reset_time()
# Reset the psfs
self._reset_psfs()
def plot(self,
idx=0,
scale='linear',
order=None,
noise=True,
traces=False,
saturation=0.8,
draw=True):
"""
Plot a TSO frame
Parameters
----------
idx: int
The frame index to plot
scale: str
Plot scale, ['linear', 'log']
order: sequence
The order to isolate
noise: bool
Plot with the noise model
traces: bool
Plot the traces used to generate the frame
saturation: float
The fraction of full well defined as saturation
draw: bool
Render the figure instead of returning it
"""
# Get the data cube
if order in [1, 2]:
tso = getattr(self, 'tso_order{}_ideal'.format(order))
else:
if noise:
tso = self.tso
else:
tso = self.tso_ideal
# Reshape data
tso.shape = self.dims3
# Set the plot args
wavecal = self.wave
title = '{} - Frame {}'.format(self.title, idx)
coeffs = locate_trace.trace_polynomial() if traces else None
# Plot the frame
fig = plotting.plot_frames(data=tso,
idx=idx,
scale=scale,
trace_coeffs=coeffs,
saturation=saturation,
title=title,
wavecal=wavecal)
if draw:
show(fig)
else:
return fig
def SOSS_psf_cube(filt='CLEAR', order=1, subarray='SUBSTRIP256', generate=False, mprocessing=True, wave_sol=None, dirname='default'):
"""
Generate/retrieve a data cube of shape (3, 2048, 76, 76) which is a
76x76 pixel psf for 2048 wavelengths for each trace order. The PSFs
are scaled to unity and rotated to reproduce the trace tilt at each
wavelength then placed on the desired subarray.
Parameters
----------
filt: str
The filter to use, ['CLEAR', 'F277W']
order: int
The trace order
subarray: str
The subarray to use, ['SUBSTRIP96', 'SUBSTRIP256', 'FULL']
generate: bool
Generate a new cube
mprocessing: bool
Use multiprocessing
wave_sol: sequence (optional)
The user provided wavelength solutions for orders 1, 2, and 3
dirname: str (optional)
The target subdirectory name for the PSFs that use a custom wavelength solution
Returns
-------
np.ndarray
An array of the SOSS psf at 2048 wavelengths for each order
"""
# Check if it's a custom wavelength solution
psf_loc = copy(PSF_DIR)
if dirname != 'default':
dirpath = os.path.join(PSF_DIR, dirname)
if not os.path.exists(dirpath):
os.system('mkdir {}'.format(dirpath))
psf_loc = psf_loc.replace('soss_psfs', 'soss_psfs/{}'.format(dirname))
if generate:
print('This takes about 2 minutes.')
# Default wavelengths
if wave_sol is None:
wavelengths = np.mean(utils.wave_solutions(subarray), axis=1)[:3 if filt == 'CLEAR' else 1]
# Or user provided
else:
if wave_sol.shape != (3, 2048):
raise TypeError("'wave_sol' input must be an array of shape (3, 2048)")
wavelengths = wave_sol
# Get trace polynomial coefficients
coeffs = locate_trace.trace_polynomial(subarray)
# Get the file
psf_file = os.path.join(PSF_DIR, 'SOSS_{}_PSF.fits'.format(filt))
# Load the SOSS psf cube
cube = fits.getdata(psf_file).swapaxes(-1, -2)
wave = fits.getdata(psf_file, ext=1)
# Initilize interpolator
psfs = interp1d(wave, cube, axis=0, kind=3)
trace_cols = np.arange(2048)
# Run datacube
for n, wavelength in enumerate(wavelengths):
# Evaluate the trace polynomial in each column to get the y-position of the trace center
trace_centers = np.polyval(coeffs[n], trace_cols)
# Don't calculate order2 or order 3 for F277W
if n > 0 and filt.lower() == 'f277w':
pass
else:
# Get the psf for each column
print('Calculating order {} SOSS psfs for {} filter...'.format(n + 1, filt))
start = time.time()
func = partial(get_SOSS_psf, filt=filt, psfs=psfs)
if mprocessing:
pool = multiprocessing.Pool(8)
raw_psfs = np.array(pool.map(func, wavelength))
pool.close()
pool.join()
del pool
else:
raw_psfs = []
for i in range(len(wavelength)):
raw_psfs.append(func(wavelength[i]))
raw_psfs = np.array(raw_psfs)
print('Finished in {} seconds.'.format(time.time()-start))
# Rotate the psfs
print('Rotating order {} SOSS psfs for {} filter...'.format(n + 1, filt))
start = time.time()
func = partial(rotate, reshape=False)
# Get the PSF tilt at each column
angles = psf_tilts(order)
if mprocessing:
pool = multiprocessing.Pool(8)
rotated_psfs = np.array(pool.starmap(func, zip(raw_psfs, angles)))
pool.close()
pool.join()
del pool
else:
rotated_psfs = []
for rp, ang in zip(raw_psfs, angles):
rotated_psfs.append(func(rp, ang))
rotated_psfs = np.array(rotated_psfs)
print('Finished in {} seconds.'.format(time.time()-start))
# Scale psfs to 1
rotated_psfs = np.abs(rotated_psfs)
scale = np.nansum(rotated_psfs, axis=(1, 2))[:, None, None]
rotated_psfs = rotated_psfs / scale
# Split it into 4 chunks to be below Github file size limit
chunks = rotated_psfs.reshape(4, 512, 76, 76)
for N, chunk in enumerate(chunks):
idx0 = N * 512
idx1 = idx0 + 512
centers = trace_centers[idx0:idx1]
# Interpolate the psfs onto the subarray
print('Interpolating chunk {}/4 for order {} SOSS psfs for {} filter onto subarray...'.format(N + 1, n + 1, filt))
start = time.time()
func = put_psf_on_subarray
if mprocessing:
pool = multiprocessing.Pool(8)
data = zip(chunk, centers)
subarray_psfs = pool.starmap(func, data)
pool.close()
pool.join()
del pool
else:
subarray_psfs = []
for ch, ce in zip(chunk, centers):
subarray_psfs.append(func(ch, ce))
print('Finished in {} seconds.'.format(time.time()-start))
# Get the filepath
file = os.path.join(psf_loc, 'SOSS_{}_PSF_order{}_{}.npy'.format(filt, n+1, N+1))
# Delete the file if it exists
if os.path.isfile(file):
os.system('rm {}'.format(file))
# Write the data
np.save(file, np.array(subarray_psfs))
print('Data saved to', file)
else:
if PSF_DIR is None:
print("No PSF files detected. Using all ones.")
subarr = 256 if subarray == 'SUBSTRIP256' else 96 if subarray == 'SUBSTRIP256' else 2048
return np.ones((2048, subarr, 76))
else:
print("Using SOSS PSF files located at {}".format(psf_loc))
# Get the chunked data and concatenate
full_data = []
for chunk in [1, 2, 3, 4]:
file = os.path.join(psf_loc, 'SOSS_{}_PSF_order{}_{}.npy'.format(filt, order, chunk))
full_data.append(np.load(file))
return np.concatenate(full_data, axis=0)
def calculate_psf_tilts():
"""
Calculate the tilt of the psf at the center of each column
using all binned pixels in the given wavelength calibration file
for both orders and save to file
"""
for order in [1, 2, 3]:
# Get the file
psf_file = os.path.join(PSF_DIR, 'SOSS_PSF_tilt_order{}.npy'.format(order))
# Dimensions
subarray = 'SUBSTRIP256'
X = range(2048)
Y = range(256)
# Get the wave map
wave_map = utils.wave_solutions(subarray, order).astype(float)
# Get the y-coordinate of the trace polynomial in this column
# (center of the trace)
coeffs = locate_trace.trace_polynomial(subarray=subarray, order=order)
trace = np.polyval(coeffs, X)
# Interpolate to get the wavelength value at the center
wave = interp2d(X, Y, wave_map)
# Get the wavelength of the trace center in each column
trace_wave = []
for x, y in zip(X, trace):
trace_wave.append(wave(x, y)[0])
# For each column wavelength (defined by the wavelength at
# the trace center) define an isowavelength contour
angles = []
for n, x in enumerate(X):
w = trace_wave[x]
# Edge cases
try:
w0 = trace_wave[x-1]
except IndexError:
w0 = 0
try:
w1 = trace_wave[x+1]
except IndexError:
w1 = 10
# Define the width of the wavelength bin as half-way
# between neighboring points
dw0 = np.mean([w0, w])
dw1 = np.mean([w1, w])
# Get the coordinates of all the pixels in that range
yy, xx = np.where(np.logical_and(wave_map >= dw0, wave_map < dw1))
# Find the angle between the vertical and the tilted wavelength bin
if len(xx) >= 1:
angle = get_angle([xx[-1], yy[-1]], [x, trace[x]])
else:
angle = 0
# Don't flip them upside down
angle = angle % 180
# Add to the array
angles.append(angle)
# Save the file
np.save(psf_file, np.array(angles))
print('Angles saved to', psf_file)
def test_coeffs(self):
"""Check the traces are drawn"""
coeffs = lt.trace_polynomial()
fig = plt.plot_frames(self.frames, scale='log', trace_coeffs=coeffs)
self.assertEqual(str(type(fig)), "<class 'bokeh.models.layouts.Column'>")
