def test_flat_stays_flat(self): recs = [{"flat1": 1, "flat2": 2}, {"flat3": 3, "flat2": 4}] result = nested_to_record(recs) expected = recs assert result == expected
def test_flat_stays_flat(self): recs = [dict(flat1=1, flat2=2), dict(flat1=3, flat2=4)] result = nested_to_record(recs) expected = recs assert result == expected
def test_with_max_level(self, max_level, expected, max_level_test_input_data): # GH23843: Enhanced JSON normalize output = nested_to_record(max_level_test_input_data, max_level=max_level) assert output == expected
def convert_image(infile, outfile=None, artifacts=None, compressor=None, chunk_shape=(-1, -1, 1, 1), nofile=False): """ Convert legacy CASA or FITS format Image to xarray Image Dataset and zarr storage format This function requires CASA6 casatools module. Parameters ---------- infile : str Input image filename (.image or .fits format) outfile : str Output zarr filename. If None, will use infile name with .img.zarr extension artifacts : list of str List of other image artifacts to include if present with infile. Default None uses ``['mask','model','pb','psf','residual','sumwt','weight']`` compressor : numcodecs.blosc.Blosc The blosc compressor to use when saving the converted data to disk using zarr. If None the zstd compression algorithm used with compression level 2. chunk_shape: 4-D tuple of ints Shape of desired chunking in the form of (x, y, channels, polarization), use -1 for entire axis in one chunk. Default is (-1, -1, 1, 1) Note: chunk size is the product of the four numbers (up to the actual size of the dimension) nofile : bool Allows legacy Image to be directly read without file conversion. If set to true, no output file will be written and entire Image will be held in memory. Requires ~4x the memory of the Image size. Default is False Returns ------- xarray.core.dataset.Dataset new xarray Datasets of Image data contents """ from casatools import image as ia import numpy as np import os import time from itertools import cycle from pandas.io.json._normalize import nested_to_record import xarray from xarray import Dataset as xd from xarray import DataArray as xa from numcodecs import Blosc print("converting Image...") infile = os.path.expanduser(infile) prefix = infile[:infile.rindex('.')] suffix = infile[infile.rindex('.') + 1:] # sanitize to avoid KeyError when calling imtypes later while suffix.endswith('/'): suffix = suffix[:-1] if outfile == None: outfile = prefix + '.img.zarr' else: outfile = os.path.expanduser(outfile) if not nofile: tmp = os.system("rm -fr " + outfile) begin = time.time() if compressor is None: compressor = Blosc(cname='zstd', clevel=2, shuffle=0) IA = ia() # all image artifacts will go in same zarr file and share common dimensions if possible # check for meta data compatibility # store necessary coordinate conversion data if artifacts == None: imtypes = [ 'image.pbcor', 'mask', 'model', 'pb', 'psf', 'residual', 'sumwt', 'weight' ] if suffix not in imtypes: imtypes = [suffix] + imtypes else: imtypes = [suffix] + artifacts meta, tm, diftypes, difmeta, xds = {}, {}, [], [], [] for imtype in imtypes: if os.path.exists(prefix + '.' + imtype): rc = IA.open(prefix + '.' + imtype) summary = IA.summary( list=False) # imhead would be better but chokes on big images ims = tuple(IA.shape()) # image shape coord_names = [ ss.replace(' ', '_').lower().replace('stokes', 'pol').replace( 'frequency', 'chan') for ss in summary['axisnames'] ] # compute world coordinates for spherical dimensions # the only way to know is to check the units for angular types (i.e. radians) sphr_dims = [ dd for dd in range(len(ims)) if summary['axisunits'][dd] == 'rad' ] coord_idxs = np.mgrid[[ range(ims[dd]) if dd in sphr_dims else range(1) for dd in range(len(ims)) ]] coord_idxs = coord_idxs.reshape(len(ims), -1) coord_world = IA.coordsys().toworldmany( coord_idxs.astype(float))['numeric'] coord_world = coord_world[sphr_dims].reshape( (len(sphr_dims), ) + tuple(np.array(ims)[sphr_dims])) spi = ['d' + str(dd) for dd in sphr_dims] coords = dict([(coord_names[dd], (spi, coord_world[di])) for di, dd in enumerate(sphr_dims)]) # compute world coordinates for cartesian dimensions cart_dims = [dd for dd in range(len(ims)) if dd not in sphr_dims] coord_idxs = np.mgrid[[ range(ims[dd]) if dd in cart_dims else range(1) for dd in range(len(ims)) ]] coord_idxs = coord_idxs.reshape(len(ims), -1) coord_world = IA.coordsys().toworldmany( coord_idxs.astype(float))['numeric'] coord_world = coord_world[cart_dims].reshape( (len(cart_dims), ) + tuple(np.array(ims)[cart_dims])) for dd, cs in enumerate(list(coord_world)): spi = tuple([ slice(None) if di == dd else slice(1) for di in range(cs.ndim) ]) coords.update(dict([(coord_names[cart_dims[dd]], cs[spi][0])])) # store metadata for later tm['coords'] = coords tm['dsize'] = np.array(summary['shape']) tm['dims'] = [ coord_names[di] if di in cart_dims else 'd' + str(di) for di in range(len(ims)) ] # store rest of image meta data as attributes omits = [ 'axisnames', 'incr', 'hasmask', 'masks', 'defaultmask', 'ndim', 'refpix', 'refval', 'shape', 'tileshape', 'messages' ] nested = [ kk for kk in summary.keys() if isinstance(summary[kk], dict) ] tm['attrs'] = dict([(kk.lower(), summary[kk]) for kk in summary.keys() if kk not in omits + nested]) tm['attrs'].update( dict([(kk, list(nested_to_record(summary[kk], sep='.').items())) for kk in nested])) # parse messages for additional keys, drop duplicate info omits = [ 'image_name', 'image_type', 'image_quantity', 'pixel_mask(s)', 'region(s)', 'image_units' ] for msg in summary['messages']: line = [ tuple(kk.split(':')) for kk in msg.lower().split('\n') if ': ' in kk ] line = [(kk[0].strip().replace(' ', '_'), kk[1].strip()) for kk in line] line = [ll for ll in line if ll[0] not in omits] tm['attrs'].update(dict(line)) # save metadata from first image product (the image itself) # compare later image products to see if dimensions match up # Note: only checking image dimensions, NOT COORDINATE VALUES!! if meta == {}: meta = dict(tm) elif (np.any(meta['dsize'] != np.array(summary['shape'])) & (imtype != 'sumwt')): diftypes += [imtype] difmeta += [tm] imtypes = [_ for _ in imtypes if _ != imtype] rc = IA.close() else: imtypes = [_ for _ in imtypes if _ != imtype] print('compatible components: ', imtypes) print('separate components: ', diftypes) # process all image artifacts with compatible metadata to same zarr file # partition by channel, read each image artifact for each channel dsize, chan_dim = meta['dsize'], meta['dims'].index('chan') pt1, pt2 = [-1 for _ in range(len(dsize))], [-1 for _ in range(len(dsize))] if chunk_shape[3] <= 0: chunk_shape[3] = dsize[chan_dim] chan_batch = dsize[chan_dim] if nofile else chunk_shape[3] for chan in range(0, dsize[chan_dim], chan_batch): print('processing channel ' + str(chan + 1) + ' of ' + str(dsize[chan_dim]), end='\r') pt1[chan_dim], pt2[chan_dim] = chan, chan + chan_batch - 1 chunk_coords = dict(meta['coords']) # only want one freq channel coord chunk_coords['chan'] = np.arange(chan_batch) + chan xdas = {} for imtype in imtypes: rc = IA.open(prefix + '.' + imtype) # extract pixel data imchunk = IA.getchunk(pt1, pt2) if imtype == 'fits': imtype = 'image' if imtype == 'mask': xdas['deconvolve'] = xa(imchunk.astype(bool), dims=meta['dims']) elif imtype == 'sumwt': xdas[imtype] = xa(imchunk.reshape(imchunk.shape[2], 1), dims=['pol', 'chan']) else: xdas[imtype] = xa(imchunk, dims=meta['dims']) # extract mask summary = IA.summary(list=False) if len(summary['masks']) > 0: imchunk = IA.getchunk(pt1, pt2, getmask=True) xdas['mask'] = xa(imchunk.astype(bool), dims=meta['dims']) rc = IA.close() chunking = dict([(meta['dims'][ii], chunk_shape[ii]) for ii in range(len(meta['dims'])) if chunk_shape[ii] > 0]) xds = xd(xdas, coords=chunk_coords, attrs=meta['attrs']).chunk(chunking) # for everyone's sanity, lets make sure the dimensions are ordered the same way as visibility data if ('pol' in xds.dims): xds = xds.transpose(xds.image.dims[0], xds.image.dims[1], 'chan', 'pol') if (chan == 0) and (not nofile): # xds = xd(xdas, coords=chunk_coords, attrs=nested_to_record(meta['attrs'], sep='_')) encoding = dict( zip(list(xds.data_vars), cycle([{ 'compressor': compressor }]))) xds.to_zarr(outfile, mode='w', encoding=encoding) elif not nofile: xds.to_zarr(outfile, mode='a', append_dim='chan') print("processed image size " + str(dsize) + " in " + str(np.float32(time.time() - begin)) + " seconds") if not nofile: xds = xarray.open_zarr(outfile) return xds