diff --git a/pyproject.toml b/pyproject.toml index c4ca0ad..5dc41a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,7 @@ build-backend = "setuptools.build_meta" name = "scope-astr" authors = [ {name = "Arjun Savel", email = "asavel@umd.edu"}, + {name = "Arjun Savel"}, {name= "Megan Bedell"}, {name= "Eliza M.-R. Kempton"}, {name= "Peter Smith"}, diff --git a/src/scope/_version.py b/src/scope/_version.py new file mode 100644 index 0000000..4556cfd --- /dev/null +++ b/src/scope/_version.py @@ -0,0 +1,17 @@ +# file generated by setuptools_scm +# don't change, don't track in version control +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import Tuple, Union + + VERSION_TUPLE = Tuple[Union[int, str], ...] +else: + VERSION_TUPLE = object + +version: str +__version__: str +__version_tuple__: VERSION_TUPLE +version_tuple: VERSION_TUPLE + +__version__ = version = "0.4.2b0" +__version_tuple__ = version_tuple = (0, 4, 2) diff --git a/src/scope/input.txt b/src/scope/input.txt index a758926..3d462d6 100644 --- a/src/scope/input.txt +++ b/src/scope/input.txt @@ -21,6 +21,7 @@ seed 42 # seed for the random number generator w planet_spectrum_path /home/asavel/scratch/retrieval_stuff/fm_wasp_77ab_noisy/fm_wasp_77ab_noisy/spectrum_emission_LTT-9779b.pic star_spectrum_path /home/asavel/scratch/scope/src/scope/data/PHOENIX_5605_4.33.txt data_cube_path /home/asavel/scratch/scope/src/scope/data/data_RAW_20201214_wl_algn_03.pic +snr_path /home/asavel/scratch/scope/src/scope/data/snr_IGRINS_LTT-9779b.json # only required for CRIRES+ simulations # Astrophysical Parameters. Can specify DATABASE to pull based on name for parameters with [DB] Rp DATABASE # planetary radius in Jupiter radii. [DB] @@ -40,6 +41,7 @@ lambda_misalign 0 # misalignment angle of the planet's or inc 90.0 # inclination of the planet's orbit with respect to the line of sight, in degrees. only matters if include_rm is set to True. and observation is set to transmission. # Instrument Parameters +spectrograph IGRINS # name of the spectrograph to simulate. currently spectrographs are ``IGRINS`` and ``CRIRES+``. blaze True # whether to include a blaze function or not. wav_error False # whether to include wavelength solution errors or not. order_dep_throughput True # whether to include order-dependent throughput variations. @@ -53,7 +55,7 @@ phase_end 0.45 # phase of the end of the observations. 0 n_exposures 10 # number of exposures to simulate. sets the phases of the exposures. star True # whether to include the star in the simulation. In general, you'd like to! telluric True # whether to include tellurics in the simulation. In general, you'd like to! -SNR 300 # signal-to-noise ratio of the observations, per pixel. I.e., what sets the photon counts at the detector. +SNR 300 # signal-to-noise ratio of the observations, per pixel. I.e., what sets the photon counts at the detector. Ignored for CRIRES+ simulations. tell_type data-driven # type of telluric simulation. supported modes are ``ATRAN`` and ``data-driven``. time_dep_tell False # whether the tellurics are time-dependent or not. diff --git a/src/scope/input_output.py b/src/scope/input_output.py index ea527fc..1a3da61 100644 --- a/src/scope/input_output.py +++ b/src/scope/input_output.py @@ -228,7 +228,12 @@ def write_input_file(data, output_file_path="input.txt"): # Define the order and categories of parameters categories = { - "Filepaths": ["planet_spectrum_path", "star_spectrum_path", "data_cube_path"], + "Filepaths": [ + "planet_spectrum_path", + "star_spectrum_path", + "data_cube_path", + "snr_path", + ], "Astrophysical Parameters": ["Rp", "Rp_solar", "Rstar", "kp", "v_rot", "v_sys"], "Instrument Parameters": ["SNR"], "Observation Parameters": [ diff --git a/src/scope/noise.py b/src/scope/noise.py index 1bb0101..760b62e 100644 --- a/src/scope/noise.py +++ b/src/scope/noise.py @@ -112,6 +112,26 @@ def add_igrins_noise(flux_cube_model, wl_grid, SNR): return add_quadratic_noise(flux_cube_model, wl_grid, SNR, IGRINS=True) +def add_crires_plus_noise(flux_cube_model, wl_grid, SNR): + """ + Adds CRIRES+ noise to the flux cube model. + """ + noisy_flux = np.zeros_like(flux_cube_model) + for exposure in range(flux_cube_model.shape[1]): + for order in range(flux_cube_model.shape[0]): + flux_level = ( + SNR**2 + * flux_cube_model[order][exposure] + / np.nanmax(flux_cube_model[order][exposure]) + ) + noisy_flux[order][exposure] = np.random.poisson(flux_level) + + # a few quick checks to make sure that nothing has gone wrong with adding noise + noisy_flux[order][exposure][noisy_flux[order][exposure] < 0.0] = 0.0 + noisy_flux[order][exposure][~np.isfinite(noisy_flux[order][exposure])] = 0.0 + return noisy_flux + + def add_custom_noise(SNR): """ Adds custom noise to the flux cube model. @@ -142,6 +162,7 @@ def add_noise_cube(flux_cube_model, wl_grid, SNR, noise_model="constant", **kwar noise_models = { "constant": add_constant_noise, "IGRINS": add_igrins_noise, + "CRIRES+": add_crires_plus_noise, "custom_quadratic": add_quadratic_noise, "custom": add_custom_noise, } diff --git a/src/scope/run_simulation.py b/src/scope/run_simulation.py index ffe2943..3c6cc65 100644 --- a/src/scope/run_simulation.py +++ b/src/scope/run_simulation.py @@ -55,6 +55,7 @@ def make_data( out_of_transit_dur=0.1, v_sys_measured=0.0, vary_throughput=True, + instrument="IGRINS", ): """ Creates a simulated HRCCS dataset. Main function. @@ -136,6 +137,11 @@ def make_data( flux_cube[:, i, :] *= throughput_factor if tellurics: + if instrument != "IGRINS": + raise NotImplementedError( + "Only IGRINS data is currently supported for telluric simulations." + "Please set tellurics=False, and at minimum ensure that your SNR input includes telluric losses." + ) flux_cube = add_tellurics( wlgrid, flux_cube, @@ -166,13 +172,13 @@ def make_data( flux_cube = detrend_cube(flux_cube, n_order, n_exposure) flux_cube[np.isnan(flux_cube)] = 0.0 - if SNR > 0: # 0 means don't add noise! + if np.any(SNR > 0): # 0 means don't add noise! if order_dep_throughput: - noise_model = "IGRINS" + noise_model = instrument else: noise_model = "constant" + print(f"Adding noise with model {noise_model}") flux_cube = add_noise_cube(flux_cube, wlgrid, SNR, noise_model=noise_model) - flux_cube = detrend_cube(flux_cube, n_order, n_exposure) if wav_error: @@ -184,7 +190,6 @@ def make_data( flux_cube = detrend_cube(flux_cube, n_order, n_exposure) flux_cube[np.isnan(flux_cube)] = 0.0 flux_cube_nopca = flux_cube.copy() - if observation == "transmission" and divide_out_of_transit: # generate the out of transit baseline n_exposures_baseline = ( @@ -218,7 +223,11 @@ def make_data( # add blaze out_of_transit_flux = add_blaze_function( - wlgrid, out_of_transit_flux, n_order, n_exposures_baseline + wlgrid, + out_of_transit_flux, + n_order, + n_exposures_baseline, + instrument=instrument, ) out_of_transit_flux[flux_cube < 0.0] = 0.0 out_of_transit_flux[np.isnan(flux_cube)] = 0.0 @@ -250,10 +259,11 @@ def make_data( flux_cube[j][i] -= np.mean(flux_cube[j][i]) if np.all(pca_noise_matrix == 0): - print("was all zero") pca_noise_matrix = np.ones_like(pca_noise_matrix) + if tellurics: return pca_noise_matrix, flux_cube, flux_cube_nopca, just_tellurics + return ( pca_noise_matrix, flux_cube, @@ -420,6 +430,11 @@ def simulate_observation( inc=90.0, seed=42, vary_throughput=True, + instrument="IGRINS", + snr_path=None, + planet_name="yourfirstplanet", + n_kp=200, + n_vsys=200, **kwargs, ): """ @@ -457,12 +472,19 @@ def simulate_observation( phases = np.linspace(phase_start, phase_end, n_exposures) Rp_solar = Rp * rjup_rsun # convert from jupiter radii to solar radii - Kp_array = np.linspace(kp - 100, kp + 100, 200) - v_sys_array = np.arange(v_sys - 100, v_sys + 100) - n_order, n_pixel = (44, 1848) # todo: generalize. - mike_wave, mike_cube = pickle.load(open(data_cube_path, "rb"), encoding="latin1") + Kp_array = np.linspace(kp - 100, kp + 100, n_kp) + v_sys_array = np.linspace(v_sys - 100, v_sys + 100, n_vsys) + + if instrument == "IGRINS": + n_order, n_pixel = (44, 1848) + mike_wave, mike_cube = pickle.load( + open(data_cube_path, "rb"), encoding="latin1" + ) - wl_cube_model = mike_wave.copy().astype(np.float64) + wl_cube_model = mike_wave.copy().astype(np.float64) + elif instrument == "CRIRES+": + # need to read in the filepath + n_order, n_pixel, wl_cube_model, SNR = read_crires_data(snr_path) wl_model, Fp, Fstar = np.load(planet_spectrum_path, allow_pickle=True) @@ -473,7 +495,7 @@ def simulate_observation( Fp_conv_rot = np.convolve(Fp, rot_ker, mode="same") # instrument profile convolution - instrument_kernel = get_instrument_kernel() + instrument_kernel = get_instrument_kernel(instrument) Fp_conv = np.convolve(Fp_conv_rot, instrument_kernel, mode="same") star_wave, star_flux = np.loadtxt( @@ -490,7 +512,7 @@ def simulate_observation( star_wave, star_flux, wl_model, instrument_kernel, smooth=False ) - lls, ccfs = np.zeros((200, 200)), np.zeros((200, 200)) + lls, ccfs = np.zeros((n_kp, n_vsys)), np.zeros((n_kp, n_vsys)) # redoing the grid. how close does PCA get to a tellurics-free signal detection? A_noplanet, flux_cube, flux_cube_nopca, just_tellurics = make_data( @@ -523,6 +545,7 @@ def simulate_observation( divide_out_of_transit=False, out_of_transit_dur=0.1, v_sys_measured=v_sys, + instrument=instrument, ) run_name = f"{n_princ_comp}_NPC_{blaze}_blaze_{star}_star_{telluric}_telluric_{SNR}_SNR_{tell_type}_{time_dep_tell}_{wav_error}_{order_dep_throughput}" @@ -555,7 +578,6 @@ def simulate_observation( star=star, observation=observation, v_sys_measured=v_sys, - vary_throughput=vary_throughput, ) lls[l, k], ccfs[l, k] = res diff --git a/src/scope/tests/conftest.py b/src/scope/tests/conftest.py index 2d5a501..196b7cf 100644 --- a/src/scope/tests/conftest.py +++ b/src/scope/tests/conftest.py @@ -28,7 +28,7 @@ def test_inputs(): rot_ker = get_rot_ker(v_rot, wl_model) Fp_conv_rot = np.convolve(Fp, rot_ker, mode="same") # instrument profile convolution - instrument_kernel = get_instrument_kernel() + instrument_kernel = get_instrument_kernel("IGRINS") Fp_conv = np.convolve(Fp_conv_rot, instrument_kernel, mode="same") star_spectrum_path = os.path.join(test_data_path, "PHOENIX_5605_4.33.txt") star_wave, star_flux = np.loadtxt( diff --git a/src/scope/tests/test_run_simulation.py b/src/scope/tests/test_run_simulation.py index 12c4eef..24926a2 100644 --- a/src/scope/tests/test_run_simulation.py +++ b/src/scope/tests/test_run_simulation.py @@ -516,3 +516,111 @@ def test_nopca_fluxlevels_deterministic( # the flux levels must be higher when there's no tellurics. np.testing.assert_array_equal(flux_cube_nopca, flux_cube_nopca_b) + + +def test_igrins_emission(test_inputs): + data_cube_path = os.path.join(test_data_path, "data_RAW_20201214_wl_algn_03.pic") + planet_spectrum_path = os.path.join(test_data_path, "best_fit_spectrum.npy") + snr_path = os.path.join(test_data_path, "output1.json") + star_spectrum_path = os.path.join(test_data_path, "PHOENIX_5605_4.33.txt") + simulate_observation( + # problem: this is all 1s! + planet_spectrum_path=planet_spectrum_path, + star_spectrum_path=star_spectrum_path, + data_cube_path=data_cube_path, + phase_start=0.4, + phase_end=0.46, + n_exposures=5, + observation="emission", + blaze=True, + n_princ_comp=4, + star=True, + SNR=250, + telluric=True, + tell_type="data-driven", + time_dep_tell=False, + wav_error=False, + rv_semiamp_orbit=0.3229, + order_dep_throughput=True, + Rp=1.21, # Jupiter radii, + Rstar=0.955, # solar radii + kp=192.02, # planetary orbital velocity, km/s + v_rot=4.5, + scale=1.0, + v_sys=0.0, + modelname="yourfirstsimulation", + planet_name="GJ1214b", + divide_out_of_transit=False, + out_of_transit_dur=0.1, + include_rm=False, + instrument="IGRINS", + v_rot_star=3.0, + a=0.033, # + lambda_misalign=0.0, + inc=90.0, + n_kp=20, + n_vsys=20, + seed=42, + vary_throughput=True, + snr_path=snr_path, + ) + + # check that there are nonzeros, basically that it runs + filetype = glob("src/scope/output/yourfirstsimulation/*ccf*") + ccfs = np.loadtxt(filetype[0]) + assert np.sum(ccfs) != 0 + + +def test_crires_emission(test_inputs): + data_cube_path = os.path.join(test_data_path, "data_RAW_20201214_wl_algn_03.pic") + planet_spectrum_path = os.path.join( + test_data_path, "best_fit_spectrum_tran_gj1214_steam_nir.pic" + ) + snr_path = os.path.join(test_data_path, "output1.json") + star_spectrum_path = os.path.join(test_data_path, "PHOENIX_5605_4.33.txt") + simulate_observation( + # problem: this is all 1s! + planet_spectrum_path=planet_spectrum_path, + star_spectrum_path=star_spectrum_path, + data_cube_path=data_cube_path, + phase_start=0.4, + phase_end=0.46, + n_exposures=5, + observation="transmission", + blaze=False, + n_princ_comp=4, + star=True, + SNR=250, + telluric=False, + tell_type="data-driven", + time_dep_tell=False, + wav_error=False, + rv_semiamp_orbit=0.3229, + order_dep_throughput=True, + Rp=1.21, # Jupiter radii, + Rstar=0.955, # solar radii + kp=192.02, # planetary orbital velocity, km/s + v_rot=4.5, + scale=1.0, + v_sys=0.0, + modelname="yoursecondsimulation", + planet_name="GJ1214b", + divide_out_of_transit=False, + out_of_transit_dur=0.1, + include_rm=False, + instrument="CRIRES+", + v_rot_star=3.0, + a=0.033, # + lambda_misalign=0.0, + inc=90.0, + n_kp=20, + n_vsys=20, + seed=42, + vary_throughput=True, + snr_path=snr_path, + ) + + # check that there are nonzeros, basically that it runs + filetype = glob("src/scope/output/yoursecondsimulation/*ccf*") + ccfs = np.loadtxt(filetype[0]) + assert np.sum(ccfs) != 0 diff --git a/src/scope/tests/test_utils.py b/src/scope/tests/test_utils.py index 5522db4..ac6d9a7 100644 --- a/src/scope/tests/test_utils.py +++ b/src/scope/tests/test_utils.py @@ -131,7 +131,9 @@ def test_same_after_shift(self): interped_flux = np.interp(eval_wave, template_wave, template_flux) v = 1e-6 # m/s. should not change much - shifted_flux = calc_doppler_shift(eval_wave, template_wave, template_flux, v) + shifted_wave, shifted_flux = calc_doppler_shift( + eval_wave, template_wave, template_flux, v + ) # test that shifted flux and interpred flux are very similar assert ( np.testing.assert_allclose(shifted_flux, interped_flux, rtol=1e-2) == None @@ -148,7 +150,9 @@ def test_diff_after_shift(self): interped_flux = np.interp(eval_wave, template_wave, template_flux) v = 2e3 # m/s. should not change much - shifted_flux = calc_doppler_shift(eval_wave, template_wave, template_flux, v) + shifted_wave, shifted_flux = calc_doppler_shift( + eval_wave, template_wave, template_flux, v + ) # test that shifted flux and interpred flux are very similar np.testing.assert_raises( AssertionError, np.testing.assert_array_equal, shifted_flux, interped_flux @@ -165,7 +169,9 @@ def test_gaussian_shifted_blue(self): interped_flux = np.interp(eval_wave, template_wave, template_flux) v = -5e4 # m/s. should not change much - shifted_flux = calc_doppler_shift(eval_wave, template_wave, template_flux, v) + shifted_wave, shifted_flux = calc_doppler_shift( + eval_wave, template_wave, template_flux, v + ) wav_max_shifted = eval_wave[np.argmax(shifted_flux)] wav_max = eval_wave[np.argmax(interped_flux)] assert wav_max_shifted < wav_max @@ -180,7 +186,9 @@ def test_gaussian_shifted_red(self): template_flux = np.exp(-0.5 * ((template_wave - 1.5e-6) / 0.01) ** 2) v = 5e4 # m/s. should not change much - shifted_flux = calc_doppler_shift(eval_wave, template_wave, template_flux, v) + shifted_wave, shifted_flux = calc_doppler_shift( + eval_wave, template_wave, template_flux, v + ) interped_flux = np.interp(eval_wave, template_wave, template_flux) wav_max_shifted = eval_wave[np.argmax(shifted_flux)] wav_max = eval_wave[np.argmax(interped_flux)] diff --git a/src/scope/utils.py b/src/scope/utils.py index 0ae2b30..2c814ac 100644 --- a/src/scope/utils.py +++ b/src/scope/utils.py @@ -2,6 +2,7 @@ Utility functions for simulating HRCCS data. """ +import json import os import pickle @@ -24,6 +25,44 @@ end_clip = 100 +def read_crires_data(data_path): + """ + Reads in CRIRES data. + + Inputs + ------ + :data_path: (str) path to the data + + Outputs + ------- + n_orders: (int) number of orders + n_pixel: (int) number of pixels + wl_cube_model: (array) wavelength cube model + snrs: (array) signal-to-noise ratios + """ + with open(data_path, "r") as file: + data = json.load(file) + + n_orders = 0 # an integer :) + for i in range(len(data["data"]["orders"])): + order_len = len(data["data"]["orders"][i]["detectors"]) + n_orders += order_len + + n_wavs = len(data["data"]["orders"][i]["detectors"][0]["wavelength"]) + + wl_grid = np.zeros((n_orders, n_wavs)) + snr_grid = np.zeros((n_orders, n_wavs)) + + for i in range(len(data["data"]["orders"])): + order_len = len(data["data"]["orders"][i]["detectors"]) + for j in range(order_len): + wl_grid[i * order_len + j] = data["data"]["orders"][i]["detectors"][j][ + "wavelength" + ] + + return n_orders, n_wavs, wl_grid * 1e6, snr_grid + + @njit def doppler_shift_planet_star( model_flux_cube, @@ -63,7 +102,6 @@ def doppler_shift_planet_star( if not reprocessing: model_flux_cube[exposure,] *= flux_star * Rstar**2 elif observation == "emission": - model_flux_cube[exposure,] = flux_planet * (Rp_solar * rsun) ** 2 elif ( observation == "transmission" @@ -88,12 +126,12 @@ def save_results(outdir, run_name, lls, ccfs): def make_outdir(outdir): try: - os.mkdir(outdir) + os.makedirs(outdir) except FileExistsError: print("Directory already exists. Continuing!") -def get_instrument_kernel(resolution_ratio=250000 / 45000, kernel_size=41): +def get_instrument_kernel(instrument, model_resolution=250000, kernel_size=41): """ Creates a Gaussian kernel for instrument response using an alternative implementation. @@ -109,6 +147,12 @@ def get_instrument_kernel(resolution_ratio=250000 / 45000, kernel_size=41): numpy.ndarray Normalized Gaussian kernel """ + instrument_resolution_dict = { + "IGRINS": 45000, + "CRIRES+": 145000, + } + instrument_resolution = instrument_resolution_dict[instrument] + resolution_ratio = instrument_resolution / model_resolution # Ensure kernel size is odd if kernel_size % 2 == 0: raise ValueError("Kernel size must be odd") @@ -410,14 +454,13 @@ def calc_rvs(v_sys, v_sys_measured, Kp, Kstar, phases): """ v_sys_tot = v_sys + v_sys_measured # total offset - rv_planet = ( - v_sys_tot + Kp * np.sin(2.0 * np.pi * phases) + rv_planet = v_sys_tot + Kp * np.sin( + 2.0 * np.pi * phases ) # input in km/s, convert to m/s - - rv_star = ( - v_sys_measured - Kstar * np.sin(2.0 * np.pi * phases) - ) # measured in m/s. note opposite sign! + rv_star = v_sys_measured - Kstar * np.sin( + 2.0 * np.pi * phases + ) # measured in m/s. note opposite sign! return rv_planet * 1e3, rv_star * 1e3 @@ -504,7 +547,9 @@ def change_wavelength_solution(wl_cube_model, flux_cube_model, doppler_shifts): return flux_cube_model -def add_blaze_function(wl_cube_model, flux_cube_model, n_order, n_exp): +def add_blaze_function( + wl_cube_model, flux_cube_model, n_order, n_exp, instrument="IGRINS" +): """ Adds the blaze function to the model. @@ -520,6 +565,10 @@ def add_blaze_function(wl_cube_model, flux_cube_model, n_order, n_exp): :flux_cube_model: (array) flux cube model with blaze function included. """ # read in...have to somehow match the telluric spectra + if instrument != "IGRINS": + raise NotImplementedError( + "Only the IGRINS blaze function is currently supported." + ) with open(abs_path + "/data/K_blaze_spectra.pic", "rb") as f: K_blaze_cube = pickle.load(f)