diff --git a/exoctk/contam_visibility/field_simulator.py b/exoctk/contam_visibility/field_simulator.py index 472e461d1..8dd35359c 100755 --- a/exoctk/contam_visibility/field_simulator.py +++ b/exoctk/contam_visibility/field_simulator.py @@ -700,9 +700,6 @@ def calc_v3pa(V3PA, stars, aperture, data=None, tilt=0, plot=False, verbose=Fals targframes = [np.zeros((subY, subX))] * n_traces starframe = np.zeros((subY, subX)) - # Get order 0 - order0 = get_order0(aperture.AperName) * 1.5e8 # Scaling factor based on observations - # Get trace masks trace_masks = NIRCam_DHS_trace_mask(aperture.AperName) if 'NRCA5' in aperture.AperName else NIRISS_SOSS_trace_mask(aperture.AperName) @@ -726,6 +723,9 @@ def calc_v3pa(V3PA, stars, aperture, data=None, tilt=0, plot=False, verbose=Fals # Add all orders to the same frame (if it is a STAR) else: + # Get correct order 0 + order0 = get_order0(aperture.AperName, star['Teff'], stype=star['type'], verbose=verbose) * 1.5e3 # Scaling factor based on observations + # Scale the order 0 image and add it to the starframe scale0 = copy(order0) * star['fluxscale'] * aper['empirical_scale'][0] starframe = add_array_at_position(starframe, scale0, int(star['xord0'] - aper['subarr_x'][0]), int(star['yord0'] - aper['subarr_y'][1]), centered=True) @@ -1197,7 +1197,7 @@ def contam_slider_plot(contam_results, threshold=0.05, plot=False): return layout -def get_order0(aperture): +def get_order0(aperture, teff, stype='STAR', verbose=False): """Get the order 0 image for the given aperture Parameters @@ -1210,15 +1210,34 @@ def get_order0(aperture): np.ndarray The 2D order 0 image """ - # Get file - # TODO: Add order 0 files for other modes - filename = 'NIS_order0.npy' if 'NIS' in aperture else 'NIS_order0.npy' + if 'DHS' in aperture: + trace = np.zeros((50,50)) + else: - # Get the path to the trace files - trace_path = os.path.join(os.environ['EXOCTK_DATA'], f'exoctk_contam/order0/{filename}') + if stype == 'STAR': + + # Get the path to the trace files + traces_path = os.path.join(os.environ['EXOCTK_DATA'], f'exoctk_contam/order0/NIS_order0_*.npy') + + # Glob the file names + trace_files = glob.glob(traces_path) + + # Get closest Teff + teffs = np.array([int(os.path.basename(file).split('_')[-1][:-4]) for file in trace_files]) + trace_file = trace_files[np.argmin((teffs - teff)**2)] + if verbose: + print(f'Fetching {aperture} {teffs[np.argmin((teffs - teff)**2)]}K trace from {trace_file}') + + # Make frame + trace = np.load(trace_file) + + else: - # Make frame - trace = np.load(trace_path) + # Get stand-in for galaxy order 0 + gal_path = os.path.join(os.environ['EXOCTK_DATA'], f'exoctk_contam/order0/NIS_gal_order0.npy') + if verbose: + print('Fetching {} galaxy trace from {}'.format(aperture, gal_path)) + trace = np.load(gal_path) return trace