Merge 6e4908c into 31b0b51

spectral_cube/io/casa_image.py

@@ -355,6 +355,120 @@ def load_casa_image(filename, skipdata=False,
     from .core import normalize_cube_stokes
     return normalize_cube_stokes(cube, target_cls=target_cls)
 
+def casa_image_array_reader(imagename):
+    """
+    Read a CASA image (a folder containing a ``table.f0_TSM0`` file) into a
+    numpy array.
+    """
+    from casatools import table
+    tb = table()
+
+    # load the metadata from the image table
+    tb.open(imagename)
+    dminfo = tb.getdminfo()
+    tb.close()
+
+    from pprint import pprint
+    pprint(dminfo)
+
+    # chunkshape definse how the chunks (array subsets) are written to disk
+    chunkshape = tuple(dminfo['*1']['SPEC']['DEFAULTTILESHAPE'])
+    chunksize = np.product(chunkshape)
+    # the total size defines the final output array size
+    totalshape = dminfo['*1']['SPEC']['HYPERCUBES']['*1']['CubeShape']
+    # the ratio between these tells you how many chunks must be combined
+    # to create a final stack
+    stacks = totalshape / chunkshape
+    nchunks = np.product(totalshape) // np.product(chunkshape)
+
+    img_fn = f'{imagename}/table.f0_TSM0'
+    # each of the chunks is stored in order on disk in fortran-order
+    chunks = [np.memmap(img_fn, dtype='float32', offset=ii*chunksize * 4,
+                        shape=chunkshape, order='F')
+              for ii in range(nchunks)]
+
+    # with all of the chunks stored in the above list, we then need to concatenate
+    # the resulting pieces into a final array
+    # this process was arrived at empirically, but in short:
+    # (1) stack the cubes along the last dimension first
+    # (2) then stack along each dimension until you get to the first
+    rslt = chunks
+    rstacks = list(stacks)
+    jj = 0
+    while len(rstacks) > 0:
+        rstacks.pop()
+        kk = len(stacks) - jj - 1
+        remaining_dims = rstacks
+        if len(remaining_dims) == 0:
+            assert kk == 0
+            rslt = np.concatenate(rslt, 0)
+        else:
+            cut = np.product(remaining_dims)
+            assert cut % 1 == 0
+            cut = int(cut)
+            rslt = [np.concatenate(rslt[ii::cut], kk) for ii in range(cut)]
+        jj += 1
+
+    # this alternative approach puts the chunks in their appropriate spots
+    # but I haven't figured out a way to turn them into the correct full-sized
+    # array.  You could do it by creating a full-sized array with a
+    # rightly-sized memmap, or something like that, but... that's not what
+    # we're trying to accomplish here.  I want an in-memory object that points
+    # to the right things with the right shape, not a copy in memory or on disk
+    #stacks = list(map(int, stacks))
+    #chunk_inds = np.arange(np.product(stacks)).reshape(stacks, order='F')
+
+    #def recursive_relist(x):
+    #    if isinstance(x, list) or isinstance(x, np.ndarray) and len(x) > 0:
+    #        return [recursive_relist(y) for y in x]
+    #    else:
+    #        return chunks[x]
+
+    return rslt
+
+
+
+def casa_image_dask_reader(imagename):
+    """
+    Read a CASA image (a folder containing a ``table.f0_TSM0`` file) into a
+    numpy array.
+    """
+    from casatools import table
+    tb = table()
+
+    # load the metadata from the image table
+    tb.open(imagename)
+    dminfo = tb.getdminfo()
+    tb.close()
+
+    # chunkshape definse how the chunks (array subsets) are written to disk
+    chunkshape = tuple(dminfo['*1']['SPEC']['DEFAULTTILESHAPE'])
+    chunksize = np.product(chunkshape)
+    # the total size defines the final output array size
+    totalshape = dminfo['*1']['SPEC']['HYPERCUBES']['*1']['CubeShape']
+    # the ratio between these tells you how many chunks must be combined
+    # to create a final stack
+    stacks = totalshape // chunkshape
+    nchunks = np.product(totalshape) // np.product(chunkshape)
+
+    img_fn = f'{imagename}/table.f0_TSM0'
+    # each of the chunks is stored in order on disk in fortran-order
+    chunks = [np.memmap(img_fn, dtype='float32', offset=ii*chunksize*4,
+                        shape=chunkshape, order='F').T
+              for ii in range(nchunks)]
+
+    def make_nested_list(chunks, stacks):
+        chunks = [chunks[i*stacks[0]:(i+1)*stacks[0]] for i in range(len(chunks) // stacks[0])]
+        if len(stacks) > 1:
+            return make_nested_list(chunks, stacks[1:])
+        else:
+            return chunks
+
+    chunks = make_nested_list(chunks, stacks)
+
+    return dask.array.block(chunks[0])
+
+
 
 io_registry.register_reader('casa', BaseSpectralCube, load_casa_image)
 io_registry.register_reader('casa_image', BaseSpectralCube, load_casa_image)

spectral_cube/tests/test_casafuncs.py

@@ -8,16 +8,17 @@
 from astropy import units as u
 
 from ..io.casa_masks import make_casa_mask
-from ..io.casa_image import wcs_casa2astropy
+from ..io.casa_image import wcs_casa2astropy, casa_image_array_reader
 from .. import SpectralCube, StokesSpectralCube, BooleanArrayMask, VaryingResolutionSpectralCube
 from . import path
 
 try:
     import casatools
+    from casatools import image
     casaOK = True
 except ImportError:
     try:
-        from taskinit import ia
+        from taskinit import ia as image
         casaOK = True
     except ImportError:
         print("Run in CASA environment.")
@@ -33,7 +34,7 @@ def make_casa_testimage(infile, outname):
         raise Exception("Attempted to make a CASA test image in a non-CASA "
                         "environment")
 
-    ia = casatools.image()
+    ia = image()
 
     ia.fromfits(infile=infile, outfile=outname, overwrite=True)
     ia.unlock()
@@ -63,6 +64,16 @@ def make_casa_testimage(infile, outname):
         ia.done()
 
 
+def make_casa_testimage_of_shape(shape, outfile):
+    ia = image()
+
+    size = np.product(shape)
+    im = np.arange(size).reshape(shape)
+
+    ia.fromarray(outfile=outfile, pixels=im, overwrite=True)
+    ia.close()
+
+
 @pytest.fixture
 def filename(request):
     return request.getfixturevalue(request.param)
@@ -96,6 +107,45 @@ def test_casa_read_stokes(data_advs, tmp_path):
     assert casacube.I.shape == cube.I.shape
 
 
+@pytest.mark.skipif(not casaOK, reason='CASA tests must be run in a CASA environment.')
+@pytest.mark.parametrize('filename', ('data_adv', 'data_advs', 'data_sdav',
+                                      'data_vad', 'data_vsad'),
+                         indirect=['filename'])
+def test_casa_numpyreader_read(filename, tmp_path):
+
+    cube = SpectralCube.read(filename)
+
+    make_casa_testimage(filename, tmp_path / 'casa.image')
+
+    casacube = SpectralCube.read(tmp_path / 'casa.image')
+
+    assert casacube.shape == cube.shape
+
+    arr = casa_image_array_reader(tmp_path / 'casa.image')
+
+    assert np.all(arr == casacube.unmasked_data[:])
+
+
+@pytest.mark.skipif(not casaOK, reason='CASA tests must be run in a CASA environment.')
+@pytest.mark.parametrize('shape', ((129,128,130), (513,128,128), (128,513,128),
+                                   (128,128,513), (512,64,64)),
+                         )
+def test_casa_numpyreader_read_shape(shape, tmp_path):
+
+    cube = SpectralCube.read(filename)
+
+    make_casa_testimage_of_shape(shape, tmp_path / 'casa.image')
+
+    casacube = SpectralCube.read(tmp_path / 'casa.image')
+
+    assert casacube.shape == cube.shape
+
+    arr = casa_image_array_reader(tmp_path / 'casa.image')
+
+    assert np.all(arr == casacube.unmasked_data[:])
+
+
+
 @pytest.mark.skipif(not casaOK, reason='CASA tests must be run in a CASA environment.')
 def test_casa_mask(data_adv, tmp_path):