Code refactor

pyproject.toml

@@ -20,7 +20,8 @@ dependencies = [
     'astroplan',
     'pandas',
     'click',
-    'matplotlib']
+    'matplotlib',
+    'pytest']
 license = {file="LICENSE"}
 requires-python = ">=3.8"
 classifiers = [

src/plumber/parsing.py

@@ -0,0 +1,221 @@
+#!/usr/bin/env python3
+
+"""
+Module to parse the input images and CSV files.
+"""
+
+import os
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+import astropy.units as u
+from astropy.io.registry import IORegistryError
+
+from spectral_cube import StokesSpectralCube
+
+from casatools import image
+ia = image()
+
+import logging
+logger = logging.getLogger(__name__)
+logger.setLevel('INFO')
+
+from typing import Union, List, TypeVar
+Quantity = TypeVar('astropy.units.quantity.Quantity')
+
+
+
+class FileParser:
+    """
+    A base class to parse files of different kinds. Doesn't do much on it's own.
+    """
+
+    def __init__(self):
+        self._filepath = None
+
+    @property
+    def filepath(self) -> str:
+        return self._filepath
+
+
+    @filepath.setter
+    def filepath(self, filepath:str) -> None:
+        """
+        Setter function for the filepath attribute.
+
+        Checks the existence of the filepath before setting it.
+        """
+
+        if not os.path.exists(filepath):
+            err = f"File {filepath} does not exist."
+            raise FileNotFoundError(err)
+
+        self._filepath = filepath
+
+
+
+class CSVParser(FileParser):
+    """
+    Parse the input CSV file(s) and return a list of pandas DataFrames.
+    """
+
+    def __init__(self):
+        self._dataframe_list = None
+        self._frequency_list = None
+        self._nstokes_csv = None
+
+    @property
+    def dataframe_list(self) -> List[pd.DataFrame]:
+        return self._dataframe_list
+
+    @dataframe_list.setter
+    def dataframe_list(self, value: List[pd.DataFrame]) -> None:
+        self._dataframe_list = value
+
+    @property
+    def frequency_list(self):
+        return self._frequency_list
+
+    @frequency_list.setter
+    def frequency_list(self, value: Quantity) -> None:
+        self._frequency_list = value
+
+    @property
+    def nstokes_csv(self):
+        return self._nstokes_csv
+
+    @nstokes_csv.setter
+    def nstokes_csv(self, value: float) -> None:
+        self._nstokes_csv = value
+
+
+    def csv_to_df(self, csv: str) -> pd.DataFrame:
+        """
+        Convert the input CSV into a Pandas DataFrame.
+
+        Inputs:
+        csv             Input CSV filename, str
+
+        Returns:
+        df              Pandas dataframe
+        """
+
+        self.filepath = csv
+        df = pd.read_csv(self.filepath, skipinitialspace=True)
+
+        return df
+
+
+    def get_zcoeffs(self, df: pd.DataFrame, imfreq: Quantity) -> Union[pd.DataFrame, float, int]:
+        """
+        Given the input frequency of the image, returns the Pandas dataframe with
+        the coefficients from the input dataframe corresponding to the input
+        frequency. The frequencies in the input dataframe file are expected to be in
+        MHz.
+
+        `imfreq` should be a list of floats, that are `astropy.units` quantities.
+
+        Inputs:
+        df                  Input pandas DataFrame
+        imfreq              Image frequency in MHz, astropy.units.quantity.Quantity
+
+        Returns:
+        dataframe_list      List of dataframe containing the subset of the CSV file which is
+                            closest to imfreq.
+        frequency_list      List of frequencies that were found in the CSV file,
+                            closest matching the input frequencies.
+        nstokes_csv         The number of Stokes parameters in the input CSV file.
+        """
+
+        self.dataframe_list = []
+        self.frequency_list = []
+
+        if not type(imfreq) is u.quantity.Quantity:
+            raise TypeError(f'Input {imfreq} must be an astropy.units Quantity.')
+
+        self.nstokes_csv = df['#stokes'].unique().size
+        freqs = df['freq'].unique()
+
+        # imfreq is astropy.units, so explicitly convert into MHz before compare.
+        for ifreq in imfreq:
+            idx = np.argmin(np.abs(freqs - ifreq.to(u.MHz).value))
+            zfreq = freqs[idx]
+
+            zdf = df[df['freq'] == zfreq]
+
+            self.dataframe_list.append(zdf)
+            self.frequency_list.append(zfreq)
+
+        return self.dataframe_list, self.frequency_list, self.nstokes_csv
+
+
+class ImageParser(FileParser):
+    """
+    Parse the input FITS/CASA image, and return the data and metadata.
+    """
+
+    def __init__(self):
+        self.imsize = [0,0]
+        self.imfreqs = [-1,]
+        self.is_stokes_cube = False
+
+    @staticmethod
+    def get_telescope(templateim: str) -> str:
+        """
+        Get the telescope name from the image header
+
+        Input:
+        templateim      Name of the input template image, string
+
+        Returns:
+        telescope       Name of the input telescope
+        """
+
+        ia.open(templateim)
+        csys = ia.coordsys().torecord()
+        ia.close()
+
+        try:
+            telescope = csys['telescope']
+            logger.info(f"Telescope is {telescope}.")
+        except KeyError:
+            message = f'Telescope key does not exist in the image header. ' \
+                      f'Please specify via the --telescope command line parameter.'
+            raise KeyError(message)
+
+        return telescope
+
+
+    def parse_image(self, imagename : str) -> Union[npt.NDArray[np.float], npt.NDArray[np.float], bool]:
+        """
+        Parse metadata from the image, such as imsize, central frequency, whether
+        it is a cube etc.
+
+        Inputs:
+        imagename           Name of the input image, string
+
+        Returns:
+        imsize              Size of the direction coordinates in pixels, array
+        reffreq             Reference frequency of the image in MHz, float
+        is_cube             Flag to indicate if the image is a cube, bool
+        """
+
+        try:
+            cube = StokesSpectralCube.read(imagename)
+        except IORegistryError:
+            # Probably CASA image
+            cube = StokesSpectralCube.read(imagename, format='casa_image')
+
+        if len(cube.shape) > 3 and cube.shape[-1] > 1:
+            self.is_stokes_cube = True
+        else:
+            self.is_stokes_cube = False
+
+        shape = cube.shape
+        self.imsize = [shape[0], shape[1]]
+
+        # This is always a list
+        self.imfreqs = cube.stokes_data['I'].spectral_axis
+
+        return self.imsize, self.imfreqs, self.is_stokes_cube

src/plumber/scripts/plumber.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 
 import click
-from plumber.image import parse_image
-from plumber.zernike import get_zcoeffs, zernikeBeam
+from plumber.zernike import zernikeBeam
+from plumber.parsing import CSVParser, ImageParser
+from plumber.telescope import TelescopeInfo
 #from plumber.parang_finder import parallctic_angle
 
 import logging
@@ -28,14 +29,12 @@
 @click.argument('CSV', type=click.File())
 @click.option('-a', '--padding', type=int, default=8, help='Padding factor for aperture, affects smoothness of output beam', show_default=True)
 @click.option('-d', '--dish-dia', type=float, default=None, help='Diameter of the antenna dish. If not one of VLA, ALMA, MeerKAT or GMRT, must be specified.')
-@click.option('-l', '--linear', is_flag=True, help='Specifies if the telescope has linear feeds. If not one of VLA, ALMA, MeerKAT or GMRT, must be specified')
-@click.option('-c', '--circular', is_flag=True, help='Specifies if the telescope has circular feeds. If not one of VLA, ALMA, MeerKAT or GMRT, must be specified')
 @click.option('-I', '--stokesI', is_flag=True, help='Only generate the Stokes I beam, not the full Stokes beams')
 @click.option('-P', '--parallel', is_flag=True, help='Use parallel processing (no MPI) to speed things up')
 @click.option('-p', '--parang', type=float, default=0, help='Parallactic angle at which to generate the PB', show_default=True)
 @click.option('--parang-file', type=click.Path(exists=True), help='Pass a file containing a list of parallactic angles and weights')
 @click.option('--scale', nargs=2, type=float, help='X and Y scaling factors for number of pixels', default=[None, None], show_default=True)
-def main(imagename, csv, padding, dish_dia, linear, circular, stokesi, parallel, parang, parang_file, scale):
+def main(imagename, csv, padding, dish_dia, stokesi, parallel, parang, parang_file, scale):
     """
     Given the input image and the coefficient CSV file, generate the full Stokes
     primary beam at the image centre frequency. If the input is a cube, a
@@ -55,25 +54,46 @@ def main(imagename, csv, padding, dish_dia, linear, circular, stokesi, parallel,
     The eta column in the CSV is optional.
     """
 
-    islinear = None
-    if linear is True:
-        islinear = True
+    csv_parser = CSVParser()
+    image_parser = ImageParser()
+    telescope_info = TelescopeInfo()
 
     #parang = sorted(parang)
     #if len(parang) > 2:
     #    raise ValueError(f"Either pass in a single PA or two values of PA. Currently set to {parang}")
 
-    imsize, imfreq, is_stokes_cube = parse_image(imagename)
-    zdflist, zfreqlist, nstokes = get_zcoeffs(csv, imfreq)
+    telescope = image_parser.get_telescope(imagename)
+    imsize, imfreq, is_stokes_cube = image_parser.parse_image(imagename)
+
+    print(f"Telescope is {telescope}")
+
+    # Given the telescope name, figure out it's properties
+    telescope_info.telescope_name = telescope
+    telescope_info.get_feed_basis()
+
+    if dish_dia is None:
+        telescope_info.get_dish_diameter()
+    else:
+        telescope_info.dish_diameter = dish_dia
+
+    # XXX need to implement - this is a dummy function for now
+    # XXX Need to fix the zb.initialize call below to reflect removing the telescope functionality
+    # This needs to be specifiable on the command line as well, to break MeerKAT L-Band/UHF ambiguity.
+    # And similar ambiguity for other instruments like VLA
+    telescope_info.get_band(imfreq)
+
+
+    df = csv_parser.csv_to_df(csv.name)
+    zdflist, zfreqlist, nstokes = csv_parser.get_zcoeffs(df, imfreq)
 
     logger.info(f"Image is at {imfreq[0].value/1e6:.2f} MHz. Model PB will be generated at {zfreqlist[0]:.2f} MHz")
-    #logger.warn(f"The above frequency is the first channel frequency if the input image is a spectral cube")
 
     zb = zernikeBeam()
 
     for zdf in zdflist:
-        zb.initialize(zdf, imagename, padfac=padding, dish_dia=dish_dia,
-                            islinear=islinear, stokesi=stokesi, parang=parang,
+        zb.initialize(zdf, imagename, padfac=padding, dish_dia=telescope_info.dish_diameter,
+                            telescope = telescope_info.telescope_name,
+                            islinear=telescope_info.is_linear, stokesi=stokesi, parang=parang,
                             parang_file=parang_file, parallel=parallel,
                             scale=scale)