diff --git a/flint/bptools/smoother.py b/flint/bptools/smoother.py index c412fba6..10e825e1 100644 --- a/flint/bptools/smoother.py +++ b/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) diff --git a/flint/calibrate/aocalibrate.py b/flint/calibrate/aocalibrate.py index 4b88e499..18f6deb5 100644 --- a/flint/calibrate/aocalibrate.py +++ b/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: diff --git a/flint/data/tests/SB38969.B1934-638.beam35.aocalibrate.bin b/flint/data/tests/SB38969.B1934-638.beam35.aocalibrate.bin new file mode 100644 index 00000000..926a00f3 Binary files /dev/null and b/flint/data/tests/SB38969.B1934-638.beam35.aocalibrate.bin differ diff --git a/pyproject.toml b/pyproject.toml index bc0eb79c..e9993769 100644 --- a/pyproject.toml +++ b/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'}, ] diff --git a/tests/test_aocalibrate.py b/tests/test_aocalibrate.py index fb925f12..f30e4447 100644 --- a/tests/test_aocalibrate.py +++ b/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)