From 74f9bdcaf06c0f9ebda609545b4596ae784f37fe Mon Sep 17 00:00:00 2001 From: Thomas Robitaille Date: Wed, 15 Jan 2020 13:48:51 +0000 Subject: [PATCH 1/6] Starting reorganizing I/O to use astropy unified I/O infrastructure --- spectral_cube/io/casa_image.py | 7 ++- spectral_cube/io/class_lmv.py | 10 ++- spectral_cube/io/core.py | 107 +++++++++++++++++++++++++++++++++ spectral_cube/io/fits.py | 40 ++++++++---- spectral_cube/spectral_cube.py | 59 ++---------------- 5 files changed, 155 insertions(+), 68 deletions(-) diff --git a/spectral_cube/io/casa_image.py b/spectral_cube/io/casa_image.py index 0e64cf7f9..aac5d6507 100644 --- a/spectral_cube/io/casa_image.py +++ b/spectral_cube/io/casa_image.py @@ -9,6 +9,7 @@ from astropy import units as u from astropy.wcs.wcsapi.sliced_low_level_wcs import sanitize_slices from astropy import log +from astropy.io import registry as io_registry import numpy as np from radio_beam import Beam, Beams @@ -29,7 +30,7 @@ # yield the same array in memory that we would get from astropy. -def is_casa_image(input, **kwargs): +def is_casa_image(input, *args, **kwargs): if isinstance(input, six.string_types): if input.endswith('.image'): return True @@ -345,3 +346,7 @@ def load_casa_image(filename, skipdata=False, return cube + + +io_registry.register_reader('casa', SpectralCube, load_casa_image) +io_registry.register_identifier('casa', SpectralCube, is_casa_image) diff --git a/spectral_cube/io/class_lmv.py b/spectral_cube/io/class_lmv.py index 71f9efa79..a3790a2aa 100644 --- a/spectral_cube/io/class_lmv.py +++ b/spectral_cube/io/class_lmv.py @@ -6,6 +6,8 @@ import warnings import string from astropy import log +from astropy.io import registry as io_registry +from ..spectral_cube import SpectralCube from .fits import load_fits_cube """ @@ -40,7 +42,7 @@ 7:'GLS', 8:'SFL', } _bunit_dict = {'k (tmb)': 'K'} -def is_lmv(input, **kwargs): +def is_lmv(input, *args, **kwargs): """ Determine whether input is in GILDAS CLASS lmv format """ @@ -50,6 +52,7 @@ def is_lmv(input, **kwargs): else: return False + def read_lmv(filename): """ Read an LMV cube file @@ -671,3 +674,8 @@ def read_lmv_type2(lf): data[data==bval] = np.nan return data,header + + +io_registry.register_reader('lmv', SpectralCube, load_lmv_cube) +io_registry.register_reader('class_lmv', SpectralCube, load_lmv_cube) +io_registry.register_identifier('lmv', SpectralCube, is_lmv) diff --git a/spectral_cube/io/core.py b/spectral_cube/io/core.py index 912fe69d7..4e7843f38 100644 --- a/spectral_cube/io/core.py +++ b/spectral_cube/io/core.py @@ -1,5 +1,112 @@ from __future__ import print_function, absolute_import, division +import warnings + +from astropy.io import registry + +__doctest_skip__ = ['SpectralCubeRead', 'SpectralCubeWrite'] + + +class SpectralCubeRead(registry.UnifiedReadWrite): + """ + Read and parse a dataset and return as a SpectralCube + + This function provides the SpectralCube interface to the astropy unified I/O + layer. This allows easily reading a dataset in several supported data + formats using syntax such as:: + + >>> from spectral_cube import SpectralCube + >>> cube1 = SpectralCube.read('cube.fits', format='fits') + >>> cube2 = SpectralCube.read('cube.image', format='casa') + + If the file contains Stokes axes, they will automatically be dropped. If + you want to read in all Stokes informtion, use + :meth:`~spectral_cube.StokesSpectralCube.read` instead. + + Get help on the available readers for ``SpectralCube`` using the``help()`` method:: + + >>> SpectralCube.read.help() # Get help reading SpectralCube and list supported formats + >>> SpectralCube.read.help('fits') # Get detailed help on SpectralCube FITS reader + >>> SpectralCube.read.list_formats() # Print list of available formats + + See also: http://docs.astropy.org/en/stable/io/unified.html + + Parameters + ---------- + *args : tuple, optional + Positional arguments passed through to data reader. If supplied the + first argument is typically the input filename. + format : str + File format specifier. + **kwargs : dict, optional + Keyword arguments passed through to data reader. + + Returns + ------- + cube : `SpectralCube` + SpectralCube corresponding to dataset + + Notes + ----- + """ + + def __init__(self, instance, cls): + super().__init__(instance, cls, 'read') + + def __call__(self, *args, **kwargs): + from .. import StokesSpectralCube + from ..utils import StokesWarning + cube = registry.read(self._cls, *args, **kwargs) + if isinstance(cube, StokesSpectralCube): + if hasattr(cube, 'I'): + warnings.warn("Cube is a Stokes cube, " + "returning spectral cube for I component", + StokesWarning) + return cube.I + else: + raise ValueError("Spectral cube is a Stokes cube that " + "does not have an I component") + else: + return cube + + +class SpectralCubeWrite(registry.UnifiedReadWrite): + """ + Write this SpectralCube object out in the specified format. + + This function provides the SpectralCube interface to the astropy unified + I/O layer. This allows easily writing a spectral cube in many supported + data formats using syntax such as:: + + >>> cube.write('cube.fits', format='fits') + + Get help on the available writers for ``SpectralCube`` using the``help()`` method:: + + >>> SpectralCube.write.help() # Get help writing SpectralCube and list supported formats + >>> SpectralCube.write.help('fits') # Get detailed help on SpectralCube FITS writer + >>> SpectralCube.write.list_formats() # Print list of available formats + + See also: http://docs.astropy.org/en/stable/io/unified.html + + Parameters + ---------- + *args : tuple, optional + Positional arguments passed through to data writer. If supplied the + first argument is the output filename. + format : str + File format specifier. + **kwargs : dict, optional + Keyword arguments passed through to data writer. + + Notes + ----- + """ + def __init__(self, instance, cls): + super().__init__(instance, cls, 'write') + + def __call__(self, *args, serialize_method=None, **kwargs): + registry.write(self._instance, *args, **kwargs) + def read(filename, format=None, hdu=None, **kwargs): """ diff --git a/spectral_cube/io/fits.py b/spectral_cube/io/fits.py index 48b99fe2e..def138eb8 100644 --- a/spectral_cube/io/fits.py +++ b/spectral_cube/io/fits.py @@ -4,11 +4,13 @@ import warnings from astropy.io import fits +from astropy.io import registry as io_registry import astropy.wcs from astropy import wcs from astropy.wcs import WCS from collections import OrderedDict from astropy.io.fits.hdu.hdulist import fitsopen as fits_open +from astropy.io.fits.connect import FITS_SIGNATURE import numpy as np import datetime @@ -29,20 +31,31 @@ def first(iterable): return next(iter(iterable)) -# FITS registry code - once Astropy includes a proper extensible I/O base -# class, we can use that instead. The following code takes care of -# interpreting string input (filename), HDU, and HDUList. - -def is_fits(input, **kwargs): +def is_fits(origin, filepath, fileobj, *args, **kwargs): """ - Determine whether input is in FITS format + Determine whether `origin` is a FITS file. + + Parameters + ---------- + origin : str or readable file-like object + Path or file object containing a potential FITS file. + + Returns + ------- + is_fits : bool + Returns `True` if the given file is a FITS file. """ - if isinstance(input, six.string_types): - if input.lower().endswith(('.fits', '.fits.gz', - '.fit', '.fit.gz', - '.fits.Z', '.fit.Z')): + print(origin, filepath, fileobj, args, kwargs) + if fileobj is not None: + pos = fileobj.tell() + sig = fileobj.read(30) + fileobj.seek(pos) + return sig == FITS_SIGNATURE + elif filepath is not None: + if filepath.lower().endswith(('.fits', '.fits.gz', '.fit', '.fit.gz', + '.fts', '.fts.gz')): return True - elif isinstance(input, (fits.HDUList, fits.PrimaryHDU, fits.ImageHDU)): + elif isinstance(args[0], (fits.HDUList, fits.ImageHDU, fits.PrimaryHDU)): return True else: return False @@ -215,3 +228,8 @@ def write_fits_cube(filename, cube, overwrite=False, hdulist.writeto(filename, clobber=overwrite) else: raise NotImplementedError() + + +io_registry.register_reader('fits', SpectralCube, load_fits_cube) +io_registry.register_writer('fits', SpectralCube, write_fits_cube) +io_registry.register_identifier('fits', SpectralCube, is_fits) diff --git a/spectral_cube/spectral_cube.py b/spectral_cube/spectral_cube.py index 078b8905d..710a4b712 100644 --- a/spectral_cube/spectral_cube.py +++ b/spectral_cube/spectral_cube.py @@ -25,6 +25,7 @@ from astropy import convolution from astropy import stats from astropy.constants import si +from astropy.io.registry import UnifiedReadWriteMethod import numpy as np @@ -52,6 +53,7 @@ BeamAverageWarning, NonFiniteBeamsWarning, BeamWarning) from .spectral_axis import (determine_vconv_from_ctype, get_rest_value_from_wcs, doppler_beta, doppler_gamma, doppler_z) +from .io.core import SpectralCubeRead, SpectralCubeWrite from distutils.version import LooseVersion @@ -2145,61 +2147,8 @@ def __pow__(self, value): else: return self._apply_everywhere(operator.pow, value) - - @classmethod - def read(cls, filename, format=None, hdu=None, **kwargs): - """ - Read a spectral cube from a file. - - If the file contains Stokes axes, they will automatically be dropped. - If you want to read in all Stokes informtion, use - :meth:`~spectral_cube.StokesSpectralCube.read` instead. - - Parameters - ---------- - filename : str - The file to read the cube from - format : str - The format of the file to read. (Currently limited to 'fits' and 'casa_image') - hdu : int or str - For FITS files, the HDU to read in (can be the ID or name of an - HDU). - kwargs : dict - If the format is 'fits', the kwargs are passed to - :func:`~astropy.io.fits.open`. - """ - from .io.core import read - from .stokes_spectral_cube import StokesSpectralCube - if isinstance(filename, PosixPath): - filename = str(filename) - cube = read(filename, format=format, hdu=hdu, **kwargs) - if isinstance(cube, StokesSpectralCube): - if hasattr(cube, 'I'): - warnings.warn("Cube is a Stokes cube, " - "returning spectral cube for I component", - StokesWarning) - return cube.I - else: - raise ValueError("Spectral cube is a Stokes cube that " - "does not have an I component") - else: - return cube - - def write(self, filename, overwrite=False, format=None): - """ - Write the spectral cube to a file. - - Parameters - ---------- - filename : str - The path to write the file to - format : str - The format of the file to write. (Currently limited to 'fits') - overwrite : bool - If True, overwrite ``filename`` if it exists - """ - from .io.core import write - write(filename, self, overwrite=overwrite, format=format) + read = UnifiedReadWriteMethod(SpectralCubeRead) + write = UnifiedReadWriteMethod(SpectralCubeWrite) def to_yt(self, spectral_factor=1.0, nprocs=None, **kwargs): """ From ff21c396f29ddccce160adf54da85850b850a4a7 Mon Sep 17 00:00:00 2001 From: Thomas Robitaille Date: Wed, 15 Jan 2020 15:06:13 +0000 Subject: [PATCH 2/6] More work on using unified I/O infrastructure --- spectral_cube/io/casa_image.py | 12 +- spectral_cube/io/class_lmv.py | 74 +++---- spectral_cube/io/core.py | 197 +++++++++++------- spectral_cube/io/fits.py | 55 ++++- spectral_cube/lower_dimensional_structures.py | 41 +--- spectral_cube/spectral_cube.py | 6 +- spectral_cube/stokes_spectral_cube.py | 47 +---- .../tests/test_analysis_functions.py | 6 +- 8 files changed, 218 insertions(+), 220 deletions(-) diff --git a/spectral_cube/io/casa_image.py b/spectral_cube/io/casa_image.py index aac5d6507..72bf8826b 100644 --- a/spectral_cube/io/casa_image.py +++ b/spectral_cube/io/casa_image.py @@ -16,6 +16,7 @@ import dask.array from .. import SpectralCube, StokesSpectralCube, BooleanArrayMask, LazyMask, VaryingResolutionSpectralCube +from ..spectral_cube import BaseSpectralCube from .. import cube_utils from .. utils import BeamWarning, cached from .. import wcs_utils @@ -30,11 +31,8 @@ # yield the same array in memory that we would get from astropy. -def is_casa_image(input, *args, **kwargs): - if isinstance(input, six.string_types): - if input.endswith('.image'): - return True - return False +def is_casa_image(origin, filepath, fileobj, *args, **kwargs): + return filepath is not None and filepath.lower().endswith('.image') def wcs_casa2astropy(ia, coordsys): @@ -348,5 +346,5 @@ def load_casa_image(filename, skipdata=False, return cube -io_registry.register_reader('casa', SpectralCube, load_casa_image) -io_registry.register_identifier('casa', SpectralCube, is_casa_image) +io_registry.register_reader('casa', BaseSpectralCube, load_casa_image) +io_registry.register_identifier('casa', BaseSpectralCube, is_casa_image) diff --git a/spectral_cube/io/class_lmv.py b/spectral_cube/io/class_lmv.py index a3790a2aa..4d5cf2d08 100644 --- a/spectral_cube/io/class_lmv.py +++ b/spectral_cube/io/class_lmv.py @@ -7,7 +7,7 @@ import string from astropy import log from astropy.io import registry as io_registry -from ..spectral_cube import SpectralCube +from ..spectral_cube import BaseSpectralCube from .fits import load_fits_cube """ @@ -42,18 +42,13 @@ 7:'GLS', 8:'SFL', } _bunit_dict = {'k (tmb)': 'K'} -def is_lmv(input, *args, **kwargs): +def is_lmv(origin, filepath, fileobj, *args, **kwargs): """ Determine whether input is in GILDAS CLASS lmv format """ - if isinstance(input, six.string_types): - if input.lower().endswith(('.lmv')): - return True - else: - return False - + return filepath is not None and filepath.lower().endswith('.lmv') -def read_lmv(filename): +def read_lmv(lf): """ Read an LMV cube file @@ -62,29 +57,28 @@ def read_lmv(filename): log.warning("CLASS LMV cube reading is tentatively supported. " "Please post bug reports at the first sign of danger!") - with open(filename,'rb') as lf: - # lf for "LMV File" - filetype = _read_string(lf, 12) - #!--------------------------------------------------------------------- - #! @ private - #! SYCODE system code - #! '-' IEEE - #! '.' EEEI (IBM like) - #! '_' VAX - #! IMCODE file code - #! '<' IEEE 64 bits (Little Endian, 99.9 % of recent computers) - #! '>' EEEI 64 bits (Big Endian, HPUX, IBM-RISC, and SPARC ...) - #!--------------------------------------------------------------------- - imcode = filetype[6] - if filetype[:6] != 'GILDAS' or filetype[7:] != 'IMAGE': - raise TypeError("File is not a GILDAS Image file") - - if imcode in ('<','>'): - if imcode =='>': - log.warning("Swap the endianness first...") - return read_lmv_type2(lf) - else: - return read_lmv_type1(lf) + # lf for "LMV File" + filetype = _read_string(lf, 12) + #!--------------------------------------------------------------------- + #! @ private + #! SYCODE system code + #! '-' IEEE + #! '.' EEEI (IBM like) + #! '_' VAX + #! IMCODE file code + #! '<' IEEE 64 bits (Little Endian, 99.9 % of recent computers) + #! '>' EEEI 64 bits (Big Endian, HPUX, IBM-RISC, and SPARC ...) + #!--------------------------------------------------------------------- + imcode = filetype[6] + if filetype[:6] != 'GILDAS' or filetype[7:] != 'IMAGE': + raise TypeError("File is not a GILDAS Image file") + + if imcode in ('<','>'): + if imcode =='>': + log.warning("Swap the endianness first...") + return read_lmv_type2(lf) + else: + return read_lmv_type1(lf) def read_lmv_type1(lf): header = {} @@ -247,9 +241,9 @@ def read_lmv_type1(lf): # debug #return data.reshape([naxis3,naxis2,naxis1]), header, hdr_f, hdr_s, hdr_i, hdr_d, hdr_d_2 -def read_lmv_tofits(filename): +def read_lmv_tofits(fileobj): from astropy.io import fits - data,header = read_lmv(filename) + data,header = read_lmv(fileobj) # LMV may contain extra dimensions that are improperly labeled data = data.squeeze() bad_kws = ['NAXIS4','CRVAL4','CRPIX4','CDELT4','CROTA4','CUNIT4','CTYPE4'] @@ -265,9 +259,9 @@ def read_lmv_tofits(filename): hdu = fits.PrimaryHDU(data=data, header=Header) return hdu -def load_lmv_cube(filename): - hdu = read_lmv_tofits(filename) - meta = {'filename':filename} +def load_lmv_cube(fileobj, target_cls=None): + hdu = read_lmv_tofits(fileobj) + meta = {'filename':fileobj.name} return load_fits_cube(hdu, meta=meta) @@ -676,6 +670,6 @@ def read_lmv_type2(lf): return data,header -io_registry.register_reader('lmv', SpectralCube, load_lmv_cube) -io_registry.register_reader('class_lmv', SpectralCube, load_lmv_cube) -io_registry.register_identifier('lmv', SpectralCube, is_lmv) +io_registry.register_reader('lmv', BaseSpectralCube, load_lmv_cube) +io_registry.register_reader('class_lmv', BaseSpectralCube, load_lmv_cube) +io_registry.register_identifier('lmv', BaseSpectralCube, is_lmv) diff --git a/spectral_cube/io/core.py b/spectral_cube/io/core.py index 4e7843f38..1e5584c34 100644 --- a/spectral_cube/io/core.py +++ b/spectral_cube/io/core.py @@ -1,18 +1,22 @@ from __future__ import print_function, absolute_import, division +from pathlib import PosixPath import warnings from astropy.io import registry -__doctest_skip__ = ['SpectralCubeRead', 'SpectralCubeWrite'] +__doctest_skip__ = ['SpectralCubeRead', + 'SpectralCubeWrite', + 'StokesSpectralCubeRead', + 'StokesSpectralCubeWrite', + 'LowerDimensionalObjectWrite'] class SpectralCubeRead(registry.UnifiedReadWrite): """ Read and parse a dataset and return as a SpectralCube - This function provides the SpectralCube interface to the astropy unified I/O - layer. This allows easily reading a dataset in several supported data + This allows easily reading a dataset in several supported data formats using syntax such as:: >>> from spectral_cube import SpectralCube @@ -53,30 +57,20 @@ class SpectralCubeRead(registry.UnifiedReadWrite): def __init__(self, instance, cls): super().__init__(instance, cls, 'read') - def __call__(self, *args, **kwargs): - from .. import StokesSpectralCube - from ..utils import StokesWarning - cube = registry.read(self._cls, *args, **kwargs) - if isinstance(cube, StokesSpectralCube): - if hasattr(cube, 'I'): - warnings.warn("Cube is a Stokes cube, " - "returning spectral cube for I component", - StokesWarning) - return cube.I - else: - raise ValueError("Spectral cube is a Stokes cube that " - "does not have an I component") - else: - return cube + def __call__(self, filename, *args, **kwargs): + from ..spectral_cube import BaseSpectralCube + if isinstance(filename, PosixPath): + filename = str(filename) + kwargs['target_cls'] = BaseSpectralCube + return registry.read(BaseSpectralCube, filename, *args, **kwargs) class SpectralCubeWrite(registry.UnifiedReadWrite): """ Write this SpectralCube object out in the specified format. - This function provides the SpectralCube interface to the astropy unified - I/O layer. This allows easily writing a spectral cube in many supported - data formats using syntax such as:: + This allows easily writing a spectral cube in many supported data formats + using syntax such as:: >>> cube.write('cube.fits', format='fits') @@ -108,85 +102,128 @@ def __call__(self, *args, serialize_method=None, **kwargs): registry.write(self._instance, *args, **kwargs) -def read(filename, format=None, hdu=None, **kwargs): +class StokesSpectralCubeRead(registry.UnifiedReadWrite): """ - Read a file into a :class:`SpectralCube` or :class:`StokesSpectralCube` - instance. + Read and parse a dataset and return as a StokesSpectralCube + + This allows easily reading a dataset in several supported data formats + using syntax such as:: + + >>> from spectral_cube import StokesSpectralCube + >>> cube1 = StokesSpectralCube.read('cube.fits', format='fits') + >>> cube2 = StokesSpectralCube.read('cube.image', format='casa') + + If the file contains Stokes axes, they will be read in. If you are only + interested in the unpolarized emission (I), you can use + :meth:`~spectral_cube.SpectralCube.read` instead. + + Get help on the available readers for ``StokesSpectralCube`` using the``help()`` method:: + + >>> StokesSpectralCube.read.help() # Get help reading StokesSpectralCube and list supported formats + >>> StokesSpectralCube.read.help('fits') # Get detailed help on StokesSpectralCube FITS reader + >>> StokesSpectralCube.read.list_formats() # Print list of available formats + + See also: http://docs.astropy.org/en/stable/io/unified.html Parameters ---------- - filename : str or HDU - File to read - format : str, optional - File format. - hdu : int or str - For FITS files, the HDU to read in (can be the ID or name of an - HDU). - kwargs : dict - If the format is 'fits', the kwargs are passed to - :func:`~astropy.io.fits.open`. + *args : tuple, optional + Positional arguments passed through to data reader. If supplied the + first argument is typically the input filename. + format : str + File format specifier. + **kwargs : dict, optional + Keyword arguments passed through to data reader. Returns ------- - cube : :class:`SpectralCube` or :class:`StokesSpectralCube` - The spectral cube read in + cube : `StokesSpectralCube` + StokesSpectralCube corresponding to dataset + + Notes + ----- """ - if format is None: - format = determine_format(filename) + def __init__(self, instance, cls): + super().__init__(instance, cls, 'read') - if format == 'fits': - from .fits import load_fits_cube - return load_fits_cube(filename, hdu=hdu, **kwargs) - elif format == 'casa_image': - from .casa_image import load_casa_image - return load_casa_image(filename) - elif format in ('class_lmv','lmv'): - from .class_lmv import load_lmv_cube - return load_lmv_cube(filename) - else: - raise ValueError("Format {0} not implemented. Supported formats are " - "'fits', 'casa_image', and 'lmv'.".format(format)) + def __call__(self, filename, *args, **kwargs): + from ..stokes_spectral_cube import StokesSpectralCube + if isinstance(filename, PosixPath): + filename = str(filename) + kwargs['target_cls'] = StokesSpectralCube + return registry.read(StokesSpectralCube, filename, *args, **kwargs) -def write(filename, cube, overwrite=False, format=None): +class StokesSpectralCubeWrite(registry.UnifiedReadWrite): """ - Write :class:`SpectralCube` or :class:`StokesSpectralCube` to a file. + Write this StokesSpectralCube object out in the specified format. + + This allows easily writing a spectral cube in many supported data formats + using syntax such as:: + + >>> cube.write('cube.fits', format='fits') + + Get help on the available writers for ``StokesSpectralCube`` using the``help()`` method:: + + >>> StokesSpectralCube.write.help() # Get help writing StokesSpectralCube and list supported formats + >>> StokesSpectralCube.write.help('fits') # Get detailed help on StokesSpectralCube FITS writer + >>> StokesSpectralCube.write.list_formats() # Print list of available formats + + See also: http://docs.astropy.org/en/stable/io/unified.html Parameters ---------- - filename : str - Name of the output file - cube : :class:`SpectralCube` or :class:`StokesSpectralCube` - The spectral cube to write out - overwrite : bool, optional - Whether to overwrite the output file - format : str, optional - File format. + *args : tuple, optional + Positional arguments passed through to data writer. If supplied the + first argument is the output filename. + format : str + File format specifier. + **kwargs : dict, optional + Keyword arguments passed through to data writer. + + Notes + ----- + """ + def __init__(self, instance, cls): + super().__init__(instance, cls, 'write') + + def __call__(self, *args, serialize_method=None, **kwargs): + registry.write(self._instance, *args, **kwargs) + + +class LowerDimensionalObjectWrite(registry.UnifiedReadWrite): """ + Write this object out in the specified format. - if format is None: - format = determine_format(filename) + This allows easily writing a data object in many supported data formats + using syntax such as:: - if format == 'fits': - from .fits import write_fits_cube - write_fits_cube(filename, cube, overwrite=overwrite) - else: - raise ValueError("Format {0} not implemented. The only supported format is 'fits'".format(format)) + >>> data.write('data.fits', format='fits') + Get help on the available writers using the``help()`` method, e.g.:: -def determine_format(input): + >>> LowerDimensionalObject.write.help() # Get help writing LowerDimensionalObject and list supported formats + >>> LowerDimensionalObject.write.help('fits') # Get detailed help on LowerDimensionalObject FITS writer + >>> LowerDimensionalObject.write.list_formats() # Print list of available formats - from .fits import is_fits - from .casa_image import is_casa_image - from .class_lmv import is_lmv + See also: http://docs.astropy.org/en/stable/io/unified.html - if is_fits(input): - return 'fits' - elif is_casa_image(input): - return 'casa_image' - elif is_lmv(input): - return 'lmv' - else: - raise ValueError("Could not determine format - use the `format=` " - "parameter to explicitly set the format") + Parameters + ---------- + *args : tuple, optional + Positional arguments passed through to data writer. If supplied the + first argument is the output filename. + format : str + File format specifier. + **kwargs : dict, optional + Keyword arguments passed through to data writer. + + Notes + ----- + """ + def __init__(self, instance, cls): + super().__init__(instance, cls, 'write') + + def __call__(self, *args, serialize_method=None, **kwargs): + registry.write(self._instance, *args, **kwargs) diff --git a/spectral_cube/io/fits.py b/spectral_cube/io/fits.py index def138eb8..8e71082b4 100644 --- a/spectral_cube/io/fits.py +++ b/spectral_cube/io/fits.py @@ -22,9 +22,10 @@ SPECTRAL_CUBE_VERSION = 'dev' from .. import SpectralCube, StokesSpectralCube, LazyMask, VaryingResolutionSpectralCube +from ..lower_dimensional_structures import LowerDimensionalObject from ..spectral_cube import BaseSpectralCube from .. import cube_utils -from ..utils import FITSWarning, FITSReadError +from ..utils import FITSWarning, FITSReadError, StokesWarning def first(iterable): @@ -45,7 +46,6 @@ def is_fits(origin, filepath, fileobj, *args, **kwargs): is_fits : bool Returns `True` if the given file is a FITS file. """ - print(origin, filepath, fileobj, args, kwargs) if fileobj is not None: pos = fileobj.tell() sig = fileobj.read(30) @@ -124,13 +124,16 @@ def read_data_fits(input, hdu=None, mode='denywrite', **kwargs): else: + if hasattr(input, 'read'): + mode = None + with fits_open(input, mode=mode, **kwargs) as hdulist: return read_data_fits(hdulist, hdu=hdu) return array_hdu.data, array_hdu.header, beam_table -def load_fits_cube(input, hdu=0, meta=None, **kwargs): +def load_fits_cube(input, hdu=0, meta=None, target_cls=None, **kwargs): """ Read in a cube from a FITS file using astropy. @@ -206,10 +209,22 @@ def load_fits_cube(input, hdu=0, meta=None, **kwargs): raise FITSReadError("Data should be 3- or 4-dimensional") - return cube + if target_cls is BaseSpectralCube and isinstance(cube, StokesSpectralCube): + if hasattr(cube, 'I'): + warnings.warn("Cube is a Stokes cube, " + "returning spectral cube for I component", + StokesWarning) + return cube.I + else: + raise ValueError("Spectral cube is a Stokes cube that " + "does not have an I component") + elif target_cls is StokesSpectralCube and isinstance(cube, BaseSpectralCube): + cube = StokesSpectralCube({'I': cube}) + else: + return cube -def write_fits_cube(filename, cube, overwrite=False, +def write_fits_cube(cube, filename, overwrite=False, include_origin_notes=True): """ Write a FITS cube with a WCS to a filename @@ -230,6 +245,30 @@ def write_fits_cube(filename, cube, overwrite=False, raise NotImplementedError() -io_registry.register_reader('fits', SpectralCube, load_fits_cube) -io_registry.register_writer('fits', SpectralCube, write_fits_cube) -io_registry.register_identifier('fits', SpectralCube, is_fits) +def write_fits_ldo(data, filename, overwrite=False): + # Spectra may have HDUList objects instead of HDUs because they + # have a beam table attached, so we want to try that first + # (a more elegant way to write this might be to do "self._hdu_general.write" + # and create a property `self._hdu_general` that selects the right one...) + if hasattr(data, 'hdulist'): + try: + data.hdulist.writeto(filename, overwrite=overwrite) + except TypeError: + data.hdulist.writeto(filename, clobber=overwrite) + elif hasattr(data, 'hdu'): + try: + data.hdu.writeto(filename, overwrite=overwrite) + except TypeError: + data.hdu.writeto(filename, clobber=overwrite) + + +io_registry.register_reader('fits', BaseSpectralCube, load_fits_cube) +io_registry.register_writer('fits', BaseSpectralCube, write_fits_cube) +io_registry.register_identifier('fits', BaseSpectralCube, is_fits) + +io_registry.register_reader('fits', StokesSpectralCube, load_fits_cube) +io_registry.register_writer('fits', StokesSpectralCube, write_fits_cube) +io_registry.register_identifier('fits', StokesSpectralCube, is_fits) + +io_registry.register_writer('fits', LowerDimensionalObject, write_fits_ldo) +io_registry.register_identifier('fits', LowerDimensionalObject, is_fits) diff --git a/spectral_cube/lower_dimensional_structures.py b/spectral_cube/lower_dimensional_structures.py index 0d51140a7..d35467e81 100644 --- a/spectral_cube/lower_dimensional_structures.py +++ b/spectral_cube/lower_dimensional_structures.py @@ -11,9 +11,10 @@ #from astropy import log from astropy.io.fits import Header, HDUList, PrimaryHDU, BinTableHDU, FITS_rec from radio_beam import Beam, Beams +from astropy.io.registry import UnifiedReadWriteMethod -from .io.core import determine_format from . import spectral_axis +from .io.core import LowerDimensionalObjectWrite from .utils import SliceWarning, BeamWarning, SmoothingWarning, FITSWarning from .cube_utils import convert_bunit from . import wcs_utils @@ -27,8 +28,6 @@ from . import cube_utils __all__ = ['LowerDimensionalObject', 'Projection', 'Slice', 'OneDSpectrum'] - - class LowerDimensionalObject(u.Quantity, BaseNDClass, HeaderMixinClass): """ Generic class for 1D and 2D objects. @@ -47,40 +46,10 @@ def hdu(self): return hdu - def write(self, filename, format=None, overwrite=False): - """ - Write the lower dimensional object to a file. - - Parameters - ---------- - filename : str - The path to write the file to - format : str - The kind of file to write. (Currently limited to 'fits') - overwrite : bool - If True, overwrite ``filename`` if it exists - """ - if format is None: - format = determine_format(filename) - if format == 'fits': - # Spectra may have HDUList objects instead of HDUs because they - # have a beam table attached, so we want to try that first - # (a more elegant way to write this might be to do "self._hdu_general.write" - # and create a property `self._hdu_general` that selects the right one...) - if hasattr(self, 'hdulist'): - try: - self.hdulist.writeto(filename, overwrite=overwrite) - except TypeError: - self.hdulist.writeto(filename, clobber=overwrite) - elif hasattr(self, 'hdu'): - try: - self.hdu.writeto(filename, overwrite=overwrite) - except TypeError: - self.hdu.writeto(filename, clobber=overwrite) - else: - raise ValueError("Unknown format '{0}' - the only available " - "format at this time is 'fits'") + def read(self, *args, **kwargs): + raise NotImplementedError() + write = UnifiedReadWriteMethod(LowerDimensionalObjectWrite) def __getslice__(self, start, end, increment=None): # I don't know why this is needed, but apparently one of the inherited diff --git a/spectral_cube/spectral_cube.py b/spectral_cube/spectral_cube.py index 710a4b712..25dbcaa82 100644 --- a/spectral_cube/spectral_cube.py +++ b/spectral_cube/spectral_cube.py @@ -282,6 +282,9 @@ def _new_cube_with(self, data=None, wcs=None, mask=None, meta=None, return cube + read = UnifiedReadWriteMethod(SpectralCubeRead) + write = UnifiedReadWriteMethod(SpectralCubeWrite) + @property def unit(self): """ The flux unit """ @@ -2147,9 +2150,6 @@ def __pow__(self, value): else: return self._apply_everywhere(operator.pow, value) - read = UnifiedReadWriteMethod(SpectralCubeRead) - write = UnifiedReadWriteMethod(SpectralCubeWrite) - def to_yt(self, spectral_factor=1.0, nprocs=None, **kwargs): """ Convert a spectral cube to a yt object that can be further analyzed in diff --git a/spectral_cube/stokes_spectral_cube.py b/spectral_cube/stokes_spectral_cube.py index 3ff372845..68d5dc796 100644 --- a/spectral_cube/stokes_spectral_cube.py +++ b/spectral_cube/stokes_spectral_cube.py @@ -3,6 +3,8 @@ import six import numpy as np +from astropy.io.registry import UnifiedReadWriteMethod +from .io.core import StokesSpectralCubeRead, StokesSpectralCubeWrite from .spectral_cube import SpectralCube, BaseSpectralCube from . import wcs_utils from .masks import BooleanArrayMask, is_broadcastable_and_smaller @@ -152,46 +154,5 @@ def with_spectral_unit(self, unit, **kwargs): return self._new_cube_with(stokes_data=stokes_data) - @classmethod - def read(cls, filename, format=None, hdu=None, **kwargs): - """ - Read a spectral cube from a file. - - If the file contains Stokes axes, they will be read in. If you are - only interested in the unpolarized emission (I), you can use - :meth:`~spectral_cube.SpectralCube.read` instead. - - Parameters - ---------- - filename : str - The file to read the cube from - format : str - The format of the file to read. (Currently limited to 'fits' and 'casa_image') - hdu : int or str - For FITS files, the HDU to read in (can be the ID or name of an - HDU). - - Returns - ------- - cube : :class:`SpectralCube` - """ - from .io.core import read - cube = read(filename, format=format, hdu=hdu) - if isinstance(cube, BaseSpectralCube): - cube = StokesSpectralCube({'I': cube}) - return cube - - def write(self, filename, overwrite=False, format=None): - """ - Write the spectral cube to a file. - - Parameters - ---------- - filename : str - The path to write the file to - format : str - The format of the file to write. (Currently limited to 'fits') - overwrite : bool - If True, overwrite ``filename`` if it exists - """ - raise NotImplementedError("") + read = UnifiedReadWriteMethod(StokesSpectralCubeRead) + write = UnifiedReadWriteMethod(StokesSpectralCubeWrite) diff --git a/spectral_cube/tests/test_analysis_functions.py b/spectral_cube/tests/test_analysis_functions.py index 5ddd1ce23..12443a186 100644 --- a/spectral_cube/tests/test_analysis_functions.py +++ b/spectral_cube/tests/test_analysis_functions.py @@ -7,7 +7,7 @@ from ..analysis_utilities import stack_spectra, fourier_shift from .utilities import generate_gaussian_cube, gaussian - +from ..utils import BadVelocitiesWarning def test_shift(): @@ -106,7 +106,8 @@ def test_stacking_badvels(): test_vels[12,11] = 500*u.km/u.s - with warnings.catch_warnings(record=True) as wrn: + with pytest.warns(BadVelocitiesWarning, + match='Some velocities are outside the allowed range and will be'): # Stack the spectra in the cube stacked = \ stack_spectra(test_cube, test_vels, v0=v0, @@ -114,7 +115,6 @@ def test_stacking_badvels(): xy_posns=None, num_cores=1, chunk_size=-1, progressbar=False, pad_edges=False) - assert 'Some velocities are outside the allowed range and will be' in str(wrn[-1].message) # Calculate residuals (the one bad value shouldn't have caused a problem) resid = np.abs(stacked.value - true_spectrum) From 859073d91c2819b5d09485ec2bf74fa8b9f0b877 Mon Sep 17 00:00:00 2001 From: Thomas Robitaille Date: Wed, 15 Jan 2020 15:20:12 +0000 Subject: [PATCH 3/6] Fix CASA I/O --- spectral_cube/io/casa_image.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/spectral_cube/io/casa_image.py b/spectral_cube/io/casa_image.py index 72bf8826b..42accac5d 100644 --- a/spectral_cube/io/casa_image.py +++ b/spectral_cube/io/casa_image.py @@ -18,7 +18,7 @@ from .. import SpectralCube, StokesSpectralCube, BooleanArrayMask, LazyMask, VaryingResolutionSpectralCube from ..spectral_cube import BaseSpectralCube from .. import cube_utils -from .. utils import BeamWarning, cached +from .. utils import BeamWarning, cached, StokesWarning from .. import wcs_utils # Read and write from a CASA image. This has a few @@ -173,7 +173,7 @@ def __getitem__(self, value): def load_casa_image(filename, skipdata=False, - skipvalid=False, skipcs=False, **kwargs): + skipvalid=False, skipcs=False, target_cls=None, **kwargs): """ Load a cube (into memory?) from a CASA image. By default it will transpose the cube into a 'python' order and drop degenerate axes. These options can @@ -342,9 +342,23 @@ def load_casa_image(filename, skipdata=False, raise ValueError("CASA image has {0} dimensions, and therefore " "is not readable by spectral-cube.".format(wcs.naxis)) + if target_cls is BaseSpectralCube and isinstance(cube, StokesSpectralCube): + if hasattr(cube, 'I'): + warnings.warn("Cube is a Stokes cube, " + "returning spectral cube for I component", + StokesWarning) + return cube.I + else: + raise ValueError("Spectral cube is a Stokes cube that " + "does not have an I component") + elif target_cls is StokesSpectralCube and isinstance(cube, BaseSpectralCube): + cube = StokesSpectralCube({'I': cube}) + else: + return cube return cube io_registry.register_reader('casa', BaseSpectralCube, load_casa_image) +io_registry.register_reader('casa_image', BaseSpectralCube, load_casa_image) io_registry.register_identifier('casa', BaseSpectralCube, is_casa_image) From d331773e160c1e952bcd31e62fa8f388fb042353 Mon Sep 17 00:00:00 2001 From: Thomas Robitaille Date: Wed, 15 Jan 2020 15:25:24 +0000 Subject: [PATCH 4/6] Fix installation of developer versions of dependencies --- tox.ini | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tox.ini b/tox.ini index d84fcf15a..7479f58d7 100644 --- a/tox.ini +++ b/tox.ini @@ -22,10 +22,10 @@ changedir = description = run tests with pytest deps = - dev: git+https://github.com/radio-astro-tools/pvextractor - dev: git+https://github.com/radio-astro-tools/radio-beam - dev: git+https://github.com/astropy/astropy - dev: git+https://github.com/astropy/reproject + dev: git+https://github.com/radio-astro-tools/pvextractor#egg=pvextractor + dev: git+https://github.com/radio-astro-tools/radio-beam#egg=radio-beam + dev: git+https://github.com/astropy/astropy#egg=astropy + dev: git+https://github.com/astropy/reproject#egg=reproject casa: :NRAO:casatools casa: :NRAO:casatasks extras = From 38f5ec83ef6b3b4d3ae8480bd477eaa812c9c8cf Mon Sep 17 00:00:00 2001 From: Thomas Robitaille Date: Wed, 15 Jan 2020 15:59:52 +0000 Subject: [PATCH 5/6] Avoid code/docstring duplication --- spectral_cube/__init__.py | 8 ++ spectral_cube/io/casa_image.py | 17 +-- spectral_cube/io/core.py | 254 +++++++++++++-------------------- spectral_cube/io/fits.py | 15 +- spectral_cube/tests/test_io.py | 1 - 5 files changed, 114 insertions(+), 181 deletions(-) diff --git a/spectral_cube/__init__.py b/spectral_cube/__init__.py index 45875df8b..bed90a8d4 100644 --- a/spectral_cube/__init__.py +++ b/spectral_cube/__init__.py @@ -11,6 +11,14 @@ FunctionMask) from .lower_dimensional_structures import (OneDSpectrum, Projection, Slice) +# Import the following sub-packages to make sure the I/O functions are registered +from .io import casa_image +del casa_image +from .io import class_lmv +del class_lmv +from .io import fits +del fits + __all__ = ['SpectralCube', 'VaryingResolutionSpectralCube', 'StokesSpectralCube', 'CompositeMask', 'LazyComparisonMask', 'LazyMask', 'BooleanArrayMask', 'FunctionMask', diff --git a/spectral_cube/io/casa_image.py b/spectral_cube/io/casa_image.py index 42accac5d..c69c38eec 100644 --- a/spectral_cube/io/casa_image.py +++ b/spectral_cube/io/casa_image.py @@ -342,21 +342,8 @@ def load_casa_image(filename, skipdata=False, raise ValueError("CASA image has {0} dimensions, and therefore " "is not readable by spectral-cube.".format(wcs.naxis)) - if target_cls is BaseSpectralCube and isinstance(cube, StokesSpectralCube): - if hasattr(cube, 'I'): - warnings.warn("Cube is a Stokes cube, " - "returning spectral cube for I component", - StokesWarning) - return cube.I - else: - raise ValueError("Spectral cube is a Stokes cube that " - "does not have an I component") - elif target_cls is StokesSpectralCube and isinstance(cube, BaseSpectralCube): - cube = StokesSpectralCube({'I': cube}) - else: - return cube - - return cube + from .core import normalize_cube_stokes + return normalize_cube_stokes(cube, target_cls=target_cls) io_registry.register_reader('casa', BaseSpectralCube, load_casa_image) diff --git a/spectral_cube/io/core.py b/spectral_cube/io/core.py index 1e5584c34..e7240d288 100644 --- a/spectral_cube/io/core.py +++ b/spectral_cube/io/core.py @@ -1,3 +1,10 @@ +# The read and write methods for SpectralCube, StokesSpectralCube, and +# LowerDimensionalObject are defined in this file and then added to the classes +# using UnifiedReadWriteMethod. This makes it possible to dynamically add the +# available formats to the read/write docstrings. For more information about +# the unified I/O framework from Astropy which is used to implement this, see +# http://docs.astropy.org/en/stable/io/unified.html + from __future__ import print_function, absolute_import, division from pathlib import PosixPath @@ -5,54 +12,90 @@ from astropy.io import registry +from ..utils import StokesWarning + __doctest_skip__ = ['SpectralCubeRead', 'SpectralCubeWrite', 'StokesSpectralCubeRead', 'StokesSpectralCubeWrite', 'LowerDimensionalObjectWrite'] +DOCSTRING_READ_TEMPLATE = """ +Read and parse a dataset and return as a {clsname} -class SpectralCubeRead(registry.UnifiedReadWrite): - """ - Read and parse a dataset and return as a SpectralCube +This allows easily reading a dataset in several supported data +formats using syntax such as:: + + >>> from spectral_cube import {clsname} + >>> cube1 = {clsname}.read('cube.fits', format='fits') + >>> cube2 = {clsname}.read('cube.image', format='casa') + +{notes} + +Get help on the available readers for {clsname} using the``help()`` method:: + + >>> {clsname}.read.help() # Get help reading {clsname} and list supported formats + >>> {clsname}.read.help('fits') # Get detailed help on {clsname} FITS reader + >>> {clsname}.read.list_formats() # Print list of available formats + +See also: http://docs.astropy.org/en/stable/io/unified.html + +Parameters +---------- +*args : tuple, optional + Positional arguments passed through to data reader. If supplied the + first argument is typically the input filename. +format : str + File format specifier. +**kwargs : dict, optional + Keyword arguments passed through to data reader. + +Returns +------- +cube : `{clsname}` + {clsname} corresponding to dataset - This allows easily reading a dataset in several supported data - formats using syntax such as:: +Notes +----- +""" - >>> from spectral_cube import SpectralCube - >>> cube1 = SpectralCube.read('cube.fits', format='fits') - >>> cube2 = SpectralCube.read('cube.image', format='casa') +DOCSTRING_WRITE_TEMPLATE = """ +Write this {clsname} object out in the specified format. - If the file contains Stokes axes, they will automatically be dropped. If - you want to read in all Stokes informtion, use - :meth:`~spectral_cube.StokesSpectralCube.read` instead. +This allows easily writing a dataset in many supported data formats +using syntax such as:: - Get help on the available readers for ``SpectralCube`` using the``help()`` method:: + >>> data.write('data.fits', format='fits') - >>> SpectralCube.read.help() # Get help reading SpectralCube and list supported formats - >>> SpectralCube.read.help('fits') # Get detailed help on SpectralCube FITS reader - >>> SpectralCube.read.list_formats() # Print list of available formats +Get help on the available writers for {clsname} using the``help()`` method:: - See also: http://docs.astropy.org/en/stable/io/unified.html + >>> {clsname}.write.help() # Get help writing {clsname} and list supported formats + >>> {clsname}.write.help('fits') # Get detailed help on {clsname} FITS writer + >>> {clsname}.write.list_formats() # Print list of available formats - Parameters - ---------- - *args : tuple, optional - Positional arguments passed through to data reader. If supplied the - first argument is typically the input filename. - format : str - File format specifier. - **kwargs : dict, optional - Keyword arguments passed through to data reader. +See also: http://docs.astropy.org/en/stable/io/unified.html - Returns - ------- - cube : `SpectralCube` - SpectralCube corresponding to dataset +Parameters +---------- +*args : tuple, optional + Positional arguments passed through to data writer. If supplied the + first argument is the output filename. +format : str + File format specifier. +**kwargs : dict, optional + Keyword arguments passed through to data writer. - Notes - ----- - """ +Notes +----- +""" + + +class SpectralCubeRead(registry.UnifiedReadWrite): + + __doc__ = DOCSTRING_READ_TEMPLATE.format(clsname='SpectralCube', + notes="If the file contains Stokes axes, they will automatically be dropped. If " + "you want to read in all Stokes informtion, use " + ":meth:`~spectral_cube.StokesSpectralCube.read` instead.") def __init__(self, instance, cls): super().__init__(instance, cls, 'read') @@ -66,35 +109,9 @@ def __call__(self, filename, *args, **kwargs): class SpectralCubeWrite(registry.UnifiedReadWrite): - """ - Write this SpectralCube object out in the specified format. - - This allows easily writing a spectral cube in many supported data formats - using syntax such as:: - >>> cube.write('cube.fits', format='fits') + __doc__ = DOCSTRING_WRITE_TEMPLATE.format(clsname='SpectralCube') - Get help on the available writers for ``SpectralCube`` using the``help()`` method:: - - >>> SpectralCube.write.help() # Get help writing SpectralCube and list supported formats - >>> SpectralCube.write.help('fits') # Get detailed help on SpectralCube FITS writer - >>> SpectralCube.write.list_formats() # Print list of available formats - - See also: http://docs.astropy.org/en/stable/io/unified.html - - Parameters - ---------- - *args : tuple, optional - Positional arguments passed through to data writer. If supplied the - first argument is the output filename. - format : str - File format specifier. - **kwargs : dict, optional - Keyword arguments passed through to data writer. - - Notes - ----- - """ def __init__(self, instance, cls): super().__init__(instance, cls, 'write') @@ -103,46 +120,11 @@ def __call__(self, *args, serialize_method=None, **kwargs): class StokesSpectralCubeRead(registry.UnifiedReadWrite): - """ - Read and parse a dataset and return as a StokesSpectralCube - - This allows easily reading a dataset in several supported data formats - using syntax such as:: - >>> from spectral_cube import StokesSpectralCube - >>> cube1 = StokesSpectralCube.read('cube.fits', format='fits') - >>> cube2 = StokesSpectralCube.read('cube.image', format='casa') - - If the file contains Stokes axes, they will be read in. If you are only - interested in the unpolarized emission (I), you can use - :meth:`~spectral_cube.SpectralCube.read` instead. - - Get help on the available readers for ``StokesSpectralCube`` using the``help()`` method:: - - >>> StokesSpectralCube.read.help() # Get help reading StokesSpectralCube and list supported formats - >>> StokesSpectralCube.read.help('fits') # Get detailed help on StokesSpectralCube FITS reader - >>> StokesSpectralCube.read.list_formats() # Print list of available formats - - See also: http://docs.astropy.org/en/stable/io/unified.html - - Parameters - ---------- - *args : tuple, optional - Positional arguments passed through to data reader. If supplied the - first argument is typically the input filename. - format : str - File format specifier. - **kwargs : dict, optional - Keyword arguments passed through to data reader. - - Returns - ------- - cube : `StokesSpectralCube` - StokesSpectralCube corresponding to dataset - - Notes - ----- - """ + __doc__ = DOCSTRING_READ_TEMPLATE.format(clsname='StokesSpectralCube', + notes="If the file contains Stokes axes, they will be read in. If you are only " + "interested in the unpolarized emission (I), you can use " + ":meth:`~spectral_cube.SpectralCube.read` instead.") def __init__(self, instance, cls): super().__init__(instance, cls, 'read') @@ -156,35 +138,9 @@ def __call__(self, filename, *args, **kwargs): class StokesSpectralCubeWrite(registry.UnifiedReadWrite): - """ - Write this StokesSpectralCube object out in the specified format. - - This allows easily writing a spectral cube in many supported data formats - using syntax such as:: - - >>> cube.write('cube.fits', format='fits') - - Get help on the available writers for ``StokesSpectralCube`` using the``help()`` method:: - >>> StokesSpectralCube.write.help() # Get help writing StokesSpectralCube and list supported formats - >>> StokesSpectralCube.write.help('fits') # Get detailed help on StokesSpectralCube FITS writer - >>> StokesSpectralCube.write.list_formats() # Print list of available formats + __doc__ = DOCSTRING_WRITE_TEMPLATE.format(clsname='StokesSpectralCube') - See also: http://docs.astropy.org/en/stable/io/unified.html - - Parameters - ---------- - *args : tuple, optional - Positional arguments passed through to data writer. If supplied the - first argument is the output filename. - format : str - File format specifier. - **kwargs : dict, optional - Keyword arguments passed through to data writer. - - Notes - ----- - """ def __init__(self, instance, cls): super().__init__(instance, cls, 'write') @@ -193,37 +149,31 @@ def __call__(self, *args, serialize_method=None, **kwargs): class LowerDimensionalObjectWrite(registry.UnifiedReadWrite): - """ - Write this object out in the specified format. - - This allows easily writing a data object in many supported data formats - using syntax such as:: - - >>> data.write('data.fits', format='fits') - - Get help on the available writers using the``help()`` method, e.g.:: - >>> LowerDimensionalObject.write.help() # Get help writing LowerDimensionalObject and list supported formats - >>> LowerDimensionalObject.write.help('fits') # Get detailed help on LowerDimensionalObject FITS writer - >>> LowerDimensionalObject.write.list_formats() # Print list of available formats + __doc__ = DOCSTRING_WRITE_TEMPLATE.format(clsname='LowerDimensionalObject') - See also: http://docs.astropy.org/en/stable/io/unified.html - - Parameters - ---------- - *args : tuple, optional - Positional arguments passed through to data writer. If supplied the - first argument is the output filename. - format : str - File format specifier. - **kwargs : dict, optional - Keyword arguments passed through to data writer. - - Notes - ----- - """ def __init__(self, instance, cls): super().__init__(instance, cls, 'write') def __call__(self, *args, serialize_method=None, **kwargs): registry.write(self._instance, *args, **kwargs) + + +def normalize_cube_stokes(cube, target_cls=None): + + from ..spectral_cube import BaseSpectralCube + from ..stokes_spectral_cube import StokesSpectralCube + + if target_cls is BaseSpectralCube and isinstance(cube, StokesSpectralCube): + if hasattr(cube, 'I'): + warnings.warn("Cube is a Stokes cube, " + "returning spectral cube for I component", + StokesWarning) + return cube.I + else: + raise ValueError("Spectral cube is a Stokes cube that " + "does not have an I component") + elif target_cls is StokesSpectralCube and isinstance(cube, BaseSpectralCube): + cube = StokesSpectralCube({'I': cube}) + else: + return cube diff --git a/spectral_cube/io/fits.py b/spectral_cube/io/fits.py index 8e71082b4..b728d5305 100644 --- a/spectral_cube/io/fits.py +++ b/spectral_cube/io/fits.py @@ -209,19 +209,8 @@ def load_fits_cube(input, hdu=0, meta=None, target_cls=None, **kwargs): raise FITSReadError("Data should be 3- or 4-dimensional") - if target_cls is BaseSpectralCube and isinstance(cube, StokesSpectralCube): - if hasattr(cube, 'I'): - warnings.warn("Cube is a Stokes cube, " - "returning spectral cube for I component", - StokesWarning) - return cube.I - else: - raise ValueError("Spectral cube is a Stokes cube that " - "does not have an I component") - elif target_cls is StokesSpectralCube and isinstance(cube, BaseSpectralCube): - cube = StokesSpectralCube({'I': cube}) - else: - return cube + from .core import normalize_cube_stokes + return normalize_cube_stokes(cube, target_cls=target_cls) def write_fits_cube(cube, filename, overwrite=False, diff --git a/spectral_cube/tests/test_io.py b/spectral_cube/tests/test_io.py index ae3e43a8e..4df3e0324 100644 --- a/spectral_cube/tests/test_io.py +++ b/spectral_cube/tests/test_io.py @@ -3,7 +3,6 @@ import numpy as np from astropy.io import fits as pyfits from astropy import units as u -from ..io import class_lmv, fits from .. import SpectralCube, StokesSpectralCube from ..lower_dimensional_structures import (OneDSpectrum, VaryingResolutionOneDSpectrum) From cad689fe72635b9ca246bf46e61fdebf9f126bba Mon Sep 17 00:00:00 2001 From: Thomas Robitaille Date: Wed, 15 Jan 2020 16:06:32 +0000 Subject: [PATCH 6/6] Added changelog entry --- CHANGES.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGES.rst b/CHANGES.rst index 177658099..51661be62 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -6,6 +6,8 @@ - Refactor package infrastructure to no longer use astropy-helpers. #599 +- Switch to using unified I/O infrastructure from Astropy. #600 + 0.4.5 (unreleased) ------------------ - Added support for casatools-based io in #541 and beam reading from CASA