add missed test case

HERA-Team · Jun 25, 2024 · de5c073 · de5c073
1 parent 3e9d2fe
commit de5c073
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/hera_cal/lst_stack/calibration.py b/hera_cal/lst_stack/calibration.py
@@ -341,7 +341,6 @@ def lstbin_absolute_calibration(
     smoothing_scale: float = 10e6,
     eigenval_cutoff: float = 1e-12,
     calibrate_inplace: bool = True,
-    return_gains: bool = True,
     smooth_gains: bool = True,
     use_autos_for_abscal: bool = True,
 ):
@@ -391,13 +390,12 @@ def lstbin_absolute_calibration(
         smoothing_scale : float, default=10e6
             The scale of the smoothing function used to smooth the gains. This is the width of the
             smoothing function in Hz.
+        smooth_gains : bool, default=True
+            Boolean flag to smooth the gains.
         eigenval_cutoff : float, default=1e-12
             The cutoff for the eigenvalues of the DPSS eigenvectors.
         calibrate_inplace : bool, default=True
             Boolean flag to calibrate the data in place.
-        return_gains : bool, default=True
-            Boolean flag to return the gains. If set to False, the function will return an empty
-            dictionary for the gains.
         use_autos_for_abscal : bool, default=True
             Boolean flag to use the auto-correlations for absolute calibration. If set to True,
             the auto-correlations will be used to calculate the gain amplitude if they are provided.

diff --git a/hera_cal/lst_stack/tests/test_calibration.py b/hera_cal/lst_stack/tests/test_calibration.py
@@ -251,8 +251,10 @@ def test_cross_polarized_amp_cal(self):
         auto_stack.data[1:] = auto_stack.data[0]  # All nights exactly the same
 
         amp_gains = {
-            "Jee": np.random.normal(1, 0.1, size=(20, stack.data.shape[2])),
-            "Jnn": np.random.normal(1, 0.1, size=(20, stack.data.shape[2])),
+            "Jee": np.random.normal(1, 0.1, size=(20, 1))
+            * np.ones((1, stack.data.shape[2])),
+            "Jnn": np.random.normal(1, 0.1, size=(20, 1))
+            * np.ones((1, stack.data.shape[2])),
         }
 
         model = np.mean(stack.data, axis=0)
@@ -278,7 +280,7 @@ def test_cross_polarized_amp_cal(self):
             run_amplitude_cal=True,
             run_phase_cal=False,
             calibrate_inplace=True,
-            smooth_gains=False,
+            smooth_gains=True,
         )
 
         post_cal_std = np.nanstd(np.where(stack.flags, np.nan, stack.data), axis=0)