Merge pull request #958 from HERA-Team/fix-lstcal

fix: lstcal with multiple bands
HERA-Team · Jul 3, 2024 · 3ca93b4 · 3ca93b4
2 parents 1858df5 + fa2cf89
commit 3ca93b4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 15 deletions.
diff --git a/hera_cal/lst_stack/averaging.py b/hera_cal/lst_stack/averaging.py
@@ -505,7 +505,6 @@ def average_and_inpaint_simultaneously_single_bl(
         CNinv_data_dpss = np.array([
             basis.T.dot(weighted_data) for weighted_data in mask[:, band] / noise_var[:, band] * stackd[:, band]
         ])
-        print(CNinv_data_dpss.shape, CNinv_data_dpss.shape)
         dpss_fits = np.array([
             a.dot(b) if np.all(np.isfinite(b)) else a.dot(np.zeros_like(b))
             for a, b in zip(CNinv_dpss_inv, CNinv_data_dpss)

diff --git a/hera_cal/lst_stack/calibration.py b/hera_cal/lst_stack/calibration.py
@@ -106,38 +106,35 @@ def _expand_degeneracies_to_ant_gains(
         if split_pol1 != split_pol2:
             continue
         for ant in gain_ants:
-            for bandidx, band in enumerate(inpaint_bands):
-                raw_ant_gain = amplitude_parameters[f"A_{split_pol1}"] * (
-                    phase_gains.get(
-                        (ant, split_pol1),
-                        np.ones_like(amplitude_parameters[f"A_{split_pol1}"]),
-                    )
-                )
+            raw_ant_gain = amplitude_parameters[f"A_{split_pol1}"] * (
+                phase_gains.get((ant, split_pol1), 1)
+            )
+            gg = gains[(ant, split_pol1)] = raw_ant_gain
+
+            if smooth_gains:
+                for bandidx, band in enumerate(inpaint_bands):
 
-                if smooth_gains:
                     # Rephase antenna gains
                     tau, _ = utils.fft_dly(
-                        raw_ant_gain,
+                        raw_ant_gain[:, band],
                         np.diff(stack.freq_array[band])[0],
                         wgts=np.logical_not(stack.flags[:, 0, band, polidx]),
                     )
                     rephasor = np.exp(-2.0j * np.pi * tau * stack.freq_array[band])
-                    raw_ant_gain *= rephasor
+                    raw_ant_gain[:, band] *= rephasor
 
                     basis = smoothing_functions[bandidx]
                     smooth_ant_gain = np.array(
                         [
                             np.dot(basis, _fmat.dot(_raw_gain))
                             for _fmat, _raw_gain in zip(
-                                fmats[split_pol1][bandidx], raw_ant_gain
+                                fmats[split_pol1][bandidx], raw_ant_gain[:, band]
                             )
                         ]
                     )
 
                     # Rephase antenna gains
-                    gains[(ant, split_pol1)] = smooth_ant_gain * rephasor.conj()
-                else:
-                    gains[(ant, split_pol1)] = raw_ant_gain
+                    gg[:, band] = smooth_ant_gain * rephasor.conj()
 
     return gains