add CRIRES+ functionality

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"},

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)

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.
 

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": [

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,
     }

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
 

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(

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