1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236
|
import warnings
from astropy import units as u
from astropy.io import registry as io_registry
from radio_beam import Beam, Beams
from .. import DaskSpectralCube, StokesSpectralCube, BooleanArrayMask, DaskVaryingResolutionSpectralCube
from ..spectral_cube import BaseSpectralCube
from .. import cube_utils
from .. utils import BeamWarning
from .. import wcs_utils
from casa_formats_io import getdesc, coordsys_to_astropy_wcs, image_to_dask
# Read and write from a CASA image. This has a few
# complications. First, by default CASA does not return the
# "python order" and so we either have to transpose the cube on
# read or have dueling conventions. Second, CASA often has
# degenerate stokes axes present in unpredictable places (3rd or
# 4th without a clear expectation). We need to replicate these
# when writing but don't want them in memory. By default, try to
# yield the same array in memory that we would get from astropy.
def is_casa_image(origin, filepath, fileobj, *args, **kwargs):
# See note before StringWrapper definition
from .core import StringWrapper
if filepath is None and len(args) > 0:
if isinstance(args[0], StringWrapper):
filepath = args[0].value
elif isinstance(args[0], str):
filepath = args[0]
return filepath is not None and filepath.lower().endswith('.image')
def load_casa_image(filename, skipdata=False, memmap=True,
skipvalid=False, skipcs=False, target_cls=None, use_dask=None,
target_chunksize=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
be suppressed. The object holds the coordsys object from the image in
memory.
"""
if use_dask is None:
use_dask = True
if not use_dask:
raise ValueError("Loading CASA datasets is not possible with use_dask=False")
from .core import StringWrapper
if isinstance(filename, StringWrapper):
filename = filename.value
# read in the data
if not skipdata:
data = image_to_dask(filename, memmap=memmap, target_chunksize=target_chunksize)
# CASA stores validity of data as a mask
if skipvalid:
valid = None
else:
try:
valid = image_to_dask(filename, memmap=memmap, mask=True, target_chunksize=target_chunksize)
except FileNotFoundError:
valid = None
# transpose is dealt with within the cube object
# read in coordinate system object
desc = getdesc(filename)
casa_cs = desc['_keywords_']['coords']
if 'units' in desc['_keywords_']:
unit = desc['_keywords_']['units']
else:
unit = ''
imageinfo = desc['_keywords_']['imageinfo']
if 'perplanebeams' in imageinfo:
beam_ = {'beams': imageinfo['perplanebeams']}
beam_['nStokes'] = beam_['beams'].pop('nStokes')
beam_['nChannels'] = beam_['beams'].pop('nChannels')
elif 'restoringbeam' in imageinfo:
beam_ = imageinfo['restoringbeam']
else:
beam_ = {}
wcs = coordsys_to_astropy_wcs(casa_cs)
del casa_cs
if 'major' in beam_:
beam = Beam(major=u.Quantity(beam_['major']['value'], unit=beam_['major']['unit']),
minor=u.Quantity(beam_['minor']['value'], unit=beam_['minor']['unit']),
pa=u.Quantity(beam_['positionangle']['value'], unit=beam_['positionangle']['unit']),
)
elif 'beams' in beam_:
bdict = beam_['beams']
# check if stokes1 or stokes 2 is present. if not assume stokes_params = ['I']
if 'stokes1' in desc['_keywords_']['coords']:
stokes_params = desc['_keywords_']['coords']['stokes1']['stokes']
elif 'stokes2' in desc['_keywords_']['coords']:
stokes_params = desc['_keywords_']['coords']['stokes2']['stokes']
else:
warnings.warn("No stokes information found in CASA image. Assuming Stokes I.")
stokes_params = ['I']
nbeams = len(bdict)
nchan = beam_['nChannels']
assert nbeams == nchan * beam_['nStokes']
beams = {}
for ii, stokes_name in enumerate(stokes_params):
majors = [u.Quantity(bdict[f"*{ii*nchan+chan}"]['major']['value'],
bdict[f"*{ii*nchan+chan}"]['major']['unit']) for chan in range(nchan)]
minors = [u.Quantity(bdict[f"*{ii*nchan+chan}"]['minor']['value'],
bdict[f"*{ii*nchan+chan}"]['minor']['unit']) for chan in range(nchan)]
pas = [u.Quantity(bdict[f"*{ii*nchan+chan}"]['positionangle']['value'],
bdict[f"*{ii*nchan+chan}"]['positionangle']['unit']) for chan in range(nchan)]
beams[stokes_name] = Beams(major=u.Quantity(majors),
minor=u.Quantity(minors),
pa=u.Quantity(pas))
else:
warnings.warn("No beam information found in CASA image.",
BeamWarning)
# don't need this yet
# stokes = get_casa_axis(temp_cs, wanttype="Stokes", skipdeg=False,)
# if stokes == None:
# order = np.arange(self.data.ndim)
# else:
# order = []
# for ax in np.arange(self.data.ndim+1):
# if ax == stokes:
# continue
# order.append(ax)
# self.casa_cs = ia.coordsys(order)
# This should work, but coordsys.reorder() has a bug
# on the error checking. JIRA filed. Until then the
# axes will be reversed from the original.
# if transpose == True:
# new_order = np.arange(self.data.ndim)
# new_order = new_order[-1*np.arange(self.data.ndim)-1]
# print new_order
# self.casa_cs.reorder(new_order)
meta = {'filename': filename,
'BUNIT': unit}
if wcs.naxis == 3:
data, wcs_slice = cube_utils._orient(data, wcs)
if valid is not None:
valid, _ = cube_utils._orient(valid, wcs)
mask = BooleanArrayMask(valid, wcs_slice)
else:
mask = None
if 'beam' in locals():
cube = DaskSpectralCube(data, wcs_slice, mask, meta=meta, beam=beam)
elif 'beams' in locals():
cube = DaskVaryingResolutionSpectralCube(data, wcs_slice, mask, meta=meta,
beams=beams['I'])
else:
cube = DaskSpectralCube(data, wcs_slice, mask, meta=meta)
# with #592, this is no longer true
# we've already loaded the cube into memory because of CASA
# limitations, so there's no reason to disallow operations
# cube.allow_huge_operations = True
if mask is not None:
assert cube.mask.shape == cube.shape
elif wcs.naxis == 4:
if valid is not None:
valid, _ = cube_utils._split_stokes(valid, wcs)
data, wcs = cube_utils._split_stokes(data, wcs)
mask = {}
for component in data:
data_, wcs_slice = cube_utils._orient(data[component], wcs)
if valid is not None:
valid_, _ = cube_utils._orient(valid[component], wcs)
mask[component] = BooleanArrayMask(valid_, wcs_slice)
else:
mask[component] = None
if 'beam' in locals():
data[component] = DaskSpectralCube(data_, wcs_slice, mask[component],
meta=meta, beam=beam)
elif 'beams' in locals():
data[component] = DaskVaryingResolutionSpectralCube(data_,
wcs_slice,
mask[component],
meta=meta,
beams=beams[component])
else:
data[component] = DaskSpectralCube(data_, wcs_slice, mask[component],
meta=meta)
data[component].allow_huge_operations = True
cube = StokesSpectralCube(stokes_data=data)
if mask['I'] is not None:
assert cube.I.mask.shape == cube.shape
assert wcs_utils.check_equality(cube.I.mask._wcs, cube.wcs)
else:
raise ValueError("CASA image has {0} dimensions, and therefore "
"is not readable by spectral-cube.".format(wcs.naxis))
from .core import normalize_cube_stokes
return normalize_cube_stokes(cube, target_cls=target_cls)
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)
io_registry.register_reader('casa', StokesSpectralCube, load_casa_image)
io_registry.register_reader('casa_image', StokesSpectralCube, load_casa_image)
io_registry.register_identifier('casa', StokesSpectralCube, is_casa_image)
|