diff --git a/spectral_cube/io/casa_image.py b/spectral_cube/io/casa_image.py index 5659e103f..d1f989222 100644 --- a/spectral_cube/io/casa_image.py +++ b/spectral_cube/io/casa_image.py @@ -16,7 +16,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 # Read and write from a CASA image. This has a few # complications. First, by default CASA does not return the @@ -116,7 +116,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 @@ -264,8 +264,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)