' mwispy module: mosaic ' __version__ = '2.0' __author__ = 'Shaobo Zhang' __all__ = ['mosaic'] import os, time, glob, logging import numpy as np from astropy.io import fits from math import floor, ceil ###the valid l/b range of mosaic, slightly larger than the MWISP project. valid_l1 = -11 valid_l2 = 251 valid_b1 = -12 valid_b2 = 12 ###output data type output_dtype = np.float32 def mosaic(*crange, sb='U', path=None, prefix='', suffix='.fits', \ undone=None, output='mosaic', weightcube=False, silent=False, display=False, ): ''' Mosaic fits file of DLH survey Parameters ---------- crange : six floats includes gl1, gl2, gb1, gb2, velo1, velo2. which indicate galactic coordinate range (in deg) and velocity range (in km/s). If a "undone" is provide, these ranges are ignored. sb : str Sideband, for 12CO, sb='U'; for 13CO, sb='L'; for C18O, sb='L2'. path : str or list of strs, optional A directories or a List of the directories containing datacube and rms files. If there are multiply data in different directories, the path on the top of the list will be used. Default path is the directories of data on SASMA. prefix : str or tuple of two strs, optional Prefix of datacube and rms files. if a string, datacube and rms will both use it. if a tuple of two strings, datacube will use the first one, while rms will use the second. suffix : str or tuple of two strs, optional Suffix of datacube and rms files. Similar to "prefix". undone : str, optional Prefix of the set of mosaicked files, which may miss several cells in the last run. You can set undone to the prefix of a mosaiked datacube to patch cells on it without re-mosaic the whole one. "crange" will be ignored, and mosaic will continue to use ranges in the undone datacube. e.g. >>> mosaic(undone='mosaic', output='newmosaic') output : str, optional Prefix of output file names, default name is 'mosaic'. weightcube: bool, optional Whether to use a 3-dimensional weight cube. Default is False. Set to True if part of cube is masked as nan. Such as cleaned cubes in which velocities are chopped. silent : bool, optional Whether to suppress message in the terminal. Default is False. display : bool, optional Whether to display the progress on a window. Default is False. display=True is not recommended when running on remote servers. Notes ----- All the input fits file (data and rms) must be the same as the grid define by template. Versions -------- Nov,09,2011,v1.0 Nov,28,2011,v1.1 add the fitspath and rmspath keywords. Jan,16,2012,v1.2 fix a error when mosaic region with gl > 180 Apr,12,2012,v1.3 change the output file name. output a mask fits file for the mosaic. May,09,2012,v1.4 complete the header of the result fits file. the mask file is now called 'coverage' file to fit its function. add the keywords "display" for users who do not support window. Dec,25,2013,v1.5 accept rms file created by procedure cuberms. Oct,13,2015,v1.6 remove keyword "fitspath" and "rmspath". new keyword "path" now accepts string array, procedure will automatically search datacube and rms file in these directories. Dec,09,2016,v1.7 add history to indicate the completeness of the mosaic. add keyword 'output' for the names of the output files. add keyword 'silent' to suppress any message in terminal. fix spelling mistakes. May,17,2018,v1.8 fix a minor error that certain input range does not match the output range. Nov,02,2018,v1.9 rewrite in python3 May,13,2021,v1.10 add keyword "undone" Jan,12,2023,v2.0 log file is added to the output set to help keyword "undone" to work better. correct the CDELT1 error when MWISP-2 data are involving. add keyword "prefix", "suffix" to mosaic cleaned datacubes. Dec,17,2024,v2.1 no longer prompt "Error" if velocity range of cells are inconsistent. mosaic can be done even if the velocity range of some cell datacubes were chopped differently. zero will be filled if chopped range is smaller than that of mosaic. Examples -------- >>> pa = ['./', '/share/data/mwisp/G020+00', '/share/data/mwisp/G030+00'] >>> mwispy.mosaic(29, 32.5, -1, 1.2, -30, 30, sb='L', path=pa, output='Test', silent=True) ''' ###Syntax prompt if len(crange) != 6 and undone is None: print("Syntax - mwispy.mosaic(l1, l2, b1, b2, v1, v2, sb='U', path='./', prefix='', suffix='.fits', \n\ undone=None, output='mosaic', weightcube=False, display=False, silent=False)") return ###log file config logging.basicConfig(filename=_getMosaicLogName(output, sb), filemode='w', \ level=logging.INFO, format='%(asctime)s %(message)s') def _prompt(message): if not silent: print(message) logging.info(message) #prepare paths paths=[] if path is None: #no path specified, use default paths = glob.glob('/share/data/mwisp/R*') + glob.glob('/share/data/mwisp/G*') else: #list all possible paths if type(path) is str: path=[path] for apath in path: if type(apath) is str: paths += glob.glob(apath) #prepare prefix if type(prefix) is str: prefix = (prefix, prefix) if type(suffix) is str: suffix = (suffix, suffix) #set display if display: import matplotlib.pyplot as plt plt.ion() if undone is None: #standardize crange l1, l2, b1, b2, v1, v2 = crange #modify gl in the range of -60 ~ 300 deg, since there is no observation at 300 deg for MWISP l1, l2 = l1 % 360, l2 % 360 if l1 > 300: l1 -= 360 if l2 > 300: l2 -= 360 #sort range if l1 > l2: l2, l1 = l1, l2 if b1 > b2: b2, b1 = b1, b2 if v1 > v2: v2, v1 = v1, v2 #return if out of MWISP range if l1 < -11 or l2 > 251 or b1 < -12 or b2 > 12: silent = False _prompt("Error - Galactic coordinate is out of the range of DLH survey!") _prompt(" l = 350 ~ 0 ~ 250 deg, b = -5 ~ 5 deg") return ###get shape of output FITS mscNX = int(floor(l2*120) - ceil(l1*120) + 1) if mscNX==0: mscNX = 1 mscCX = int(floor(l2*120) + 1) mscNY = int(floor(b2*120) - ceil(b1*120) + 1) mscCY = int(-ceil(b1*120) + 1) #1-base for FITS file mscCells = [] else: if os.path.abspath(undone) == os.path.abspath(output): _prompt("Error - output should be different from undone.") return try: mscHdr = fits.getheader(_getMosaicCubeName(undone, sb)) mscNX = mscHdr['NAXIS1'] mscNY = mscHdr['NAXIS2'] mscNC = mscHdr['NAXIS3'] mscCX = mscHdr['CRPIX1'] mscCY = mscHdr['CRPIX2'] l1 = _carx2l(mscHdr, mscNX-1) l2 = _carx2l(mscHdr, 0) b1 = _cary2b(mscHdr, 0) b2 = _cary2b(mscHdr, mscNY-1) v1 = _c2v(mscHdr, 0) v2 = _c2v(mscHdr, mscNC-1) mscHdr = fits.getheader(_getMosaicWeightName(undone, sb)) if mscHdr['NAXIS']>2: weightcube=True mscLog = open(_getMosaicLogName(undone, sb), 'r') mscCells = mscLog.readlines()[-1] mscCells = mscCells.split('Done Cells:')[-1] mscCells = mscCells.split() except: silent = False _prompt('Error - Unable to load undone files!') _prompt(' Make sure datacube/weight/coverage/log files are all ready.') return #related FITS file range with value*10) cellL1x10 = int(floor(l1*2)*5) cellL2x10 = int(ceil(l2*2)*5) cellsLx10 = range(cellL1x10, cellL2x10+1, 5) cellB1x10 = int(floor(b1*2)*5) cellB2x10 = int(ceil(b2*2)*5) cellsBx10 = range(cellB1x10, cellB2x10+1, 5) cellNum = len(cellsLx10) * len(cellsBx10) time_start = time.time() _prompt('Mosaic begins:') _prompt('GL from %f to %f' % (l1 % 360, l2 % 360)) _prompt('GB from %f to %f' % (b1, b2)) _prompt('V from %f to %f' % (v1, v2)) index, valid = 0, 0 for cellLx10 in cellsLx10: for cellBx10 in cellsBx10: index += 1 cellName = _getCellName(cellLx10, cellBx10) ### check if a cell is already mosaicked when resuming a undone mosaic. if undone is not None: if cellName in mscCells: _prompt('[%d/%d]%s%s: mosaic is already in the undone.' % (index, cellNum, cellName, sb)) continue ###Prepare files #search for cube file observed = False for apath in paths: cellCubeFile = _getCellCubeName(apath, prefix, cellName, sb, suffix) if os.path.exists(cellCubeFile): observed = True break if not observed: _prompt('[%d/%d]%s%s: datacube file not found.' % (index, cellNum, cellName, sb)) continue #search for rms file haveRms = False for apath in paths: cellRmsFile = _getCellRmsName(apath, prefix, cellName, sb, suffix) if os.path.exists(cellRmsFile): haveRms = True break if not haveRms: _prompt('[%d/%d]%s%s: RMS file not found.' % (index, cellNum, cellName, sb)) continue ###Load data #read cube file _prompt("[%d/%d]%s%s: mosaic using: '%s'" % (index, cellNum, cellName, sb, cellCubeFile)) try: cellCubeHDU = fits.open(cellCubeFile)[0] cellHdr = cellCubeHDU.header cellDat = cellCubeHDU.data ### Unify to 4 dimension while cellDat.ndim<4: cellDat = cellDat[np.newaxis] ### Clip velocity nc = cellHdr['NAXIS3'] cv = cellHdr['CRVAL3'] cp = cellHdr['CRPIX3'] cd = cellHdr['CDELT3'] #flip channels if necessary if cd < 0: cellDat = np.flip(cellDat, axis=-3) cd = -cd cp = nc-cp+1 ### get channel range to extract c1 = int(floor((v1*1e3 - cv) / cd + cp -1)) #convert faster than v2c c2 = int(ceil((v2*1e3 - cv) / cd + cp -1)) if c1 > nc-1 or c2 < 0: #out of range _prompt('Error - Velocity from %f to %s is out of the range!' % (v1, v2)) return ### extract channel, fill nan if c1<0 or c2>nc-1 shp = list(cellDat.shape) if c1>=0: cellDat = np.take(cellDat, range(c1, nc), axis=-3) padBefore = None else: shp[-3] = -c1 padBefore = np.empty(shp) padBefore[:] = np.nan if c2<=nc-1: cellDat = np.take(cellDat, range(0, cellDat.shape[-3]-nc+c2+1), axis=-3) padAfter = None else: shp[-3] = c2-nc+1 padAfter = np.empty(shp) padAfter[:] = np.nan if padBefore is not None: cellDat = np.concatenate((padBefore, cellDat), axis=-3) if padAfter is not None: cellDat = np.concatenate((cellDat, padAfter), axis=-3) ### revise keyword ''' if c1 < 0: c1 = 0 if c2 > nc-1: c2 = nc-1 #cellDat = cellDat[:, c1:c2+1, :, :] cellDat = np.take(cellDat, range(c1,c2+1), axis=-3) ''' cellHdr['NAXIS3'] = c2-c1+1 cellHdr['CRPIX3'] = cp-c1 cellHdr['CDELT3'] = cd ###deal with NAN if cellHdr['BITPIX'] > 0: nanDat = cellDat == np.nanmax(cellDat) else: nanDat = (cellDat == -1000) | ~np.isfinite(cellDat) cellDat[nanDat] = 0 except: _prompt('Error - Unable to load datacube file: %s' % (index, cellNum, cellName, sb, cellCubeFile)) continue #read rms file try: cellRmsHDU = fits.open(cellRmsFile)[0] cellRms = np.squeeze(cellRmsHDU.data) ###deal with NAN if cellRmsHDU.header['BITPIX'] > 0: nanRms = cellRms == np.nanmax(cellRms) else: nanRms = (cellRms == -1000) | ~np.isfinite(cellRms) cellRms[nanRms] = np.inf cellWei = 1/cellRms**2 if weightcube: cellWei = np.zeros_like(cellDat)+cellWei cellWei[nanDat] = 0 except: _prompt('Error - Unable to load rms file: %s' % (index, cellNum, cellRmsFile)) continue ###Check SHAPE !!! #cube and weight are ready, it is a valid cell valid += 1 #initiate a mosaic cube and header if valid == 1: if undone is None: ###init an empty one mscNC = cellHdr['NAXIS3'] mscDat = np.zeros((cellDat.shape[0], mscNC, mscNY, mscNX), dtype=output_dtype) if weightcube: mscWei = mscDat.copy() else: mscWei = np.zeros((mscNY, mscNX), dtype=output_dtype) mscCov = np.zeros((mscNY, mscNX), dtype=np.uint8) mscHdr = cellHdr.copy() else: ###load a mosaicked undone mscHDU = fits.open(_getMosaicCubeName(undone, sb))[0] mscHdr = mscHDU.header mscNC = mscHdr['NAXIS3'] mscDat = mscHDU.data mscDat[np.isnan(mscDat)] = 0 mscWei = fits.open(_getMosaicWeightName(undone, sb))[0].data mscWei[np.isnan(mscWei)] = 0 mscDat *= mscWei mscCov = fits.open(_getMosaicCoverageName(undone, sb))[0].data mscV1 = _c2v(mscHdr, 0) mscV2 = _c2v(mscHdr, mscHdr['NAXIS3']-1) else: #check velocity consistency ### check velocity before, but only velocity sameNC = mscHdr['NAXIS3'] == cellHdr['NAXIS3'] #consistent channel number sameV1 = abs(_c2v(cellHdr,0) - mscV1) < mscHdr['CDELT3']/1e3 #consistent velocity range sameV2 = abs(_c2v(cellHdr,cellHdr['NAXIS3']-1) - mscV2) < mscHdr['CDELT3']/1e3 ''' mscV0 = mscHdr['CRVAL3'] % abs(mscHdr['CDELT3']) cellV0 = cellHdr['CRVAL3'] % abs(cellHdr['CDELT3']) sameV0 = (mscV0 - cellV0) / mscV0 < 1e-4 sameVD = abs((mscHdr['CDELT3'] - cellHdr['CDELT3']) / mscHdr['CDELT3']) < 1e-4 ''' if not (sameNC and sameV1 and sameV2): _prompt('Error - Inconsistent velocity dimension!') continue #mosaic dat to mscDat cellDat *= cellWei ###clip cube/weight cellCX = mscCX-cellLx10*12-1 cellCY = mscCY+cellBx10*12-1 cellNX = cellHdr['NAXIS1'] cellNY = cellHdr['NAXIS2'] cellX1 = cellCX-(cellNX-1)//2 cellX2 = cellCX+(cellNX-1)//2 #subscript on the mosaic array, within cellX1:cellX2 cellY1 = cellCY-(cellNY-1)//2 cellY2 = cellCY+(cellNY-1)//2 if cellX2 < 0 or cellX1 > (mscNX-1) or cellY2 < 0 or cellY1 > (mscNY-1): continue #fits is out of the mosaic range #cut cell if on boundary of msc cutX1 = -cellX1 if cellX1<0 else 0 cutX2 = cellNX-cellX2+mscNX-2 if cellX2>mscNX-1 else cellNX-1 cutY1 = -cellY1 if cellY1<0 else 0 cutY2 = cellNY-cellY2+mscNY-2 if cellY2>mscNY-1 else cellNY-1 #position of cell on msc posX1 = max(cellX1, 0) posX2 = min(cellX2, mscNX-1) posY1 = max(cellY1, 0) posY2 = min(cellY2, mscNY-1) cellDat = cellDat[..., cutY1:cutY2+1, cutX1:cutX2+1] mscDat[..., posY1:posY2+1, posX1:posX2+1] += cellDat if weightcube: cellWei = cellWei[..., cutY1:cutY2+1, cutX1:cutX2+1] mscWei[..., posY1:posY2+1, posX1:posX2+1] += cellWei else: cellWei = cellWei[cutY1:cutY2+1, cutX1:cutX2+1] mscWei[posY1:posY2+1, posX1:posX2+1] += cellWei cellNX = 61 cellNY = 61 cellX1 = cellCX-(cellNX-1)//2 cellX2 = cellCX+(cellNX-1)//2 cellY1 = cellCY-(cellNY-1)//2 cellY2 = cellCY+(cellNY-1)//2 if cellX2 < 0 or cellX1 > (mscNX-1) or cellY2 < 0 or cellY1 > (mscNY-1): continue posX1 = max(cellX1, 0) posX2 = min(cellX2, mscNX-1) posY1 = max(cellY1, 0) posY2 = min(cellY2, mscNY-1) mscCov[posY1:posY2+1, posX1:posX2+1] = 1 if display: plt.imshow(mscWei) plt.pause(0.01) mscCells.append(cellName) if valid == 0: _prompt('Error - No file in the selected region!') return _prompt('Writing to files...') mscWei[mscWei == 0] = np.nan mscDat /= mscWei if weightcube: mscWei = np.nanmax(mscWei, axis=-3) mscHdr['NAXIS1'] = mscNX mscHdr['NAXIS2'] = mscNY mscHdr['CTYPE1'] = 'GLON-CAR' mscHdr['CTYPE2'] = 'GLAT-CAR' mscHdr['CRVAL1'] = 0. mscHdr['CRVAL2'] = 0. mscHdr['CDELT1'] = -30/3600 mscHdr['CDELT2'] = 30/3600 if 'CTYPE4' not in mscHdr: mscHdr['CTYPE4'] = ' ' mscHdr['CRVAL4'] = 1 mscHdr['CDELT4'] = 1 mscHdr['CRPIX4'] = 1 mscHdr['CROTA4'] = 0 if mscCX > 21601: mscHdr['CRPIX1'] = mscCX-43200 else: mscHdr['CRPIX1'] = mscCX mscHdr['CRPIX2'] = mscCY mscHdr['BUNIT'] = 'K (T_MB)' if undone is None: mscHdr.remove('OBJECT', ignore_missing=True) mscHdr.remove('GLAT', ignore_missing=True) mscHdr.remove('GLON', ignore_missing=True) mscHdr.remove('PV1_0', ignore_missing=True) mscHdr.remove('PV1_1', ignore_missing=True) mscHdr.remove('PV1_2', ignore_missing=True) mscHdr['HISTORY'] = 'MOSAIC: '+time.strftime("%b %d %Y %H:%M:%S", time.localtime()) else: mscHdr['HISTORY'] = 'MOSAIC PATCH: '+time.strftime("%b %d %Y %H:%M:%S", time.localtime()) mscHdr['HISTORY'] = str('L from %f to %f' % (l1, l2)) mscHdr['HISTORY'] = str('B from %f to %f' % (b1, b2)) mscHdr['HISTORY'] = str('V from %f to %f' % (v1, v2)) mscHdr['HISTORY'] = str('Mosaic Completeness: [%d/%d]' % (len(mscCells), cellNum)) #output mscHDU = fits.PrimaryHDU(mscDat, mscHdr) mscHDU.writeto(_getMosaicCubeName(output, sb), overwrite = True) mscHDU = fits.PrimaryHDU(mscWei, mscHdr) mscHDU.writeto(_getMosaicWeightName(output, sb), overwrite = True) mscHDU = fits.PrimaryHDU(mscCov, mscHdr) mscHDU.writeto(_getMosaicCoverageName(output, sb), overwrite = True) mscRms = 1/np.sqrt(mscWei) #mscRms[~np.isfinite(mscRms)] = np.nan mscHDU = fits.PrimaryHDU(mscRms, mscHdr) mscHDU.writeto(_getMosaicRmsName(output, sb), overwrite = True) time_end = time.time() _prompt('Mosaic successfully.') _prompt('Done in %f seconds.' % (time_end-time_start)) _prompt('Done Cells: '+' '.join(mscCells)) def _getCellName(gl, gb): gl = gl % 3600 return '%04d%+04d' % (gl, gb) def _getCellCubeName(path, prefix, cellName, sb, suffix): return os.path.join(path, '%s%s%s%s' % (prefix[0], cellName, sb, suffix[0])) def _getCellRmsName(path, prefix, cellName, sb, suffix): return os.path.join(path, '%s%s%s_rms%s' % (prefix[1], cellName, sb, suffix[1])) def _getMosaicLogName(output, sb): return '%s_%s.log' % (output, sb) def _getMosaicCubeName(output, sb): return '%s_%s.fits' % (output, sb) def _getMosaicWeightName(output, sb): return '%s_%s_weight.fits' % (output, sb) def _getMosaicCoverageName(output, sb): return '%s_%s_coverage.fits' % (output, sb) def _getMosaicRmsName(output, sb): return '%s_%s_rms.fits' % (output, sb) def _carx2l(hdr, x): return (x - hdr['CRPIX1'] + 1) * hdr['CDELT1'] + hdr['CRVAL1'] def _cary2b(hdr, y): return (y - hdr['CRPIX2'] + 1) * hdr['CDELT2'] + hdr['CRVAL2'] #c is 0-base; v in km/s def _c2v(hdr, c): v = (c - hdr['CRPIX3'] + 1) * hdr['CDELT3'] + hdr['CRVAL3'] return v/1e3 def _v2c(hdr, v): c = (v * 1e3 - hdr['CRVAL3']) / hdr['CDELT3'] + hdr['CRPIX3'] -1 return c if __name__ == '__main__': if 1: import random f = lambda cell:'/share/data/mwisp/G%2s0+00/%sL.fits' % (cell[:2], cell) for l in range(2235, 2245, 5): for b in range(-10, 0, 5): cell = _getCellName(l, b) print(cell, f(cell)) c1 = random.randint(-200, -50) + 10826 c2 = random.randint(50, 200) + 10826 hdu = fits.open(f(cell))[0] hdu.data = np.take(hdu.data, range(c1,c2+1), axis=-3) hdu.header['NAXIS3'] = c2-c1+1 hdu.header['CRPIX3'] = hdu.header['CRPIX3']-c1 hdu.writeto(cell+'L.fits', overwrite=True) mosaic(223,224.5,-1.5,0,-24,24, sb='L',path=['.', '/share/data/mwisp/G220+00']) #mosaic(95.5, 97, 4.5, 6, -10, 10, path='/Users/shaobo/Downloads/', output='mosaic', weightcube=True) #mosaic(95.5, 97, 4.5, 6, -120, 120, path='/Users/shaobo/Downloads/', output='clip1', weightcube=True, prefix=('CLIP','')) #mosaic(path='/Users/shaobo/Downloads/', output='clip3', undone='clip1', weightcube=True, prefix=('CLIP','')) #mosaic(214,217,-1,2,-300,300, path='/share/public/qzyanShare/mwispIfinal/longBaseClean/all_cubes/', output='test1', prefix='NC') #mosaic(214,217,-1,2,-300,300, path='/share/public/qzyanShare/mwispIfinal/longBaseClean/CO13/all_cubes', output='test1', prefix='NC', sb='L')