Add code to subtract sources

lofarimaging/hdf5util.py

@@ -25,6 +25,10 @@ Per observation, the following attributes are used:
  * extent_lonlat   Extent of image in longitude and latitude (w.r.t. WGS84 ellipsoid)
                    [lon_min, lon_max, lat_min, lat_max]
  * height          Height (w.r.t. station phase centre) of the image plane in metres
+ * subtracted      List of sources subtracted from visibilities
+
+ For sky images, the following attributes are used:
+ * subtracted      List of sources subtracted from visibilities
 """
 
 import datetime
@@ -61,7 +65,8 @@ def get_new_obsname(h5file: h5py.File):
 def write_hdf5(filename: str, xst_data: np.ndarray, visibilities: np.ndarray, sky_img: np.ndarray,
                ground_img: np.ndarray, station_name: str, subband: int, rcu_mode: int, frequency: float,
                obstime: datetime.datetime, extent: List[float], extent_lonlat: List[float],
-               height: float, bodies_lmn: Dict[str, Tuple[float]], calibration_info: Dict[str, str]):
+               height: float, bodies_lmn: Dict[str, Tuple[float]], calibration_info: Dict[str, str],
+               subtracted: List[str]):
     """
     Write an HDF5 file with all data
 
@@ -81,6 +86,7 @@ def write_hdf5(filename: str, xst_data: np.ndarray, visibilities: np.ndarray, sk
         height (float): Height of ground image (in metres)
         bodies_lmn (Dict[str, Tuple[float]]): lmn coordinates of some objects on the sky
         calibration_info (Dict[str, str]): Calibration metadata
+        subtracted (List[str]): List of sources subtracted
 
     Returns:
         None
@@ -91,7 +97,7 @@ def write_hdf5(filename: str, xst_data: np.ndarray, visibilities: np.ndarray, sk
         >>> write_hdf5("test/test.h5", xst_data, visibilities, sky_img, ground_img, "DE603", \
                        297, 3, 150e6, datetime.datetime.now(), [-150, 150, -150, 150], \
                        [11.709, 11.713, 50.978, 50.981], 1.5, {'Cas A': (0.3, 0.5, 0.2)}, \
-                       {'CalTableHeader.Calibration.Date': '20181214'})
+                       {'CalTableHeader.Calibration.Date': '20181214'}, ["Cas A", "Sun"])
     """
     short_station_name = station_name[:5]
 
@@ -110,13 +116,15 @@ def write_hdf5(filename: str, xst_data: np.ndarray, visibilities: np.ndarray, sk
         obs_group.create_dataset("calibrated_data", data=visibilities, compression="gzip")
         for key, value in calibration_info.items():
             obs_group["calibrated_data"].attrs[key] = value
-        obs_group.create_dataset("sky_img", data=sky_img, compression="gzip")
+        dataset_sky_img = obs_group.create_dataset("sky_img", data=sky_img, compression="gzip")
+        dataset_sky_img.attrs["subtracted"] = subtracted
 
         ground_img_group = obs_group.create_group("ground_images")
         dataset_ground_img = ground_img_group.create_dataset("ground_img000", data=ground_img, compression="gzip")
         dataset_ground_img.attrs["extent"] = extent
         dataset_ground_img.attrs["extent_lonlat"] = extent_lonlat
         dataset_ground_img.attrs["height"] = height
+        dataset_ground_img.attrs["subtracted"] = str(subtracted)
 
 
 def merge_hdf5(src_filename: str, dest_filename: str, obslist: List[str] = None):

lofarimaging/lofarimaging.py

@@ -1,5 +1,6 @@
 """Functions for working with LOFAR single station data"""
 
+from typing import Dict, List
 import numpy as np
 from numpy.linalg import norm, lstsq
 import numexpr as ne
@@ -7,7 +8,8 @@ import numba
 from astropy.coordinates import SkyCoord, SkyOffsetFrame, CartesianRepresentation
 
 
-__all__ = ["nearfield_imager", "sky_imager", "ground_imager", "skycoord_to_lmn", "calibrate", "simulate_sky_source"]
+__all__ = ["nearfield_imager", "sky_imager", "ground_imager", "skycoord_to_lmn", "calibrate", "simulate_sky_source",
+           "subtract_sources"]
 
 __version__ = "1.5.0"
 SPEED_OF_LIGHT = 299792458.0
@@ -180,13 +182,25 @@ def simulate_sky_source(lmn_coord: np.array, baselines: np.array, freq: float):
     return np.exp(2j * np.pi * freq * baselines.dot(np.array(lmn_coord)) / SPEED_OF_LIGHT)
 
 
-def simulate_nearfield_source(pqr_coord: np.array, baselines: np.array, freq: float):
+def subtract_sources(vis: np.array, baselines: np.array, freq: float, lmn_dict: Dict[str, np.array],
+                     sources=["Cas A", "Cyg A", "Sun"]):
     """
-    Simulate visibilities for a nearfield source
+    Subtract sky sources from visibilities
 
     Args:
-        pqr_coord (np.array): l, m, n coordinate
+        vis (np.array): visibility matrix, shape [n_ant, n_ant]
+        lmn_dict (Dict[str, np.array]): dictionary with lmn coordinates
         baselines (np.array): baseline distances in metres, shape (n_ant, n_ant)
         freq (float): Frequency in Hz
+        sources (List[str]): list with source names to subtract (should all be in lmn_dict).
+                             Default ["Cas A", "Sun"]
+
+    Returns:
+        vis (np.array): visibility matrix with sources subtracted
     """
-    pass
+    modelvis = [simulate_sky_source(lmn_dict[srcname], baselines, freq) for srcname in lmn_dict
+                if srcname in sources]
+
+    residual, _ = calibrate(vis, modelvis)
+
+    return residual

lofarimaging/singlestationutil.py

@@ -28,7 +28,7 @@ from lofarantpos.db import LofarAntennaDatabase
 import lofarantpos
 
 from .maputil import get_map, make_leaflet_map
-from .lofarimaging import nearfield_imager, sky_imager, skycoord_to_lmn
+from .lofarimaging import nearfield_imager, sky_imager, skycoord_to_lmn, subtract_sources
 from .hdf5util import write_hdf5
 
 
@@ -565,7 +565,8 @@ def make_xst_plots(xst_data: np.ndarray,
                    sky_only: bool = False,
                    opacity: float = 0.6,
                    hdf5_filename: str = None,
-                   outputpath: str = "results"):
+                   outputpath: str = "results",
+                   subtract: List[str] = None):
     """
     Create sky and ground plots for an XST file
 
@@ -584,6 +585,7 @@ 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:
@@ -594,7 +596,8 @@ def make_xst_plots(xst_data: np.ndarray,
         >>> obstime = datetime.datetime(2017, 7, 20, 9, 58, 16)
         >>> sky_fig, ground_fig, leafletmap = make_xst_plots(xst_data, "DE603", obstime, 297, \
                                                              3, caltable_dir="test/CalTables", \
-                                                             hdf5_filename="test/test.h5")
+                                                             hdf5_filename="test/test.h5", \
+                                                             subtract=["Cas A", "Sun"])
         Maximum at -6m east, 70m north of station center (lat/long 50.97998, 11.71118)
 
         >>> type(leafletmap)
@@ -650,10 +653,6 @@ def make_xst_plots(xst_data: np.ndarray,
 
     baselines = station_xyz[:, np.newaxis, :] - station_xyz[np.newaxis, :, :]
 
-    # Fourier transform
-    # visibilities = cube_xx[2,:,:]
-    sky_img = sky_imager(visibilities_stokesI, baselines, freq, npix_l, npix_m)
-
     obstime_astropy = Time(obstime)
     # Determine positions of Cas A and Cyg A
     station_earthlocation = EarthLocation.from_geocentric(*(db.phase_centres[station_name] * u.m))
@@ -679,11 +678,20 @@ def make_xst_plots(xst_data: np.ndarray,
         if body_coord.transform_to(AltAz(location=station_earthlocation, obstime=obstime_astropy)).alt > 0:
             marked_bodies_lmn[body_name] = skycoord_to_lmn(marked_bodies[body_name], zenith)
 
+    if subtract is not None:
+        visibilities_stokesI = subtract_sources(visibilities_stokesI, baselines, freq, marked_bodies_lmn, subtract)
+
+    sky_img = sky_imager(visibilities_stokesI, baselines, freq, npix_l, npix_m)
+
     marked_bodies_lmn_only3 = {k: v for (k, v) in marked_bodies_lmn.items() if k in ('Cas A', 'Cyg A', 'Sun')}
 
     # Plot the resulting sky image
     sky_fig = plt.figure(figsize=(10, 10))
 
+    if sky_vmin is None and subtract is not None:
+        # Tendency to oversubtract, we don't want to see that
+        sky_vmin = np.quantile(sky_img, 0.05)
+
     make_sky_plot(sky_img, marked_bodies_lmn_only3, title=f"Sky image for {station_name}",
                   subtitle=f"SB {subband} ({freq / 1e6:.1f} MHz), {str(obstime)[:16]}", fig=sky_fig,
                   vmin=sky_vmin, vmax=sky_vmax)
@@ -731,10 +739,9 @@ def make_xst_plots(xst_data: np.ndarray,
     [maxpixel_p, maxpixel_q, _] = pqr_to_xyz.T @ np.array([maxpixel_x, maxpixel_y, height])
     maxpixel_lon, maxpixel_lat, _ = lofargeotiff.pqr_to_longlatheight([maxpixel_p, maxpixel_q], station_name)
 
-    # Show location of maximum if not at the image border
+    # Show location of maximum
-    if 2 < maxpixel_xpix < npix_x - 2 and 2 < maxpixel_ypix < npix_y - 2:
+    print(f"Maximum at {maxpixel_x:.0f}m east, {maxpixel_y:.0f}m north of station center " +
-        print(f"Maximum at {maxpixel_x:.0f}m east, {maxpixel_y:.0f}m north of station center " +
+          f"(lat/long {maxpixel_lat:.5f}, {maxpixel_lon:.5f})")
-              f"(lat/long {maxpixel_lat:.5f}, {maxpixel_lon:.5f})")
 
     tags = {"generated_with": f"lofarimaging v{__version__}",
             "subband": subband,
@@ -752,7 +759,7 @@ def make_xst_plots(xst_data: np.ndarray,
     leaflet_map = make_leaflet_map(folium_overlay, lon_center, lat_center, lon_min, lat_min, lon_max, lat_max)
 
     write_hdf5(hdf5_filename, xst_data, visibilities, sky_img, ground_img, station_name, subband, rcu_mode,
-               freq, obstime, extent, extent_lonlat, height, marked_bodies_lmn, calibration_info)
+               freq, obstime, extent, extent_lonlat, height, marked_bodies_lmn, calibration_info, subtract)
 
     return sky_fig, ground_fig, leaflet_map