@@ -11,6 +11,22 @@
 
 http://adsabs.harvard.edu/cgi-bin/nph-data_query?bibcode=2011MNRAS.414.1656K&link_type=EJOURNAL
 
+In this code:
+
+- A single theta vector is taken to be the amplitude of a given point source and the gain (phase only) solutions for
+that point source for all solution point
+   
+
+- The E step for a specific window is the sum of the window data model and the discrepancy between the observed data and the summed (over all windows) data models.
+   
+   
+- The M step for a specific window is the optimisation of the theta vector given the window data model. This involves fitting a point source (flux only for the moment) and fitting for the gain phases.
+
+
+To run sagecal, you must provide a visibility dataset and a set of skycomponents. The output will be the model
+parameters (source flux and gaintable for all skycomponents), and the residual visibility.
+
+Sagecal should only be used once an initial phase calibration has been obtained using an isoplanatic approximation.
 """
 
 import logging
@@ -29,9 +45,11 @@
 log = logging.getLogger(__name__)
 
 
-def create_sagecal_thetas(vis, comps, **kwargs):
+def create_sagecal_thetas(vis: BlockVisibility, comps, **kwargs):
     """Create the thetas
     
+    Create the data model for each window, from the visibility and the existing components
+    
     :param comps:
     :param gt:
     :return:
@@ -48,7 +66,9 @@ def create_sagecal_thetas(vis, comps, **kwargs):
 def sagecal_fit_component(evis, theta, gain=0.1, **kwargs):
     """Fit a single component to a visibility i.e. A13
 
-    # Just do the amplitude for now
+    This is the update to the component part of the window
+
+    Just do the amplitude for now
 
     :param evis:
     :param theta:
@@ -64,6 +84,8 @@ def sagecal_fit_component(evis, theta, gain=0.1, **kwargs):
 
 def sagecal_fit_gaintable(evis, theta, gain=0.1, niter=3, tol=1e-3, **kwargs):
     """Fit a gaintable to a visibility i.e. A13
+    
+    This is the update to the gain part of the window
 
     :param evis:
     :param theta:
@@ -83,6 +105,8 @@ def sagecal_fit_gaintable(evis, theta, gain=0.1, niter=3, tol=1e-3, **kwargs):
 
 def sagecal_e_step(vis: BlockVisibility, evis_all: BlockVisibility, theta, beta=1.0, **kwargs):
     """Calculates E step in equation A12
+    
+    This is the data model for this window plus the difference between observed data and summed data models
 
     :param vis:
     :param theta:
@@ -100,6 +124,8 @@ def sagecal_e_step(vis: BlockVisibility, evis_all: BlockVisibility, theta, beta=
 
 def sagecal_e_all(vis: BlockVisibility, thetas, **kwargs):
     """Calculates E step in equation A12
+    
+    This is the sum of the data models over all windows
 
     :param vis:
     :param thetas:
@@ -116,10 +142,11 @@ def sagecal_e_all(vis: BlockVisibility, thetas, **kwargs):
     return evis
 
 
-def sagecal_m_step(evis: BlockVisibility, theta, gain=0.25,
-                   **kwargs):
+def sagecal_m_step(evis: BlockVisibility, theta, gain=0.25, **kwargs):
     """Calculates M step in equation A13
     
+    This maximises the likelihood of the theta parameters given the existing model data
+    
     :param vis:
     :param theta:
     :param kwargs:
@@ -129,28 +156,16 @@ def sagecal_m_step(evis: BlockVisibility, theta, gain=0.25,
             sagecal_fit_gaintable(evis, theta, gain=gain, **kwargs))
 
 
-def sagecal_monitor(iter, thetas, window_index):
-    """Callback for monitoring
-
-    :param iter:
-    :param thetas:
-    :param window_index:
-    :return:
-    """
-    flux = thetas[window_index][0].flux[0, 0]
-    qa = qa_gaintable(thetas[window_index][1])
-    residual = qa.data['residual']
-    print("\t Window %d, flux %s, residual %.3f" % (window_index, str(flux), residual))
-
-
-def sagecal_solve(vis, components, niter=10, tol=1e-8, monitor=sagecal_monitor, **kwargs):
+def sagecal_solve(vis, components, niter=10, tol=1e-8, gain=0.25, callback=None, **kwargs):
     """ Solve
     
+    Solve by iterating, performing E step and M step.
+    
     :param vis:
     :param components:
     :param gaintables:
     :param kwargs:
-    :return:
+    :return: The individual data models and the residual visibility
     """
     thetas = create_sagecal_thetas(vis, components, **kwargs)
 
@@ -159,10 +174,13 @@ def sagecal_solve(vis, components, niter=10, tol=1e-8, monitor=sagecal_monitor,
         evis_all = sagecal_e_all(vis, thetas, **kwargs)
         print("Iteration %d" % (iter))
         for window_index, theta in enumerate(thetas):
-            evis = sagecal_e_step(vis, evis_all, theta, **kwargs)
-            new_theta = sagecal_m_step(evis, theta, **kwargs)
+            evis = sagecal_e_step(vis, evis_all, theta, gain=gain, **kwargs)
+            new_theta = sagecal_m_step(evis, theta, gain=gain, **kwargs)
             new_thetas.append(new_theta)
 
+            if callback is not None:
+                callback(iter, thetas)
+                
             flux = new_theta[0].flux[0, 0]
             qa = qa_gaintable(new_theta[1])
             residual = qa.data['residual']
 
@@ -46,9 +46,9 @@ def solve_gaintable(vis: BlockVisibility, modelvis: BlockVisibility = None, gt=N
         assert isinstance(modelvis, BlockVisibility), modelvis
 
     if phase_only:
-        log.info('solve_gaintable: Solving for phase only')
+        log.debug('solve_gaintable: Solving for phase only')
     else:
-        log.info('solve_gaintable: Solving for complex gain')
+        log.debug('solve_gaintable: Solving for complex gain')
 
     if gt is None:
         log.debug("solve_gaintable: creating new gaintable")
@@ -158,7 +158,6 @@ def solve_antenna_gains_itsubs_scalar(gain, gwt, x, xwt, niter=30, tol=1e-8, pha
         gain = 0.5 * (gain + gainLast)
         change = numpy.max(numpy.abs(gain - gainLast))
         if change < tol:
-            print("solve_antenna_gains_itsubs_scalar: Converged in %d iterations" % iter)
             return gain, gwt, solution_residual_scalar(gain, x, xwt)
 
     return gain, gwt, solution_residual_scalar(gain, x, xwt)
 
@@ -13,6 +13,8 @@
 from astropy.wcs import FITSFixedWarning
 from astropy.wcs import WCS
 from reproject import reproject_interp
+from astropy.wcs.utils import skycoord_to_pixel
+
 
 from arl.data.data_models import Image, QA
 from arl.data.parameters import arl_path
@@ -259,12 +261,15 @@ def qa_image(im, context="") -> QA:
     return qa
 
 
-def show_image(im: Image, fig=None, title: str = '', pol=0, chan=0, cm='rainbow'):
+def show_image(im: Image, fig=None, title: str = '', pol=0, chan=0, cm='rainbow', components=None):
     """ Show an Image with coordinates using matplotlib
 
     :param im:
     :param fig:
     :param title:
+    :param pol: Polarisation
+    :param chan: Channel
+    :param components: Optional components
     :return:
     """
     import matplotlib.pyplot as plt
@@ -282,6 +287,11 @@ def show_image(im: Image, fig=None, title: str = '', pol=0, chan=0, cm='rainbow'
     plt.ylabel('DEC--SIN')
     plt.title(title)
     plt.colorbar()
+    
+    if components is not None:
+        for sc in components:
+            x, y = skycoord_to_pixel(sc.direction, im.wcs, 1, 'wcs')
+            plt.plot(x, y, marker='+')
     return fig
 
 
 
@@ -267,8 +267,7 @@ def predict_skycomponent_visibility(vis: Union[Visibility, BlockVisibility],
 
     elif isinstance(vis, BlockVisibility):
 
-        nchan = vis.nchan
-        npol = vis.npol
+        ntimes, nant, _, nchan, npol = vis.vis.shape
 
         k = numpy.array(vis.frequency) / constants.c.to('m/s').value
 
@@ -280,10 +279,12 @@ def predict_skycomponent_visibility(vis: Union[Visibility, BlockVisibility],
                 flux = convert_pol_frame(flux, comp.polarisation_frame, vis.polarisation_frame)
 
             l, m, n = skycoord_to_lmn(comp.direction, vis.phasecentre)
+            uvw = vis.uvw[..., numpy.newaxis] * k
+            phasor = numpy.ones([ntimes, nant, nant, nchan, npol], dtype='complex')
             for chan in range(nchan):
-                phasor = simulate_point(vis.uvw * k[chan], l, m)
-                for pol in range(npol):
-                    vis.data['vis'][..., chan, pol] += flux[chan, pol] * phasor[...]
+                phasor[:,:,:,chan,:] = simulate_point(uvw[...,chan], l, m)[...,numpy.newaxis]
+                
+            vis.data['vis'][..., :, :] += flux[:, :] * phasor[..., :]
 
     return vis
 
 
@@ -24,23 +24,16 @@ These are well-modeled in ARL. A example of three cycles of a major/minor cycle
 Conceptually we can think of one visibility set, which may be decomposed into different subsets in processing, and
 one image, which also may be subsetted for processing.
 
-Dask.delayed is supported in various ways:
+Dask.delayed is used for distribution:
 
     - :py:mod:`arl.graphs.graphs`: Graphs to perform various types of prediction and inversion of visibility data
     - :py:mod:`arl.graphs.generic_graphs`: Graphs to perform generic image or visibility unary operations
     - :py:mod:`arl.util.graph_support`: Graphs to support testing and simulations
     - :py:mod:`arl.pipelines.graphs`: Graphs to implement the canonical pipelines
 
-Dask.bags is supported as follows:
-
-    - :py:mod:`arl.graphs.bags`: Bags to perform various types of prediction and inversion of visibility data
-    - :py:mod:`arl.util.bag_support`: Bags to support testing and simulations
-    - :py:mod:`arl.pipelines.bags`: Bags to implement the canonical pipelines
-
 In addition there are notebooks in examples/arl.
 
     - imaging-dask: General imaging example using Dask.delayed
-    - imaging-bags: General imaging example using Dask.bag
     - imaging-pipelines: Pipelines run on small LOW observation (core only)
     - calibration-pipelines: Pipelines for self-calibration
 
 
@@ -116,6 +116,12 @@ Solvers
 .. automodule:: arl.calibration.solvers
    :members:
 
+Sagecal
++++++++
+
+.. automodule:: arl.calibration.sagecal
+   :members:
+
 Fourier transform
 -----------------