UV coverage and plot. Takes a while. Needs a lot of RAM.

ms_uvcoverage.py

#! /usr/bin/env python

from __future__ import print_function, division

import numpy as np
import sys

from astropy.constants import c; c = c.value

import matplotlib as mpl
mpl.use("Agg")

from matplotlib import pyplot as plt

mpl.rcParams["lines.linewidth"] = 1.2
mpl.rcParams["xtick.direction"] = "in"
mpl.rcParams["ytick.direction"] = "in"
mpl.rcParams["xtick.major.width"] = 1.2
mpl.rcParams["ytick.major.width"] = 1.2
mpl.rcParams["xtick.minor.width"] = 1.2
mpl.rcParams["ytick.minor.width"] = 1.2


from casacore.tables import table, taql

from argparse import ArgumentParser

def get_uv(ms, max_uv=1e9, inc=1):
    """
    """

    mset = table(ms, readonly=True, ack=False)
    filtered = taql("select * from $ms where not FLAG_ROW and ANTENNA1 <> ANTENNA2")

    spw = table(ms+"/SPECTRAL_WINDOW", readonly=True, ack=False)
    uvw = filtered.getcol("UVW")
    # data = filtered.getcol("DATA")
    # flags = filtered.getcol("FLAG")
    # data[flags] = np.nan
    bnd = spw.getcol("CHAN_FREQ")[0]
    print(bnd.shape)
    wln = c / bnd
    ref = spw.getcol("REF_FREQUENCY")[0]



    with open(ms.replace(".ms", "_uvw.txt"), "w+") as f:
        for i in range(len(uvw)):
            if uvw[i][0] != 0.0 and uvw[i][1] != 0.0:
                u = uvw[:, 0][i] / wln
                # print(u.shape)
                v = uvw[:, 1][i] / wln
                for j in range(0, len(u), inc):
                    if abs(u[j]) < max_uv and abs(v[j]) < max_uv:
                        f.write("{},{}\n".format(u[j], v[j]))

def plot_uv(uv, limits=None, inc=1, color="crimson", label=""):
    """
    """

    fig = plt.figure(figsize=(8, 8))
    ax = plt.axes([0.01, 0.01, 0.98, 0.98])

    for ms in uv:

        uvw = np.genfromtxt(ms.replace(".ms", "_uvw.txt"), delimiter=",", names="u,v", dtype=np.float16)
        ax.plot(uvw["u"][::inc], uvw["v"][::inc], ls="", marker=".", ms=2., c=color, zorder=0)
        ax.plot(-uvw["u"][::inc], -uvw["v"][::inc], ls="", marker=".", ms=2., c=color, zorder=0)

    ax.tick_params(axis="x", which="both", bottom=True, top=True, length=7., zorder=-1)
    ax.tick_params(axis="y", which="both", left=True, right=True, length=7., zorder=-1)
    # ax.tick_params(which="minor", length=5., zorder=-1)

    ax.axes.xaxis.set_ticklabels([])
    ax.axes.yaxis.set_ticklabels([])

    if limits is not None:
        ax.set_xlim(limits)
        ax.set_ylim(limits)

    ax.text(0.5, 0.05, label, ha="center", va="center",
            fontsize=24., transform=ax.transAxes)


    fig.savefig(uv[0].replace(".", "_").replace("_txt", ".png"), bbox_inches="tight", dpi=300)


def main():
    """
    """

    ps = ArgumentParser()
    ps.add_argument("ms", nargs="*")
    ps.add_argument("-m", "--max-uv", "--max_uv", dest="max_uv", default=1.e9,
        type=float)
    ps.add_argument("-i", "--inc", default=1, type=int)
    ps.add_argument("--plot", action="store_true")
    ps.add_argument("--limits", default=None, nargs=2, type=float)
    ps.add_argument("--color", default="crimson")
    ps.add_argument("--label", default="")
    ps.add_argument("--both", action="store_true")


    args = ps.parse_args()

    if not args.plot or args.both:
        args.ms = args.ms[0]  # only do the first one
        get_uv(args.ms, args.max_uv, args.inc)  
        args.inc = 1
        args.ms = [args.ms]

    if args.plot:
        plot_uv(args.ms, args.limits, args.inc, args.color, args.label)



if __name__ == "__main__":
    main()