From 61049f4a1b2337a30537033b311e9b95f72fa31e Mon Sep 17 00:00:00 2001 From: arjunsavel Date: Wed, 21 Aug 2024 17:58:31 -0400 Subject: [PATCH 1/2] mild refactor! --- src/scope/ccf.py | 4 +-- src/scope/run_simulation.py | 54 +++++++++---------------------------- src/scope/tests/conftest.py | 1 - src/scope/utils.py | 14 ++++------ 4 files changed, 19 insertions(+), 54 deletions(-) diff --git a/src/scope/ccf.py b/src/scope/ccf.py index 594120f..273f026 100644 --- a/src/scope/ccf.py +++ b/src/scope/ccf.py @@ -37,13 +37,13 @@ def calc_ccf(model_flux, data_arr_slice, n_pixel): cross_variance = (model_vector * data_arr_slice).sum(axis=1) # now need to sum - CCF = (cross_variance / jnp.sqrt(variance_data * variance_model)).sum() + ccf = (cross_variance / jnp.sqrt(variance_data * variance_model)).sum() logl = ( -0.5 * n_pixel * jnp.log(variance_data + variance_model - 2.0 * cross_variance) ).sum() - return logl, CCF + return logl, ccf calc_ccf_map = jax.vmap(calc_ccf, in_axes=(0, None, None), out_axes=0) diff --git a/src/scope/run_simulation.py b/src/scope/run_simulation.py index 9f00c18..19b5f85 100644 --- a/src/scope/run_simulation.py +++ b/src/scope/run_simulation.py @@ -35,7 +35,6 @@ def make_data( Rstar, do_pca=True, blaze=False, - do_airmass_detrending=False, tellurics=False, n_princ_comp=4, v_sys=0, @@ -65,7 +64,6 @@ def make_data( :wlgrid: (array) wavelength grid for the data. :do_pca: (bool) if True, do PCA on the data. :blaze: (bool) if True, include the blaze function in the data. - :do_airmass_detrending: (bool) if True, do airmass detrending on the data. :tellurics: (bool) if True, include tellurics in the data. :n_princ_comp: (int) number of principal components to use. :v_sys: (float) systemic velocity of the planet. @@ -80,7 +78,7 @@ def make_data( Outputs ------- - :A_noplanet: (array) the data cube with only the larger-scale trends. + :pca_noise_matrix: (array) the data cube with only the larger-scale trends. :flux_cube: (array) the data cube with the larger-scale trends removed. :fTemp_nopca: (array) the data cube with all components. :just_tellurics: (array) the telluric model that's multiplied to the dataset. @@ -98,7 +96,7 @@ def make_data( flux_cube = np.zeros( (n_order, n_exposure, n_pixel) ) # will store planet and star signal - A_noplanet = np.zeros((n_order, n_exposure, n_pixel)) + pca_noise_matrix = np.zeros((n_order, n_exposure, n_pixel)) for order in range(n_order): wlgrid_order = np.copy(wlgrid[order,]) # Cropped wavelengths flux_cube[order] = doppler_shift_planet_star( @@ -238,40 +236,24 @@ def make_data( for j in range(n_order): flux_cube[j] -= np.mean(flux_cube[j]) flux_cube[j] /= np.std(flux_cube[j]) - flux_cube[j], A_noplanet[j] = perform_pca( + flux_cube[j], pca_noise_matrix[j] = perform_pca( flux_cube[j], n_princ_comp, return_noplanet=True ) # todo: think about the svd # todo: redo all analysis centering on 0? - elif do_airmass_detrending: - zenith_angles = np.loadtxt("data/zenith_angles_w77ab.txt") - airm = 1 / np.cos(np.radians(zenith_angles)) # todo: check units - polyDeg = 2.0 # Setting the degree of the polynomial fit - xvec = airm - fAirDetr = np.zeros((n_order, n_exposure, n_pixel)) - # todo: think about looping - for io in range(n_order): - # Now looping over columns - for i in range(n_pixel): - yvec = flux_cube[io, :, i].copy() - fit = np.poly1d(np.polyfit(xvec, yvec, polyDeg))(xvec) - fAirDetr[io, :, i] = flux_cube[io, :, i] / fit - A_noplanet[io, :, i] = fit - flux_cube = fAirDetr - flux_cube[~np.isfinite(flux_cube)] = 0.0 else: for j in range(n_order): for i in range(n_exposure): flux_cube[j][i] -= np.mean(flux_cube[j][i]) # todo: check vars - if np.all(A_noplanet == 0): + if np.all(pca_noise_matrix == 0): print("was all zero") - A_noplanet = np.ones_like(A_noplanet) + pca_noise_matrix = np.ones_like(pca_noise_matrix) if tellurics: - return A_noplanet, flux_cube, flux_cube_nopca, just_tellurics + return pca_noise_matrix, flux_cube, flux_cube_nopca, just_tellurics return ( - A_noplanet, + pca_noise_matrix, flux_cube, flux_cube_nopca, np.ones_like(flux_cube), @@ -399,22 +381,11 @@ def calc_log_likelihood( # process the model same as the "data"! model_flux_cube *= A_noplanet[order] model_flux_cube, _ = perform_pca(model_flux_cube, n_princ_comp, False) - I = np.ones(n_pixel) - for i in range(n_exposure): - gVec = np.copy(model_flux_cube[i,]) - gVec -= (gVec.dot(I)) / float(n_pixel) # mean subtracting here... - sg2 = (gVec.dot(gVec)) / float(n_pixel) - - fVec = np.copy( - flux_cube[order][i,] - ) # already mean-subtracted! needed for the previous PCA! or...can I do the PCA without normalizing first? TODO - # fVec-=(fVec.dot(I))/float(Npix) - sf2 = (fVec.dot(fVec)) / float(n_pixel) - - R = (fVec.dot(gVec)) / float(n_pixel) # cross-covariance - CC = R / np.sqrt(sf2 * sg2) # cross-correlation - CCF += CC - logL += -0.5 * n_pixel * np.log(sf2 + sg2 - 2.0 * R) + # I = np.ones(n_pixel) + + logl, ccf = calc_ccf(model_flux_cube, flux_cube[order], n_pixel) + CCF += ccf + logL += logl # # todo: airmass detrending reprocessing @@ -538,7 +509,6 @@ def simulate_observation( Rp_solar, Rstar, do_pca=True, - do_airmass_detrending=True, blaze=blaze, n_princ_comp=n_princ_comp, tellurics=telluric, diff --git a/src/scope/tests/conftest.py b/src/scope/tests/conftest.py index c66d69b..2d5a501 100644 --- a/src/scope/tests/conftest.py +++ b/src/scope/tests/conftest.py @@ -68,7 +68,6 @@ def test_baseline_outouts(test_inputs): Rstar, do_pca=True, blaze=True, - do_airmass_detrending=False, tellurics=True, n_princ_comp=4, v_sys=0, diff --git a/src/scope/utils.py b/src/scope/utils.py index 541da8a..5398398 100644 --- a/src/scope/utils.py +++ b/src/scope/utils.py @@ -198,15 +198,11 @@ def calc_doppler_shift(eval_wave, template_wave, template_flux, v): ------- :flux_shifted: shifted flux grid """ - # beta = v / const_c - # delta_lam = eval_wave * beta - # shifted_wave = eval_wave + delta_lam - # shifted_flux = np.interp(shifted_wave, template_wave, template_flux) - # return shifted_flux - dl_l = v / const_c - wShift = eval_wave * (1.0 - dl_l) - flux_shifted = np.interp(wShift, template_wave, template_flux) - return flux_shifted + beta = v / const_c + delta_lam = eval_wave * beta + shifted_wave = eval_wave - delta_lam + shifted_flux = np.interp(shifted_wave, template_wave, template_flux) + return shifted_flux def calc_crossing_time( From e46b4729a4945583a207eec759f6b6f5939bfc7e Mon Sep 17 00:00:00 2001 From: arjunsavel Date: Wed, 21 Aug 2024 18:03:46 -0400 Subject: [PATCH 2/2] remove airmass detrending option --- src/scope/run_simulation.py | 1 - src/scope/tests/test_run_simulation.py | 4 ---- src/scope/utils.py | 12 ++++++------ 3 files changed, 6 insertions(+), 11 deletions(-) diff --git a/src/scope/run_simulation.py b/src/scope/run_simulation.py index 19b5f85..4859b88 100644 --- a/src/scope/run_simulation.py +++ b/src/scope/run_simulation.py @@ -10,7 +10,6 @@ from scope.broadening import * from scope.ccf import * -from scope.grid import * from scope.input_output import parse_input_file, write_input_file from scope.noise import * from scope.tellurics import * diff --git a/src/scope/tests/test_run_simulation.py b/src/scope/tests/test_run_simulation.py index 74b8886..12c4eef 100644 --- a/src/scope/tests/test_run_simulation.py +++ b/src/scope/tests/test_run_simulation.py @@ -46,7 +46,6 @@ def test_noblaze_outputs(test_inputs): Rstar, do_pca=True, blaze=False, - do_airmass_detrending=False, tellurics=True, n_princ_comp=4, v_sys=0, @@ -105,7 +104,6 @@ def test_notell_outputs(test_inputs): Rstar, do_pca=True, blaze=True, - do_airmass_detrending=False, tellurics=False, n_princ_comp=4, v_sys=0, @@ -164,7 +162,6 @@ def test_noisy_outputs(test_inputs): Rstar, do_pca=True, blaze=True, - do_airmass_detrending=False, tellurics=True, n_princ_comp=4, v_sys=0, @@ -223,7 +220,6 @@ def test_baseline_outputs_take2(test_inputs): Rstar, do_pca=True, blaze=True, - do_airmass_detrending=False, tellurics=True, n_princ_comp=4, v_sys=0, diff --git a/src/scope/utils.py b/src/scope/utils.py index 5398398..0fe04ee 100644 --- a/src/scope/utils.py +++ b/src/scope/utils.py @@ -12,6 +12,12 @@ from scope.constants import * +abs_path = os.path.dirname(__file__) + +np.random.seed(42) +start_clip = 200 +end_clip = 100 + @njit def doppler_shift_planet_star( @@ -141,12 +147,6 @@ def calc_limb_darkening(u1, u2, a, b, Rstar, ph, LD): return I -abs_path = os.path.dirname(__file__) - -np.random.seed(42) -start_clip = 200 -end_clip = 100 - # todo: download atran scripts # todo: fit wavelength solution stuff # todo: plot the maps