From ae7535c5156f87c72f9da840a511e66ff64aac72 Mon Sep 17 00:00:00 2001
From: arjunsavel <asavel@gmail.com>
Date: Sun, 18 Aug 2024 12:29:08 -0400
Subject: [PATCH] test doppler shift

---
 src/scope/tests/test_utils.py | 68 +++++++++++++++++++++++++++++++++++
 src/scope/utils.py            |  1 +
 2 files changed, 69 insertions(+)

diff --git a/src/scope/tests/test_utils.py b/src/scope/tests/test_utils.py
index 96fef84..9054d4c 100644
--- a/src/scope/tests/test_utils.py
+++ b/src/scope/tests/test_utils.py
@@ -117,3 +117,71 @@ def test_perform_pca_variance_removed_compared_input_diffcompoents(self):
         assert np.std(scaled_cube10) < np.std(
             scaled_cube5
         )  # the variance should be contained in the main components
+
+
+class TestDopplerShift(unittest.TestCase):
+    def test_same_after_shift(self):
+        """
+        shouldn't globally change the properties. should really just shift things.
+        """
+        eval_wave = np.linspace(1.1, 1.8, 10000) * 1e-6  # meters.
+        template_wave = np.linspace(1, 2, 10000) * 1e-6
+
+        template_flux = np.random.random(10000)
+        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)
+        # test that shifted flux and interpred flux are very similar
+        assert (
+            np.testing.assert_allclose(shifted_flux, interped_flux, rtol=1e-2) == None
+        )
+
+    def test_diff_after_shift(self):
+        """
+        shouldn't globally change the properties. should really just shift things.
+        """
+        eval_wave = np.linspace(1.1, 1.8, 10000) * 1e-6  # meters.
+        template_wave = np.linspace(1, 2, 10000) * 1e-6
+
+        template_flux = np.random.random(10000)
+        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)
+        # test that shifted flux and interpred flux are very similar
+        np.testing.assert_raises(
+            AssertionError, np.testing.assert_array_equal, shifted_flux, interped_flux
+        )
+
+    def test_gaussian_shifted_red(self):
+        """
+        shouldn't globally change the properties. should really just shift things.
+        """
+        eval_wave = np.linspace(1.1, 1.8, 10000) * 1e-6
+        template_wave = np.linspace(1, 2, 10000) * 1e-6
+        # make a gaussian centered at 1.5
+        template_flux = np.exp(-0.5 * ((template_wave - 1.5e-6) / 0.01) ** 2)
+        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)
+        wav_max_shifted = eval_wave[np.argmax(shifted_flux)]
+        wav_max = eval_wave[np.argmax(interped_flux)]
+        assert wav_max_shifted > wav_max
+
+    def test_gaussian_shifted_blue(self):
+        """
+        shouldn't globally change the properties. should really just shift things.
+        """
+        eval_wave = np.linspace(1.1, 1.8, 10000) * 1e-6
+        template_wave = np.linspace(1, 2, 10000) * 1e-6
+        # make a gaussian centered at 1.5
+        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)
+        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)]
+        assert wav_max_shifted < wav_max
diff --git a/src/scope/utils.py b/src/scope/utils.py
index 7946b9f..b865411 100644
--- a/src/scope/utils.py
+++ b/src/scope/utils.py
@@ -186,6 +186,7 @@ def perform_pca(input_matrix, n_princ_comp, return_noplanet=False):
 def calc_doppler_shift(eval_wave, template_wave, template_flux, v):
     """
     Doppler shifts a spectrum. Evaluates the flux at a different grid.
+    convention: negative v is redshift.
 
     Inputs
     ------