import numpy as np
from astropy.io import fits, ascii
import matplotlib.pyplot as plt
import warnings
try:
from ..utils import get_baseDir
except ImportError as err1:
warnings.warn("Could not import baseDir. Will save diagnostics to . ")
def get_baseDir():
return "."
import os
import pdb
from copy import deepcopy
from scipy.signal import savgol_filter, savgol_coeffs
from astropy.convolution import convolve
import celerite2
[docs]
def do_even_odd(thisAmp):
"""
Do an even-odd correction for a given amplifier
If only one amplifier is used, it can be the whole image
"""
even_odd_model = np.zeros_like(thisAmp)
even_offset = np.nanmedian(thisAmp[:,0::2])
even_odd_model[:,0::2] = even_offset
odd_offset = np.nanmedian(thisAmp[:,1::2])
even_odd_model[:,1::2] = odd_offset
return thisAmp - even_odd_model, even_odd_model
[docs]
def col_by_col(thisAmp):
"""
Do a column-by-column slow read correction
instead of an even/odd correction
"""
colMedian = np.nanmedian(thisAmp,axis=0)
slowread_model = np.tile(colMedian,[thisAmp.shape[0],1])
return thisAmp - slowread_model, slowread_model
def make_time2D(img,amplifiers=4):
nY, nX = img.shape
if amplifiers == 4:
numColsPerAmp = 512
elif amplifiers == 1:
numColsPerAmp = nX
else:
raise Exception("Unrecognized number of amps")
time1Damp_over = np.arange((numColsPerAmp + 12) * nY)
time2Damp_over = time_wrap(time1Damp_over,nY,(numColsPerAmp+12))
time2Damp = time2Damp_over[:,0:numColsPerAmp]
if amplifiers == 1:
time2D = time2Damp
elif amplifiers == 4:
time2D = np.zeros([nY,2048]) * np.nan
time2D[:,0:512] = time2Damp
time2D[:,512:1024] = time2Damp[:,::-1]
time2D[:,1024:1536] = time2Damp
time2D[:,1536:2048] = time2Damp[:,::-1]
else:
raise Exception("Unrecognized number of amps")
return time2D * 10e-6
def time_wrap(img1D,rows,columns,amp=0):
#return np.reshape(img1D,[rows,columns],order='F')
img2D = np.reshape(img1D,[rows,columns],order='C')
if (amp == 0) | (amp == 2):
return img2D
else:
return img2D[:,::-1]
def time_unwrap(img2D,amp=0):
if (amp == 0) | (amp == 2):
return img2D.ravel()
else:
return img2D[:,::-1].ravel()
[docs]
def show_gp_psd():
"""
Show the power spectral density function for the Gaussian process
"""
#minfreq, maxfreq = 1.0 / np.max(times), 1.0 / (times[1] - times[0])
minfreq, maxfreq = 1e-4, 1e2
freq = np.logspace(np.log10(minfreq),np.log10(maxfreq), 9000)
omega = 2 * np.pi * freq
gp_comb, gp_list = make_gp()
for ind,oneGP in enumerate(gp_list):
plt.loglog(freq, oneGP.kernel.get_psd(omega),label="SHOTerm {}".format(ind+1))
plt.loglog(freq, gp_comb.kernel.get_psd(omega),label="Combined")
plt.plot(freq,1./freq * 7.,label='1/f power law',linestyle='dashed')
plt.legend()
plt.ylim(1e-3,1e5)
plt.xlabel("Frequency (1/ms)")
plt.ylabel("Power")
plt.show()
[docs]
def make_gp(sigma_w=7.):
"""
Make a Gaussian Process Kernel for 1/f corrections
Parameters
-----------
sigma_w: float
White noise error in DN
"""
terms_list = []
gp_list = []
rho_list = [1000, 100, 10, 1, 0.1, 0.01]
for ind,rho in enumerate(rho_list):
oneTerm = celerite2.terms.SHOTerm(sigma=sigma_w, rho=rho, tau=0.1 * rho)
if ind == 0:
kernel_comb = oneTerm
else:
kernel_comb = kernel_comb + oneTerm
gp = celerite2.GaussianProcess(oneTerm, mean=0)
gp_list.append(gp)
gp_comb = celerite2.GaussianProcess(kernel_comb, mean=0.0)
return gp_comb, gp_list
def gp_predict(t,f,gp_model,sigma_w=7.,nsmoothKern=501):
pts_finite = np.isfinite(f)
pts_non_outlier = (np.abs(f - np.nanmedian(f)) < 100.)
pts = pts_finite & pts_non_outlier
gp_model.compute(t[pts],yerr=sigma_w * np.ones(np.sum(pts)))
mu_gp = gp_model.predict(f[pts], t=t, return_var=False)
kernel_for_gp_smoothing = savgol_coeffs(nsmoothKern,3)
gp_smoothed = convolve(mu_gp, kernel_for_gp_smoothing, boundary='extend')
#pdb.set_trace()
#plt.plot(t,f,'.'); plt.plot(t,gp_smoothed); plt.ylim(-100,100); plt.show()
return gp_smoothed
[docs]
def do_backsub(img,photObj=None,amplifiers=4,saveDiagnostics=False,
evenOdd=True,activePixMask=None,backgMask=None,
grismr=False,returnFastSlow=False,
colByCol=False,smoothSlowDir=None,
badRowsAllowed=0,
GROEBA=False,GPnsmoothKern=None,
showGP1D=False):
"""
Do background subtraction amplifier-by-amplifier, row-by-row around the sources
Parameters
---------------
img: numpy 2D array
The input uncorrected image
photObj: (optional) a tshirt photometry pipeline object
If supplied, it will use the background apertures to mask out sources
amplifiers: int
How many outputamplifiers are used? 4 for NIRCam stripe mode and 1 is for 1 output amplifier
evenOdd: bool
Remove the even and odd offsets before doing row-by-row medians?
colByCol: bool
Do column-by-column subtraction. This will supercede the even/odd correction.
saveDiagnostics: bool
Save diagnostic files?
activePixMask: numpy 2D array or None
Mask of reference pixels to ignore (optionally). Pixels that are False
will be ignored in the background estimation, and pixels that are
true will be kept in the background estimation. If refpixMask is None
no extra points will be masked
badRowsAllowed: int
Number of bad rows to accept and still do fast-read correction
(ie. if not enough background and/or ref pixels,
no fast-read correction can be done)
backgMask: numpy 2D array or None
Mask for the background. Pixels that are False will be ignored in
row-by-row background estimation. Pixels that are true will
be kept for the background estimation.
grismr: bool
Is this NIRCam GRISMR data? Special treatment is needed for NIRCam
GRISMR, where the spectra run through multiple amplifiers
returnFastSlow: bool
Return both the fast and slow read models?
smoothSlowDir: None or float
Length of the smoothing kernel to apply along the fast read direction
If None, no smoothing is applied
GROEBA: bool
Use Gaussian-Process row-by-row, odd/even by amplifier subtraction?
This will use a Gaussian process to do 1/f noise interpolation
GPnsmoothKern: int or None
(only if GROEBA==True). Specify the window length for a Savgol smoothing
filter for the GP. Otherwise it is chosen automatically
showGP1D: bool
(only if GROEBA==True). Plot the GP prediction in 1D
"""
## Npix Threshold
## this many pixels must be in a row in order to do an median
## otherwise, it attempts to interpolate from other amplifiers
## this only applies when noutputs=4
Npix_threshold = 3
## Start by including all pixels
useMask = np.ones_like(img,dtype=bool)
## make a mask with Tshirt object
if photObj is not None:
y, x = np.mgrid[0:img.shape[0],0:img.shape[1]]
xcen = photObj.srcApertures.positions[:,0]
ycen = photObj.srcApertures.positions[:,1]
useMask = np.ones_like(img,dtype=bool)
for srcInd in np.arange(photObj.nsrc):
r = np.sqrt((x - xcen)**2 + (y - ycen)**2)
srcPts = r < photObj.param['backStart']
useMask[srcPts] = False
if backgMask is not None:
useMask = useMask & backgMask
## only include active pixels
if activePixMask is None:
activeMask = np.ones_like(img,dtype=bool)
else:
activeMask = activePixMask
useMask = useMask & activeMask
## let nanmedian do the work of masking
maskedPts = (useMask == False)
masked_img = deepcopy(img)
masked_img[maskedPts] = np.nan
outimg = np.zeros_like(img)
slowread_model = np.zeros_like(img)
fastread_model = np.zeros_like(img)
if GROEBA == True:
sigma_w = 7.
gp_comb, gp_list = make_gp(sigma_w=sigma_w)
time2D = make_time2D(img,amplifiers=amplifiers)
rowsAmp, totcols = img.shape
if amplifiers == 4:
colsAmp = 512
else:
colsAmp = totcols
## mask to keep track of which amplifiers have enough pixels to use
## in NIRCam GRISMR spectroscopy, the spectra cut across rows
## so these amplifiers' 1/f noise has to interpolated from other
## amps
amp_check_mask = np.ones(amplifiers,dtype=bool)
if amplifiers == 4:
## list where the amplifiers are
ampStarts = [0,512,1024,1536]
ampEnds = [512,1024,1536,2048]
ampWidth = 512
elif amplifiers == 1:
ampStarts = [0]
ampEnds = [colsAmp]
ampWidth = colsAmp
else:
raise Exception("{} amplifiers not implemented".format(amplifiers))
## loop through amps
for amp in np.arange(amplifiers):
if colByCol == True:
thisAmp, colbycol_model = col_by_col(masked_img[:,ampStarts[amp]:ampEnds[amp]])
slowread_model[:,ampStarts[amp]:ampEnds[amp]] = colbycol_model
elif evenOdd == True:
thisAmp, even_odd_model = do_even_odd(masked_img[:,ampStarts[amp]:ampEnds[amp]])
slowread_model[:,ampStarts[amp]:ampEnds[amp]] = even_odd_model
else:
thisAmp = masked_img[:,ampStarts[amp]:ampEnds[amp]]
## even if not doing an even/odd correction, still do an overall median
slowread_model[:,ampStarts[amp]:ampEnds[amp]] = np.nanmedian(thisAmp)
thisAmp = thisAmp - slowread_model[:,ampStarts[amp]:ampEnds[amp]]
## check if the mask leaves too few pixels in this amplifier
## in that case
pxPerRow = np.sum(np.isfinite(thisAmp),axis=1)
bad_rows = np.sum(pxPerRow <= Npix_threshold)
if bad_rows <= badRowsAllowed:
if GROEBA == True:
t = time_unwrap(time2D[:,ampStarts[amp]:ampEnds[amp]],amp=amp)
f = time_unwrap(thisAmp,amp=amp)
if GPnsmoothKern is None:
if np.median(pxPerRow) < 6:
## Smooth more for small numbers of pixels like when using refpix
GPnsmoothKern = 15001
else:
GPnsmoothKern = 2001
mu_gp1D = gp_predict(t,f,gp_comb,sigma_w=sigma_w,nsmoothKern=GPnsmoothKern)
if showGP1D == True:
plt.plot(t,f); plt.plot(t,mu_gp1D)
plt.show()
pdb.set_trace()
fastread_model[:,ampStarts[amp]:ampEnds[amp]] = time_wrap(mu_gp1D,
columns=colsAmp,
rows=rowsAmp,
amp=amp)
else:
medVals = np.nanmedian(thisAmp,axis=1)
if smoothSlowDir is not None:
medVals = savgol_filter(medVals,smoothSlowDir,3)
## tile this to make a model across the fast-read direction
fastread_model[:,ampStarts[amp]:ampEnds[amp]] = np.tile(medVals, [ampWidth,1]).T
if np.median(pxPerRow) < 6:
amp_check_mask[amp] = False ## For amplifiers where there are few pixels,
## count this as "unchecked" to average other amplifier's information
else:
amp_check_mask[amp] = True
else:
amp_check_mask[amp] = False
fastread_model[:,ampStarts[amp]:ampEnds[amp]] = 0.0
## If there is at least 1 good amp and 1 bad amp,
## the good amp can be used to
## estimate the 1/f noise in "bad" amps where some rows have
## no background pixels that can used
ngood_amps = np.sum(amp_check_mask)
## only do this if we have >=1 good and >=1 bad amp
if (ngood_amps >= 1) & (ngood_amps < amplifiers):
rowModel = np.nanmean(fastread_model,axis=1)
tiled_avg = np.tile(rowModel,[ampWidth,1]).T
for amp in np.arange(4):
if amp_check_mask[amp] == False:
fastread_model[:,ampStarts[amp]:ampEnds[amp]] = tiled_avg
## put the results in the model image
modelimg = slowread_model + fastread_model
outimg = img - modelimg
if saveDiagnostics == True:
if photObj is None:
outPrefix = 'unnamed'
else:
outPrefix = photObj.dataFileDescrip
diag_dir = os.path.join(get_baseDir(),'diagnostics','rowamp_sub')
descrips = ['orig','mask','slowread_model','slowread_sub',
'fastread_model','model','subtracted']
diag_imgs = [img,masked_img,slowread_model,img-slowread_model,
fastread_model,modelimg,outimg]
for outInd,outDescrip in enumerate(descrips):
out_path = os.path.join(diag_dir,"{}_{}.fits".format(outPrefix,outDescrip))
outHDU = fits.PrimaryHDU(diag_imgs[outInd])
outHDU.writeto(out_path,overwrite=True)
if returnFastSlow == True:
return outimg, slowread_model, fastread_model
else:
return outimg, modelimg