diff --git a/.DS_Store b/.DS_Store index 08fcf5c..bda3bdc 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/changelog.md b/changelog.md index c1a047c..808220d 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,20 @@ # Changelog All notable changes to this project will be documented in this file. +### [1.2.0] -- 2023-11-21 +#### Added +- Best fitting transit models are saved by default. +- Limb-darkening coefficients are automatically calculated at desired binning during light curve fitting. +- Streamline process for fitting (free or prior) or fixing LD coefficients during fits. +- Incorporate PASTASOSS to calculate SOSS wavelength solution pased on pupil wheel position. +- Option to specify background estimation region. +- Improvements to bad pixel interpolation. +- extra_functions.py for potentially helpful functions not directly used by the pipeline. +- Misc. bug fixes. + +#### Removed +- Stability calculations via the cross-correlation method. + ### [1.1.7] -- 2023-07-28 #### Added - Streamline installation process. diff --git a/setup.py b/setup.py index 1f8ba6f..3355259 100644 --- a/setup.py +++ b/setup.py @@ -4,7 +4,7 @@ from setuptools import setup setup(name='supreme_spoon', - version='1.1.7', + version='1.2.0', license='MIT', author='Michael Radica', author_email='michael.radica@umontreal.ca', @@ -16,7 +16,7 @@ install_requires=['applesoss==2.0.2', 'astropy', 'bottleneck', 'corner', 'jwst==1.8.5', 'matplotlib', 'more_itertools', 'numpy', 'pandas', 'opencv-python', 'ray', 'scikit-learn', - 'scipy', 'tqdm', 'pyyaml'], + 'scipy', 'tqdm', 'pastasoss', 'pyyaml'], extras_require={'stage4': ['supreme_spoon', 'cython', 'dynesty==1.2.2', 'exotic_ld', 'h5py', 'juliet'], 'webbpsf': ['supreme_spoon', 'webbpsf>=1.1.1']}, diff --git a/supreme_spoon/.DS_Store b/supreme_spoon/.DS_Store index b05bdd5..02d331c 100644 Binary files a/supreme_spoon/.DS_Store and b/supreme_spoon/.DS_Store differ diff --git a/supreme_spoon/extra_functions.py b/supreme_spoon/extra_functions.py new file mode 100644 index 0000000..804a965 --- /dev/null +++ b/supreme_spoon/extra_functions.py @@ -0,0 +1,249 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Mon Nov 20 15:08 2023 + +@author: MCR + +Extra functions to be used alongside the main pipeline. +""" + +from astropy.io import fits +import numpy as np + + +def get_throughput_from_photom_file(photom_path): + """Calculate the throughput based on the photom reference file. + Function from Loïc, and is apparently the proper way to get the + throughput? Gives different results from contents of spectra reference + file. + + Parameters + ---------- + photom_path : str + Path to photom reference file + + Returns + ------- + w1 : np.array(float) + Order 1 wavelength axis. + w2 : np.array(float) + Order 2 wavelength axis. + thpt1 : np.array(float) + Order 1 throughput values + thpt2 : np.array(float) + Order 2 throughput values + """ + + # From the photom ref file, get conversion factors + wavelength/pixel grid. + photom = fits.open(photom_path) + w1 = photom[1].data['wavelength'][0] + ii = np.where((photom[1].data['wavelength'][0] >= 0.84649785) & + (photom[1].data['wavelength'][0] <= 2.83358154)) + w1 = w1[ii] + scale1 = photom[1].data['relresponse'][0][ii] + + w2 = photom[1].data['wavelength'][1] + ii = np.where((photom[1].data['wavelength'][1] >= 0.49996997) & + (photom[1].data['wavelength'][1] <= 1.40884607)) + w2 = w2[ii] + scale2 = photom[1].data['relresponse'][1][ii] + + # Calculate throughput from conversion factor. + thpt1 = 1 / (scale1 * 3e8 / w1 ** 2) + thpt2 = 1 / (scale2 * 3e8 / w2 ** 2) + + return w1, w2, thpt1, thpt2 + + +# ====== Deprecated CCF Stability Functions --- May Have Later Uses?? ===== +def soss_stability_xy(cube, nsteps=501, axis='x', nthreads=4, + smoothing_scale=None): + """Perform a CCF analysis to track the movement of the SOSS trace + relative to the median stack over the course of a TSO. + Parameters + ---------- + cube : array-like[float] + Data cube. Should be 3D (ints, dimy, dimx). + nsteps : int + Number of CCF steps to test. + axis : str + Axis over which to calculate the CCF - either 'x', or 'y'. + nthreads : int + Number of CPUs for multiprocessing. + smoothing_scale : int + Length scale over which to smooth results. + Returns + ------- + ccf : array-like[float] + The cross-correlation results. + """ + # Initialize ray with specified number of threads. + ray.shutdown() + ray.init(num_cpus=nthreads) + # Subtract integration-wise median from cube for CCF. + cube = cube - np.nanmedian(cube, axis=(1, 2))[:, None, None] + # Calculate median stack. + deep = bn.nanmedian(cube, axis=0) + # Divide total data cube into segments and run each segment in parallel + # with ray. + ii = 0 + all_fits = [] + nints = np.shape(cube)[0] + seglen = nints // nthreads + for i in range(nthreads): + if i == nthreads - 1: + cube_seg = cube[ii:] + else: + cube_seg = cube[ii:ii + seglen] + all_fits.append(soss_stability_xy_run.remote(cube_seg, deep, seg_no=i+1, + nsteps=nsteps, axis=axis)) + ii += seglen + # Run the CCFs. + ray_results = ray.get(all_fits) + # Stack all the CCF results into a single array. + maxvals = [] + for i in range(nthreads): + if i == 0: + maxvals = ray_results[i] + else: + maxvals = np.concatenate([maxvals, ray_results[i]]) + # Smooth results if requested. + if smoothing_scale is not None: + ccf = median_filter(np.linspace(-0.01, 0.01, nsteps)[maxvals], + smoothing_scale) + else: + ccf = np.linspace(-0.01, 0.01, nsteps)[maxvals] + ccf = ccf.reshape(nints) + return ccf + + +@ray.remote +def soss_stability_xy_run(cube_sub, med, seg_no, nsteps=501, axis='x'): + """Wrapper to perform CCF calculations in parallel with ray. + """ + # Get data dimensions. + nints, dimy, dimx = np.shape(cube_sub) + # Get integration numbers to show progress prints. + marks = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] + locs = np.nanpercentile(np.arange(nints), marks) + # Initialize CCF variables. + ccf = np.zeros((nints, nsteps)) + f = interp2d(np.arange(dimx), np.arange(dimy), med, kind='cubic') + # Perform cross-correlation over desired axis. + loc = 0 + for i in range(nints): + # Progress print. + if i >= int(locs[loc]): + fancyprint('Slice {}: {}% complete.'.format(seg_no, marks[loc])) + loc += 1 + for j, jj in enumerate(np.linspace(-0.01, 0.01, nsteps)): + if axis == 'x': + interp = f(np.arange(dimx) + jj, np.arange(dimy)) + elif axis == 'y': + interp = f(np.arange(dimx), np.arange(dimy) + jj) + else: + raise ValueError('Unknown axis: {}'.format(axis)) + ccf[i, j] = np.nansum(cube_sub[i] * interp) + # Determine the peak of the CCF for each integration to get the + # best-fitting shift. + maxvals = [] + for i in range(nints): + maxvals.append(np.where(ccf[i] == np.max(ccf[i]))[0]) + maxvals = np.array(maxvals) + maxvals = maxvals.reshape(maxvals.shape[0]) + return maxvals + + +def soss_stability_fwhm(cube, ycens_o1, nthreads=4, smoothing_scale=None): + """Estimate the FWHM of the trace over the course of a TSO by fitting a + Gaussian to each detector column. + Parameters + ---------- + cube : array-like[float] + Data cube. Should be 3D (ints, dimy, dimx). + ycens_o1 : arrray-like[float] + Y-centroid positions of the order 1 trace. Should have length dimx. + nthreads : int + Number of CPUs for multiprocessing. + smoothing_scale : int + Length scale over which to smooth results. + Returns + ------- + fwhm : array-like[float] + FWHM estimates for each column at every integration. + """ + # Initialize ray with specified number of threads. + ray.shutdown() + ray.init(num_cpus=nthreads) + # Divide total data cube into segments and run each segment in parallel + # with ray. + ii = 0 + all_fits = [] + nints = np.shape(cube)[0] + seglen = nints // nthreads + for i in range(nthreads): + if i == nthreads - 1: + cube_seg = cube[ii:] + else: + cube_seg = cube[ii:ii + seglen] + all_fits.append(soss_stability_fwhm_run.remote(cube_seg, ycens_o1, + seg_no=i+1)) + ii += seglen + # Run the CCFs. + ray_results = ray.get(all_fits) + # Stack all the CCF results into a single array. + fwhm = [] + for i in range(nthreads): + if i == 0: + fwhm = ray_results[i] + else: + fwhm = np.concatenate([fwhm, ray_results[i]]) + # Set median of trend to zero. + fwhm -= np.median(fwhm) + # Smooth the trend. + if smoothing_scale is None: + smoothing_scale = int(0.2*nints) + fwhm = median_filter(fwhm, smoothing_scale) + return fwhm + + +@ray.remote +def soss_stability_fwhm_run(cube, ycens_o1, seg_no): + """Wrapper to perform FWHM calculations in parallel with ray. + """ + def gauss(x, *p): + amp, mu, sigma = p + return amp * np.exp(-(x - mu) ** 2 / (2. * sigma ** 2)) + # Get data dimensions. + nints, dimy, dimx = np.shape(cube) + # Get integration numbers to show progress prints. + marks = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] + locs = np.nanpercentile(np.arange(nints), marks) + # Initialize storage array for widths. + fwhm = np.zeros((nints, dimx-254)) + # Fit a Gaussian to the PSF in each detector column. + loc = 0 + for j in range(nints): + # Progress print. + if j >= int(locs[loc]): + fancyprint('Slice {}: {}% complete.'.format(seg_no, marks[loc])) + loc += 1 + # Cut out first 250 columns as there is order 2 contmination. + for i in range(250, dimx-4): + p0 = [1., ycens_o1[i], 1.] + data = np.copy(cube[j, :, i]) + # Replace any NaN values with a median. + if np.isnan(data).any(): + ii = np.where(np.isnan(data)) + data[ii] = np.nanmedian(data) + # Fit a Gaussian to the profile, and save the FWHM. + try: + coeff, var_matrix = curve_fit(gauss, np.arange(dimy), data, + p0=p0) + fwhm[j, i-250] = coeff[2] * 2.355 + except RuntimeError: + fwhm[j, i-250] = np.nan + # Get median FWHM per integration. + fwhm = np.nanmedian(fwhm, axis=1) + return fwhm diff --git a/supreme_spoon/fit_lightcurves.py b/supreme_spoon/fit_lightcurves.py index 3362d05..de3a9c0 100644 --- a/supreme_spoon/fit_lightcurves.py +++ b/supreme_spoon/fit_lightcurves.py @@ -8,6 +8,8 @@ Juliet light curve fitting script. """ +import warnings + from astropy.io import fits import copy from datetime import datetime @@ -86,7 +88,7 @@ 'ecc_p1': r'$e$', 'omega_p1': r'$\Omega$', 'a_p1': r'$a/R_*$', 'sigma_w_SOSS': r'$\sigma_w$', 'theta0_SOSS': r'$\theta_0$', 'theta1_SOSS': r'$\theta_1$', - 'theta2_SOSS': r'$\theta_2$', + 'theta2_SOSS': r'$\theta_2$', 'theta3_SOSS': r'$\theta_3$', 'GP_sigma_SOSS': r'$GP \sigma$', 'GP_rho_SOSS': r'$GP \rho$', 'GP_S0_SOSS': r'$GP S0$', 'GO_omega0_SOSS': r'$GP \Omega_0$', 'GP_Q_SOSS': r'$GP Q$', @@ -183,14 +185,39 @@ priors[param] = {} priors[param]['distribution'] = dist priors[param]['hyperparameters'] = hyperp - # For transit fits, interpolate LD coefficients from stellar models. - if config['occultation_type'] == 'transit': - if order == 1 and config['ldcoef_file_o1'] is not None: - q1, q2 = stage4.read_ld_coefs(config['ldcoef_file_o1'], wave_low, - wave_up) - if order == 2 and config['ldcoef_file_o2'] is not None: - q1, q2 = stage4.read_ld_coefs(config['ldcoef_file_o2'], wave_low, - wave_up) + + # For transit fits, calculate LD coefficients from stellar models. + if config['occultation_type'] == 'transit' and config['ld_fit_type'] != 'free': + calculate = True + # First check if LD coefficient files have been provided. + if config['ldcoef_file_o{}'.format(order)] is not None: + fancyprint('Reading limb-darkening coefficient file.') + try: + q1, q2 = stage4.read_ld_coefs(config['ldcoef_file_o{}'.format(order)], + wave_low, wave_up) + except ValueError: + msg = 'LD coefficient file could not be correctly parsed. ' \ + 'Falling back onto LD calculation.' + warnings.warn(msg) + calculate = True + calculate = False + if calculate is True: + # Calculate LD coefficients on specified wavelength grid. + fancyprint('Calculating limb-darkening coefficients.') + m_h, logg, teff = config['m_h'], config['logg'], config['teff'] + msg = 'All stellar parameters must be provided to calculate ' \ + 'limb-darkening coefficients.' + assert np.all(np.array([m_h, logg, teff]) != None), msg + c1, c2 = stage4.gen_ld_coefs(config['spectrace_ref'], wave_low, + wave_up, order, m_h, logg, teff, + config['ld_data_path']) + q1, q2 = juliet.reverse_q_coeffs('quadratic', c1, c2) + # Save calculated coefficients. + target = fits.getheader(config['infile'], 0)['TARGET'] + outdir_ld = outdir + 'speclightcurve{}/'.format(fit_suffix) + utils.verify_path(outdir_ld) + utils.save_ld_priors(wave, c1, c2, order, target, m_h, teff, logg, + outdir=outdir_ld) # Pack fitting arrays and priors into dictionaries. data_dict, prior_dict = {}, {} @@ -212,16 +239,23 @@ prior_dict[thisbin] = copy.deepcopy(priors) # For transit only; update the LD prior for this bin if available. if config['occultation_type'] == 'transit': - # Set prior width to 0.2 around the model value - based on - # findings of Patel & Espinoza 2022. - if config['ldcoef_file_o1'] is not None: + if config['ld_fit_type'] == 'prior': + # Set prior width to 0.2 around the model value - based on + # findings of Patel & Espinoza 2022. if np.isfinite(q1[wavebin]): prior_dict[thisbin]['q1_SOSS']['distribution'] = 'truncatednormal' prior_dict[thisbin]['q1_SOSS']['hyperparameters'] = [q1[wavebin], 0.2, 0.0, 1.0] - if config['ldcoef_file_o2'] is not None: if np.isfinite(q2[wavebin]): prior_dict[thisbin]['q2_SOSS']['distribution'] = 'truncatednormal' prior_dict[thisbin]['q2_SOSS']['hyperparameters'] = [q2[wavebin], 0.2, 0.0, 1.0] + elif config['ld_fit_type'] == 'fixed': + # Fix LD to model values. + if np.isfinite(q1[wavebin]): + prior_dict[thisbin]['q1_SOSS']['distribution'] = 'fixed' + prior_dict[thisbin]['q1_SOSS']['hyperparameters'] = q1[wavebin] + if np.isfinite(q2[wavebin]): + prior_dict[thisbin]['q2_SOSS']['distribution'] = 'fixed' + prior_dict[thisbin]['q2_SOSS']['hyperparameters'] = q2[wavebin] # === Do the Fit === # Fit each light curve @@ -235,7 +269,7 @@ # make summary plots if necessary. fancyprint('Summarizing fit results.') data = np.ones((nints, nbins)) * np.nan - models = np.ones((nints, nbins)) * np.nan + models = np.ones((2, nints, nbins)) * np.nan residuals = np.ones((nints, nbins)) * np.nan order_results = {'dppm': [], 'dppm_err': [], 'wave': wave, 'wave_err': np.mean([wave - wave_low, wave_up - wave], @@ -326,16 +360,21 @@ outpdf=outpdf) data[:, i] = norm_flux[:, i] - models[:, i] = transit_model + models[0, :, i] = transit_model + if systematics is not None: + models[1, :, i] = systematics residuals[:, i] = norm_flux[:, i] - transit_model except: pass results_dict['order {}'.format(order)] = order_results + # Save best-fitting light curve models. + np.save(outdir + 'speclightcurve{0}/' + '_models_order{1}.npy'.format(fit_suffix, order), models) # Plot 2D lightcurves. if doplot is True: plotting.make_2d_lightcurve_plot(wave, data, outpdf=outpdf, title='Normalized Lightcurves') - plotting.make_2d_lightcurve_plot(wave, models, outpdf=outpdf, + plotting.make_2d_lightcurve_plot(wave, models[0], outpdf=outpdf, title='Model Lightcurves') plotting.make_2d_lightcurve_plot(wave, residuals, outpdf=outpdf, title='Residuals') diff --git a/supreme_spoon/fit_lightcurves.yaml b/supreme_spoon/fit_lightcurves.yaml index 2f03495..a2bd41a 100644 --- a/supreme_spoon/fit_lightcurves.yaml +++ b/supreme_spoon/fit_lightcurves.yaml @@ -2,7 +2,7 @@ # # This is the configuration file for fit_lightcurves.py. # -# ===== Fit Metadata ===== +# ====== Fit Metadata ====== # Name tag for output file directory. output_tag : '' # File containing light curves to fit. @@ -28,7 +28,7 @@ npix : None # Planet identifier. planet_letter : 'b' -# ===== Fit Priors + Parameters ===== +# ====== Fit Priors + Parameters ====== # Fitting priors in juliet format. # For eclipse fits, must not pass q1_soss and q2_soss, and must pass t_secondary_p1 and fp_p1. params : ['P_p1', 't0_p1', 'p_p1', 'b_p1', @@ -41,9 +41,7 @@ hyperps : [3.42525650, 2459751.821681146, [0.01, 0.9], 0.748, [0., 1.], [0., 1.], 0.0, 90., 8.82, 1.0, 0, [0.1, 10000]] -# Paths to files containing model limb-darkening coefficients. -ldcoef_file_o1 : None -ldcoef_file_o2 : None +# === Detrending Paramaters === # Path to file containing linear detrending parameters. lm_file : None # Key names for detrending parametrers. @@ -52,4 +50,24 @@ lm_parameters : ['x'] gp_file : None # Key name for GP training parametrer. gp_parameter : '' + +# === Parameters for Limb-Darkening -- Transit Only === +# Options for limb-darkening fitting. One of 'fixed', 'prior', or 'free'. +ld_fit_type : 'prior' +# Stellar Metallicity +m_h : None +# Star log Gravity +logg : None +# Star Effective Temperature +teff : None +# Path to ExoTiC-LD Reference Files +# See ExoTiC-LD documentation for more info: https://exotic-ld.readthedocs.io/en/latest/views/installation.html. +ld_data_path : +# Path to JWST spectrace Reference File +# Will be in crds_cache with file name like jwst_niriss_spectrace_XXXX.fits +spectrace_ref : './crds_cache/references/jwst/niriss/jwst_niriss_spectrace_0023.fits' +# Paths to files containing model limb-darkening coefficients. +# If provided, will take precedence over calculated values. Must be in format compatible with read_ld_coefs. +ldcoef_file_o1 : None +ldcoef_file_o2 : None # ======================= END LIGHT CURVE FIT CONFIG FILE =========================== \ No newline at end of file diff --git a/supreme_spoon/plotting.py b/supreme_spoon/plotting.py index 0b9c1bd..57a02d2 100644 --- a/supreme_spoon/plotting.py +++ b/supreme_spoon/plotting.py @@ -792,6 +792,9 @@ def nine_panel_plot(data, text=None, outfile=None, show_plot=True, **kwargs): else: max_percentile = kwargs['max_percentile'] vmax = np.nanpercentile(data[frame], max_percentile) + while vmax <= vmin: + max_percentile += 5 + vmax = np.nanpercentile(data[frame], max_percentile) else: vmax = kwargs['vmax'] ax.imshow(data[frame], aspect='auto', origin='lower', vmin=vmin, diff --git a/supreme_spoon/run_DMS.py b/supreme_spoon/run_DMS.py index 6c29ff7..42405fa 100644 --- a/supreme_spoon/run_DMS.py +++ b/supreme_spoon/run_DMS.py @@ -66,6 +66,8 @@ # Open some of the input files. background_model = np.load(config['background_file']) +if np.ndim(background_model) == 3: + background_model = background_model[0] if config['smoothed_wlc'] is not None: config['smoothed_wlc'] = np.load(config['smoothed_wlc']) if config['centroids'] is not None: @@ -93,23 +95,24 @@ # ===== Run Stage 2 ===== if 2 in config['run_stages']: + if config['soss_width'] <= 40: + mask_width = 45 + else: + mask_width = config['soss_width'] + 5 stage2_results = run_stage2(stage1_results, background_model=background_model, baseline_ints=config['baseline_ints'], smoothed_wlc=config['smoothed_wlc'], save_results=config['save_results'], force_redo=config['force_redo'], - calculate_stability_ccf=config['calculate_stability_ccf'], - stability_params_ccf=config['stability_params_ccf'], - nthreads=config['nthreads'], - calculate_stability_pca=config['calculate_stability_pca'], + calculate_stability=config['calculate_stability'], pca_components=config['pca_components'], output_tag=config['output_tag'], smoothing_scale=config['smoothing_scale'], skip_steps=config['stage2_skip'], generate_lc=config['generate_lc'], generate_tracemask=config['generate_tracemask'], - mask_width=config['mask_width'], + mask_width=mask_width, pixel_flags=config['outlier_maps'], generate_order0_mask=config['generate_order0_mask'], f277w=config['f277w'], diff --git a/supreme_spoon/run_DMS.yaml b/supreme_spoon/run_DMS.yaml index 1a4e441..119d44c 100644 --- a/supreme_spoon/run_DMS.yaml +++ b/supreme_spoon/run_DMS.yaml @@ -11,7 +11,8 @@ input_dir : 'DMS_uncal/' input_filetag : 'uncal' # ===== Stage 1 Input Files & Parameters ===== -# Background model. +# Background model. Default model available here: +# https://jwst-docs.stsci.edu/jwst-near-infrared-imager-and-slitless-spectrograph/niriss-observing-strategies/niriss-soss-recommended-strategies background_file : 'model_background256.npy' # For 1/f correction; outlier pixel maps (optional). outlier_maps : None @@ -31,20 +32,9 @@ stage1_kwargs : {} # ===== Stage 2 Input Files & Parameters ===== # If True, create a mask of the target trace orders. generate_tracemask : True -# Size of box to mask around the trace - should be wider than extraction box. -# Only necessary if generate_tracemask is True. -mask_width : 45 # If True, calculate the stability of the SOSS trace over the course of the -# TSO using a CCF analysis. -calculate_stability_ccf : False -# Parameters for which to calcuate the stability: 'x', 'y', 'FWHM, or 'ALL'. Only -# necessary if calculate_stability_ccf is True. -stability_params_ccf : 'ALL' -# Number of CPUs for CCF stability parameter calculation multiprocessing. -nthreads : 4 -# If True, calculate the stability of the SOSS trace over the course of the -# TSO using a principle component analysis. -calculate_stability_pca : True +# TSO using principle component analysis. +calculate_stability : True # Number of principle components to extract. Only neceesary if calculate_stability_pca # is True. pca_components : 10 @@ -64,9 +54,9 @@ stage2_kwargs : {} # ===== Stage 3 Input Files & Parameters ===== # Extraction method: box or atoca. -extract_method : 'atoca' +extract_method : 'box' # Box width to extract around the trace center. -soss_width : 30 +soss_width : 40 # Specprofile reference file for ATOCA (optional). specprofile : None # File with centroids for all three orders for box extraction (optional). diff --git a/supreme_spoon/stage2.py b/supreme_spoon/stage2.py index 2466b16..b4a2d2e 100644 --- a/supreme_spoon/stage2.py +++ b/supreme_spoon/stage2.py @@ -15,11 +15,8 @@ import numpy as np import os import pandas as pd -import ray from sklearn.decomposition import PCA -from scipy.interpolate import interp2d from scipy.ndimage import median_filter -from scipy.optimize import curve_fit from tqdm import tqdm import warnings @@ -157,18 +154,19 @@ def run(self, save_results=True, force_redo=False, do_plot=False, else: with warnings.catch_warnings(): warnings.filterwarnings('ignore') - scale1, scale2 = None, None - if 'scale1' in kwargs.keys(): - scale1 = kwargs['scale1'] - if 'scale2' in kwargs.keys(): - scale2 = kwargs['scale2'] + scale, background_coords = None, None + if 'scale' in kwargs.keys(): + scale = kwargs['scale'] + if 'background_coords' in kwargs.keys(): + background_coords = kwargs['background_coords'] step_results = backgroundstep(self.datafiles, self.background_model, output_dir=self.output_dir, save_results=save_results, fileroots=self.fileroots, fileroot_noseg=self.fileroot_noseg, - scale1=scale1, scale2=scale2) + scale=scale, + background_coords=background_coords) results, background_models = step_results # Do step plot if requested. @@ -297,8 +295,7 @@ def __init__(self, input_data, deepframe, output_dir='./'): def run(self, generate_tracemask=True, mask_width=45, pixel_flags=None, generate_order0_mask=True, f277w=None, - calculate_stability_ccf=True, stability_params='ALL', nthreads=4, - calculate_stability_pca=True, pca_components=10, + calculate_stability=True, pca_components=10, save_results=True, force_redo=False, generate_lc=True, baseline_ints=None, smoothing_scale=None, do_plot=False, show_plot=False): @@ -332,10 +329,7 @@ def run(self, generate_tracemask=True, mask_width=45, pixel_flags=None, pixel_flags = new_pixel_flags step_results = tracingstep(self.datafiles, self.deepframe, - calculate_stability_ccf=calculate_stability_ccf, - stability_params=stability_params, - nthreads=nthreads, - calculate_stability_pca=calculate_stability_pca, + calculate_stability=calculate_stability, pca_components=pca_components, generate_tracemask=generate_tracemask, mask_width=mask_width, @@ -355,13 +349,13 @@ def run(self, generate_tracemask=True, mask_width=45, pixel_flags=None, def backgroundstep(datafiles, background_model, output_dir='./', save_results=True, fileroots=None, fileroot_noseg='', - scale1=None, scale2=None): + scale=None, background_coords=None): """Background subtraction must be carefully treated with SOSS observations. Due to the extent of the PSF wings, there are very few, if any, non-illuminated pixels to serve as a sky region. Furthermore, the zodi background has a unique stepped shape, which would render a constant background subtraction ill-advised. Therefore, a background subtracton is - performed by scaling a model background to the countns level of a median + performed by scaling a model background to the counts level of a median stack of the exposure. This scaled model background is then subtracted from each integration. @@ -380,14 +374,13 @@ def backgroundstep(datafiles, background_model, output_dir='./', Root names for output files. fileroot_noseg : str Root name with no segment information. - scale1 : float, None + scale : float, array-like[float], None Scaling value to apply to background model to match data. Will take - precedence over calculated scaling value. If only scale1 is provided, - this will multiply the entire frame. If scale2 is also provided, this - will be the "pre-stp" scaling. - scale2 : float, None - "Post-step" scaling value. scale1 must also be passed if this - parameter is not None. + precedence over calculated scaling value. If applied at group level, + length of scaling array must equal ngroup. + background_coords : array-like[int], None + Region of frame to use to estimate the background. Must be 1D: + [x_low, x_up, y_low, y_up]. Returns ------- @@ -416,39 +409,41 @@ def backgroundstep(datafiles, background_model, output_dir='./', # Make median stack of all integrations to use for background scaling. # This is to limit the influence of cosmic rays, which can greatly effect # the background scaling factor calculated for an individual inegration. - fancyprint('Generating a deep stack using all integrations.') - deepstack = utils.make_deepstack(cube) - # If applied at the integration level, reshape deepstack to 3D. - if np.ndim(deepstack) != 3: - dimy, dimx = np.shape(deepstack) - deepstack = deepstack.reshape(1, dimy, dimx) - ngroup, dimy, dimx = np.shape(deepstack) + fancyprint('Generating a median stack using all integrations.') + stack = utils.make_deepstack(cube) + # If applied at the integration level, reshape median stack to 3D. + if np.ndim(stack) != 3: + dimy, dimx = np.shape(stack) + stack = stack.reshape(1, dimy, dimx) + ngroup, dimy, dimx = np.shape(stack) + # Ensure if user-defined scalings are provided that there is one per group. + if scale is not None: + scale = np.atleast_1d(scale) + assert len(scale) == ngroup fancyprint('Calculating background model scaling.') - model_scaled = np.zeros_like(deepstack) - if scale1 is None: - fancyprint('Scale factor(s):') - else: - fancyprint(' Using user-defined background scaling(s):') - if scale2 is not None: - fancyprint('Pre-step scale factor: {:.5f}'.format(scale1)) - fancyprint('Post-step scale factor: {:.5f}'.format(scale2)) - else: - fancyprint('Background scale factor: {:.5f}'.format(scale1)) + model_scaled = np.zeros_like(stack) first_time = True for i in range(ngroup): - if scale1 is None: - # Calculate the scaling of the model background to the median - # stack. - if dimy == 96: - # Use area in bottom left corner of detector for SUBSTRIP96. - xl, xu = 5, 21 - yl, yu = 5, 401 + if scale is None: + if background_coords is None: + # If region to estimate background is not provided, use a + # default region. + if dimy == 96: + # Use area in bottom left corner for SUBSTRIP96. + xl, xu = 5, 21 + yl, yu = 5, 401 + else: + # Use area in the top left corner for SUBSTRIP256 + xl, xu = 210, 250 + yl, yu = 250, 500 else: - # Use area in the top left corner of detector for SUBSTRIP256 - xl, xu = 210, 250 - yl, yu = 250, 500 - bkg_ratio = deepstack[i, xl:xu, yl:yu] / background_model[xl:xu, yl:yu] + # Use user-defined background scaling region. + assert len(background_coords) == 4 + # Convert to int if not already. + background_coords = np.array(background_coords).astype(int) + xl, xu, yl, yu = background_coords + bkg_ratio = stack[i, xl:xu, yl:yu] / background_model[xl:xu, yl:yu] # Instead of a straight median, use the median of the 2nd quartile # to limit the effect of any remaining illuminated pixels. q1 = np.nanpercentile(bkg_ratio, 25) @@ -457,22 +452,14 @@ def backgroundstep(datafiles, background_model, output_dir='./', scale_factor = np.nanmedian(bkg_ratio[ii]) if scale_factor < 0: scale_factor = 0 - fancyprint('Background scale factor: {1:.5f}'.format(i+1, scale_factor)) + fancyprint('Using calculated background scale factor: ' + '{:.5f}'.format(scale_factor)) model_scaled[i] = background_model * scale_factor - elif scale1 is not None and scale2 is None: - # If using a user specified scaling for the whole frame. - model_scaled[i] = background_model * scale1 else: - # If using seperate pre- and post- step scalings. - # Locate the step position using the gradient of the background. - grad_bkg = np.gradient(background_model, axis=1) - # The boundary between the left and right scalings - # actually happens before the step (-4). - step_pos = np.argmax(grad_bkg[:, 10:-10], axis=1) + 10 - 4 - # Seperately scale both sides of the step. - for j in range(dimy): - model_scaled[i, j, :(step_pos[j])] = background_model[j, :(step_pos[j])] * scale1 - model_scaled[i, j, (step_pos[j]):] = background_model[j, (step_pos[j]):] * scale2 + # If using a user specified scaling for the whole frame. + fancyprint('Using user-defined background scaling: ' + '{:.5f}'.format(scale[i])) + model_scaled[i] = background_model * scale[i] # Loop over all segments in the exposure and subtract the background from # each of them. @@ -580,11 +567,11 @@ def badpixstep(datafiles, baseline_ints, smoothed_wlc=None, thresh=15, hotpix = np.zeros_like(deepframe_itl) nanpix = np.zeros_like(deepframe_itl) otherpix = np.zeros_like(deepframe_itl) - nan, hot, other = 0, 0, 0 + nan, hot, other, neg = 0, 0, 0, 0 nint, dimy, dimx = np.shape(newdata) # Loop over whole deepstack and flag deviant pixels. - for i in tqdm(range(4, dimx-4)): - for j in range(dimy-4): + for i in tqdm(range(4, dimx - 4)): + for j in range(dimy - 4): # If the pixel is known to be hot, add it to list to interpolate. if hot_pix[j, i]: hotpix[j, i] = 1 @@ -608,9 +595,12 @@ def badpixstep(datafiles, baseline_ints, smoothed_wlc=None, thresh=15, nanpix[j, i] = 1 nan += 1 elif deepframe_itl[j, i] < 0: + # Interpolate if bright, set to zero if dark. if med >= np.nanpercentile(deepframe_itl, 10): nanpix[j, i] = 1 - nan += 1 + else: + deepframe_itl[j, i] = 0 + neg += 1 elif np.abs(deepframe_itl[j, i] - med) >= (thresh * std): otherpix[j, i] = 1 other += 1 @@ -619,8 +609,8 @@ def badpixstep(datafiles, baseline_ints, smoothed_wlc=None, thresh=15, badpix = hotpix.astype(bool) | nanpix.astype(bool) | otherpix.astype(bool) badpix = badpix.astype(int) - fancyprint('{0} hot, {1} negative, and {2} other deviant pixels ' - 'identified.'.format(hot, nan, other)) + fancyprint('{0} hot, {1} nan, {2} negative, and {3} deviant pixels ' + 'identified.'.format(hot, nan, neg, other)) # Replace the flagged pixels in the median integration. newdeep, deepdq = utils.do_replacement(deepframe_itl, badpix, dq=np.ones_like(deepframe_itl), @@ -652,12 +642,12 @@ def badpixstep(datafiles, baseline_ints, smoothed_wlc=None, thresh=15, deepdq = ~deepdq.astype(bool) newdq[:, deepdq] = 0 - # Generate a final corrected deep frame for the baseline integrations. - deepframe_fnl = utils.make_deepstack(newdata[baseline_ints]) + # Generate a temporary corrected deep frame. + deepframe_tmp = utils.make_deepstack(newdata[baseline_ints]) # Final check along the time axis for outlier pixels. std_dev = bn.nanstd(newdata, axis=0) - scale = np.abs(newdata - deepframe_fnl)/std_dev + scale = np.abs(newdata - deepframe_tmp) / std_dev ii = np.where(scale > 5) mask = np.zeros_like(cube).astype(bool) mask[ii] = True @@ -668,6 +658,11 @@ def badpixstep(datafiles, baseline_ints, smoothed_wlc=None, thresh=15, newdata[ii] = newdeep[ii] ii = np.where(np.isnan(err_cube)) err_cube[ii] = np.nanmedian(err_cube) + # And replace any negatives with zeros + newdata[newdata < 0] = 0 + + # Make a final, corrected deepframe for the baseline intergations. + deepframe_fnl = utils.make_deepstack(newdata[baseline_ints]) results = [] current_int = 0 @@ -715,14 +710,12 @@ def badpixstep(datafiles, baseline_ints, smoothed_wlc=None, thresh=15, return results, deepframe_fnl -def tracingstep(datafiles, deepframe=None, calculate_stability_ccf=True, - stability_params='ALL', nthreads=4, - calculate_stability_pca=True, pca_components=10, - generate_tracemask=True, mask_width=45, pixel_flags=None, - generate_order0_mask=False, f277w=None, generate_lc=True, - baseline_ints=None, smoothing_scale=None, output_dir='./', - save_results=True, fileroot_noseg='', do_plot=False, - show_plot=False): +def tracingstep(datafiles, deepframe=None, calculate_stability=True, + pca_components=10, generate_tracemask=True, mask_width=45, + pixel_flags=None, generate_order0_mask=False, f277w=None, + generate_lc=True, baseline_ints=None, smoothing_scale=None, + output_dir='./', save_results=True, fileroot_noseg='', + do_plot=False, show_plot=False): """A multipurpose step to perform some initial analysis of the 2D dataframes and produce products which can be useful in further reduction iterations. The five functionalities are detailed below: @@ -743,15 +736,7 @@ def tracingstep(datafiles, deepframe=None, calculate_stability_ccf=True, deepframe : str, array-like[float], None Path to median stack file, or the median stack itself. Should be 2D (dimy, dimx). If None is passed, one will be generated. - calculate_stability_ccf : bool - If True, calculate the stabilty of the SOSS trace over the TSO using a - CCF method. - stability_params : str, array-like[str] - List of parameters for which to calculate the stability. Any of: 'x', - 'y', and/or 'FWHM', or 'ALL' for all three. - nthreads : int - Number of CPUs for CCF stability parameter calculation multiprocessing. - calculate_stability_pca : bool + calculate_stability : bool If True, calculate the stabilty of the SOSS trace over the TSO using a PCA method. pca_components : int @@ -806,7 +791,7 @@ def tracingstep(datafiles, deepframe=None, calculate_stability_ccf=True, datafiles = np.atleast_1d(datafiles) # If no deepframe is passed, construct one. Also generate a datacube for # later white light curve or stability calculations. - if deepframe is None or np.any([generate_lc, calculate_stability_ccf, calculate_stability_pca]) == True: + if deepframe is None or np.any([generate_lc, calculate_stability]) == True: # Construct datacube from the data files. for i, file in enumerate(datafiles): if isinstance(file, str): @@ -935,53 +920,13 @@ def tracingstep(datafiles, deepframe=None, calculate_stability_ccf=True, fancyprint('Order 0 mask saved to {}'.format(outfile)) # ===== PART 4: Calculate the trace stability ===== - # === CCF Method === - # If requested, calculate the change in position of the trace, as well as - # its FWHM over the course of the TSO. These quantities may be useful for - # lightcurve detrending. - if calculate_stability_ccf is True: - fancyprint('Calculating trace stability using the CCF method... ' - 'This might take a while.') - assert save_results is True, 'save_results must be True to run ' \ - 'soss_stability_ccf' - if stability_params == 'ALL': - stability_params = ['x', 'y', 'FWHM'] - - # Calculate the stability of the requested parameters. - stability_results = {} - if 'x' in stability_params: - fancyprint('Getting trace X-positions...') - ccf_x = soss_stability_xy(cube, axis='x', nthreads=nthreads) - stability_results['X'] = ccf_x - if 'y' in stability_params: - fancyprint('Getting trace Y-positions...') - ccf_y = soss_stability_xy(cube, axis='y', nthreads=nthreads) - stability_results['Y'] = ccf_y - if 'FWHM' in stability_params: - fancyprint('Getting trace FWHM values...') - fwhm = soss_stability_fwhm(cube, y1, nthreads=nthreads) - stability_results['FWHM'] = fwhm - - # Save stability results. - suffix = 'soss_stability.csv' - if os.path.exists(output_dir + fileroot_noseg + suffix): - old_data = pd.read_csv(output_dir + fileroot_noseg + suffix, - comment='#') - for key in stability_results.keys(): - old_data[key] = stability_results[key] - os.remove(output_dir + fileroot_noseg + suffix) - old_data.to_csv(output_dir + fileroot_noseg + suffix, index=False) - else: - df = pd.DataFrame(data=stability_results) - df.to_csv(output_dir + fileroot_noseg + suffix, index=False) - - # === PCA Method === - # If requested, calculate the stability of the SOSS trace using PCA. - if calculate_stability_pca is True: + # If requested, calculate the stability of the SOSS trace over the course + # of the TSO using PCA. + if calculate_stability is True: fancyprint('Calculating trace stability using the PCA method...' ' This might take a while.') assert save_results is True, 'save_results must be True to run ' \ - 'soss_stability_pca' + 'soss_stability.' # Calculate the trace stability using PCA. outfile = output_dir + 'soss_stability_pca.pdf' @@ -1074,228 +1019,6 @@ def make_order0_mask_from_f277w(f277w, thresh_std=1, thresh_size=10): return mask -def soss_stability_xy(cube, nsteps=501, axis='x', nthreads=4, - smoothing_scale=None): - """Perform a CCF analysis to track the movement of the SOSS trace - relative to the median stack over the course of a TSO. - - Parameters - ---------- - cube : array-like[float] - Data cube. Should be 3D (ints, dimy, dimx). - nsteps : int - Number of CCF steps to test. - axis : str - Axis over which to calculate the CCF - either 'x', or 'y'. - nthreads : int - Number of CPUs for multiprocessing. - smoothing_scale : int - Length scale over which to smooth results. - - Returns - ------- - ccf : array-like[float] - The cross-correlation results. - """ - - # Initialize ray with specified number of threads. - ray.shutdown() - ray.init(num_cpus=nthreads) - - # Subtract integration-wise median from cube for CCF. - cube = cube - np.nanmedian(cube, axis=(1, 2))[:, None, None] - # Calculate median stack. - deep = bn.nanmedian(cube, axis=0) - - # Divide total data cube into segments and run each segment in parallel - # with ray. - ii = 0 - all_fits = [] - nints = np.shape(cube)[0] - seglen = nints // nthreads - for i in range(nthreads): - if i == nthreads - 1: - cube_seg = cube[ii:] - else: - cube_seg = cube[ii:ii + seglen] - - all_fits.append(soss_stability_xy_run.remote(cube_seg, deep, seg_no=i+1, - nsteps=nsteps, axis=axis)) - ii += seglen - - # Run the CCFs. - ray_results = ray.get(all_fits) - - # Stack all the CCF results into a single array. - maxvals = [] - for i in range(nthreads): - if i == 0: - maxvals = ray_results[i] - else: - maxvals = np.concatenate([maxvals, ray_results[i]]) - - # Smooth results if requested. - if smoothing_scale is not None: - ccf = median_filter(np.linspace(-0.01, 0.01, nsteps)[maxvals], - smoothing_scale) - else: - ccf = np.linspace(-0.01, 0.01, nsteps)[maxvals] - ccf = ccf.reshape(nints) - - return ccf - - -@ray.remote -def soss_stability_xy_run(cube_sub, med, seg_no, nsteps=501, axis='x'): - """Wrapper to perform CCF calculations in parallel with ray. - """ - - # Get data dimensions. - nints, dimy, dimx = np.shape(cube_sub) - - # Get integration numbers to show progress prints. - marks = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] - locs = np.nanpercentile(np.arange(nints), marks) - - # Initialize CCF variables. - ccf = np.zeros((nints, nsteps)) - f = interp2d(np.arange(dimx), np.arange(dimy), med, kind='cubic') - # Perform cross-correlation over desired axis. - loc = 0 - for i in range(nints): - # Progress print. - if i >= int(locs[loc]): - fancyprint('Slice {}: {}% complete.'.format(seg_no, marks[loc])) - loc += 1 - for j, jj in enumerate(np.linspace(-0.01, 0.01, nsteps)): - if axis == 'x': - interp = f(np.arange(dimx) + jj, np.arange(dimy)) - elif axis == 'y': - interp = f(np.arange(dimx), np.arange(dimy) + jj) - else: - raise ValueError('Unknown axis: {}'.format(axis)) - ccf[i, j] = np.nansum(cube_sub[i] * interp) - - # Determine the peak of the CCF for each integration to get the - # best-fitting shift. - maxvals = [] - for i in range(nints): - maxvals.append(np.where(ccf[i] == np.max(ccf[i]))[0]) - maxvals = np.array(maxvals) - maxvals = maxvals.reshape(maxvals.shape[0]) - - return maxvals - - -def soss_stability_fwhm(cube, ycens_o1, nthreads=4, smoothing_scale=None): - """Estimate the FWHM of the trace over the course of a TSO by fitting a - Gaussian to each detector column. - - Parameters - ---------- - cube : array-like[float] - Data cube. Should be 3D (ints, dimy, dimx). - ycens_o1 : arrray-like[float] - Y-centroid positions of the order 1 trace. Should have length dimx. - nthreads : int - Number of CPUs for multiprocessing. - smoothing_scale : int - Length scale over which to smooth results. - - Returns - ------- - fwhm : array-like[float] - FWHM estimates for each column at every integration. - """ - - # Initialize ray with specified number of threads. - ray.shutdown() - ray.init(num_cpus=nthreads) - - # Divide total data cube into segments and run each segment in parallel - # with ray. - ii = 0 - all_fits = [] - nints = np.shape(cube)[0] - seglen = nints // nthreads - for i in range(nthreads): - if i == nthreads - 1: - cube_seg = cube[ii:] - else: - cube_seg = cube[ii:ii + seglen] - - all_fits.append(soss_stability_fwhm_run.remote(cube_seg, ycens_o1, - seg_no=i+1)) - ii += seglen - - # Run the CCFs. - ray_results = ray.get(all_fits) - - # Stack all the CCF results into a single array. - fwhm = [] - for i in range(nthreads): - if i == 0: - fwhm = ray_results[i] - else: - fwhm = np.concatenate([fwhm, ray_results[i]]) - - # Set median of trend to zero. - fwhm -= np.median(fwhm) - # Smooth the trend. - if smoothing_scale is None: - smoothing_scale = int(0.2*nints) - fwhm = median_filter(fwhm, smoothing_scale) - - return fwhm - - -@ray.remote -def soss_stability_fwhm_run(cube, ycens_o1, seg_no): - """Wrapper to perform FWHM calculations in parallel with ray. - """ - - def gauss(x, *p): - amp, mu, sigma = p - return amp * np.exp(-(x - mu) ** 2 / (2. * sigma ** 2)) - - # Get data dimensions. - nints, dimy, dimx = np.shape(cube) - - # Get integration numbers to show progress prints. - marks = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] - locs = np.nanpercentile(np.arange(nints), marks) - - # Initialize storage array for widths. - fwhm = np.zeros((nints, dimx-254)) - # Fit a Gaussian to the PSF in each detector column. - loc = 0 - for j in range(nints): - # Progress print. - if j >= int(locs[loc]): - fancyprint('Slice {}: {}% complete.'.format(seg_no, marks[loc])) - loc += 1 - # Cut out first 250 columns as there is order 2 contmination. - for i in range(250, dimx-4): - p0 = [1., ycens_o1[i], 1.] - data = np.copy(cube[j, :, i]) - # Replace any NaN values with a median. - if np.isnan(data).any(): - ii = np.where(np.isnan(data)) - data[ii] = np.nanmedian(data) - # Fit a Gaussian to the profile, and save the FWHM. - try: - coeff, var_matrix = curve_fit(gauss, np.arange(dimy), data, - p0=p0) - fwhm[j, i-250] = coeff[2] * 2.355 - except RuntimeError: - fwhm[j, i-250] = np.nan - - # Get median FWHM per integration. - fwhm = np.nanmedian(fwhm, axis=1) - - return fwhm - - def soss_stability_pca(cube, n_components=10, outfile=None, do_plot=False, show_plot=False): """Calculate the stability of the SOSS trace over the course of a TSO @@ -1353,13 +1076,12 @@ def soss_stability_pca(cube, n_components=10, outfile=None, do_plot=False, def run_stage2(results, background_model, baseline_ints, smoothed_wlc=None, - save_results=True, force_redo=False, - calculate_stability_ccf=True, stability_params_ccf='ALL', - nthreads=4, calculate_stability_pca=True, pca_components=10, - root_dir='./', output_tag='', smoothing_scale=None, - skip_steps=None, generate_lc=True, generate_tracemask=True, - mask_width=45, pixel_flags=None, generate_order0_mask=True, - f277w=None, do_plot=False, show_plot=False, **kwargs): + save_results=True, force_redo=False, calculate_stability=True, + pca_components=10, root_dir='./', output_tag='', + smoothing_scale=None, skip_steps=None, generate_lc=True, + generate_tracemask=True, mask_width=45, pixel_flags=None, + generate_order0_mask=True, f277w=None, do_plot=False, + show_plot=False, **kwargs): """Run the supreme-SPOON Stage 2 pipeline: spectroscopic processing, using a combination of official STScI DMS and custom steps. Documentation for the official DMS steps can be found here: @@ -1379,17 +1101,9 @@ def run_stage2(results, background_model, baseline_ints, smoothed_wlc=None, If True, save results of each step to file. force_redo : bool If True, redo steps even if outputs files are already present. - calculate_stability_ccf : bool - If True, calculate the stability of the SOSS trace over the course of - the TSO using a CCF method. - stability_params_ccf : str, array-like[str] - List of parameters for which to calculate the stability. Any of: 'x', - 'y', and/or 'FWHM', or 'ALL' for all three. - nthreads : int - Number of CPUs for stability parameter calculation multiprocessing. - calculate_stability_pca : bool + calculate_stability : bool If True, calculate the stability of the SOSS trace over the course of - the TSO using a CCF method. + the TSO using a PCA method. pca_components : int Number of PCA components to calculate. root_dir : str @@ -1505,10 +1219,7 @@ def run_stage2(results, background_model, baseline_ints, smoothed_wlc=None, # Custom DMS step. if 'TracingStep' not in skip_steps: step = TracingStep(results, deepframe=deepframe, output_dir=outdir) - step_results = step.run(calculate_stability_ccf=calculate_stability_ccf, - stability_params=stability_params_ccf, - nthreads=nthreads, - calculate_stability_pca=calculate_stability_pca, + step_results = step.run(calculate_stability=calculate_stability, pca_components=pca_components, generate_tracemask=generate_tracemask, mask_width=mask_width, pixel_flags=pixel_flags, diff --git a/supreme_spoon/stage3.py b/supreme_spoon/stage3.py index 35ad8e8..aca5b58 100644 --- a/supreme_spoon/stage3.py +++ b/supreme_spoon/stage3.py @@ -13,6 +13,7 @@ import glob import numpy as np import pandas as pd +import pastasoss from scipy.interpolate import interp1d from scipy.signal import butter, filtfilt import os @@ -165,12 +166,13 @@ def run(self, soss_width=25, specprofile=None, centroids=None, # Save the final extraction parameters. extract_params = {'soss_width': soss_width, 'method': self.extract_method} - # Get timestamps. + # Get timestamps and pupil wheel position. for i, datafile in enumerate(self.datafiles): with utils.open_filetype(datafile) as file: this_time = file.int_times['int_mid_BJD_TDB'] if i == 0: times = this_time + pwcpos = file.meta.instrument.pupil_position else: times = np.concatenate([times, this_time]) @@ -183,6 +185,7 @@ def run(self, soss_width=25, specprofile=None, centroids=None, spectra = format_extracted_spectra(results, times, extract_params, self.pl_name, st_teff, st_logg, st_met, throughput=thpt, + pwcpos=pwcpos, output_dir=self.output_dir, save_results=save_results) @@ -392,7 +395,7 @@ def highpass_filter(signal, order=3, freq=0.05): def format_extracted_spectra(datafiles, times, extract_params, target_name, st_teff=None, st_logg=None, st_met=None, - throughput=None, output_dir='./', + throughput=None, pwcpos=None, output_dir='./', save_results=True): """Unpack the outputs of the 1D extraction and format them into lightcurves at the native detector resolution. @@ -419,6 +422,8 @@ def format_extracted_spectra(datafiles, times, extract_params, target_name, Stellar metallicity as [Fe/H]. throughput : str Path to JWST spectrace reference file. + pwcpos : float + Filter wheel position. Returns ------- @@ -429,10 +434,8 @@ def format_extracted_spectra(datafiles, times, extract_params, target_name, fancyprint('Formatting extracted 1d spectra.') # Box extract outputs will just be a tuple of arrays. if isinstance(datafiles, tuple): - wave1d_o1 = datafiles[0] flux_o1 = datafiles[1] ferr_o1 = datafiles[2] - wave1d_o2 = datafiles[3] flux_o2 = datafiles[4] ferr_o2 = datafiles[5] # Whereas ATOCA extract outputs are in the atoca extract1dstep format. @@ -443,25 +446,33 @@ def format_extracted_spectra(datafiles, times, extract_params, target_name, for i, file in enumerate(datafiles): segment = utils.unpack_atoca_spectra(file) if i == 0: - wave2d_o1 = segment[1]['WAVELENGTH'] flux_o1 = segment[1]['FLUX'] ferr_o1 = segment[1]['FLUX_ERROR'] - wave2d_o2 = segment[2]['WAVELENGTH'] flux_o2 = segment[2]['FLUX'] ferr_o2 = segment[2]['FLUX_ERROR'] else: - wave2d_o1 = np.concatenate([wave2d_o1, segment[1]['WAVELENGTH']]) flux_o1 = np.concatenate([flux_o1, segment[1]['FLUX']]) ferr_o1 = np.concatenate([ferr_o1, segment[1]['FLUX_ERROR']]) - wave2d_o2 = np.concatenate([wave2d_o2, segment[2]['WAVELENGTH']]) flux_o2 = np.concatenate([flux_o2, segment[2]['FLUX']]) ferr_o2 = np.concatenate([ferr_o2, segment[2]['FLUX_ERROR']]) - # Create 1D wavelength axes from the 2D wavelength solution. - wave1d_o1, wave1d_o2 = wave2d_o1[0], wave2d_o2[0] - # Refine wavelength solution. + # Get wavelength solution. + # First, use PASTASOSS to predict wavelength solution from pupil wheel + # position. + wave1d_o1 = pastasoss.get_soss_traces(pwcpos=pwcpos, order='1', + interp=True).wavelength + soln_o2 = pastasoss.get_soss_traces(pwcpos=pwcpos, order='2', + interp=True) + xpos_o2, wave1d_o2 = soln_o2.x.astype(int), soln_o2.wavelength + # Trim extracted quantities to match shapes of pastasoss quantities. + flux_o1 = flux_o1[:, 4:-4] + ferr_o1 = ferr_o1[:, 4:-4] + flux_o2 = flux_o2[:, xpos_o2] + ferr_o2 = ferr_o2[:, xpos_o2] + + # Now cross-correlate with stellar model. # If one or more of the stellar parameters are not provided, use the - # default wavelength solution. + # wavelength solution from pastasoss. if None in [st_teff, st_logg, st_met]: fancyprint('Stellar parameters not provided. ' 'Using default wavelength solution.', msg_type='WARNING') @@ -501,6 +512,12 @@ def format_extracted_spectra(datafiles, times, extract_params, target_name, wave1d_o1 += wave_shift wave1d_o2 += wave_shift + # Restrict order 2 to only the useful bit (0.6 - 0.85µm). + ii = np.where((wave1d_o2 >= 0.6) & (wave1d_o2 < 0.85))[0] + wave1d_o2 = wave1d_o2[ii] + flux_o2 = flux_o2[:, ii] + ferr_o2 = ferr_o2[:, ii] + # Clip remaining 3-sigma outliers. flux_o1_clip = utils.sigma_clip_lightcurves(flux_o1, ferr_o1) flux_o2_clip = utils.sigma_clip_lightcurves(flux_o2, ferr_o2) @@ -534,8 +551,8 @@ def format_extracted_spectra(datafiles, times, extract_params, target_name, def run_stage3(results, save_results=True, root_dir='./', force_redo=False, - extract_method='atoca', specprofile=None, centroids=None, - soss_width=25, st_teff=None, st_logg=None, st_met=None, + extract_method='box', specprofile=None, centroids=None, + soss_width=40, st_teff=None, st_logg=None, st_met=None, planet_letter='b', output_tag='', do_plot=False, show_plot=False): """Run the supreme-SPOON Stage 3 pipeline: 1D spectral extraction, using diff --git a/supreme_spoon/stage4.py b/supreme_spoon/stage4.py index 755bf91..82fa9f9 100644 --- a/supreme_spoon/stage4.py +++ b/supreme_spoon/stage4.py @@ -18,7 +18,6 @@ from tqdm import tqdm from jwst import datamodels -from jwst.pipeline import calwebb_spec2 from supreme_spoon.utils import fancyprint @@ -345,15 +344,15 @@ def fit_lightcurves(data_dict, prior_dict, order, output_dir, fit_suffix, return results -def gen_ld_coefs(datafile, wavebin_low, wavebin_up, order, m_h, logg, teff, +def gen_ld_coefs(spectrace_ref, wavebin_low, wavebin_up, order, m_h, logg, teff, ld_data_path, model_type='stagger'): """Generate estimates of quadratic limb-darkening coefficients using the ExoTiC-LD package. Parameters ---------- - datafile : str - Path to extract1d output file. + spectrace_ref : str + Path to spectrace reference file. wavebin_low : array-like[float] Lower edge of wavelength bins. wavebin_up: array-like[float] @@ -385,8 +384,6 @@ def gen_ld_coefs(datafile, wavebin_low, wavebin_up, order, m_h, logg, teff, sld = StellarLimbDarkening(m_h, teff, logg, model_type, ld_data_path) # Load the most up to date throughput info for SOSS - step = calwebb_spec2.extract_1d_step.Extract1dStep() - spectrace_ref = step.get_reference_file(datafile, 'spectrace') spec_trace = datamodels.SpecTraceModel(spectrace_ref) wavelengths = spec_trace.trace[order-1].data['WAVELENGTH'] * 10000 throughputs = spec_trace.trace[order-1].data['THROUGHPUT'] @@ -413,7 +410,8 @@ def gen_ld_coefs(datafile, wavebin_low, wavebin_up, order, m_h, logg, teff, def read_ld_coefs(filename, wavebin_low, wavebin_up, ld_model='quadratic'): """Unpack limb darkening coefficients and interpolate to the wavelength - grid of data being fit. + grid of data being fit. File must be comma separated with three columns: + wavelength c1 and c2. Parameters ---------- @@ -446,6 +444,17 @@ def read_ld_coefs(filename, wavebin_low, wavebin_up, ld_model='quadratic'): waves = waves[ii] q1s, q2s = q1s[ii], q2s[ii] + # Check that coeffs in file span the wavelength range of the observations + # with at least the same resolution. + if np.min(waves) < np.min(wavebin_low) or np.max(waves) > np.max(wavebin_up): + msg = 'LD coefficient file does not span the full wavelength range ' \ + 'of the observations.' + raise ValueError(msg) + if len(waves) < len(wavebin_low): + msg = 'LD coefficient file has a coarser wavelength grid than ' \ + 'the observations.' + raise ValueError(msg) + prior_q1, prior_q2 = [], [] # Loop over all fitting bins. Calculate mean of model LD coefs within that # range. diff --git a/supreme_spoon/utils.py b/supreme_spoon/utils.py index 3a5297b..44aa1e3 100644 --- a/supreme_spoon/utils.py +++ b/supreme_spoon/utils.py @@ -644,59 +644,6 @@ def outlier_resistant_variance(data): return var -def pack_ld_priors(wave, c1, c2, order, target, m_h, teff, logg, outdir): - """Write model limb darkening parameters to a file to be used as priors - for light curve fitting. - - Parameters - ---------- - wave : array-like[float] - Wavelength axis. - c1 : array-like[float] - c1 parameter for two-parameter limb darkening law. - c2 : array-like[float] - c2 parameter for two-parameter limb darkening law. - order : int - SOSS order. - target : str - Name of the target. - m_h : float - Host star metallicity. - teff : float - Host star effective temperature. - logg : float - Host star gravity. - outdir : str - Directory to which to save file. - """ - - # Create dictionary with model LD info. - dd = {'wave': wave, 'c1': c1, 'c2': c2} - df = pd.DataFrame(data=dd) - # Remove old LD file if one exists. - filename = target+'_order' + str(order) + '_exotic-ld_quadratic.csv' - if os.path.exists(outdir + filename): - os.remove(outdir + filename) - # Add header info. - f = open(outdir + filename, 'a') - f.write('# Target: {}\n'.format(target)) - f.write('# Instrument: NIRISS/SOSS\n') - f.write('# Order: {}\n'.format(order)) - f.write('# Author: {}\n'.format(os.environ.get('USER'))) - f.write('# Date: {}\n'.format(datetime.utcnow().replace(microsecond=0).isoformat())) - f.write('# Stellar M/H: {}\n'.format(m_h)) - f.write('# Stellar log g: {}\n'.format(logg)) - f.write('# Stellar Teff: {}\n'.format(teff)) - f.write('# Algorithm: ExoTiC-LD\n') - f.write('# Limb Darkening Model: quadratic\n') - f.write('# Column wave: Central wavelength of bin (micron)\n') - f.write('# Column c1: Quadratic Coefficient 1\n') - f.write('# Column c2: Quadratic Coefficient 2\n') - f.write('#\n') - df.to_csv(f, index=False) - f.close() - - def parse_config(config_file): """Parse a yaml config file. @@ -825,6 +772,59 @@ def save_extracted_spectra(filename, wl1, wu1, f1, e1, wl2, wu2, f2, e2, t, return param_dict +def save_ld_priors(wave, c1, c2, order, target, m_h, teff, logg, outdir): + """Write model limb darkening parameters to a file to be used as priors + for light curve fitting. + + Parameters + ---------- + wave : array-like[float] + Wavelength axis. + c1 : array-like[float] + c1 parameter for two-parameter limb darkening law. + c2 : array-like[float] + c2 parameter for two-parameter limb darkening law. + order : int + SOSS order. + target : str + Name of the target. + m_h : float + Host star metallicity. + teff : float + Host star effective temperature. + logg : float + Host star gravity. + outdir : str + Directory to which to save file. + """ + + # Create dictionary with model LD info. + dd = {'wave': wave, 'c1': c1, 'c2': c2} + df = pd.DataFrame(data=dd) + # Remove old LD file if one exists. + filename = target+'_order' + str(order) + '_exotic-ld_quadratic.csv' + if os.path.exists(outdir + filename): + os.remove(outdir + filename) + # Add header info. + f = open(outdir + filename, 'a') + f.write('# Target: {}\n'.format(target)) + f.write('# Instrument: NIRISS/SOSS\n') + f.write('# Order: {}\n'.format(order)) + f.write('# Author: {}\n'.format(os.environ.get('USER'))) + f.write('# Date: {}\n'.format(datetime.utcnow().replace(microsecond=0).isoformat())) + f.write('# Stellar M/H: {}\n'.format(m_h)) + f.write('# Stellar log g: {}\n'.format(logg)) + f.write('# Stellar Teff: {}\n'.format(teff)) + f.write('# Algorithm: ExoTiC-LD\n') + f.write('# Limb Darkening Model: quadratic\n') + f.write('# Column wave: Central wavelength of bin (micron)\n') + f.write('# Column c1: Quadratic Coefficient 1\n') + f.write('# Column c2: Quadratic Coefficient 2\n') + f.write('#\n') + df.to_csv(f, index=False) + f.close() + + def sigma_clip_lightcurves(flux, ferr, thresh=3, window=10): """Sigma clip outliers in time from final lightcurves.