def read_ms(vis):
''' Read a CASA ms file and return a dictionary of amplitude, phase, uvdistance,
uvangle, frequency (GHz) and time (MJD). Currently only returns the XX IF channel.
vis Name of the visibility (ms) folder
'''
ms.open(vis)
spwinfo = ms.getspectralwindowinfo()
nspw = len(spwinfo.keys())
for i in range(nspw):
print('Working on spw', i)
ms.selectinit(datadescid=0, reset=True)
ms.selectinit(datadescid=i)
if i == 0:
spw = ms.getdata(['amplitude', 'phase', 'u', 'v', 'axis_info'], ifraxis=True)
xxamp = spw['amplitude']
xxpha = spw['phase']
fghz = spw['axis_info']['freq_axis']['chan_freq'][:, 0] / 1e9
band = np.ones_like(fghz) * i
mjd = spw['axis_info']['time_axis']['MJDseconds'] / 86400.
uvdist = np.sqrt(spw['u'] ** 2 + spw['v'] ** 2)
uvang = np.angle(spw['u'] + 1j * spw['v'])
else:
spw = ms.getdata(['amplitude', 'phase', 'axis_info'], ifraxis=True)
xxamp = np.concatenate((xxamp, spw['amplitude']), 1)
xxpha = np.concatenate((xxpha, spw['phase']), 1)
fg = spw['axis_info']['freq_axis']['chan_freq'][:, 0] / 1e9
fghz = np.concatenate((fghz, fg))
band = np.concatenate((band, np.ones_like(fg) * i))
ms.close()
return {'amp': xxamp, 'phase': xxpha, 'fghz': fghz, 'band': band, 'mjd': mjd, 'uvdist': uvdist, 'uvangle': uvang}
def getspwfromfreq(vis, freqrange):
from taskinit import ms
ms.open(vis)
axisInfo = ms.getdata(["axis_info"], ifraxis=True)
spwInfo = ms.getspectralwindowinfo()
freqInfo = axisInfo["axis_info"]["freq_axis"]["chan_freq"].swapaxes(0, 1) / 1e9
freqInfo_ravel = freqInfo.ravel()
timeInfo = axisInfo["axis_info"]["time_axis"]['MJDseconds']
mstimran = ms.range(["time"])
ms.close()
freq0, freq1 = freqrange.split(' ')[0].split('~')
freq0, freq1 = float(freq0), float(freq1)
for ll in [freq0, freq1]:
if not freqInfo_ravel[0] <= ll <= freqInfo_ravel[-1]:
raise ValueError('Selected frequency out of range!!!')
freqIdx0 = np.where(freqInfo == freq0)
freqIdx1 = np.where(freqInfo == freq1)
sz_freqInfo = freqInfo.shape
ms_spw = ['{}'.format(ll) for ll in xrange(freqIdx0[0], freqIdx1[0] + 1)]
if len(ms_spw) == 1:
ms_chan = ['{}~{}'.format(freqIdx0[1][0], freqIdx1[1][0])]
else:
ms_chan = ['{}~{}'.format(freqIdx0[1][0], sz_freqInfo[1] - 1)] + ['0~{}'.format(sz_freqInfo[1] - 1) for ll in
xrange(freqIdx0[0] + 1, freqIdx1[0])]
ms_chan.append('0~{}'.format(freqIdx1[1][0]))
spw = ','.join('{}:{}'.format(t[0], t[1]) for t in zip(ms_spw, ms_chan))
return spw
def insertdiskmodel(vis, sizescale=1.0, fdens=None, dsize=None, xmlfile='SOLDISK.xml', writediskinfoonly=False,
active=False, overwrite=True):
# Apply size scale adjustment (default is no adjustment)
for i in range(len(dsize)):
num, unit = dsize[i].split('arc')
dsize[i] = str(float(num) * sizescale)[:6] + 'arc' + unit
msfile = vis
ms.open(msfile)
spwinfo = ms.getspectralwindowinfo()
nspw = len(spwinfo.keys())
ms.close()
diskimdir = 'diskim/'
if not os.path.exists(diskimdir):
os.makedirs(diskimdir)
frq = []
spws = range(nspw)
for sp in spws:
spw = spwinfo[str(sp)]
frq.append('{:.4f}GHz'.format((spw['RefFreq'] + spw['TotalWidth'] / 2.0) / 1e9))
frq = np.array(frq)
writediskxml(dsize, fdens, frq, xmlfile=xmlfile)
if not writediskinfoonly:
tb.open(msfile + '/FIELD')
phadir = tb.getcol('PHASE_DIR').flatten()
tb.close()
ra = phadir[0]
dec = phadir[1]
direction = 'J2000 ' + str(ra) + 'rad ' + str(dec) + 'rad'
diskim = []
for sp in tqdm(spws, desc='Generating {} disk models'.format(nspw), ascii=True):
diskim.append(
mk_diskmodel(outname=diskimdir + 'disk{:02d}_'.format(sp), bdwidth=spwinfo[str(sp)],
direction=direction,
reffreq=frq[sp],
flux=fdens[sp], eqradius=dsize[sp], polradius=dsize[sp], overwrite=overwrite))
if not active:
delmod(msfile, otf=True, scr=True)
for sp in tqdm(spws, desc='Inserting disk model', ascii=True):
ft(vis=msfile, spw=str(sp), field='', model=str(diskim[sp]), nterms=1,
reffreq="", complist="", incremental=True, usescratch=True)
return msfile, diskim
def WEI_plot(dofile=True,
vis=None,
timerange=None,
spw='',
aiafits='',
imagehead='',
workdir='',
spwCol=3,
phasecenter='J2000 11h00m48 06d14m60'):
if vis[-1] == '/':
vis = vis[:-1]
ms.open(vis)
spwInfo = ms.getspectralwindowinfo()
ms.close()
tb.open(vis)
starttim = Time(tb.getcell('TIME', 0) / 24. / 3600., format='mjd')
endtim = Time(tb.getcell('TIME',
tb.nrows() - 1) / 24. / 3600.,
format='mjd')
tb.close()
tb.open(vis + '/SPECTRAL_WINDOW')
reffreqs = tb.getcol('REF_FREQUENCY')
bdwds = tb.getcol('TOTAL_BANDWIDTH')
cfreqs = reffreqs + bdwds / 2.
tb.close()
sbeam = 35.
#get timerange from vis file
if not timerange:
timerange = '{0}~{1}'.format(
starttim.iso.replace('-', '/').replace(' ', '/'),
endtim.iso.replace('-', '/').replace(' ', '/'))
nspw = len(spwInfo)
#draw plot aia
fig = plt.figure(figsize=(12, 7), dpi=100)
gs1 = gridspec.GridSpec(4, 3)
gs1.update(left=0.08, right=0.32, wspace=0.05)
ax1 = plt.subplot(gs1[11])
aiamap = smap.Map(aiafits)
aiamap.plot(axes=ax1)
#do clean spwCol by spwCol
for cur_spwCol in range(0, np.floor_divide(nspw, spwCol)):
if ((cur_spwCol + 1) * spwCol) < nspw:
cur_spwRange = str(cur_spwCol * spwCol + 1) + '~' + str(
(cur_spwCol + 1) * spwCol)
else:
cur_spwRange = str(cur_spwCol * spwCol + 1) + '~' + '31'
imagename = imagehead + cur_spwRange + 'SPWs'
cur_eovsaFits = imagename + '.fits'
if cur_spwCol < 6:
cur_mask = '/srg/ywei/data/eovsa/mask/sep_6mask_' + str(
cur_spwCol + 1) + '.rgn'
else:
cur_mask = '/srg/ywei/data/eovsa/mask/sep_6/mask_6.rgn'
if dofile:
#clean(vis=vis, spw=cur_spwRange, timerange=timerange, imagename=imagename, imsize=[256,256], niter=100, cell=['2arcsec'] )
#clean(vis=vis, spw=cur_spwRange, timerange=timerange, imagename=imagename, imsize=[512,512], niter=1000, cell=['1arcsec'],stokes='XX', gain=0.05,weighting='briggs', mode='mfs',imagermode='csclean',psfmode='clark',robust=0.0,restoringbeam = ['10.0arcsec'], mask=cur_mask,pbcor=True)
clean(vis=vis,
spw=cur_spwRange,
timerange=timerange,
imagename=imagename,
imsize=[512, 512],
niter=1000,
cell=['1arcsec'],
stokes='XX',
gain=0.05,
weighting='briggs',
mode='mfs',
imagermode='csclean',
psfmode='clark',
robust=0.0,
restoringbeam=['10.0arcsec'],
mask='',
pbcor=True)
print 'fits name =' + str(cur_eovsaFits)
hf.imreg(vis=vis,
imagefile=imagename + '.image',
fitsfile=imagename + '.fits',
timerange=timerange)
#plot eovsa
cur_emap = smap.Map(cur_eovsaFits)
cur_axname = plt.subplot(gs1[cur_spwCol + 1])
(npol, nf, nx, ny) = cur_emap.data.shape
print 'shape = ' + str(cur_emap.data.shape)
if npol != 1:
print 'To be determined'
else:
cur_emap.data = cur_emap.data.reshape((512, 512))
cur_emap.plot_settings['cmap'] = plt.get_cmap('jet')
cur_emap.plot(axes=cur_axname)
def mk_qlook_image(vis, ncpu=10, twidth=12, stokes='I,V', antenna='', imagedir=None, spws=[], toTb=True, overwrite=True, doslfcal=False,
phasecenter='', c_external=True):
vis = [vis]
subdir = ['/']
for idx, f in enumerate(vis):
if f[-1] == '/':
vis[idx] = f[:-1]
if not imagedir:
imagedir = './'
imres = {'Succeeded': [], 'BeginTime': [], 'EndTime': [], 'ImageName': [], 'Spw': [], 'Vis': [], 'Freq': []}
msfile = vis[0]
ms.open(msfile)
metadata = ms.metadata()
observatory = metadata.observatorynames()[0]
# axisInfo = ms.getdata(["axis_info"], ifraxis=True)
spwInfo = ms.getspectralwindowinfo()
# freqInfo = axisInfo["axis_info"]["freq_axis"]["chan_freq"].swapaxes(0, 1) / 1e9
# freqInfo_ravel = freqInfo.ravel()
ms.close()
if not spws:
if observatory == 'EVLA':
spws = ['0', '1', '2', '3', '4', '5', '6', '7']
if observatory == 'EOVSA':
spws = ['1~5', '6~10', '11~15', '16~25']
if observatory == 'EOVSA':
print 'Provide stokes: ' + str(stokes) + '. However EOVSA has linear feeds. Force stokes to be IV'
stokes = 'I,V'
msfilebs = os.path.basename(msfile)
imdir = imagedir + subdir[0]
if not os.path.exists(imdir):
os.makedirs(imdir)
if doslfcal:
slfcalms = './' + msfilebs + '.rr'
split(msfile, outputvis=slfcalms, datacolumn='corrected', correlation='RR')
for spw in spws:
spwran = [s.zfill(2) for s in spw.split('~')]
# freqran = [
# (int(s) * spwInfo['0']['TotalWidth'] + spwInfo['0']['RefFreq'] + spwInfo['0']['TotalWidth'] / 2.0) / 1.0e9
# for s in spw.split('~')]
spw_ = spw.split('~')
if len(spw_) == 2:
freqran = [(spwInfo['{}'.format(s)]['RefFreq'] + spwInfo['{}'.format(s)]['TotalWidth'] / 2.0) / 1.0e9 for s in spw.split('~')]
elif len(spw_) == 1:
s = spw_[0]
freqran = np.array([0, spwInfo['{}'.format(s)]['TotalWidth']]) + spwInfo['{}'.format(s)]['RefFreq']
freqran = freqran / 1.0e9
freqran = list(freqran)
else:
raise ValueError("Keyword 'spw' in wrong format")
cfreq = np.mean(freqran)
bmsz = max(30. / cfreq, 30.)
uvrange = '<3km'
if cfreq < 10.:
imsize = 512
cell = ['5arcsec']
else:
imsize = 1024
cell = ['2.5arcsec']
if len(spwran) == 2:
spwstr = spwran[0] + '~' + spwran[1]
else:
spwstr = spwran[0]
restoringbeam = ['{0:.1f}arcsec'.format(bmsz)]
imagesuffix = '.spw' + spwstr.replace('~', '-')
# if cfreq > 10.:
# antenna = antenna + ';!0&1;!0&2' # deselect the shortest baselines
sto = stokes.replace(',', '')
if c_external:
cleanscript = os.path.join(imdir, 'ptclean_external.py')
resfile = os.path.join(imdir, os.path.basename(msfile) + '.res.npz')
os.system('rm -rf {}'.format(cleanscript))
inpdict = {'vis': msfile, 'imageprefix': imdir, 'imagesuffix': imagesuffix, 'twidth': twidth, 'uvrange': uvrange, 'spw': spw,
'ncpu': ncpu, 'niter': 1000, 'gain': 0.05, 'antenna': antenna, 'imsize': imsize, 'cell': cell, 'stokes': sto, 'doreg': True,
'overwrite': overwrite, 'toTb': toTb, 'restoringbeam': restoringbeam, 'uvtaper': True, 'outertaper': ['30arcsec'],
'phasecenter': phasecenter}
for key, val in inpdict.items():
if type(val) is str:
inpdict[key] = '"{}"'.format(val)
fi = open(cleanscript, 'wb')
fi.write('from ptclean_cli import ptclean_cli as ptclean \n')
fi.write('import numpy as np \n')
fi.write(
'res = ptclean(vis={i[vis]},imageprefix={i[imageprefix]},imagesuffix={i[imagesuffix]},twidth={i[twidth]},uvrange={i[uvrange]},spw={i[spw]},ncpu={i[ncpu]},niter={i[niter]},gain={i[gain]},antenna={i[antenna]},imsize={i[imsize]},cell={i[cell]},stokes={i[stokes]},doreg={i[doreg]},overwrite={i[overwrite]},toTb={i[toTb]},restoringbeam={i[restoringbeam]},uvtaper={i[uvtaper]},outertaper={i[outertaper]},phasecenter={i[phasecenter]}) \n'.format(
i=inpdict))
fi.write('np.savez("{}",res=res) \n'.format(resfile))
fi.close()
os.system('casa --nologger -c {}'.format(cleanscript))
res = np.load(resfile)
res = res['res'].item()
else:
res = ptclean(vis=msfile, imageprefix=imdir, imagesuffix=imagesuffix, twidth=twidth, uvrange=uvrange, spw=spw, ncpu=ncpu, niter=1000,
gain=0.05, antenna=antenna, imsize=imsize, cell=cell, stokes=sto, doreg=True, overwrite=overwrite, toTb=toTb,
restoringbeam=restoringbeam, uvtaper=True, outertaper=['30arcsec'], phasecenter=phasecenter)
if res:
imres['Succeeded'] += res['Succeeded']
imres['BeginTime'] += res['BeginTime']
imres['EndTime'] += res['EndTime']
imres['ImageName'] += res['ImageName']
imres['Spw'] += [spwstr] * len(res['ImageName'])
imres['Vis'] += [msfile] * len(res['ImageName'])
imres['Freq'] += [freqran] * len(res['ImageName'])
else:
return None
# save it for debugging purposes
np.savez(os.path.join(imagedir, '{}.imres.npz'.format(os.path.basename(msfile))), imres=imres)
return imres
def svplot(vis, timerange=None, spw='', workdir='./', specfile=None, bl=None, uvrange=None,
stokes='RR,LL', dmin=None, dmax=None,
goestime=None, reftime=None,
xycen=None, fov=[500.,500.], xyrange=None, restoringbeam=[''], robust=0.0,
niter=500, imsize=[512], cell=['5.0arcsec'],interactive=False,
usemsphacenter=True, imagefile=None, fitsfile=None, plotaia=True,
aiawave=171, aiafits=None, savefig=False, mkmovie=False, overwrite=True, ncpu=10, twidth=1, verbose=True):
'''
Required inputs:
vis: calibrated CASA measurement set
Important optional inputs:
timerange: timerange for clean. Standard CASA time selection format.
If not provided, use the entire range (*BE CAREFUL, COULD BE VERY SLOW*)
spw: spectral window selection following the CASA syntax.
Examples: spw='1:2~60' (spw id 1, channel range 2-60); spw='*:1.2~1.3GHz' (selects all channels within 1.2-1.3 GHz; note the *)
specfile: supply dynamic spectrum save file (from suncasa.utils.dspec2.get_dspec()). Otherwise
generate a median dynamic spectrum on the fly
Optional inputs:
bl: baseline to generate dynamic spectrum
uvrange: uvrange to select baselines for generating dynamic spectrum
stokes: polarization of the clean image, can be 'RR,LL' or 'I,V'
dmin,dmax: color bar parameter
goestime: goes plot time, example ['2016/02/18 18:00:00','2016/02/18 23:00:00']
rhessisav: rhessi savefile
reftime: reftime for the image
xycen: center of the image in helioprojective coordinates (HPLN/HPLT), in arcseconds. Example: [900, -150.]
fov: field of view in arcsecs. Example: [500., 500.]
xyrange: field of view in solar XY coordinates. Format: [[x1,x2],[y1,y2]]. Example: [[900., 1200.],[0,300]]
***NOTE: THIS PARAMETER OVERWRITES XYCEN AND FOV***
aiawave: wave length of aia file in a
imagefile: if imagefile provided, use it. Otherwise do clean and generate a new one.
fitsfile: if fitsfile provided, use it. Otherwise generate a new one
savefig: whether to save the figure
Example:
'''
if xycen:
xc, yc = xycen
xlen, ylen = fov
if parse_version(sunpy.__version__)>parse_version('0.8.0'):
xyrange = [[xc - xlen / 2.0, yc - ylen / 2.0], [xc + xlen / 2.0, yc + ylen / 2.0]]
else:
xyrange = [[xc - xlen / 2.0, xc + xlen / 2.0], [yc - ylen / 2.0, yc + ylen / 2.0]]
stokes_allowed = ['RR,LL', 'I,V', 'RRLL', 'IV']
if not stokes in stokes_allowed:
print 'wrong stokes parameter ' + str(stokes) + '. Allowed values are ' + ', '.join(stokes_allowed)
return -1
if stokes == 'RRLL':
stokes = 'RR,LL'
if stokes == 'IV':
stokes = 'I,V'
if vis[-1] == '/':
vis = vis[:-1]
if not os.path.exists(vis):
print 'input measurement not exist'
return -1
if aiafits is None:
aiafits = ''
# split the data
# generating dynamic spectrum
if not os.path.exists(workdir):
os.makedirs(workdir)
if specfile:
try:
specdata = np.load(specfile)
except:
print('Provided dynamic spectrum file not numpy npz. Generating one from the visibility data')
specfile = os.path.join(workdir, os.path.basename(vis) + '.dspec.npz')
dspec_external(vis, workdir=workdir, specfile=specfile)
specdata = np.load(specfile) # specdata = ds.get_dspec(vis, domedian=True, verbose=True)
else:
print('Dynamic spectrum file not provided; Generating one from the visibility data')
# specdata = ds.get_dspec(vis, domedian=True, verbose=True)
specfile = os.path.join(workdir, os.path.basename(vis) + '.dspec.npz')
dspec_external(vis, workdir=workdir, specfile=specfile)
specdata = np.load(specfile)
tb.open(vis)
starttim = Time(tb.getcell('TIME', 0) / 24. / 3600., format='mjd')
endtim = Time(tb.getcell('TIME', tb.nrows() - 1) / 24. / 3600., format='mjd')
tb.close()
datstr = starttim.iso[:10]
if timerange is None or timerange == '':
starttim1 = starttim
endtim1 = endtim
timerange = '{0}~{1}'.format(starttim.iso.replace('-', '/').replace(' ', '/'), endtim.iso.replace('-', '/').replace(' ', '/'))
else:
try:
(tstart, tend) = timerange.split('~')
if tstart[2] == ':':
starttim1 = Time(datstr + 'T' + tstart)
endtim1 = Time(datstr + 'T' + tend)
timerange = '{0}/{1}~{0}/{2}'.format(datstr.replace('-', '/'), tstart, tend)
else:
starttim1 = Time(qa.quantity(tstart, 'd')['value'], format='mjd')
endtim1 = Time(qa.quantity(tend, 'd')['value'], format='mjd')
except ValueError:
print "keyword 'timerange' in wrong format"
midtime_mjd = (starttim1.mjd + endtim1.mjd) / 2.
if vis.endswith('/'):
vis = vis[:-1]
visname = os.path.basename(vis)
bt = starttim1.plot_date
et = endtim1.plot_date
# find out min and max frequency for plotting in dynamic spectrum
ms.open(vis)
metadata = ms.metadata()
observatory = metadata.observatorynames()[0]
spwInfo = ms.getspectralwindowinfo()
nspw = len(spwInfo)
if not spw:
spw = '0~' + str(nspw - 1)
staql = {'timerange': timerange, 'spw': spw}
if ms.msselect(staql, onlyparse=True):
ndx = ms.msselectedindices()
chan_sel = ndx['channel']
nspw = chan_sel.shape[0]
bspw = chan_sel[0, 0]
bchan = chan_sel[0, 1]
espw = chan_sel[-1, 0]
echan = chan_sel[-1, 2]
bfreq = spwInfo[str(bspw)]['Chan1Freq'] + spwInfo[str(bspw)]['ChanWidth'] * bchan
efreq = spwInfo[str(espw)]['Chan1Freq'] + spwInfo[str(espw)]['ChanWidth'] * echan
bfreqghz = bfreq / 1e9
efreqghz = efreq / 1e9
if verbose:
print 'selected timerange {}'.format(timerange)
print 'selected frequency range {0:6.3f} to {1:6.3f} GHz'.format(bfreqghz, efreqghz)
else:
print "spw or timerange selection failed. Aborting..."
ms.close()
return -1
ms.close()
if observatory == 'EOVSA':
print 'Provide stokes: ' + str(stokes) + '. However EOVSA has linear feeds. Force stokes to be IV'
stokes = 'I,V'
if mkmovie:
plt.ioff()
# fig = plt.figure(figsize=(12, 7.5), dpi=100)
if fitsfile:
pass
else:
if not imagefile:
# from ptclean_cli import ptclean_cli as ptclean
eph = hf.read_horizons(t0=Time(midtime_mjd, format='mjd'))
if observatory == 'EOVSA' or (not usemsphacenter):
phasecenter = ''
else:
phasecenter = 'J2000 ' + str(eph['ra'][0])[:15] + 'rad ' + str(eph['dec'][0])[:15] + 'rad'
print 'use phasecenter: ' + phasecenter
qlookfitsdir = os.path.join(workdir, 'qlookfits/')
qlookfigdir = os.path.join(workdir, 'qlookimgs/')
imresfile = os.path.join(qlookfitsdir, '{}.imres.npz'.format(os.path.basename(vis)))
if overwrite:
imres = mk_qlook_image(vis, twidth=twidth, ncpu=ncpu, imagedir=qlookfitsdir, phasecenter=phasecenter, stokes=stokes,
c_external=True)
else:
if os.path.exists(imresfile):
imres = np.load(imresfile)
imres = imres['imres'].item()
else:
print('Image results file not found; Creating new images.')
imres = mk_qlook_image(vis, twidth=twidth, ncpu=ncpu, imagedir=qlookfitsdir, phasecenter=phasecenter, stokes=stokes,
c_external=True)
if not os.path.exists(qlookfigdir):
os.makedirs(qlookfigdir)
plt_qlook_image(imres, figdir=qlookfigdir, specdata=specdata, verbose=True, stokes=stokes, fov=xyrange)
else:
spec = specdata['spec']
(npol, nbl, nfreq, ntim) = spec.shape
tidx = range(ntim)
fidx = range(nfreq)
tim = specdata['tim']
freq = specdata['freq']
freqghz = freq / 1e9
spec_tim = Time(specdata['tim'] / 3600. / 24., format='mjd')
timstrr = spec_tim.plot_date
plt.ion()
fig = plt.figure(figsize=(12, 7), dpi=100)
gs1 = gridspec.GridSpec(3, 1)
gs1.update(left=0.08, right=0.32, wspace=0.05)
gs2 = gridspec.GridSpec(2, 2)
gs2.update(left=0.38, right=0.98, hspace=0.02, wspace=0.02)
spec_1 = np.absolute(spec[0, 0, :, :])
spec_2 = np.absolute(spec[1, 0, :, :])
if observatory == 'EVLA':
# circular feeds
polstr = ['RR', 'LL']
if observatory == 'EOVSA' or observatory == 'ALMA':
# linear feeds
polstr = ['XX', 'YY']
print 'plot the dynamic spectrum in pol ' + ' & '.join(polstr)
ax1 = plt.subplot(gs1[0])
ax1.pcolormesh(timstrr, freqghz, spec_1, cmap='jet', vmin=dmin, vmax=dmax)
ax1.set_xlim(timstrr[tidx[0]], timstrr[tidx[-1]])
ax1.xaxis_date()
ax1.xaxis.set_major_formatter(DateFormatter("%H:%M:%S"))
# ax1.set_xticklabels(['']*10)
ax1.set_ylim(freqghz[fidx[0]], freqghz[fidx[-1]])
ax1.set_ylabel('Frequency (GHz)', fontsize=10)
ax1.set_title(observatory + ' ' + datstr + ' ' + polstr[0] + ' & ' + polstr[1], fontsize=12)
ax1.set_autoscale_on(False)
ax1.add_patch(patches.Rectangle((bt, bfreqghz), et - bt, efreqghz - bfreqghz, ec='w', fill=False))
ax1.plot([(bt + et) / 2.], [(bfreqghz + efreqghz) / 2.], '*w', ms=12)
for tick in ax1.get_xticklabels():
tick.set_fontsize(8)
for tick in ax1.get_yticklabels():
tick.set_fontsize(8)
ax2 = plt.subplot(gs1[1])
ax2.pcolormesh(timstrr, freqghz, spec_2, cmap='jet', vmin=dmin, vmax=dmax)
ax2.set_xlim(timstrr[tidx[0]], timstrr[tidx[-1]])
ax2.xaxis_date()
ax2.xaxis.set_major_formatter(DateFormatter("%H:%M:%S"))
ax2.set_ylim(freqghz[fidx[0]], freqghz[fidx[-1]])
ax2.set_ylabel('Frequency (GHz)', fontsize=10)
for tick in ax2.get_xticklabels():
tick.set_fontsize(8)
for tick in ax2.get_yticklabels():
tick.set_fontsize(8)
ax2.set_autoscale_on(False)
ax2.add_patch(patches.Rectangle((bt, bfreqghz), et - bt, efreqghz - bfreqghz, ec='w', fill=False))
ax2.plot([(bt + et) / 2.], [(bfreqghz + efreqghz) / 2.], '*w', ms=12)
# Second part: GOES plot
if goestime:
btgoes = goestime[0]
etgoes = goestime[1]
else:
datstrg = datstr.replace('-', '/')
btgoes = datstrg + ' ' + qa.time(qa.quantity(tim[0] - 1800, 's'), form='clean', prec=9)[0]
etgoes = datstrg + ' ' + qa.time(qa.quantity(tim[tidx[-1] - 1] + 1800, 's'), form='clean', prec=9)[0]
if verbose:
print 'Acquire GOES soft X-ray data in from ' + btgoes + ' to ' + etgoes
ax3 = plt.subplot(gs1[2])
try:
from sunpy import lightcurve as lc
from sunpy.time import TimeRange
goest = lc.GOESLightCurve.create(TimeRange(btgoes, etgoes))
except:
goesscript = os.path.join(workdir, 'goes.py')
goesdatafile = os.path.join(workdir, 'goes.dat')
os.system('rm -rf {}'.format(goesscript))
fi = open(goesscript, 'wb')
fi.write('import os \n')
fi.write('from sunpy.time import TimeRange \n')
fi.write('from sunpy import lightcurve as lc \n')
fi.write('import pickle \n')
fi.write('goesplottim = TimeRange("{0}", "{1}") \n'.format(btgoes, etgoes))
fi.write('goes = lc.GOESLightCurve.create(goesplottim) \n')
fi.write('fi2 = open("{}", "wb") \n'.format(goesdatafile))
fi.write('pickle.dump(goes, fi2) \n')
fi.write('fi2.close()')
fi.close()
try:
os.system('python {}'.format(goesscript))
os.system('rm -rf {}'.format(goesscript))
except NameError:
print "Bad input names"
except ValueError:
print "Bad input values"
except:
print "Unexpected error:", sys.exc_info()[0]
print "Error in generating GOES light curves. Proceed without GOES..."
if os.path.exists(goesdatafile):
fi1 = file(goesdatafile, 'rb')
goest = pickle.load(fi1)
fi1.close()
try:
dates = mpl.dates.date2num(parse_time(goest.data.index))
goesdif = np.diff(goest.data['xrsb'])
gmax = np.nanmax(goesdif)
gmin = np.nanmin(goesdif)
ran = gmax - gmin
db = 2.8 / ran
goesdifp = goesdif * db + gmin + (-6)
ax3.plot_date(dates, np.log10(goest.data['xrsb']), '-', label='1.0--8.0 $\AA$', color='red', lw=2)
ax3.plot_date(dates[0:-1], goesdifp, '-', label='derivate', color='blue', lw=0.4)
ax3.set_ylim([-7, -3])
ax3.set_yticks([-7, -6, -5, -4, -3])
ax3.set_yticklabels([r'$10^{-7}$', r'$10^{-6}$', r'$10^{-5}$', r'$10^{-4}$', r'$10^{-3}$'])
ax3.set_title('Goes Soft X-ray', fontsize=12)
ax3.set_ylabel('Watts m$^{-2}$')
ax3.set_xlabel(datetime.datetime.isoformat(goest.data.index[0])[0:10])
ax3.axvspan(dates[899], dates[dates.size - 899], alpha=0.2)
ax2 = ax3.twinx()
# ax2.set_yscale("log")
ax2.set_ylim([-7, -3])
ax2.set_yticks([-7, -6, -5, -4, -3])
ax2.set_yticklabels(['B', 'C', 'M', 'X', ''])
ax3.yaxis.grid(True, 'major')
ax3.xaxis.grid(False, 'major')
ax3.legend(prop={'size': 6})
formatter = mpl.dates.DateFormatter('%H:%M')
ax3.xaxis.set_major_formatter(formatter)
ax3.fmt_xdata = mpl.dates.DateFormatter('%H:%M')
except:
print 'Error in downloading GOES soft X-ray data. Proceeding with out soft X-ray plot.'
# third part
# start to download the fits files
if plotaia:
if not aiafits:
newlist = []
items = glob.glob('*.fits')
for names in items:
str1 = starttim1.iso[:4] + '_' + starttim1.iso[5:7] + '_' + starttim1.iso[8:10] + 't' + starttim1.iso[
11:13] + '_' + starttim1.iso[14:16]
str2 = str(aiawave)
if names.endswith(".fits"):
if names.find(str1) != -1 and names.find(str2) != -1:
newlist.append(names)
newlist.append('0')
if newlist and os.path.exists(newlist[0]):
aiafits = newlist[0]
else:
print 'downloading the aiafits file'
wave1 = aiawave - 3
wave2 = aiawave + 3
t1 = Time(starttim1.mjd - 0.02 / 24., format='mjd')
t2 = Time(endtim1.mjd + 0.02 / 24., format='mjd')
try:
from sunpy.net import vso
client = vso.VSOClient()
qr = client.query(vso.attrs.Time(t1.iso, t2.iso), vso.attrs.Instrument('aia'), vso.attrs.Wave(wave1 * u.AA, wave2 * u.AA))
res = client.get(qr, path='{file}')
except:
SdoDownloadscript = os.path.join(workdir, 'SdoDownload.py')
os.system('rm -rf {}'.format(SdoDownloadscript))
fi = open(SdoDownloadscript, 'wb')
fi.write('from sunpy.net import vso \n')
fi.write('from astropy import units as u \n')
fi.write('client = vso.VSOClient() \n')
fi.write(
"qr = client.query(vso.attrs.Time('{0}', '{1}'), vso.attrs.Instrument('aia'), vso.attrs.Wave({2} * u.AA, {3} * u.AA)) \n".format(
t1.iso, t2.iso, wave1, wave2))
fi.write("res = client.get(qr, path='{file}') \n")
fi.close()
try:
os.system('python {}'.format(SdoDownloadscript))
except NameError:
print "Bad input names"
except ValueError:
print "Bad input values"
except:
print "Unexpected error:", sys.exc_info()[0]
print "Error in Downloading AIA fits files. Proceed without AIA..."
# Here something is needed to check whether it has finished downloading the fits files or not
if not aiafits:
newlist = []
items = glob.glob('*.fits')
for nm in items:
str1 = starttim1.iso[:4] + '_' + starttim1.iso[5:7] + '_' + starttim1.iso[8:10] + 't' + starttim1.iso[
11:13] + '_' + starttim1.iso[14:16]
str2 = str(aiawave)
if nm.find(str1) != -1 and nm.find(str2) != -1:
newlist.append(nm)
if newlist:
aiafits = newlist[0]
print 'AIA fits ' + aiafits + ' selected'
else:
print 'no AIA fits files found. Proceed without AIA'
try:
aiamap = smap.Map(aiafits)
except:
print 'error in reading aiafits. Proceed without AIA'
# RCP or I
ax4 = plt.subplot(gs2[0, 0])
ax5 = plt.subplot(gs2[1, 0])
# LCP or V
ax6 = plt.subplot(gs2[0, 1])
ax7 = plt.subplot(gs2[1, 1])
if fitsfile:
pass
else:
if not imagefile:
eph = hf.read_horizons(t0=Time(midtime_mjd, format='mjd'))
if observatory == 'EOVSA' or (not usemsphacenter):
print 'This is EOVSA data'
# use RA and DEC from FIELD ID 0
tb.open(vis+'/FIELD')
phadir = tb.getcol('PHASE_DIR').flatten()
tb.close()
ra0 = phadir[0]
dec0 = phadir[1]
if stokes == 'RRLL' or stokes == 'RR,LL':
print 'Provide stokes: ' + str(stokes) + '. However EOVSA has linear feeds. Force stokes to be IV'
stokes = 'I,V'
else:
ra0 = eph['ra'][0]
dec0 = eph['dec'][0]
if not xycen:
# use solar disk center as default
phasecenter = 'J2000 ' + str(ra0) + 'rad ' + str(dec0) + 'rad'
else:
x0 = np.radians(xycen[0]/3600.)
y0 = np.radians(xycen[1]/3600.)
p0 = np.radians(eph['p0'][0]) # p angle in radians
raoff = -((x0) * np.cos(p0) - y0 * np.sin(p0))/np.cos(eph['dec'][0])
decoff = (x0) * np.sin(p0) + y0 * np.cos(p0)
newra = ra0 + raoff
newdec = dec0 + decoff
phasecenter = 'J2000 ' + str(newra) + 'rad ' + str(newdec) + 'rad'
imagename = os.path.join(workdir, visname + '.outim')
if os.path.exists(imagename + '.image') or os.path.exists(imagename + '.flux'):
os.system('rm -rf ' + imagename + '.*')
sto = stokes.replace(',', '')
print 'do clean for ' + timerange + ' in spw ' + spw + ' stokes ' + sto
print 'Original phasecenter: '+ str(ra0) + str(dec0)
print 'use phasecenter: ' + phasecenter
clean(vis=vis, imagename=imagename, selectdata=True, spw=spw, timerange=timerange, stokes=sto,
niter=niter, interactive=interactive, npercycle=50, imsize=imsize, cell=cell, restoringbeam=restoringbeam,
weighting='briggs', robust=robust, phasecenter=phasecenter)
os.system('rm -rf ' + imagename + '.psf')
os.system('rm -rf ' + imagename + '.flux')
os.system('rm -rf ' + imagename + '.model')
os.system('rm -rf ' + imagename + '.mask')
os.system('rm -rf ' + imagename + '.residual')
imagefile = imagename + '.image'
fitsfile = imagefile + '.fits'
hf.imreg(vis=vis, ephem=eph, imagefile=imagefile, timerange=timerange, reftime=reftime, fitsfile=fitsfile, verbose=True, overwrite=True)
print 'fits file ' + fitsfile + ' selected'
ax4.cla()
ax5.cla()
ax6.cla()
ax7.cla()
rfits = fitsfile
try:
hdulist = fits.open(rfits)
hdu = hdulist[0]
(npol, nf, nx, ny) = hdu.data.shape
rmap = smap.Map(hdu.data[0, 0, :, :], hdu.header)
except:
print 'radio fits file not recognized by sunpy.map. Aborting...'
return -1
if npol > 1:
rmap1 = smap.Map(hdu.data[0, 0, :, :], hdu.header)
rmap2 = smap.Map(hdu.data[1, 0, :, :], hdu.header)
XX, YY = np.meshgrid(np.arange(rmap.data.shape[1]), np.arange(rmap.data.shape[0]))
try:
rmapx, rmapy = rmap.pixel_to_data(XX * u.pix, YY * u.pix)
except:
rmapxy = rmap.pixel_to_data(XX * u.pix, YY * u.pix)
rmapx = rmapxy.Tx
rmapy = rmapxy.Ty
if not xyrange:
if xycen:
x0 = xycen[0] * u.arcsec
y0 = xycen[1] * u.arcsec
if not xycen:
row, col = rmap1.data.shape
positon = np.nanargmax(rmap1.data)
m, n = divmod(positon, col)
x0 = rmap1.xrange[0] + rmap1.scale[1] * (n + 0.5) * u.pix
y0 = rmap1.yrange[0] + rmap1.scale[0] * (m + 0.5) * u.pix
if len(fov) == 1:
fov=[fov]*2
sz_x = fov[0] * u.arcsec
sz_y = fov[1] * u.arcsec
x1 = x0 - sz_x/2.
x2 = x0 + sz_x/2.
y1 = y0 - sz_y/2.
y2 = y0 + sz_y/2.
xyrange = [[x1.value, x2.value], [y1.value, y2.value]]
else:
sz_x = xyrange[0][1] - xyrange[0][0]
sz_y = xyrange[1][1] - xyrange[1][0]
clevels1 = np.linspace(0.2, 0.9, 5)
if stokes.split(',')[1] == 'V':
clevels2 = np.array([0.8, -0.6, -0.4, -0.2, 0.2, 0.4, 0.6, 0.8])
else:
clevels2 = np.linspace(0.2, 0.9, 5)
if 'aiamap' in vars():
aiamap.plot_settings['cmap'] = plt.get_cmap('binary')
if rmap:
title = 'AIA {0:.0f} + {1} {2:6.3f} GHz'.format(aiamap.wavelength.value, observatory, (bfreqghz + efreqghz) / 2.0)
else:
title = 'AIA {0:.0f}'.format(aiamap.wavelength.value)
aiamap.plot(axes=ax4)
ax4.set_title(title + ' ' + stokes.split(',')[0], fontsize=12)
aiamap.draw_limb()
aiamap.draw_grid()
aiamap.draw_rectangle((xyrange[0][0], xyrange[1][0]) * u.arcsec, sz_x, sz_y)
aiamap.plot(axes=ax6)
ax6.set_title(title + ' ' + stokes.split(',')[1], fontsize=12)
aiamap.draw_limb()
aiamap.draw_grid()
aiamap.draw_rectangle((xyrange[0][0], xyrange[1][0]) * u.arcsec, sz_x, sz_y)
if rmap:
ax4.contour(rmapx.value, rmapy.value, rmap1.data, levels=clevels1 * np.nanmax(rmap1.data), cmap=cm.jet)
ax6.contour(rmapx.value, rmapy.value, rmap2.data, levels=clevels2 * np.nanmax(rmap2.data), cmap=cm.RdBu)
ax4.text(0.02, 0.02, 'AIA {0:.0f} '.format(aiamap.wavelength.value) + aiamap.date.strftime('%H:%M:%S'),
verticalalignment='bottom', horizontalalignment='left', transform=ax4.transAxes, color='k',
fontsize=10)
ax6.text(0.02, 0.02, 'AIA {0:.0f} '.format(aiamap.wavelength.value) + aiamap.date.strftime('%H:%M:%S'),
verticalalignment='bottom', horizontalalignment='left', transform=ax6.transAxes, color='k',
fontsize=10)
else:
title = '{0} {1:6.3f} GHz'.format(observatory, (bfreqghz + efreqghz) / 2.0)
rmap1.plot(axes=ax4, cmap=cm.jet)
ax4.set_title(title + ' ' + stokes.split(',')[0], fontsize=12)
rmap1.draw_limb()
rmap1.draw_grid()
rmap1.draw_rectangle((xyrange[0][0], xyrange[1][0]) * u.arcsec, sz_x, sz_y)
rmap2.plot(axes=ax6, cmap=cm.RdBu)
ax6.set_title(title + ' ' + stokes.split(',')[1], fontsize=12)
rmap2.draw_limb()
rmap2.draw_grid()
# ax4.contour(rmapx.value, rmapy.value, rmap1.data, levels=np.linspace(0.2, 0.9, 5) * np.nanmax(rmap1.data),
# cmap=cm.gray)
# ax6.contour(rmapx.value, rmapy.value, rmap2.data, levels=np.linspace(0.2, 0.9, 5) * np.nanmax(rmap2.data),
# cmap=cm.gray)
rmap2.draw_rectangle((xyrange[0][0], xyrange[1][0]) * u.arcsec, sz_x, sz_y)
ax4.set_xlim(-1200, 1200)
ax4.set_ylim(-1200, 1200)
ax6.set_xlim(-1200, 1200)
ax6.set_ylim(-1200, 1200)
try:
subrmap1 = rmap1.submap(xyrange[0] * u.arcsec, xyrange[1] * u.arcsec)
subrmap2 = rmap2.submap(xyrange[0] * u.arcsec, xyrange[1] * u.arcsec)
except:
bl = SkyCoord(xyrange[0][0] * u.arcsec, xyrange[1][0] * u.arcsec, frame=rmap1.coordinate_frame)
tr = SkyCoord(xyrange[0][1] * u.arcsec, xyrange[1][1] * u.arcsec, frame=rmap1.coordinate_frame)
subrmap1 = rmap1.submap(bl, tr)
subrmap2 = rmap2.submap(bl, tr)
XX, YY = np.meshgrid(np.arange(subrmap1.data.shape[1]), np.arange(subrmap1.data.shape[0]))
try:
subrmapx, subrmapy = subrmap1.pixel_to_data(XX * u.pix, YY * u.pix)
except:
subrmapxy = subrmap1.pixel_to_data(XX * u.pix, YY * u.pix)
subrmapx = subrmapxy.Tx
subrmapy = subrmapxy.Ty
if 'aiamap' in vars():
try:
subaiamap = aiamap.submap(xyrange[0] * u.arcsec, xyrange[1] * u.arcsec)
except:
bl = SkyCoord(xyrange[0][0] * u.arcsec, xyrange[1][0] * u.arcsec, frame=aiamap.coordinate_frame)
tr = SkyCoord(xyrange[0][1] * u.arcsec, xyrange[1][1] * u.arcsec, frame=aiamap.coordinate_frame)
subaiamap = aiamap.submap(bl, tr)
subaiamap.plot(axes=ax5, title='')
subaiamap.draw_limb()
subaiamap.draw_grid()
subaiamap.plot(axes=ax7, title='')
subaiamap.draw_limb()
subaiamap.draw_grid()
ax5.contour(subrmapx.value, subrmapy.value, subrmap1.data, levels=clevels1 * np.nanmax(subrmap1.data), cmap=cm.jet)
ax7.contour(subrmapx.value, subrmapy.value, subrmap2.data, levels=clevels2 * np.nanmax(subrmap2.data),
cmap=cm.RdBu) # subaiamap.draw_rectangle((fov[0][0], fov[1][0]) * u.arcsec, 400 * u.arcsec, 400 * u.arcsec)
else:
subrmap1.plot(axes=ax5, cmap=cm.jet, title='')
subrmap1.draw_limb()
subrmap1.draw_grid()
subrmap2.plot(axes=ax7, cmap=cm.RdBu, title='')
subrmap2.draw_limb()
subrmap2.draw_grid() # ax5.contour(subrmapx.value, subrmapy.value, subrmap1.data, # levels=clevels1 * np.nanmax(subrmap1.data), cmap=cm.gray) # ax7.contour(subrmapx.value, subrmapy.value, subrmap2.data, # levels=clevels2 * np.nanmax(subrmap2.data), cmap=cm.gray) # subrmap1.draw_rectangle((fov[0][0], fov[1][0]) * u.arcsec, 400 * u.arcsec, 400 * u.arcsec) # subrmap2.draw_rectangle((fov[0][0], fov[1][0]) * u.arcsec, 400 * u.arcsec, 400 * u.arcsec)
ax5.set_xlim(xyrange[0])
ax5.set_ylim(xyrange[1])
ax5.text(0.02, 0.02, observatory + ' ' + rmap.date.strftime('%H:%M:%S.%f')[:-3], verticalalignment='bottom',
horizontalalignment='left', transform=ax5.transAxes, color='k', fontsize=10)
ax7.set_xlim(xyrange[0])
ax7.set_ylim(xyrange[1])
ax7.text(0.02, 0.02, observatory + ' ' + rmap.date.strftime('%H:%M:%S.%f')[:-3], verticalalignment='bottom',
horizontalalignment='left', transform=ax7.transAxes, color='k', fontsize=10)
fig.show()
def get_dspec(vis=None, savespec=True, specfile=None, bl='', uvrange='', field='', scan='', datacolumn='data',
domedian=False, timeran=None, spw=None, timebin='0s', regridfreq=False, fillnan=None, verbose=False):
# from split_cli import split_cli as split
if vis.endswith('/'):
vis = vis[:-1]
msfile = vis
if not spw:
spw = ''
if not timeran:
timeran = ''
if not bl:
bl = ''
if domedian:
if not uvrange:
uvrange = '0.2~0.8km'
bl = ''
else:
uvrange = ''
# Open the ms and plot dynamic spectrum
if verbose:
print('Splitting selected data...')
vis_spl = './tmpms.splitted'
if os.path.exists(vis_spl):
os.system('rm -rf ' + vis_spl)
# split(vis=msfile, outputvis=vis_spl, timerange=timeran, antenna=bl, field=field, scan=scan, spw=spw,
# uvrange=uvrange, timebin=timebin, datacolumn=datacolumn)
ms.open(msfile, nomodify=False)
ms.split(outputms=vis_spl, whichcol=datacolumn, time=timeran, spw=spw, baseline=bl, field=field, scan=scan,
uvrange=uvrange, timebin=timebin)
ms.close()
if verbose:
print('Regridding into a single spectral window...')
# print('Reading data spw by spw')
try:
tb.open(vis_spl + '/POLARIZATION')
corrtype = tb.getcell('CORR_TYPE', 0)
pols = [stokesenum[p] for p in corrtype]
tb.close()
except:
pols = []
if regridfreq:
ms.open(vis_spl, nomodify=False)
ms.cvel(outframe='LSRK', mode='frequency', interp='nearest')
ms.selectinit(datadescid=0, reset=True)
data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
specamp = data['amplitude']
freq = data['axis_info']['freq_axis']['chan_freq']
else:
ms.open(vis_spl)
ms.selectinit(datadescid=0, reset=True)
spwinfo = ms.getspectralwindowinfo()
specamp = []
freq = []
time = []
for descid in range(len(spwinfo.keys())):
ms.selectinit(datadescid=0, reset=True)
ms.selectinit(datadescid=descid)
data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
specamp_ = data['amplitude']
freq_ = data['axis_info']['freq_axis']['chan_freq']
time_ = data['time']
if fillnan is not None:
flag_ = ms.getdata(['flag', 'time', 'axis_info'], ifraxis=True)['flag']
if type(fillnan) in [int, float, long]:
specamp_[flag_] = float(fillnan)
else:
specamp_[flag_] = 0.0
specamp.append(specamp_)
freq.append(freq_)
time.append(time_)
specamp = np.concatenate(specamp, axis=1)
freq = np.concatenate(freq, axis=0)
ms.selectinit(datadescid=0, reset=True)
ms.close()
os.system('rm -rf ' + vis_spl)
(npol, nfreq, nbl, ntim) = specamp.shape
freq = freq.reshape(nfreq)
if verbose:
print('npol, nfreq, nbl, ntime:', (npol, nfreq, nbl, ntim))
spec = np.swapaxes(specamp, 2, 1)
tim = data['time']
if domedian:
if verbose:
print('doing median of all the baselines')
# mask zero values before median
spec_masked = np.ma.masked_where(spec < 1e-9, spec)
spec_med = np.ma.filled(np.ma.median(spec_masked, axis=1), fill_value=0.)
nbl = 1
ospec = spec_med.reshape((npol, nbl, nfreq, ntim))
else:
ospec = spec
# Save the dynamic spectral data
if savespec:
if not specfile:
specfile = msfile + '.dspec.npz'
if os.path.exists(specfile):
os.system('rm -rf ' + specfile)
np.savez(specfile, spec=ospec, tim=tim, freq=freq,
timeran=timeran, spw=spw, bl=bl, uvrange=uvrange, pol=pols)
if verbose:
print('Median dynamic spectrum saved as: ' + specfile)
return {'spec': ospec, 'tim': tim, 'freq': freq, 'timeran': timeran, 'spw': spw, 'bl': bl, 'uvrange': uvrange,
'pol': pols}
def get_dspec(vis=None, savespec=True, specfile=None, bl='', uvrange='', field='', scan='', datacolumn='data',
domedian=False, timeran=None, spw=None, timebin='0s', regridfreq=False, fillnan=None, verbose=False,
usetbtool=False):
# from split_cli import split_cli as split
if vis.endswith('/'):
vis = vis[:-1]
msfile = vis
if not spw:
spw = ''
if not timeran:
timeran = ''
if domedian:
if not uvrange:
uvrange = '0.2~0.8km'
# bl = ''
else:
uvrange = ''
if not bl:
bl = ''
else:
uvrange = ''
# Open the ms and plot dynamic spectrum
if verbose:
print('Splitting selected data...')
if usetbtool:
try:
tb.open(vis + '/POLARIZATION')
corrtype = tb.getcell('CORR_TYPE', 0)
pols = [stokesenum[p] for p in corrtype]
tb.close()
except:
pols = []
antmask = []
if uvrange is not '' or bl is not '':
ms.open(vis)
ms.selectinit(datadescid=0)
mdata = ms.metadata()
antlist = mdata.antennaids()
mdata.done()
staql = {'uvdist': uvrange, 'baseline': bl, 'spw': spw, 'field': field, 'scan': scan, 'timerange': timeran}
### todo the selection only works for uvrange and bl. To make the selection of other items works,
## I need to make mask for other items.
a = ms.msselect(staql)
mdata = ms.metadata()
baselines = mdata.baselines()
for lidx, l in enumerate(antlist):
antmask.append(baselines[l][antlist[lidx:]])
antmask = np.hstack(antmask)
mdata.done()
ms.close()
tb.open(vis)
spwtb = tbtool()
spwtb.open(vis + '/SPECTRAL_WINDOW')
ptb = tbtool()
ptb.open(vis + '/POLARIZATION')
ms.open(vis)
spwlist = []
mdata = ms.metadata()
nspw = mdata.nspw()
nbl = mdata.nbaselines() + mdata.nantennas()
nscans = mdata.nscans()
spw_nfrq = [] # List of number of frequencies in each spw
for i in range(nspw):
spw_nfrq.append(mdata.nchan(i))
spw_nfrq = np.array(spw_nfrq)
nf = np.sum(spw_nfrq)
smry = mdata.summary()
scan_ntimes = [] # List of number of times in each scan
for iscan in range(nscans):
scan_ntimes.append(
smry['observationID=0']['arrayID=0']['scan=' + str(iscan)]['fieldID=0']['nrows'] / nspw / nbl)
scan_ntimes = np.array(scan_ntimes)
scan_ntimes_integer = scan_ntimes.astype(np.int)
if len(np.where(scan_ntimes % scan_ntimes_integer != 0)[0]) != 0:
# if True:
scan_ntimes = [] # List of number of times in each scan
for iscan in range(nscans):
scan_ntimes.append(
len(smry['observationID=0']['arrayID=0']['scan=' + str(iscan)]['fieldID=0'].keys()) - 6)
scan_ntimes = np.array(scan_ntimes)
else:
scan_ntimes = scan_ntimes_integer
nt = np.sum(scan_ntimes)
times = tb.getcol('TIME')
if times[nbl] - times[0] != 0:
# This is frequency/scan sort order
order = 'f'
elif times[nbl * nspw - 1] - times[0] != 0:
# This is time sort order
order = 't'
npol = ptb.getcol('NUM_CORR', 0, 1)[0]
ptb.close()
freq = np.zeros(nf, float)
times = np.zeros(nt, float)
if order == 't':
specamp = np.zeros((npol, nf, nbl, nt), np.complex)
for j in range(nt):
fptr = 0
# Loop over spw
for i in range(nspw):
# Get channel frequencies for this spw (annoyingly comes out as shape (nf, 1)
cfrq = spwtb.getcol('CHAN_FREQ', i, 1)[:, 0]
if j == 0:
# Only need this the first time through
spwlist += [i] * len(cfrq)
if i == 0:
times[j] = tb.getcol('TIME', nbl * (i + nspw * j), 1) # Get the time
spec_ = tb.getcol('DATA', nbl * (i + nspw * j), nbl) # Get complex data for this spw
flag = tb.getcol('FLAG', nbl * (i + nspw * j), nbl) # Get flags for this spw
nfrq = len(cfrq)
# Apply flags
if type(fillnan) in [int, float]:
spec_[flag] = float(fillnan)
else:
spec_[flag] = 0.0
# Insert data for this spw into larger array
specamp[:, fptr:fptr + nfrq, :, j] = spec_
freq[fptr:fptr + nfrq] = cfrq
fptr += nfrq
else:
specf = np.zeros((npol, nf, nt, nbl), np.complex) # Array indexes are swapped
iptr = 0
for j in range(nscans):
# Loop over scans
for i in range(nspw):
# Loop over spectral windows
s = scan_ntimes[j]
f = spw_nfrq[i]
s1 = np.sum(scan_ntimes[:j]) # Start time index
s2 = np.sum(scan_ntimes[:j + 1]) # End time index
f1 = np.sum(spw_nfrq[:i]) # Start freq index
f2 = np.sum(spw_nfrq[:i + 1]) # End freq index
spec_ = tb.getcol('DATA', iptr, nbl * s)
flag = tb.getcol('FLAG', iptr, nbl * s)
if j == 0:
cfrq = spwtb.getcol('CHAN_FREQ', i, 1)[:, 0]
freq[f1:f2] = cfrq
spwlist += [i] * len(cfrq)
times[s1:s2] = tb.getcol('TIME', iptr, nbl * s).reshape(s, nbl)[:, 0] # Get the times
iptr += nbl * s
# Apply flags
if type(fillnan) in [int, float]:
spec_[flag] = float(fillnan)
else:
spec_[flag] = 0.0
# Insert data for this spw into larger array
specf[:, f1:f2, s1:s2] = spec_.reshape(npol, f, s, nbl)
# Swap the array indexes back to the desired order
specamp = np.swapaxes(specf, 2, 3)
tb.close()
spwtb.close()
ms.close()
if len(antmask) > 0:
specamp = specamp[:, :, np.where(antmask)[0], :]
(npol, nfreq, nbl, ntim) = specamp.shape
tim = times
else:
# Open the ms and plot dynamic spectrum
if verbose:
print('Splitting selected data...')
vis_spl = './tmpms.splitted'
if os.path.exists(vis_spl):
os.system('rm -rf ' + vis_spl)
# split(vis=msfile, outputvis=vis_spl, timerange=timeran, antenna=bl, field=field, scan=scan, spw=spw,
# uvrange=uvrange, timebin=timebin, datacolumn=datacolumn)
try:
from split_cli import split_cli as split
split(vis=msfile, outputvis=vis_spl, datacolumn=datacolumn, timerange=timeran, spw=spw, antenna=bl,
field=field,
scan=scan, uvrange=uvrange, timebin=timebin)
except:
ms.open(msfile, nomodify=True)
ms.split(outputms=vis_spl, whichcol=datacolumn, time=timeran, spw=spw, baseline=bl, field=field, scan=scan,
uvrange=uvrange, timebin=timebin)
ms.close()
if verbose:
print('Regridding into a single spectral window...')
# print('Reading data spw by spw')
try:
tb.open(vis_spl + '/POLARIZATION')
corrtype = tb.getcell('CORR_TYPE', 0)
pols = [stokesenum[p] for p in corrtype]
tb.close()
except:
pols = []
if regridfreq:
ms.open(vis_spl, nomodify=False)
ms.cvel(outframe='LSRK', mode='frequency', interp='nearest')
ms.selectinit(datadescid=0, reset=True)
data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
specamp = data['amplitude']
freq = data['axis_info']['freq_axis']['chan_freq']
else:
ms.open(vis_spl)
ms.selectinit(datadescid=0, reset=True)
spwinfo = ms.getspectralwindowinfo()
specamp = []
freq = []
time = []
for descid in range(len(spwinfo.keys())):
ms.selectinit(datadescid=0, reset=True)
ms.selectinit(datadescid=descid)
data = ms.getdata(['amplitude', 'time', 'axis_info'], ifraxis=True)
specamp_ = data['amplitude']
freq_ = data['axis_info']['freq_axis']['chan_freq']
time_ = data['time']
if fillnan is not None:
flag_ = ms.getdata(['flag', 'time', 'axis_info'], ifraxis=True)['flag']
if type(fillnan) in [int, float, long]:
specamp_[flag_] = float(fillnan)
else:
specamp_[flag_] = 0.0
specamp.append(specamp_)
freq.append(freq_)
time.append(time_)
specamp = np.concatenate(specamp, axis=1)
freq = np.concatenate(freq, axis=0)
ms.selectinit(datadescid=0, reset=True)
ms.close()
os.system('rm -rf ' + vis_spl)
(npol, nfreq, nbl, ntim) = specamp.shape
freq = freq.reshape(nfreq)
tim = data['time']
if verbose:
print('npol, nfreq, nbl, ntime:', (npol, nfreq, nbl, ntim))
spec = np.swapaxes(specamp, 2, 1)
if domedian:
if verbose:
print('doing median of all the baselines')
# mask zero values before median
# spec_masked = np.ma.masked_where(spec < 1e-9, spec)
# spec_masked2 = np.ma.masked_invalid(spec)
# spec_masked = np.ma.masked_array(spec, mask=np.logical_or(spec_masked.mask, spec_masked2.mask))
# spec_med = np.ma.filled(np.ma.median(spec_masked, axis=1), fill_value=0.)
spec = np.abs(spec)
spec_med = np.nanmedian(spec, axis=1)
nbl = 1
ospec = spec_med.reshape((npol, nbl, nfreq, ntim))
else:
ospec = spec
# Save the dynamic spectral data
if savespec:
if not specfile:
specfile = msfile + '.dspec.npz'
if os.path.exists(specfile):
os.system('rm -rf ' + specfile)
np.savez(specfile, spec=ospec, tim=tim, freq=freq,
timeran=timeran, spw=spw, bl=bl, uvrange=uvrange, pol=pols)
if verbose:
print('Median dynamic spectrum saved as: ' + specfile)
return {'spec': ospec, 'tim': tim, 'freq': freq, 'timeran': timeran, 'spw': spw, 'bl': bl, 'uvrange': uvrange,
'pol': pols}
def insertdiskmodel(vis, sizescale=1.0, fdens=None, dsize=None, overwrite=True, xmlfile='SOLDISK.xml'):
if fdens is None:
# Default flux density for solar minimum
fdens = np.array([891282, 954570, 1173229, 1245433, 1373730, 1506802,
1613253, 1702751, 1800721, 1946756, 2096020, 2243951,
2367362, 2525968, 2699795, 2861604, 3054829, 3220450,
3404182, 3602625, 3794312, 3962926, 4164667, 4360683,
4575677, 4767210, 4972824, 5211717, 5444632, 5648266,
5926634, 6144249, 6339863, 6598018, 6802707, 7016012,
7258929, 7454951, 7742816, 7948976, 8203206, 8411834,
8656720, 8908130, 9087766, 9410760, 9571365, 9827078,
10023598, 8896671])
if dsize is None:
# Default solar disk radius for solar minimum
dsize = np.array(['1228.0arcsec', '1194.0arcsec', '1165.0arcsec', '1139.0arcsec', '1117.0arcsec',
'1097.0arcsec', '1080.0arcsec', '1065.0arcsec', '1053.0arcsec', '1042.0arcsec',
'1033.0arcsec', '1025.0arcsec', '1018.0arcsec', '1012.0arcsec', '1008.0arcsec',
'1003.0arcsec', '1000.0arcsec', '997.0arcsec', '994.0arcsec', '992.0arcsec',
'990.0arcsec', '988.0arcsec', '986.0arcsec', '985.0arcsec', '983.0arcsec', '982.0arcsec',
'980.0arcsec', '979.0arcsec', '978.0arcsec', '976.0arcsec', '975.0arcsec', '974.0arcsec',
'972.0arcsec', '971.0arcsec', '970.0arcsec', '969.0arcsec', '968.0arcsec', '967.0arcsec',
'966.0arcsec', '965.0arcsec', '964.0arcsec', '964.0arcsec', '963.0arcsec', '962.0arcsec',
'962.0arcsec', '961.0arcsec', '960.0arcsec', '959.0arcsec', '957.0arcsec', '956.0arcsec'])
# Apply size scale adjustment (default is no adjustment)
for i in range(len(dsize)):
num, unit = dsize[i].split('arc')
dsize[i] = str(float(num) * sizescale)[:6] + 'arc' + unit
msfile = vis
diskim = []
ms.open(msfile)
spwinfo = ms.getspectralwindowinfo()
nspw = len(spwinfo.keys())
ms.close()
diskimdir = 'diskim/'
if not os.path.exists(diskimdir):
os.makedirs(diskimdir)
frq = []
for sp in range(nspw):
spw = spwinfo[str(sp)]
frq.append('{:.4f}GHz'.format((spw['RefFreq'] + spw['TotalWidth'] / 2.0) / 1e9))
frq = np.array(frq)
writediskxml(dsize, fdens, frq, xmlfile=xmlfile)
tb.open(msfile + '/FIELD')
phadir = tb.getcol('PHASE_DIR').flatten()
tb.close()
ra = phadir[0]
dec = phadir[1]
direction = 'J2000 ' + str(ra) + 'rad ' + str(dec) + 'rad'
for sp in tqdm(range(nspw), desc='Generating {} disk models'.format(nspw), ascii=True):
diskim.append(
diskmodel(outname=diskimdir + 'disk{:02d}_'.format(sp), bdwidth=spwinfo[str(sp)], direction=direction,
reffreq=frq[sp],
flux=fdens[sp], eqradius=dsize[sp], polradius=dsize[sp], overwrite=overwrite))
delmod(msfile, otf=True, scr=True)
mstl.clearflagrow(msfile, mode='clear')
for sp in tqdm(range(nspw), desc='Inserting disk model', ascii=True):
ft(vis=msfile, spw=str(sp), field='', model=str(diskim[sp]), nterms=1,
reffreq="", complist="", incremental=False, usescratch=True)
uvsub(vis=msfile)
# tb.open(os.path.join(msfile, 'DATA_DESCRIPTION'), nomodify=False)
#
# tabdesc = {'DISK_SIZE': {'comment': 'Size of solar disk',
# 'dataManagerGroup': 'StandardStMan',
# 'dataManagerType': 'StandardStMan',
# 'keywords': {},
# 'maxlen': 0,
# 'option': 0,
# 'valueType': 'double'},
# 'DISK_FLUX': {'comment': 'Flux [Jy] of solar disk',
# 'dataManagerGroup': 'StandardStMan',
# 'dataManagerType': 'StandardStMan',
# 'keywords': {},
# 'maxlen': 0,
# 'option': 0,
# 'valueType': 'double'},
# '_define_hypercolumn_': {},
# '_keywords_': {},
# '_private_keywords_': {}}
#
# tb.addcols(tabdesc)
# tb.close()
# tabdesc = {'DISK_SIZE': {'comment': 'Size of solar disk',
# 'dataManagerGroup': 'StandardStMan',
# 'dataManagerType': 'StandardStMan',
# 'keywords': {},
# 'maxlen': 0,
# 'option': 0,
# 'valueType': 'double'},
# 'DISK_FLUX': {'comment': 'Flux [Jy] of solar disk',
# 'dataManagerGroup': 'StandardStMan',
# 'dataManagerType': 'StandardStMan',
# 'keywords': {},
# 'maxlen': 0,
# 'option': 0,
# 'valueType': 'double'},
# '_define_hypercolumn_': {},
# '_keywords_': {},
# '_private_keywords_': {}}
# tb.create(os.path.join(msfile, 'SOLDISK'),tabdesc)
# tb.addrows(50) # Add the rows _before_ filling the columns.
# tb.close()
return msfile