Merge pull request #40 from tjgalvin/nansmoother

Handle completely flagged antenna when smoothing

flint/bptools/smoother.py

@@ -76,6 +76,9 @@ def smooth_data(
     # where ever it might be. Trust nothing you sea dog.
     data = data.copy()
 
+    if np.all(~np.isfinite(data)):
+        return data
+
     if apply_median_filter:
         data = median_filter(input=data, size=window_size)
 

flint/calibrate/aocalibrate.py

@@ -668,6 +668,7 @@ def flag_aosolutions(
     flag_ant_xyyx_mean_gain: bool = False,
     zero_cross_terms: bool = False,
     smooth_solutions: bool = False,
+    plot_solutions_throughout: bool = True,
 ) -> Path:
     """Will open a previously solved ao-calibrate solutions file and flag additional channels and antennae.
 
@@ -686,6 +687,7 @@ def flag_aosolutions(
         flag_ant_xyyx_mean_gain (bool, optional): Whether to flag antennas based on the mean gain ratio of the XY and YX amplitude gains. Defaults to False.
         zero_cross_terms (bool, optional): Set the XY and YX terms of each Jones to be 0. Defaults to False.
         smooth_solutions (blool, optional): Smooth the complex gain solutions after flaggined. Defaults to False.
+        plot_solutions_throughout (bool, Optional): If True, the solutions will be plotted at different stages of processing. Defaults to True.
 
     Returns:
         Path: Path to the updated solutions file. This is out_solutions_path if provided, otherwise solutions_path
@@ -717,7 +719,6 @@ def flag_aosolutions(
         for pol in (0, 3):
             logger.info(f"Processing {pols[pol]} polarisation")
             ref_ant_gains = bandpass[time, ref_ant, :, pol]
-            # TODO: A better way of selecting the most appropriate reference antenna is needed.
             if np.sum(np.isfinite(ref_ant_gains)) == 0:
                 raise ValueError(f"The ref_ant={ref_ant} is completely bad. ")
 
@@ -813,7 +814,8 @@ def flag_aosolutions(
         ms_path=solutions_path, include_preflagger=True, include_smoother=False
     )
     solutions.save(output_path=out_solutions_path)
-    plot_solutions(solutions=out_solutions_path, ref_ant=ref_ant)
+    if plot_solutions_throughout:
+        plot_solutions(solutions=out_solutions_path, ref_ant=ref_ant)
 
     if smooth_solutions:
         logger.info("Smoothing the bandpass solutions. ")
@@ -827,7 +829,8 @@ def flag_aosolutions(
             ms_path=solutions_path, include_preflagger=True, include_smoother=True
         )
         solutions.save(output_path=out_solutions_path)
-        plot_solutions(solutions=out_solutions_path, ref_ant=None)
+        if plot_solutions_throughout:
+            plot_solutions(solutions=out_solutions_path, ref_ant=None)
 
     total_flagged = np.sum(~np.isfinite(bandpass)) / np.prod(bandpass.shape)
     if total_flagged > 0.8:

pyproject.toml

@@ -10,6 +10,7 @@ include = [
     {path = 'flint/data/source_counts/de_zotti_1p4.txt'},
     {path = 'flint/data/source_counts/SKADS_1p4GHz.fits'},
     {path = 'flint/data/aoflagger/ATCA.lua'},
+    {path = 'flint/data/aoflagger/ASKAP.lua'},
     {path = 'flint/data/tests'},
 ]
 

tests/test_aocalibrate.py

@@ -7,8 +7,17 @@
 import pkg_resources
 import pytest
 
-from flint.bptools.smoother import divide_bandpass_by_ref_ant
-from flint.calibrate.aocalibrate import AOSolutions, plot_solutions, select_refant
+from flint.bptools.smoother import (
+    divide_bandpass_by_ref_ant,
+    smooth_bandpass_complex_gains,
+    smooth_data,
+)
+from flint.calibrate.aocalibrate import (
+    AOSolutions,
+    flag_aosolutions,
+    plot_solutions,
+    select_refant,
+)
 
 
 @pytest.fixture
@@ -26,6 +35,28 @@ def ao_sols(tmpdir):
     return out_ao_sols
 
 
+@pytest.fixture
+def ao_sols_known_bad(tmpdir):
+    # The file contains a binary solutions file that failed previously.
+    # It was fixed by testing for all nans in the `flint.bptools.smoother.smooth_data`
+    # function.
+    ao_sols = Path(
+        pkg_resources.resource_filename(
+            "flint", "data/tests/SB38969.B1934-638.beam35.aocalibrate.bin"
+        )
+    )
+
+    out_ao_sols = Path(tmpdir) / ao_sols.name
+
+    shutil.copyfile(ao_sols, out_ao_sols)
+
+    return out_ao_sols
+
+
+def test_known_bad_sols(ao_sols_known_bad):
+    flag_aosolutions(solutions_path=ao_sols_known_bad, plot_solutions_throughout=False)
+
+
 def test_load_aosols(ao_sols):
     ao = AOSolutions.load(path=ao_sols)
 
@@ -40,6 +71,41 @@ def test_aosols_bandpass_plot(ao_sols):
     plot_solutions(solutions=ao_sols, ref_ant=None)
 
 
+def test_aosols_all_nans_smooth_data(ao_sols):
+    ao = AOSolutions.load(ao_sols)
+    ao.bandpass[0, 20, :, :] = np.nan
+    assert np.all(~np.isfinite(ao.bandpass[0, 20, :, 0]))
+
+    smoothed = smooth_data(
+        data=ao.bandpass[0, 20, :, 0].real, window_size=16, polynomial_order=4
+    )
+    assert np.all(~np.isfinite(smoothed))
+
+
+def test_smooth_bandpass_complex_gains_nans(ao_sols):
+    ao = AOSolutions.load(ao_sols)
+    ao.bandpass[0, 20, :, :] = np.nan
+    assert np.all(~np.isfinite(ao.bandpass[0, 20, :, 0]))
+
+    smoothed = smooth_bandpass_complex_gains(
+        complex_gains=ao.bandpass[0], window_size=16, polynomial_order=4
+    )
+    assert np.all(~np.isfinite(smoothed[20, :, 0]))
+
+
+def test_smooth_bandpass_complex_gains_nans_with_refant(ao_sols):
+    ao = AOSolutions.load(ao_sols)
+    ao.bandpass[0, 20, :, :] = np.nan
+    assert np.all(~np.isfinite(ao.bandpass[0, 20, :, 0]))
+
+    ref = divide_bandpass_by_ref_ant(complex_gains=ao.bandpass[0], ref_ant=0)
+
+    smoothed = smooth_bandpass_complex_gains(
+        complex_gains=ref, window_size=16, polynomial_order=4
+    )
+    assert np.all(~np.isfinite(smoothed[20, :, 0]))
+
+
 def test_aosols_bandpass_ref_nu_rank_error(ao_sols):
     ao = AOSolutions.load(path=ao_sols)