From 918d42a06a9ab734467808456362d959373913d5 Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Tue, 16 Jun 2020 15:59:23 +0200
Subject: [PATCH 1/6] Add self cal algorithm and example script

---
 lofarimaging/lofarimaging.py      | 74 ++++++++++++++++++++++++++--
 lofarimaging/singlestationutil.py | 80 ++++++++++++++++++-------------
 self_cal_demo.py                  | 66 +++++++++++++++++++++++++
 3 files changed, 184 insertions(+), 36 deletions(-)
 create mode 100644 self_cal_demo.py

diff --git a/lofarimaging/lofarimaging.py b/lofarimaging/lofarimaging.py
index 354e619..f5f335d 100644
--- a/lofarimaging/lofarimaging.py
+++ b/lofarimaging/lofarimaging.py
@@ -1,12 +1,14 @@
 """Functions for working with LOFAR single station data"""
 
+import logging
 from typing import Dict, List
-import numpy as np
-from numpy.linalg import norm, lstsq
-import numexpr as ne
+
 import numba
+import numexpr as ne
+import numpy as np
+import scipy.optimize
 from astropy.coordinates import SkyCoord, SkyOffsetFrame, CartesianRepresentation
-
+from numpy.linalg import norm, lstsq
 
 __all__ = ["nearfield_imager", "sky_imager", "ground_imager", "skycoord_to_lmn", "calibrate", "simulate_sky_source",
            "subtract_sources"]
@@ -170,6 +172,70 @@ def calibrate(vis, modelvis, maxiter=30, amplitudeonly=True):
     return residual, gains
 
 
+def compute_calibrated_model(vis, model_vis, maxiter=30):
+    n_ant = vis.shape[0]
+    gain_phase = np.zeros((n_ant), dtype=complex)
+
+    def gains_difference(gain_phase, vis, model_vis):
+        gains = np.exp(1.j * gain_phase)
+        gains_mat = np.diag(gains)
+        gains_mat_star = np.diag(np.conj(gains))
+        predicted = gains_mat_star @ model_vis @ gains_mat
+        residual = np.sum(np.abs(vis - predicted))
+        return residual
+
+    result = scipy.optimize.minimize(gains_difference, gain_phase, args=(vis, model_vis), options={'maxiter': maxiter})
+    gain_phase = result.x
+    gains_mat = np.diag(np.exp(1.j * gain_phase))
+
+    calibrated = np.conj(gains_mat) @ model_vis @ gains_mat
+
+    return calibrated, gains_difference(gain_phase, vis, model_vis)
+
+
+def compute_pointing_matrix(sources_positions, baselines, frequency):
+    n_baselines = baselines.shape[0]
+    n_sources = sources_positions.shape[0]
+
+    pointing_matrix = np.zeros(shape=(n_baselines, n_baselines, n_sources), dtype=complex)
+
+    for q in range(n_sources):
+        pointing_matrix[:, :, q] = simulate_sky_source(sources_positions[q, :], baselines, frequency)
+    return pointing_matrix
+
+
+def estimate_sources_flux(visibilities, pointing_matrix):
+    def residual_flux(signal, vis, point_matrix):
+        residual = vis - point_matrix @ signal
+        residual = np.sum(np.abs(residual))
+        return residual
+
+    n_antennas = visibilities.shape[0]
+    n_sources = pointing_matrix.shape[2]
+    lin_visibilities = visibilities.reshape((n_antennas * n_antennas))
+    lin_pointing_matrix = pointing_matrix.reshape((n_antennas * n_antennas, n_sources))
+    fluxes, *_ = np.linalg.lstsq(lin_pointing_matrix, lin_visibilities)
+
+    result = scipy.optimize.minimize(residual_flux, fluxes, args=(visibilities, pointing_matrix))
+    return result.x
+
+
+def estimate_model_visibilities(sources_positions, visibilities, baselines, frequency):
+    pointing_matrix = compute_pointing_matrix(sources_positions, baselines, frequency)
+    fluxes = estimate_sources_flux(visibilities, pointing_matrix)
+    model = pointing_matrix @ fluxes
+
+    return model
+
+
+def self_cal(visibilities, expected_sources, baselines, frequency, iterations=10):
+    model = estimate_model_visibilities(expected_sources, visibilities, baselines, frequency)
+    for i in range(iterations):
+        model, residual = compute_calibrated_model(visibilities, model, maxiter=50)
+        logging.debug('residual after iteration %d: %s', i, residual)
+    return model
+
+
 def simulate_sky_source(lmn_coord: np.array, baselines: np.array, freq: float):
     """
     Simulate visibilities for a sky source
diff --git a/lofarimaging/singlestationutil.py b/lofarimaging/singlestationutil.py
index 0b57b43..995f4b2 100644
--- a/lofarimaging/singlestationutil.py
+++ b/lofarimaging/singlestationutil.py
@@ -1,36 +1,32 @@
 """Functions for working with LOFAR single station data"""
 
-import os
 import datetime
+import os
 from typing import List, Dict, Tuple, Union
 
-import numpy as np
-from packaging import version
-import tqdm
+import astropy.units as u
 import h5py
-
-import matplotlib.pyplot as plt
+import lofarantpos
+import lofargeotiff
 import matplotlib.animation
-from matplotlib.ticker import FormatStrFormatter
-from matplotlib import cm
-from matplotlib.figure import Figure
-from matplotlib.colors import ListedColormap, Normalize
-from matplotlib.patches import Circle
 import matplotlib.axes as maxes
-from mpl_toolkits.axes_grid1 import make_axes_locatable
-
+import matplotlib.pyplot as plt
+import numpy as np
+import tqdm
 from astropy.coordinates import SkyCoord, GCRS, EarthLocation, AltAz, get_sun, get_moon
-import astropy.units as u
 from astropy.time import Time
-
-import lofargeotiff
 from lofarantpos.db import LofarAntennaDatabase
-import lofarantpos
+from matplotlib import cm
+from matplotlib.colors import ListedColormap, Normalize
+from matplotlib.figure import Figure
+from matplotlib.patches import Circle
+from matplotlib.ticker import FormatStrFormatter
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from packaging import version
 
-from .maputil import get_map, make_leaflet_map
-from .lofarimaging import nearfield_imager, sky_imager, skycoord_to_lmn, subtract_sources
 from .hdf5util import write_hdf5
-
+from .lofarimaging import nearfield_imager, sky_imager, skycoord_to_lmn, subtract_sources
+from .maputil import get_map, make_leaflet_map
 
 __all__ = ["sb_from_freq", "freq_from_sb", "find_caltable", "read_caltable",
            "rcus_in_station", "read_acm_cube", "get_station_pqr", "get_station_type",
@@ -331,7 +327,7 @@ def get_station_pqr(station_name: str, rcu_mode: Union[str, int], db):
             'intl': GENERIC_INT_201512, 'remote': GENERIC_REMOTE_201512, 'core': GENERIC_CORE_201512
         }
         selected_dipoles = selected_dipole_config[station_type] + \
-            np.arange(len(selected_dipole_config[station_type])) * 16
+                           np.arange(len(selected_dipole_config[station_type])) * 16
         station_pqr = db.hba_dipole_pqr(full_station_name)[selected_dipoles]
     else:
         raise RuntimeError("Station name did not contain LBA or HBA, could not load antenna positions")
@@ -340,7 +336,7 @@ def get_station_pqr(station_name: str, rcu_mode: Union[str, int], db):
 
 
 def make_ground_plot(image: np.ndarray, background_map: np.ndarray, extent: List[float], title: str = "Ground plot",
-        subtitle: str = "", opacity: float = 0.6, fig: Figure = None, draw_contours: bool = True, **kwargs) \
+                     subtitle: str = "", opacity: float = 0.6, fig: Figure = None, draw_contours: bool = True, **kwargs) \
         -> Tuple[Figure, np.ndarray]:
     """
     Make a ground plot of an array with data
@@ -585,7 +581,6 @@ def make_xst_plots(xst_data: np.ndarray,
         opacity: Opacity for map overlay. Defaults to 0.6.
         hdf5_filename: Filename where hdf5 results can be written. Defaults to outputpath + '/results.h5'
         outputpath: Directory where results can be saved. Defaults to 'results'
-        subtract: List of sources to subtract. Defaults to None
 
 
     Returns:
@@ -662,14 +657,14 @@ def make_xst_plots(xst_data: np.ndarray,
     marked_bodies = {
         'Cas A': SkyCoord(ra=350.85 * u.deg, dec=58.815 * u.deg),
         'Cyg A': SkyCoord(ra=299.86815191 * u.deg, dec=40.73391574 * u.deg),
-        'Per A': SkyCoord(ra=49.95066567*u.deg, dec=41.51169838 * u.deg),
-        'Her A': SkyCoord(ra=252.78343333*u.deg, dec=4.99303056*u.deg),
-        'Cen A': SkyCoord(ra=201.36506288*u.deg, dec=-43.01911267*u.deg),
-        'Vir A': SkyCoord(ra=187.70593076*u.deg, dec=12.39112329*u.deg),
-        '3C295': SkyCoord(ra=212.83527917*u.deg, dec=52.20264444*u.deg),
+        'Per A': SkyCoord(ra=49.95066567 * u.deg, dec=41.51169838 * u.deg),
+        'Her A': SkyCoord(ra=252.78343333 * u.deg, dec=4.99303056 * u.deg),
+        'Cen A': SkyCoord(ra=201.36506288 * u.deg, dec=-43.01911267 * u.deg),
+        'Vir A': SkyCoord(ra=187.70593076 * u.deg, dec=12.39112329 * u.deg),
+        '3C295': SkyCoord(ra=212.83527917 * u.deg, dec=52.20264444 * u.deg),
         'Moon': get_moon(obstime_astropy, location=station_earthlocation).transform_to(GCRS),
         'Sun': get_sun(obstime_astropy),
-        '3C196': SkyCoord(ra=123.40023371*u.deg, dec=48.21739888*u.deg)
+        '3C196': SkyCoord(ra=123.40023371 * u.deg, dec=48.21739888 * u.deg)
     }
 
     marked_bodies_lmn = {}
@@ -752,7 +747,7 @@ def make_xst_plots(xst_data: np.ndarray,
             "pixels_per_metre": pixels_per_metre}
     tags.update(calibration_info)
     lon_min, lon_max, lat_min, lat_max = extent_lonlat
-    lofargeotiff.write_geotiff(ground_img[::-1,:], os.path.join(outputpath, f"{fname}_nearfield_calibrated.tiff"),
+    lofargeotiff.write_geotiff(ground_img[::-1, :], os.path.join(outputpath, f"{fname}_nearfield_calibrated.tiff"),
                                (lon_min, lat_max), (lon_max, lat_min), as_pqr=False,
                                stationname=station_name, obsdate=obstime, tags=tags)
 
@@ -769,7 +764,7 @@ def make_sky_movie(moviefilename: str, h5file: h5py.File, obsnums: List[str], vm
     """
     Make movie of a list of observations
     """
-    fig = plt.figure(figsize=(10,10))
+    fig = plt.figure(figsize=(10, 10))
     for obsnum in tqdm.tqdm(obsnums):
         obs_h5 = h5file[obsnum]
         skydata_h5 = obs_h5["sky_img"]
@@ -787,7 +782,28 @@ def make_sky_movie(moviefilename: str, h5file: h5py.File, obsnums: List[str], vm
 
     # Thanks to Maaijke Mevius for making this animation work!
     ims = fig.get_children()[1:]
-    ims = [ims[i:i+2] for i in range(0, len(ims), 2)]
+    ims = [ims[i:i + 2] for i in range(0, len(ims), 2)]
     ani = matplotlib.animation.ArtistAnimation(fig, ims, interval=30, blit=False, repeat_delay=1000)
     writer = matplotlib.animation.writers['ffmpeg'](fps=5, bitrate=800)
     ani.save(moviefilename, writer=writer, dpi=fig.dpi)
+
+
+def compute_baselines(station_name, rcu_mode):
+    # Setup the database
+    db = LofarAntennaDatabase()
+
+    station_pqr = get_station_pqr(station_name, rcu_mode, db)
+
+    station_name = get_full_station_name(station_name, rcu_mode)
+
+    # Rotate station_pqr to a north-oriented xyz frame, where y points North, in a plane through the station.
+    rotation = db.rotation_from_north(station_name)
+
+    pqr_to_xyz = np.array([[np.cos(-rotation), -np.sin(-rotation), 0],
+                           [np.sin(-rotation), np.cos(-rotation), 0],
+                           [0, 0, 1]])
+
+    station_xyz = (pqr_to_xyz @ station_pqr.T).T
+
+    baselines = station_xyz[:, np.newaxis, :] - station_xyz[np.newaxis, :, :]
+    return baselines
diff --git a/self_cal_demo.py b/self_cal_demo.py
new file mode 100644
index 0000000..24082ac
--- /dev/null
+++ b/self_cal_demo.py
@@ -0,0 +1,66 @@
+import h5py
+import numpy
+import lofarimaging.hdf5util as h5utils
+from lofarimaging.singlestationutil import compute_baselines, make_sky_plot
+from lofarimaging.lofarimaging import compute_pointing_matrix, compute_calibrated_model, sky_imager,\
+    self_cal, estimate_model_visibilities
+import matplotlib.pyplot as plt
+
+test_data_path = '/home/mmancini/Documents/Projects/SingleStationImager/test_dataset.h5'
+test_station = 'CS103'
+test_data = h5py.File(test_data_path, 'r')
+obs_names = h5utils.get_obsnums(test_data, station_name=test_station)
+an_obs = test_data[obs_names[10]]
+
+a_frequency = an_obs.attrs['frequency']
+a_source_lmn = an_obs.attrs['source_lmn']
+a_source_name = an_obs.attrs['source_names']
+a_rcu_mode = an_obs.attrs['rcu_mode']
+a_dataset = an_obs['calibrated_data']
+
+a_dataset_xx = a_dataset[::2, ::2]
+a_dataset_yy = a_dataset[1::2, 1::2]
+
+a_dataset_i = a_dataset_xx + a_dataset_yy
+a_dataset_q = a_dataset_xx - a_dataset_yy
+a_dataset_u = 2 * (a_dataset_xx * a_dataset_yy.conj()).real
+a_dataset_v = -2 * (a_dataset_xx * a_dataset_yy.conj()).imag
+
+baselines = compute_baselines(test_station, a_rcu_mode)
+sky_image = sky_imager(a_dataset_i, baselines, a_frequency, 100, 100)
+
+model = estimate_model_visibilities(a_source_lmn, a_dataset_i, baselines, a_frequency)
+
+model_image = sky_imager(model, baselines, a_frequency, 100, 100)
+
+pointing_matrix = compute_pointing_matrix(a_source_lmn, baselines, a_frequency)
+calibrated_model, _ = compute_calibrated_model(a_dataset_i, model, maxiter=50)
+self_cal_model = self_cal(a_dataset_i, a_source_lmn, baselines, a_frequency)
+
+sky_without_calibrated_model = sky_imager(a_dataset_i - self_cal_model, baselines, a_frequency, 100, 100)
+calibrated_model_image = sky_imager(calibrated_model, baselines, a_frequency, 100, 100)
+sky_without_model = sky_imager(a_dataset_i - model, baselines, a_frequency, 100, 100)
+
+
+f, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
+min, max = numpy.min(sky_image), numpy.max(sky_image)
+
+ax1.scatter(a_source_lmn[:, 0], a_source_lmn[:, 1], marker='o', color='r', s=100, facecolors='none')
+im = ax1.imshow(sky_image[::-1,:], vmin=min, vmax=max, extent=(1, -1, -1, 1))
+ax1.set_title('initial sky')
+
+ax2.scatter(a_source_lmn[:, 0], a_source_lmn[:, 1], marker='o', color='r', s=100, facecolors='none')
+ax2.imshow(sky_without_model[::-1, :], extent=(1, -1, -1, 1), vmin=min, vmax=max)
+ax2.set_title('sky without model')
+
+ax3.scatter(a_source_lmn[:, 0], a_source_lmn[:, 1], marker='o', color='r', s=100, facecolors='none')
+ax3.imshow(calibrated_model_image[::-1, :], extent=(1, -1, -1, 1))
+ax3.set_title('calibrated model image')
+
+ax4.imshow(sky_without_calibrated_model[::-1, :], extent=(1, -1, -1, 1), vmin=min, vmax=max)
+to_plot = {name: pos for pos, name in zip(a_source_lmn, a_source_name)}
+ax4.set_title('sky without calibrated model')
+
+make_sky_plot(sky_image, to_plot)
+
+plt.show()

From 35cfee3304a1bc80373bd02dce8cd103fa155d18 Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Tue, 16 Jun 2020 16:30:14 +0200
Subject: [PATCH 2/6] Add missing dependency

---
 requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/requirements.txt b/requirements.txt
index e456e46..9669e00 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,3 +13,4 @@ Pillow
 opcua
 h5py
 tqdm
+scipy
\ No newline at end of file

From 9844a41573703f2b8e8e429f5c88d48cb5a0247e Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Tue, 16 Jun 2020 17:53:34 +0200
Subject: [PATCH 3/6] Add test and doc strings

---
 lofarimaging/lofarimaging.py | 67 +++++++++++++++++++++++++++++--
 self_cal_demo.py             | 66 -------------------------------
 test/test_single_station.py  | 77 ++++++++++++++++++++++++++++++++++++
 3 files changed, 141 insertions(+), 69 deletions(-)
 delete mode 100644 self_cal_demo.py
 create mode 100644 test/test_single_station.py

diff --git a/lofarimaging/lofarimaging.py b/lofarimaging/lofarimaging.py
index f5f335d..f904e71 100644
--- a/lofarimaging/lofarimaging.py
+++ b/lofarimaging/lofarimaging.py
@@ -11,7 +11,8 @@
 from numpy.linalg import norm, lstsq
 
 __all__ = ["nearfield_imager", "sky_imager", "ground_imager", "skycoord_to_lmn", "calibrate", "simulate_sky_source",
-           "subtract_sources"]
+           "subtract_sources", "compute_calibrated_model", "compute_pointing_matrix", "estimate_model_visibilities",
+           "estimate_sources_flux", "self_cal"]
 
 __version__ = "1.5.0"
 SPEED_OF_LIGHT = 299792458.0
@@ -173,6 +174,18 @@ def calibrate(vis, modelvis, maxiter=30, amplitudeonly=True):
 
 
 def compute_calibrated_model(vis, model_vis, maxiter=30):
+    """
+    Calibrate the gains phases and apply to the model
+
+    Args:
+        vis: visibility matrix, shape [n_st, n_st]
+        model_vis: model visibility matrix, shape [n_st, n_st]
+        maxiter: max iterations (default 30)
+    Returns:
+        calibrated: model visibilities, shape [n_st, n_st]
+        residual: amplitude difference between model and actual visibilities, float
+    """
+
     n_ant = vis.shape[0]
     gain_phase = np.zeros((n_ant), dtype=complex)
 
@@ -194,6 +207,17 @@ def gains_difference(gain_phase, vis, model_vis):
 
 
 def compute_pointing_matrix(sources_positions, baselines, frequency):
+    """
+    Compute the pointing matrix used to compute the visibilities from
+    the sources position
+
+    Args:
+        sources_positions: matrix of the lmn sources position, shape [n_dir, 3]
+        baselines: baseline matrix, shape [n_st, n_st]
+        frequency: frequency
+    Returns:
+        pointing_matrix: pointing matrix, shape [n_st, n_st, n_dir]
+    """
     n_baselines = baselines.shape[0]
     n_sources = sources_positions.shape[0]
 
@@ -205,6 +229,16 @@ def compute_pointing_matrix(sources_positions, baselines, frequency):
 
 
 def estimate_sources_flux(visibilities, pointing_matrix):
+    """
+    Estimate the flux of the sources given the pointing matrix and the measured visibilities
+
+    Args:
+        visibilities: visibilities, shape [n_st, n_st]
+        pointing_matrix: pointing matrix, shape [n_st, n_st, n_dir]
+    Returns:
+        fluxes: amplitude of the source at the given direction, shape [n_st, n_st, n_dir]
+    """
+
     def residual_flux(signal, vis, point_matrix):
         residual = vis - point_matrix @ signal
         residual = np.sum(np.abs(residual))
@@ -217,10 +251,22 @@ def residual_flux(signal, vis, point_matrix):
     fluxes, *_ = np.linalg.lstsq(lin_pointing_matrix, lin_visibilities)
 
     result = scipy.optimize.minimize(residual_flux, fluxes, args=(visibilities, pointing_matrix))
-    return result.x
+    fluxes = result.x
+    return fluxes
 
 
 def estimate_model_visibilities(sources_positions, visibilities, baselines, frequency):
+    """
+    Estimate the visibilities of the models given the sources position and the measured visibilities
+
+    Args:
+        sources_positions: lmn coords of the sources, shape [n_dir, 3]
+        visibilities: visibilities, shape [n_st, n_st]
+        baselines: baselines matrix in m, shape [n_st, n_st, 3]
+        frequency: frequency
+    Returns:
+        model: the model visibilities, shape [n_st, n_st]
+    """
     pointing_matrix = compute_pointing_matrix(sources_positions, baselines, frequency)
     fluxes = estimate_sources_flux(visibilities, pointing_matrix)
     model = pointing_matrix @ fluxes
@@ -229,6 +275,21 @@ def estimate_model_visibilities(sources_positions, visibilities, baselines, freq
 
 
 def self_cal(visibilities, expected_sources, baselines, frequency, iterations=10):
+    """
+    Compute the gain phase comparing the model and the actual visibilities returning
+    the model with the gain phase correction applied.
+
+    TODO introduce a valid stop condition and make the iteration parameters a max_iter parameter
+
+    Args:
+        visibilities: visibilities, shape [n_st, n_st]
+        expected_sources: lmn coords of the sources, shape [n_dir, 3]
+        baselines: baselines matrix in m, shape [n_st, n_st, 3]
+        frequency: frequency
+        iterations: number of iterations to perform
+    Returns:
+        model: the model visibilities with the applied gains, shape [n_st, n_st]
+    """
     model = estimate_model_visibilities(expected_sources, visibilities, baselines, frequency)
     for i in range(iterations):
         model, residual = compute_calibrated_model(visibilities, model, maxiter=50)
@@ -242,7 +303,7 @@ def simulate_sky_source(lmn_coord: np.array, baselines: np.array, freq: float):
 
     Args:
         lmn_coord (np.array): l, m, n coordinate
-        baselines (np.array): baseline distances in metres, shape (n_ant, n_ant)
+        baselines (np.array): baseline distances in metres, shape (n_ant, n_ant, 3)
         freq (float): Frequency in Hz
     """
     return np.exp(2j * np.pi * freq * baselines.dot(np.array(lmn_coord)) / SPEED_OF_LIGHT)
diff --git a/self_cal_demo.py b/self_cal_demo.py
deleted file mode 100644
index 24082ac..0000000
--- a/self_cal_demo.py
+++ /dev/null
@@ -1,66 +0,0 @@
-import h5py
-import numpy
-import lofarimaging.hdf5util as h5utils
-from lofarimaging.singlestationutil import compute_baselines, make_sky_plot
-from lofarimaging.lofarimaging import compute_pointing_matrix, compute_calibrated_model, sky_imager,\
-    self_cal, estimate_model_visibilities
-import matplotlib.pyplot as plt
-
-test_data_path = '/home/mmancini/Documents/Projects/SingleStationImager/test_dataset.h5'
-test_station = 'CS103'
-test_data = h5py.File(test_data_path, 'r')
-obs_names = h5utils.get_obsnums(test_data, station_name=test_station)
-an_obs = test_data[obs_names[10]]
-
-a_frequency = an_obs.attrs['frequency']
-a_source_lmn = an_obs.attrs['source_lmn']
-a_source_name = an_obs.attrs['source_names']
-a_rcu_mode = an_obs.attrs['rcu_mode']
-a_dataset = an_obs['calibrated_data']
-
-a_dataset_xx = a_dataset[::2, ::2]
-a_dataset_yy = a_dataset[1::2, 1::2]
-
-a_dataset_i = a_dataset_xx + a_dataset_yy
-a_dataset_q = a_dataset_xx - a_dataset_yy
-a_dataset_u = 2 * (a_dataset_xx * a_dataset_yy.conj()).real
-a_dataset_v = -2 * (a_dataset_xx * a_dataset_yy.conj()).imag
-
-baselines = compute_baselines(test_station, a_rcu_mode)
-sky_image = sky_imager(a_dataset_i, baselines, a_frequency, 100, 100)
-
-model = estimate_model_visibilities(a_source_lmn, a_dataset_i, baselines, a_frequency)
-
-model_image = sky_imager(model, baselines, a_frequency, 100, 100)
-
-pointing_matrix = compute_pointing_matrix(a_source_lmn, baselines, a_frequency)
-calibrated_model, _ = compute_calibrated_model(a_dataset_i, model, maxiter=50)
-self_cal_model = self_cal(a_dataset_i, a_source_lmn, baselines, a_frequency)
-
-sky_without_calibrated_model = sky_imager(a_dataset_i - self_cal_model, baselines, a_frequency, 100, 100)
-calibrated_model_image = sky_imager(calibrated_model, baselines, a_frequency, 100, 100)
-sky_without_model = sky_imager(a_dataset_i - model, baselines, a_frequency, 100, 100)
-
-
-f, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4)
-min, max = numpy.min(sky_image), numpy.max(sky_image)
-
-ax1.scatter(a_source_lmn[:, 0], a_source_lmn[:, 1], marker='o', color='r', s=100, facecolors='none')
-im = ax1.imshow(sky_image[::-1,:], vmin=min, vmax=max, extent=(1, -1, -1, 1))
-ax1.set_title('initial sky')
-
-ax2.scatter(a_source_lmn[:, 0], a_source_lmn[:, 1], marker='o', color='r', s=100, facecolors='none')
-ax2.imshow(sky_without_model[::-1, :], extent=(1, -1, -1, 1), vmin=min, vmax=max)
-ax2.set_title('sky without model')
-
-ax3.scatter(a_source_lmn[:, 0], a_source_lmn[:, 1], marker='o', color='r', s=100, facecolors='none')
-ax3.imshow(calibrated_model_image[::-1, :], extent=(1, -1, -1, 1))
-ax3.set_title('calibrated model image')
-
-ax4.imshow(sky_without_calibrated_model[::-1, :], extent=(1, -1, -1, 1), vmin=min, vmax=max)
-to_plot = {name: pos for pos, name in zip(a_source_lmn, a_source_name)}
-ax4.set_title('sky without calibrated model')
-
-make_sky_plot(sky_image, to_plot)
-
-plt.show()
diff --git a/test/test_single_station.py b/test/test_single_station.py
new file mode 100644
index 0000000..49b7ace
--- /dev/null
+++ b/test/test_single_station.py
@@ -0,0 +1,77 @@
+import numpy
+from numpy.testing import assert_almost_equal
+
+from lofarimaging.lofarimaging import compute_calibrated_model, \
+    compute_pointing_matrix, simulate_sky_source, \
+    estimate_sources_flux, \
+    estimate_model_visibilities, \
+    self_cal
+
+
+def test_compute_calibrated_model():
+    # ----TEST DATA
+    vis = numpy.diag(numpy.ones((10)))
+    model_vis = numpy.diag(numpy.ones(10))
+    # ----TEST
+    calibrated, residual = compute_calibrated_model(vis, model_vis)
+    assert_almost_equal(calibrated, vis)
+    assert residual < 1.e-5
+
+
+def test_simulate_sky_source():
+    # ----TEST DATA
+    source_position = [0, 0, 0]
+    baselines = numpy.zeros((2, 2, 3))
+    baselines[0, 1, :] = [1, 0, 0]
+    baselines[1, 0, :] = [1, 0, 0]
+    # ----TEST
+    visibilities = simulate_sky_source(source_position, baselines, 1)
+    expected = numpy.ones((2, 2))
+    assert_almost_equal(visibilities, expected)
+
+
+def test_compute_pointing_matrix():
+    # ----TEST DATA
+    source_position = numpy.zeros((1, 3))
+    baselines = numpy.zeros((2, 2, 3))
+    baselines[0, 1, :] = [1, 0, 0]
+    baselines[1, 0, :] = [1, 0, 0]
+    expected = numpy.ones((2, 2, 1))
+    # ----TEST
+    pointing_matrix = compute_pointing_matrix(source_position, baselines, 1)
+    assert_almost_equal(pointing_matrix, expected)
+
+
+def test_estimate_sources_flux():
+    # ----TEST DATA
+    source_position = numpy.zeros((1, 3))
+    pointing_matrix = numpy.ones((2, 2, 1))
+    visibilities = numpy.ones((2, 2))
+    # ----TEST
+    source_fluxes = estimate_sources_flux(visibilities, pointing_matrix)
+    expected = numpy.ones((1))
+    assert_almost_equal(source_fluxes, expected)
+
+
+def test_estimate_model_visibilities():
+    # ----TEST DATA
+    source_position = numpy.zeros((1, 3))
+    visibilities = numpy.ones((2, 2))
+    baselines = numpy.zeros((2, 2, 3))
+    baselines[0, 1, :] = [1, 0, 0]
+    baselines[1, 0, :] = [1, 0, 0]
+    # ----TEST
+    model_visibilities = estimate_model_visibilities(source_position, visibilities, baselines, 1)
+    assert_almost_equal(visibilities, model_visibilities)
+
+
+def test_self_cal():
+    # ----TEST DATA
+    source_position = numpy.zeros((1, 3))
+    visibilities = numpy.ones((2, 2))
+    baselines = numpy.zeros((2, 2, 3))
+    baselines[0, 1, :] = [1, 0, 0]
+    baselines[1, 0, :] = [1, 0, 0]
+    # ----TEST
+    calibrated_model = self_cal(visibilities, source_position, baselines, 1)
+    assert_almost_equal(calibrated_model, visibilities)

From 3b0dd89cba614d045d01ce91f7ec4edc722b08b3 Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Tue, 16 Jun 2020 20:26:50 +0200
Subject: [PATCH 4/6] Fix error in lofarimaging notebook

---
 lofarimaging.ipynb | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/lofarimaging.ipynb b/lofarimaging.ipynb
index 20b1cca..5098bb0 100644
--- a/lofarimaging.ipynb
+++ b/lofarimaging.ipynb
@@ -170,8 +170,7 @@
      "output_type": "stream",
      "text": [
       "Searching for available files in ./test\n",
-      "0: ./test/20200525_084652_mode_5_xst_sb100.dat\n",
-      "1: ./test/20170720_095816_mode_3_xst_sb297.dat\n"
+      "0: ./test/20170720_095816_mode_3_xst_sb297.dat\n"
      ]
     }
    ],
@@ -198,7 +197,7 @@
    ],
    "source": [
     "# Select a file\n",
-    "xst_filename = files[1]\n",
+    "xst_filename = files[0]\n",
     "\n",
     "print(\"File selected:\", xst_filename)"
    ]

From 470e19b3b006f93d62e2c1b7eefa63214a6bcd34 Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Tue, 16 Jun 2020 22:19:58 +0200
Subject: [PATCH 5/6] Add exit criteria and return grains vector

---
 lofarimaging/lofarimaging.py | 30 ++++++++++++++++++++----------
 test/test_single_station.py  |  6 +++---
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/lofarimaging/lofarimaging.py b/lofarimaging/lofarimaging.py
index f904e71..b8a30cb 100644
--- a/lofarimaging/lofarimaging.py
+++ b/lofarimaging/lofarimaging.py
@@ -183,6 +183,7 @@ def compute_calibrated_model(vis, model_vis, maxiter=30):
         maxiter: max iterations (default 30)
     Returns:
         calibrated: model visibilities, shape [n_st, n_st]
+        gains: the antenna grains, shape [n_st]
         residual: amplitude difference between model and actual visibilities, float
     """
 
@@ -199,11 +200,11 @@ def gains_difference(gain_phase, vis, model_vis):
 
     result = scipy.optimize.minimize(gains_difference, gain_phase, args=(vis, model_vis), options={'maxiter': maxiter})
     gain_phase = result.x
-    gains_mat = np.diag(np.exp(1.j * gain_phase))
-
+    gains = np.exp(1.j * gain_phase)
+    gains_mat = np.diag(gains)
     calibrated = np.conj(gains_mat) @ model_vis @ gains_mat
 
-    return calibrated, gains_difference(gain_phase, vis, model_vis)
+    return calibrated, gains, gains_difference(gain_phase, vis, model_vis)
 
 
 def compute_pointing_matrix(sources_positions, baselines, frequency):
@@ -274,27 +275,36 @@ def estimate_model_visibilities(sources_positions, visibilities, baselines, freq
     return model
 
 
-def self_cal(visibilities, expected_sources, baselines, frequency, iterations=10):
+def self_cal(visibilities, expected_sources, baselines, frequency, max_iterations=10, tollerance=1.e-4):
     """
     Compute the gain phase comparing the model and the actual visibilities returning
     the model with the gain phase correction applied.
 
-    TODO introduce a valid stop condition and make the iteration parameters a max_iter parameter
 
     Args:
         visibilities: visibilities, shape [n_st, n_st]
         expected_sources: lmn coords of the sources, shape [n_dir, 3]
         baselines: baselines matrix in m, shape [n_st, n_st, 3]
         frequency: frequency
-        iterations: number of iterations to perform
+        max_iterations: maximum number of iterations to perform
     Returns:
         model: the model visibilities with the applied gains, shape [n_st, n_st]
+        gains: the antenna gains, shape [n_st]
+
     """
     model = estimate_model_visibilities(expected_sources, visibilities, baselines, frequency)
-    for i in range(iterations):
-        model, residual = compute_calibrated_model(visibilities, model, maxiter=50)
-        logging.debug('residual after iteration %d: %s', i, residual)
-    return model
+    model, gains,  residual = compute_calibrated_model(visibilities, model, maxiter=50)
+    for i in range(1, max_iterations):
+        model, gains, next_residual = compute_calibrated_model(visibilities, model, maxiter=50)
+        print(i, residual, next_residual)
+        if next_residual == 0:
+            break
+        elif abs(1 - (residual / next_residual)) >= tollerance:
+            residual = next_residual
+        else:
+            break
+
+    return model, gains
 
 
 def simulate_sky_source(lmn_coord: np.array, baselines: np.array, freq: float):
diff --git a/test/test_single_station.py b/test/test_single_station.py
index 49b7ace..12747ea 100644
--- a/test/test_single_station.py
+++ b/test/test_single_station.py
@@ -13,7 +13,7 @@ def test_compute_calibrated_model():
     vis = numpy.diag(numpy.ones((10)))
     model_vis = numpy.diag(numpy.ones(10))
     # ----TEST
-    calibrated, residual = compute_calibrated_model(vis, model_vis)
+    calibrated, gains, residual = compute_calibrated_model(vis, model_vis)
     assert_almost_equal(calibrated, vis)
     assert residual < 1.e-5
 
@@ -65,7 +65,7 @@ def test_estimate_model_visibilities():
     assert_almost_equal(visibilities, model_visibilities)
 
 
-def test_self_cal():
+def test_self_cal_zero_residual():
     # ----TEST DATA
     source_position = numpy.zeros((1, 3))
     visibilities = numpy.ones((2, 2))
@@ -73,5 +73,5 @@ def test_self_cal():
     baselines[0, 1, :] = [1, 0, 0]
     baselines[1, 0, :] = [1, 0, 0]
     # ----TEST
-    calibrated_model = self_cal(visibilities, source_position, baselines, 1)
+    calibrated_model, _ = self_cal(visibilities, source_position, baselines, 1)
     assert_almost_equal(calibrated_model, visibilities)

From 3c177be3596e5da890fbc3f124e83e657400fc56 Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Wed, 17 Jun 2020 10:13:42 +0200
Subject: [PATCH 6/6] Fix typos and refactors if branch

---
 lofarimaging/lofarimaging.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/lofarimaging/lofarimaging.py b/lofarimaging/lofarimaging.py
index b8a30cb..59a61e6 100644
--- a/lofarimaging/lofarimaging.py
+++ b/lofarimaging/lofarimaging.py
@@ -183,7 +183,7 @@ def compute_calibrated_model(vis, model_vis, maxiter=30):
         maxiter: max iterations (default 30)
     Returns:
         calibrated: model visibilities, shape [n_st, n_st]
-        gains: the antenna grains, shape [n_st]
+        gains: the antenna gains, shape [n_st]
         residual: amplitude difference between model and actual visibilities, float
     """
 
@@ -275,7 +275,7 @@ def estimate_model_visibilities(sources_positions, visibilities, baselines, freq
     return model
 
 
-def self_cal(visibilities, expected_sources, baselines, frequency, max_iterations=10, tollerance=1.e-4):
+def self_cal(visibilities, expected_sources, baselines, frequency, max_iterations=10, tolerance=1.e-4):
     """
     Compute the gain phase comparing the model and the actual visibilities returning
     the model with the gain phase correction applied.
@@ -296,10 +296,7 @@ def self_cal(visibilities, expected_sources, baselines, frequency, max_iteration
     model, gains,  residual = compute_calibrated_model(visibilities, model, maxiter=50)
     for i in range(1, max_iterations):
         model, gains, next_residual = compute_calibrated_model(visibilities, model, maxiter=50)
-        print(i, residual, next_residual)
-        if next_residual == 0:
-            break
-        elif abs(1 - (residual / next_residual)) >= tollerance:
+        if next_residual != 0 and abs(1 - (residual / next_residual)) >= tolerance:
             residual = next_residual
         else:
             break