Refactor helper-classes using attrs, improve & u-test extract_lsm.

This began as a basic refactoring of SkySource / SkyRegion using attrs, which
cuts down on verbosity / duplication and has the added bonus of providing prettyprinting.

This led to work on ensuring the extract_lsm script actually works through unit-testing,
including several minor improvements to consistency of output:
* Only the catalog is output to stdout, so piping will give the same results as specifying a filename.
* An empty catalog (with headers) is output even if there are no matches.

Author: timstaley

src/fastimgproto/scripts/extract_lsm.py

@@ -6,21 +6,22 @@
 from fastimgproto.skymodel.helpers import SkyRegion
 from fastimgproto.skymodel.extraction import (
     sumss_file_to_dataframe,
-    lsm_extract
+    lsm_extract,
+    SumssSrc,
 )
 import click
 import sys
 import csv
 
+
 def write_catalog(src_list, filehandle):
-    if src_list:
-        fieldnames = src_list[0].to_ordereddict().keys()
-        dw = csv.DictWriter(filehandle,
-                            fieldnames=fieldnames,
-                            delimiter='\t')
-        dw.writeheader()
-        for src in src_list:
-            dw.writerow(src.to_ordereddict())
+    dw = csv.DictWriter(filehandle,
+                        fieldnames=SumssSrc._list_dictkeys(),
+                        delimiter='\t')
+    dw.writeheader()
+    for src in src_list:
+        dw.writerow(src.to_ordereddict())
+
 
 @click.command()
 @click.argument('ra', type=click.FLOAT)
@@ -42,6 +43,10 @@ def cli(ra, dec, radius, outfile, vcat, catalog_file):
 
         fastimg_extract_lsm --vcat var.csv -- 189.2 -45.6 1.5 lsm.csv
 
+    Note: arguments (RA, dec, radius, [outfile]) are separated from options by the
+    separator `--`, this avoids mistakenly trying to parse a negative
+    declination as an option flag.
+
     \b
     Args:
     - ra (float): RA of centre (decimal degrees, J2000)
@@ -64,8 +69,8 @@ def cli(ra, dec, radius, outfile, vcat, catalog_file):
 
     """
 
-    field_of_view = SkyRegion(SkyCoord(ra*u.deg, dec*u.deg),
-                              Angle(radius*u.deg))
+    field_of_view = SkyRegion(SkyCoord(ra * u.deg, dec * u.deg),
+                              Angle(radius * u.deg))
 
     if catalog_file is None:
         catalog_file = download_file(
@@ -80,9 +85,9 @@ def cli(ra, dec, radius, outfile, vcat, catalog_file):
     if vcat:
         write_catalog(variable_cat, vcat)
 
-    click.echo("{} sources matched".format(len(lsm_cat)))
-    click.echo("Of which {} variable.".format(len(variable_cat)))
-    sys.exit(0)
+    click.echo("{} sources matched, of which {} variable".format(
+        len(lsm_cat), len(variable_cat)),
+        err=True)
 
 
 if __name__ == '__main__':

src/fastimgproto/skymodel/extraction.py

@@ -36,7 +36,7 @@
 
 class SumssSrc(SkySource):
     """
-    Represents an entry from the SUMSS catalog.
+    Represents a subset of data from a row/source in the SUMSS catalog.
 
     (Sydney University Molonglo Southern Sky Survey)
     """
@@ -51,23 +51,37 @@ def __init__(self, row):
         self.peak_flux_err = row.peak_flux_err_mjy * u.mJy
         self.variable = row.variable
         self.position_err = PositionError(
-            ra_err=Angle(row.ra_err_arcsec * u.arcsecond),
-            dec_err=Angle(row.dec_err_arcsec * u.arcsecond)
+            ra=Angle(row.ra_err_arcsec * u.arcsecond),
+            dec=Angle(row.dec_err_arcsec * u.arcsecond)
         )
 
         super(SumssSrc, self).__init__(
             position=position, flux=self.peak_flux, variable=self.variable)
-
+    
+    class DictKeys(object):
+        ra = 'ra'
+        dec = 'dec'
+        ra_err = 'ra_err'
+        dec_err = 'dec_err'
+        peak_flux = 'peak_flux_mjy'
+        peak_flux_err = 'peak_flux_err_mjy'
+        variable = 'variable'
+
+    @classmethod
+    def _list_dictkeys(cls):
+        innards = vars(cls.DictKeys)
+        return  [innards[k] for k in innards if not k.startswith('__')]
 
     def to_ordereddict(self):
         od = OrderedDict()
-        od['ra'] = self.position.ra.deg
-        od['dec'] = self.position.dec.deg
-        od['ra_err'] = self.position_err.ra.deg
-        od['dec_err'] = self.position_err.dec.deg
-        od['peak_flux_mjy'] = self.peak_flux.to(u.mJy).value
-        od['peak_flux_err_mjy'] = self.peak_flux_err.to(u.mJy).value
-        od['variable'] = self.variable
+        keys = SumssSrc.DictKeys
+        od[keys.ra] = self.position.ra.deg
+        od[keys.dec] = self.position.dec.deg
+        od[keys.ra_err] = self.position_err.ra.deg
+        od[keys.dec_err] = self.position_err.dec.deg
+        od[keys.peak_flux] = self.peak_flux.to(u.mJy).value
+        od[keys.peak_flux_err] = self.peak_flux_err.to(u.mJy).value
+        od[keys.variable] = self.variable
         return od
 
 

src/fastimgproto/skymodel/helpers.py

@@ -4,49 +4,42 @@
 
 from astropy.coordinates import Angle, SkyCoord
 import astropy.units as u
+import attr.validators
+from attr import attrs, attrib
 
 
+@attrs
 class SkyRegion(object):
     """
     Defines a circular region of the sky.
     """
-
-    def __init__(self, centre, radius):
-        assert isinstance(centre, SkyCoord)
-        assert isinstance(radius, Angle)
-        self.centre = centre
-        self.radius = radius
+    centre = attrib(validator=attr.validators.instance_of(SkyCoord))
+    radius = attrib(validator=attr.validators.instance_of(Angle))
 
 
+@attrs
 class PositionError(object):
-    def __init__(self, ra_err, dec_err):
-        assert isinstance(ra_err, Angle)
-        assert isinstance(dec_err, Angle)
-        self.ra = ra_err
-        self.dec = dec_err
+    """
+    Represent positional uncertainty.
 
-    def __str__(self):
-        return "<PositionError: (ra, dec) in deg ({}, {})>".format(
-            self.ra.deg, self.dec.deg)
+    (Mainly used for representing entries in the SUMSS catalog.)
+    """
+    ra = attrib(validator=attr.validators.instance_of(Angle))
+    dec = attrib(validator=attr.validators.instance_of(Angle))
 
 
+@attrs
 class SkySource(object):
     """
     Basic point source w/ flux modelled at a single frequency
 
-    Args:
+    Attributes:
         position (astropy.coordinates.SkyCoord): Sky-coordinates of source.
         flux (astropy.units.Quantity): Source flux at measured frequency.
         frequency (astropy.units.Quantity): Measurement frequency.
         variable (bool): 'Known variable' flag.
     """
-    def __init__(self, position, flux,
-                 frequency=2.5*u.GHz, variable=False):
-        assert isinstance(position, SkyCoord)
-        #This will raise if the flux has wrong units or no units:
-        flux.to(u.Jy)
-        frequency.to(u.Hz)
-        self.position = position
-        self.flux = flux
-        self.variable = variable
-        self.frequency = frequency
+    position = attrib(validator=attr.validators.instance_of(SkyCoord))
+    flux = attrib(convert=lambda x: x.to(u.Jy))
+    frequency = attrib(default=2.5 * u.GHz, convert=lambda x: x.to(u.GHz))
+    variable = attrib(default=False)

tests/test_scripts/test_extract_lsm.py

@@ -0,0 +1,35 @@
+from __future__ import print_function
+import json
+import numpy as np
+
+from click.testing import CliRunner
+from fastimgproto.resources.testdata import simple_vis_npz_filepath
+from fastimgproto.scripts.extract_lsm import cli as extract_lsm_cli
+from fastimgproto.skymodel.extraction import SumssSrc
+import csv
+
+
+def test_extract_lsm():
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        catalog_output_filename = 'foo.csv'
+        args = '-- 189.2 -45.6 0.2'.split()
+        # Invoke with output to stdout
+        result = runner.invoke(extract_lsm_cli,
+                               args=args)
+        assert result.exit_code == 0
+
+        # And output to TSV file:
+        args.append(catalog_output_filename)
+        result = runner.invoke(extract_lsm_cli,
+                               args=args)
+        print(result.output)
+
+        with open(catalog_output_filename, 'rb') as tsvfile:
+            dr=csv.DictReader(tsvfile, delimiter='\t')
+            rows = [r for r in dr]
+
+        assert len(rows)==1
+
+        for key in SumssSrc._list_dictkeys():
+            assert key in rows[0].keys()