diff --git a/pyproject.toml b/pyproject.toml index bd0e4a7..899af01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,8 @@ dependencies = [ 'astroplan', 'pandas', 'click', - 'matplotlib'] + 'matplotlib', + 'pytest'] license = {file="LICENSE"} requires-python = ">=3.8" classifiers = [ diff --git a/src/plumber/image.py b/src/plumber/image.py deleted file mode 100644 index 4bf7bb2..0000000 --- a/src/plumber/image.py +++ /dev/null @@ -1,48 +0,0 @@ -#!/usr/bin/env python3 - -import spectral_cube -print(spectral_cube.__path__) - -from spectral_cube import SpectralCube, StokesSpectralCube -from casatools import image -ia = image() - -from astropy.io.registry import IORegistryError - -import logging -logger = logging.getLogger(__name__) -logging.basicConfig(format="%(asctime)-15s %(levelname)s: %(message)s", - level=logging.INFO) - -def parse_image(imagename): - """ - 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: - is_stokes_cube = True - else: - is_stokes_cube = False - - shape = cube.shape - imsize = [shape[0], shape[1]] - - # This is always a list - imfreqs = cube.stokes_data['I'].spectral_axis - - return imsize, imfreqs, is_stokes_cube diff --git a/src/plumber/parsing.py b/src/plumber/parsing.py new file mode 100644 index 0000000..4d1c8a0 --- /dev/null +++ b/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 diff --git a/src/plumber/scripts/plumber.py b/src/plumber/scripts/plumber.py index c0d7369..4d0574c 100644 --- a/src/plumber/scripts/plumber.py +++ b/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) diff --git a/src/plumber/telescope.py b/src/plumber/telescope.py new file mode 100644 index 0000000..391075e --- /dev/null +++ b/src/plumber/telescope.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 + +""" +Module that holds functionality related to telescopes. The class and functions +to return the metadata relating to the different instruments, and the class and +functions to setup the aperture illumination, and obtain the model PB. +""" + +import numpy as np +import numpy.typing as npt +from typing import Union + +import logging +logger = logging.getLogger(__name__) +logger.setLevel('INFO') + +class TelescopeInfo(): + """ + Given the telescope name, return various parameters related to the + instrument. + """ + + def __init__(self): + self._telescope_name = '' + self.known_telescopes = ['MeerKAT', 'VLA', 'JVLA', 'EVLA', 'ALMA', 'uGMRT', 'GMRT'] + self.is_known = True + self.is_linear = False + self.dish_diameter = [-1, -1] + self.band = '' + + + @property + def telescope_name(self) -> str: + return self._telescope_name + + @telescope_name.setter + def telescope_name(self, name:str) : + # If telescope name is unknown, print warning and set flag + if not any([name.lower() in tname.lower() for tname in self.known_telescopes]): + self.is_known = False + logger.warning(f'Telescope name {name} is unknown. Using default parameters everywhere. ' + 'Please pass in specific values through the command line if necessary.') + + # XXX Set the telescope name to a standard string here, based on the input + # XXX Makes parsing it again much easier. + self._telescope_name = name + + + def get_feed_basis(self) -> None: + """ + Get the feed basis from the telescope name. + """ + + # XXX : Need to take into acount uGMRT linear feeds at L-Band and VLA + # linear feeds at P Band + + if self.telescope_name == '': + raise ValueError('Please set the telescope name to determine the band.') + + if not self.is_known: + logger.warning(f"Using default linear basis for unknown telescope {self.telescope_name}") + self.is_linear = False + return + + if "vla" in self.telescope_name.lower(): + self.is_linear = False + elif "meerkat" in self.telescope_name.lower(): + self.is_linear = True + elif "alma" in self.telescope_name.lower(): + self.is_linear = True + elif "gmrt" in self.telescope_name.lower(): + self.is_linear = False + else: # This should never happen + raise ValueError("Unable to determine feed basis, something went wrong.") + + basis = "linear" if self.is_linear else "circular" + logger.info(f"Using {basis} basis for telescope {self.telescope_name}") + return basis + + + def get_dish_diameter(self) -> None: + """ + Get the dish diameter from the telescope name. + """ + + if self.telescope_name == '': + raise ValueError('Please set the telescope name to determine the band.') + + if not self.is_known: + logger.warning(f"Using default dish diameter of 1m for unknown telescope {self.telescope_name}") + self.dish_diameter = [1,1] + return + + #if self.dish_diameter is not None: + # logger.info(f"Overriding automatic determination of telescope dish " + # "diameter. Using the input of {self.dish_diameter} m") + # return + + if 'vla' in self.telescope_name.lower(): + self.dish_diameter = [25, 25] + elif 'meerkat' in self.telescope_name.lower(): + self.dish_diameter = [13.5, 13.5] + elif 'alma' in self.telescope_name.lower(): + self.dish_diameter = [12, 12] + elif 'gmrt' in self.telescope_name.lower(): + self.dish_diameter = [45, 45] + else: # This should never happen. + raise ValueError("Unable to determine telescope type. Something went wrong.") + + return self.dish_diameter + logger.info(f"Using dish diameter of ({self.dish_diameter[0]}m, {self.dish_diameter[1]}m) for telescope {self.telescope_name}.") + + + def get_band(self, inp_frequency: Union[float,npt.NDArray[float]]) -> None: + """ + Given the input frequency/frequencies, determine the band of the telescope. + + This is necessary, because the PB models can be different for different bands, even at overlapping frequencies. + This sets the TelescopeInfo().band attribute. + + Inputs: + inp_frequency The input frequency, astropy.quantity.Quantity + + Returns: + None + """ + + if self.telescope_name == '': + raise ValueError('Please set the telescope name to determine the band.') + + if self.telescope_name == 'MeerKAT': + pass diff --git a/src/plumber/zernike.py b/src/plumber/zernike.py index 502b06a..f99ef3e 100644 --- a/src/plumber/zernike.py +++ b/src/plumber/zernike.py @@ -25,55 +25,16 @@ logger.setLevel('INFO') from plumber.misc import wipe_file, make_unique -from plumber.image import parse_image +from plumber.parsing import ImageParser from casatasks import immath, imregrid from casatools import image ia = image() -# Use up to 4 concunrrent processes +# Use up to 4 concunrrent processes -- one for each Stokes NCPU = min(multiprocessing.cpu_count(), 4) -def get_zcoeffs(csv: str, imfreq: float) -> Union[pd.Dataframe, float, int]: - """ - Given the input frequency of the image, returns the Pandas dataframe with - the coefficients corresponding to the input frequency. The frequencies in - the input CSV file are expected to be in MHz. - - Inputs: - csv Input CSV filename, string - imfreq Image frequency in MHz, float - - Returns: - zdf Dataframe containing the subset of the CSV file which is closest - to imfreq. - zfreq Frequency of the coefficients in MHz, float - nstokes Number of stokes in the CSV, integer - """ - - zdflist = [] - freqlist = [] - - df = pd.read_csv(csv, skipinitialspace=True) - nstokes = df['#stokes'].unique().size - freqs = df['freq'].unique() - - #imfreq is an array astropy.units, so convert to right unit and grab - #numerical value - for ifreq in imfreq: - idx = np.argmin(np.abs(freqs - ifreq.to(u.MHz).value)) - zfreq = freqs[idx] - - zdf = df[df['freq'] == zfreq] - - zdflist.append(zdf) - freqlist.append(zfreq) - - return zdflist, freqlist, nstokes - - - class zernikeBeam(): """ Generate the model beam on a 2D pixel grid. @@ -102,7 +63,7 @@ def __init__(self): self.oversamp = 20 - def initialize(self, df, templateim, padfac=8, dish_dia=[], islinear=None, stokesi=False, parang=None, parang_file=None, parallel=False, scale=None): + def initialize(self, df, templateim, padfac=8, dish_dia=[], telescope=None, islinear=None, stokesi=False, parang=None, parang_file=None, parallel=False, scale=None): """ Initialize the class with an input DataFrame, and optionally padding factor for the FFT. @@ -119,13 +80,11 @@ def initialize(self, df, templateim, padfac=8, dish_dia=[], islinear=None, stoke """ self.df = df - self.telescope = self.get_telescope(templateim) + self.telescope = telescope self.eta = [1., 1.] self.dish_dia = dish_dia - self.get_dish_diameter() self.islinear = islinear - self.get_feed_basis() self.padfac = int(padfac) self.parallel = parallel @@ -145,79 +104,6 @@ def initialize(self, df, templateim, padfac=8, dish_dia=[], islinear=None, stoke self.get_npix_aperture(templateim) - def get_telescope(self, 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() - - telescope = csys['telescope'] - logger.debug(f"Telescope is {telescope}.") - - return csys['telescope'] - - - def get_feed_basis(self) -> None: - """ - Get the feed basis, from the telescope name. - """ - - if self.islinear is not None: - outstr = 'circular' if self.islinear == False else 'linear' - logger.info(f"Overriding automatic determination of feed basis, using user supplied value of {outstr}") - - if "vla" in self.telescope.lower(): - self.islinear = False - elif "meerkat" in self.telescope.lower(): - self.islinear = True - elif "alma" in self.telescope.lower(): - self.islinear = True - elif "gmrt" in self.telescope.lower() and self.freq < 9e2: - self.islinear = False - elif "gmrt" in self.telescope.lower() and self.freq >= 9e2: - self.islinear = True - else: - raise ValueError("Unable to determine feed basis, unknown telescope. " - "Please pass in the feed basis via the islinear paramter to " - ".initialize()") - - - def get_dish_diameter(self) -> None: - """ - Get the dish diameter for different instruments. - - Returns: - None, sets self.dish_dia in place - """ - - if self.dish_dia is not None: - logger.info(f"Overriding automatic determination of telescope dish " - "diameter. Using the input of {self.dish_dia} m") - return - - if 'vla' in self.telescope.lower(): - self.dish_dia = [25, 25] - elif 'meerkat' in self.telescope.lower(): - #self.dish_dia = 13.5 - self.dish_dia = [13.5, 13.5] - elif 'alma' in self.telescope.lower(): - self.dish_dia = [12, 12] - elif 'gmrt' in self.telescope.lower(): - self.dish_dia = [45, 45] - else: - raise ValueError('Unknown telescope type. Please initialize the ' - 'class with the dish_diameter value in metres.') - - logger.info(f"Using dish diameter of ({self.dish_dia[0]}m, {self.dish_dia[1]}m) for telescope {self.telescope}.") @@ -263,7 +149,8 @@ def get_npix_aperture(self, templateim: str) -> Union[int, float]: cdelt The pixel delta in lambda, float """ - imsize, imfreq, is_stokes_cube = parse_image(templateim) + image_parser = ImageParser() + imsize, imfreq, is_stokes_cube = image_parser.parse_image(templateim) freq = imfreq[0].value ia.open(templateim) diff --git a/tests/test_metadata.py b/tests/test_metadata.py new file mode 100644 index 0000000..d0e193f --- /dev/null +++ b/tests/test_metadata.py @@ -0,0 +1,57 @@ +#! /usr/bin/env python + +""" +Module to test all the metadata-related functionality in plumber. + +Currently tests the TelescopeInfo class, but will be expanded to support other +metadata classes. +""" + +from plumber.telescope import TelescopeInfo +import numpy as np + + +def test_meerkat_info(): + ti = TelescopeInfo() + ti.telescope_name = 'MeerKAT' + + feed_basis = ti.get_feed_basis() + dish_dia = ti.get_dish_diameter() + + print(f"Feed basis {feed_basis}, Dish_dia {dish_dia}") + + np.testing.assert_string_equal(feed_basis, "linear") + np.testing.assert_array_equal(dish_dia, [13.5,13.5]) + + +def test_vla_info(): + ti = TelescopeInfo() + ti.telescope_name = 'VLA' + + feed_basis = ti.get_feed_basis() + dish_dia = ti.get_dish_diameter() + + np.testing.assert_string_equal(feed_basis, "circular") + np.testing.assert_array_equal(dish_dia, [25,25]) + + +def test_alma_info(): + ti = TelescopeInfo() + ti.telescope_name = 'ALMA' + + feed_basis = ti.get_feed_basis() + dish_dia = ti.get_dish_diameter() + + np.testing.assert_string_equal(feed_basis, "linear") + np.testing.assert_array_equal(dish_dia, [12,12]) + + +def test_gmrt_info(): + ti = TelescopeInfo() + ti.telescope_name = 'GMRT' + + feed_basis = ti.get_feed_basis() + dish_dia = ti.get_dish_diameter() + + np.testing.assert_string_equal(feed_basis, "circular") + np.testing.assert_array_equal(dish_dia, [45,45]) diff --git a/tests/test_parsing.py b/tests/test_parsing.py new file mode 100644 index 0000000..37e29e4 --- /dev/null +++ b/tests/test_parsing.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 + +from plumber.parsing import CSVParser + +import pdb +import os +import numpy as np +import pandas as pd + +import astropy.units as u + +import pytest + +mock_csv = """ +#stokes,freq,ind,real,imag,etax,etay +9,1000,0,1,1,1,1 +10,1000,0,1,1,1,1 +11,1000,0,1,1,1,1 +12,1000,0,1,1,1,1 +""" + +def test_csv_to_df(): + """ + Test that the CSV to DataFrame parsing works correctly. + """ + + tmpfile = 'tmp_parsing.txt' + with open(tmpfile, 'w') as fptr: + fptr.writelines(mock_csv) + + parser = CSVParser() + df = parser.csv_to_df(tmpfile) + + os.remove(tmpfile) + + assert type(df) is pd.DataFrame + + +def test_set_filepath(): + """ + Test that the filepath method does the right thing. + """ + + tmpfile = 'tmp_parsing.txt' + with open(tmpfile, 'w') as fptr: + fptr.writelines(mock_csv) + + parser = CSVParser() + parser.filepath = tmpfile + + assert type(parser.filepath) is str + + +def test_set_filepath_exception(): + """ + Test that the filepath method throws an exception when there is no file + passed in. + """ + + parser = CSVParser() + + with pytest.raises(FileNotFoundError): + parser.filepath = mock_csv + + +def test_imfreq_units_exception_scalar(): + """ + Test that the imfreq exception raising works. + """ + + tmpfile = 'tmp_parsing.txt' + with open(tmpfile, 'w') as fptr: + fptr.writelines(mock_csv) + + parser = CSVParser() + df = parser.csv_to_df(tmpfile) + + os.remove(tmpfile) + + # imfreq is not astropy.unit so should throw exception + with pytest.raises(TypeError): + outdf, outfreq, nstokes = parser.get_zcoeffs(df, 1000) + + +def test_imfreq_units_exception_vector(): + """ + Test that the imfreq exception raising works. + """ + + tmpfile = 'tmp_parsing.txt' + with open(tmpfile, 'w') as fptr: + fptr.writelines(mock_csv) + + parser = CSVParser() + df = parser.csv_to_df(tmpfile) + + os.remove(tmpfile) + + # imfreq is not astropy.unit so should throw exception + with pytest.raises(TypeError): + outdf, outfreq, nstokes = parser.get_zcoeffs(df, [1000,2000]) + + + +def test_imfreq_units_vector(): + """ + Test that the imfreq works for vectors + """ + + tmpfile = 'tmp_parsing.txt' + with open(tmpfile, 'w') as fptr: + fptr.writelines(mock_csv) + + parser = CSVParser() + df = parser.csv_to_df(tmpfile) + + os.remove(tmpfile) + + freqs = np.arange(1000, 2000, 500) * u.MHz + outdf, outfreq, nstokes = parser.get_zcoeffs(df, freqs) + #pdb.set_trace() + + assert outfreq == [1000, 1000] + assert nstokes == 4 + assert len(outdf) == 2 + assert outdf[0].equals(df) diff --git a/tests/test_plumber.py b/tests/test_plumber.py index 87b1df2..8dc3073 100644 --- a/tests/test_plumber.py +++ b/tests/test_plumber.py @@ -1,4 +1,42 @@ #!/usr/bin/env python3 -def test_dummy(): - assert 1 == 1 +import shutil +import numpy as np +import os + +from plumber.scripts.plumber import main as plumber_main + +from casatools import image +ia = image() + +def get_imdat(imname): + ia.open(imname) + dat = np.squeeze(ia.getchunk()) + ia.close() + + return dat + + +def test_StokesI_generation(): + """ + Given a fixed CSV file, compare against a known beam and confirm the PB + generation is working. + """ + + template_image = './plumber-data/images/I_MeerKAT_1577.832MHz_template.im' + csv_path = './plumber-data/csv/MeerKAT_coeffs_for_testing.csv' + + try: + # Need to pass args as a list since plumnber.main is wrapped with Click + plumber_main([f'{template_image}', f'{csv_path}', '--stokesI']) + except SystemExit: + pass + + imname = 'I_MeerKAT_1577.832MHz.im' + + imdat_gen = get_imdat(imname) + imdat_template = get_imdat(template_image) + + shutil.rmtree('I_MeerKAT_1577.832MHz.im') + + np.testing.assert_allclose(imdat_gen, imdat_template, rtol=1e-06, atol=1e-06) diff --git a/tests/test_sky.py b/tests/test_sky.py index 94f9b81..3596559 100644 --- a/tests/test_sky.py +++ b/tests/test_sky.py @@ -3,7 +3,7 @@ from plumber.sky import ParallacticAngle import numpy as np -def test_parang_casa(input_ms='../../plumber-data/measurement_sets/J1939-6342_MeerKAT_1400MHz.ms'): +def test_parang_casa(input_ms='./plumber-data/measurement_sets/J1939-6342_MeerKAT_1400MHz.ms'): parang_casa = ParallacticAngle(ms=input_ms, use_astropy=False) print(parang_casa.parangs[0]) print(parang_casa.parangs[-1]) @@ -11,7 +11,7 @@ def test_parang_casa(input_ms='../../plumber-data/measurement_sets/J1939-6342_Me assert np.allclose(parang_casa.parangs[0], 262.576893) assert np.allclose(parang_casa.parangs[-1], 263.558544) -def test_parang_astropy(input_ms='../../plumber-data/measurement_sets/J1939-6342_MeerKAT_1400MHz.ms'): +def test_parang_astropy(input_ms='./plumber-data/measurement_sets/J1939-6342_MeerKAT_1400MHz.ms'): parang_astro = ParallacticAngle(ms=input_ms, use_astropy=True) print(parang_astro.parangs[0]) print(parang_astro.parangs[-1])