From b73e4a694a79e3d7cbed94d9f7bdbeec31caa34e Mon Sep 17 00:00:00 2001 From: Arjun Savel Date: Mon, 3 Mar 2025 17:18:58 -0500 Subject: [PATCH] add command line fun --- src/scope/input.txt | 22 ++++++++-------- src/scope/input_output.py | 52 ++++++++++++++++++++++++++++++++++++- src/scope/noise.py | 1 + src/scope/run_simulation.py | 40 +++++++++++++++++++++------- src/scope/utils.py | 4 +++ 5 files changed, 98 insertions(+), 21 deletions(-) diff --git a/src/scope/input.txt b/src/scope/input.txt index a758926..aee6550 100644 --- a/src/scope/input.txt +++ b/src/scope/input.txt @@ -10,23 +10,23 @@ #·······································: Created: 2024-08-15 Author: Arjun Savel! -Planet name: LTT-9779 b +Planet name: WASP-166 b # Simulation setup -modelname LTT-9779b_dayside # name of the model. used for saving the output files. -seed 42 # seed for the random number generator when generating noisy data. set to None for a random seed. +modelname wasp_166_150seed_ld # name of the model. used for saving the output files. +seed 150 # seed for the random number generator when generating noisy data. set to None for a random seed. # Filepaths -planet_spectrum_path /home/asavel/scratch/retrieval_stuff/fm_wasp_77ab_noisy/fm_wasp_77ab_noisy/spectrum_emission_LTT-9779b.pic +planet_spectrum_path /home/asavel/scratch/retrieval_stuff/fm_wasp_77ab_noisy/fm_wasp_77ab_noisy/spectrum_tran_WASP-166b_0.3_c_to_o.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 # Astrophysical Parameters. Can specify DATABASE to pull based on name for parameters with [DB] Rp DATABASE # planetary radius in Jupiter radii. [DB] Rstar DATABASE # stellar radius, in solar radii. [DB] -kp NULL # expected planetary orbital velocity assuming circular orbit, in km/s. input NULL if you'd like to calculate from orbital parameters. -v_sys DATABASE # systemic velocity, in km/s. [DB] +kp 128.1 # expected planetary orbital velocity assuming circular orbit, in km/s. input NULL if you'd like to calculate from orbital parameters. +v_sys 0.0 # systemic velocity, in km/s. [DB] v_rot NULL # equatorial rotational velocity. input NULL if you'd like to calculate from orbital parameters. P_rot DATABASE # orbital period of the planet, in days. [DB] a DATABASE # semi-major axis of the planet, in AU. [DB] @@ -47,13 +47,13 @@ vary_throughput True # whether to include throughput variations. # Observation Parameters -observation emission # type of observation to perform. supported observations are ``emission`` and ``transmission``. -phase_start 0.27 # phase of the beginning of the observations. 0 is center of transit, 0.5 is secondary eclipse. If DATABASE, just the transit duration. [DB] -phase_end 0.45 # phase of the end of the observations. 0 is center of transit, 0.5 is secondary eclipse. If DATABASE, just the transit duration. [DB] -n_exposures 10 # number of exposures to simulate. sets the phases of the exposures. +observation transmission # type of observation to perform. supported observations are ``emission`` and ``transmission``. +phase_start 0.986222200994206 # phase of the beginning of the observations. 0 is center of transit, 0.5 is secondary eclipse. If DATABASE, just the transit duration. [DB] +phase_end 1.0137777990057941 # phase of the end of the observations. 0 is center of transit, 0.5 is secondary eclipse. If DATABASE, just the transit duration. [DB] +n_exposures 17 # 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 250 # signal-to-noise ratio of the observations, per pixel. I.e., what sets the photon counts at the detector. 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..c72ae6b 100644 --- a/src/scope/input_output.py +++ b/src/scope/input_output.py @@ -4,6 +4,7 @@ import io import os +import argparse from datetime import datetime import pandas as pd @@ -88,11 +89,12 @@ def coerce_nulls(data, key, value): if value == "NULL": data[key] = np.nan + return data def coerce_integers(data, key, value): - integer_fields = ["n_exposures", "n_princ_comp"] + integer_fields = ["n_exposures", "n_princ_comp", "seed"] if key in integer_fields: data[key] = int(value) else: @@ -288,3 +290,51 @@ def write_input_file(data, output_file_path="input.txt"): f.write("\n") print(f"Input file written to {output_file_path}") + + + +def parse_arguments(): + parser = argparse.ArgumentParser(description='Simulate observation') + + # Required parameters + parser.add_argument('--planet_spectrum_path', type=str, default=".", help='Path to planet spectrum') + parser.add_argument('--star_spectrum_path', type=str, default=".", help='Path to star spectrum') + parser.add_argument('--data_cube_path', type=str, default=".", help='Path to data cube') + + # Optional parameters with their defaults matching your function + parser.add_argument('--phase_start', type=float, default=0, help='Start phase of the simulated observations') + parser.add_argument('--phase_end', type=float, default=1, help='End phase of the simulated observations') + parser.add_argument('--n_exposures', type=int, default=10, help='Number of exposures') + parser.add_argument('--observation', type=str, default="emission", help='Observation type') + parser.add_argument('--blaze', type=bool, default=True, help='Blaze flag') + parser.add_argument('--n_princ_comp', type=int, default=4, help='Number of principal components') + parser.add_argument('--star', type=bool, default=True, help='Star flag') + parser.add_argument('--SNR', type=float, default=250, help='Signal to noise ratio') + parser.add_argument('--telluric', type=bool, default=True, help='Telluric flag') + parser.add_argument('--tell_type', type=str, default="data-driven", help='Telluric type') + parser.add_argument('--time_dep_tell', type=bool, default=False, help='Time dependent telluric') + parser.add_argument('--wav_error', type=bool, default=False, help='Wavelength error flag') + parser.add_argument('--rv_semiamp_orbit', type=float, default=0.3229, help='RV semi-amplitude orbit') + parser.add_argument('--order_dep_throughput', type=bool, default=True, help='Order dependent throughput') + parser.add_argument('--Rp', type=float, default=1.21, help='Planet radius (Jupiter radii)') + parser.add_argument('--Rstar', type=float, default=0.955, help='Star radius (solar radii)') + parser.add_argument('--kp', type=float, default=192.02, help='Planetary orbital velocity (km/s)') + parser.add_argument('--v_rot', type=float, default=4.5, help='Rotation velocity') + parser.add_argument('--scale', type=float, default=1.0, help='Scale factor') + parser.add_argument('--v_sys', type=float, default=0.0, help='Systemic velocity') + parser.add_argument('--modelname', type=str, default="yourfirstsimulation", help='Model name') + parser.add_argument('--divide_out_of_transit', type=bool, default=False, help='Divide out of transit') + parser.add_argument('--out_of_transit_dur', type=float, default=0.1, help='Out of transit duration') + parser.add_argument('--include_rm', type=bool, default=False, help='Include RM effect') + parser.add_argument('--v_rot_star', type=float, default=3.0, help='Star rotation velocity') + parser.add_argument('--a', type=float, default=0.033, help='Semi-major axis') + parser.add_argument('--lambda_misalign', type=float, default=0.0, help='Misalignment angle') + parser.add_argument('--inc', type=float, default=90.0, help='Inclination') + parser.add_argument('--seed', type=int, default=42, help='Random seed') + parser.add_argument('--LD', type=bool, default=True, help='Limb darkening') + parser.add_argument('--vary_throughput', type=bool, default=True, help='Vary throughput') + + # For file input option + parser.add_argument('--input_file', type=str, default="input.txt", help='Input file with parameters') + + return parser.parse_args() diff --git a/src/scope/noise.py b/src/scope/noise.py index 1bb0101..b9e11d9 100644 --- a/src/scope/noise.py +++ b/src/scope/noise.py @@ -81,6 +81,7 @@ def add_quadratic_noise(flux_cube_model, wl_grid, SNR, IGRINS=False, **kwargs): * 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 diff --git a/src/scope/run_simulation.py b/src/scope/run_simulation.py index ffe2943..6729653 100644 --- a/src/scope/run_simulation.py +++ b/src/scope/run_simulation.py @@ -11,7 +11,7 @@ from scope.broadening import * from scope.ccf import * -from scope.input_output import parse_input_file, write_input_file +from scope.input_output import parse_input_file, write_input_file, parse_arguments from scope.noise import * from scope.tellurics import * from scope.utils import * @@ -85,12 +85,13 @@ def make_data( :just_tellurics: (array) the telluric model that's multiplied to the dataset. """ + print(seed) np.random.seed(seed) rv_planet, rv_star = calc_rvs( v_sys, v_sys_measured, Kp, rv_semiamp_orbit, phases ) # measured in m/s now - + flux_cube = np.zeros( (n_order, n_exposure, n_pixel) ) # will store planet and star signal @@ -118,7 +119,7 @@ def make_data( star, observation, ) - + throughput_baselines = np.loadtxt(abs_path + "/data/throughputs.txt") flux_cube = detrend_cube(flux_cube, n_order, n_exposure) @@ -171,6 +172,7 @@ def make_data( noise_model = "IGRINS" else: noise_model = "constant" + flux_cube = add_noise_cube(flux_cube, wlgrid, SNR, noise_model=noise_model) flux_cube = detrend_cube(flux_cube, n_order, n_exposure) @@ -262,7 +264,7 @@ def make_data( ) # returning CCF and logL values -# @njit +#@njit def calc_log_likelihood( v_sys, Kp, @@ -419,6 +421,7 @@ def simulate_observation( lambda_misalign=0.0, inc=90.0, seed=42, + LD=True, vary_throughput=True, **kwargs, ): @@ -458,6 +461,8 @@ 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) + print(v_sys) + print('that was vsys') 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") @@ -475,7 +480,8 @@ def simulate_observation( # instrument profile convolution instrument_kernel = get_instrument_kernel() Fp_conv = np.convolve(Fp_conv_rot, instrument_kernel, mode="same") - + print('here spectrum') + print(Fp_conv) star_wave, star_flux = np.loadtxt( star_spectrum_path ).T # Phoenix stellar model packing @@ -520,9 +526,10 @@ def simulate_observation( wav_error=wav_error, order_dep_throughput=order_dep_throughput, observation=observation, - divide_out_of_transit=False, + divide_out_of_transit=divide_out_of_transit, out_of_transit_dur=0.1, v_sys_measured=v_sys, + LD=LD, ) 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 +562,8 @@ def simulate_observation( star=star, observation=observation, v_sys_measured=v_sys, - vary_throughput=vary_throughput, + LD=LD + ) lls[l, k], ccfs[l, k] = res @@ -563,6 +571,20 @@ def simulate_observation( if __name__ == "__main__": - file = "input.txt" - inputs = parse_input_file(file) + args = parse_arguments() + # First, parse the input file to get base parameters + inputs = parse_input_file(args.input_file) + + # Then override with any command line arguments that were explicitly provided + args_dict = vars(args) + for key, value in args_dict.items(): + # Skip the input_file parameter + if key == 'input_file': + continue + + # Only override if the argument was explicitly provided (not None) + if value is not None: + inputs[key] = value + + # Call the simulation function with the merged parameters simulate_observation(**inputs) diff --git a/src/scope/utils.py b/src/scope/utils.py index 0ae2b30..8c94d4a 100644 --- a/src/scope/utils.py +++ b/src/scope/utils.py @@ -52,6 +52,7 @@ def doppler_shift_planet_star( wlgrid_order, wl_model, Fp_conv, rv_planet[exposure] ) flux_planet *= scale # apply scale factor + _, flux_star = calc_doppler_shift( wlgrid_order, wl_model, Fstar_conv, rv_star[exposure] ) @@ -69,7 +70,9 @@ def doppler_shift_planet_star( observation == "transmission" ): # in transmission, after we "divide out" (with PCA) the star and tellurics, we're left with Fp. I = calc_limb_darkening(u1, u2, a, b, Rstar, phases[exposure], LD) + model_flux_cube[exposure,] = 1.0 - flux_planet * I + if not reprocessing: model_flux_cube[exposure,] *= flux_star return model_flux_cube @@ -418,6 +421,7 @@ def calc_rvs(v_sys, v_sys_measured, Kp, Kstar, phases): 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