File: casa_image.py

package info (click to toggle)
spectral-cube 0.6.6-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 2,136 kB
  • sloc: python: 13,236; makefile: 154
file content (236 lines) | stat: -rw-r--r-- 9,280 bytes parent folder | download | duplicates (2)
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)