<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Comparing Arecibo, VLA, and Effelsberg constraints on FRB 121102 burst spectra
Build broad band spectra of FRB for Law et al 2017
Note that this notebook assumes
Step1:  Useful functions
Step2:  Read coherently dedispersed Arecibo dynamic spectrum
Step3:  Define python names for Arecibo header info
Step4:  Read dynamic spectrum and average down a bit
Define time/freq binning and barycentric time grid for Arecibo data
Step5:  Arecibo burst SNR
Step6:  crude integrated flux measurement
Step7:  Make nice Arecibo burst spectrum like the VLA ones
Step10:  VLA
Set calibration table and set useful functions for rtpipe
Step11:  Read data with rtpipe and phase it to FRB 121102
Step12:  Select VLA data for comparison to Arecibo
Step13:  Calculate dm delay and topocentric correction for VLA
Step14:  Regrid VLA and AO data to same fixed image grid
Step15:  Plot it
Step16:  Calculate residual sweep and compare to DM model
Step17:  Burst spectra with and without detections
Compare Effelsberg, VLA, Arecibo measured burst fluxes and limits
Step18:  Calculate implied spectral index and limits
<ASSISTANT_TASK:>
Python Code:
%matplotlib inline
import numpy as np
import pylab as pl
import astropy.io.fits as fits
import rtpipe
import rtlib_cython as rtlib
import astropy.units as units
import astropy.coordinates as coord
from astropy.time import Time
# confirm version is is earlier than 1.54 if using old dm scale
print(rtpipe.__version__)
Explanation: Comparing Arecibo, VLA, and Effelsberg constraints on FRB 121102 burst spectra
Build broad band spectra of FRB for Law et al 2017
Note that this notebook assumes:
- rtpipe <= v1.54
- sdmpy<=1.35
This is due to an error in the DM scale that must be preserved to properly redetect the FRB. Starting at v1.55, the DM scale factor changed from 4.2e-3 to 4.1488e-3. 
The DM scale error is 1.2%, so DM=560.5 (in v 1.54) is actually 567.4 (in v1.55 and in other codes and in reality).
End of explanation
dmdelay_new = lambda dm, f0, f1: 4.1488e-3*dm*(1./f0**2 - 1./f1**2)  # inconsistent with rtpipe
dmdelay = lambda dm, f0, f1: 4.2e-3*dm*(1./f0**2 - 1./f1**2)  # consistent with rtpipe, but requires scaling by 1.2%
Explanation: Useful functions
End of explanation
#name = 'puppi_57648_C0531+33_0048_3756.99.ar.paz'
name = 'puppi_57648_C0531+33_0048_3757.00.ar.paz.pT'
fits.info(name)
hdu = fits.open(name)
hdu0, hdu1, hdu2, hdu3, hdu4 = hdu[0], hdu[1], hdu[2], hdu[3], hdu[4]
hdu0.header
Explanation: Read coherently dedispersed Arecibo dynamic spectrum
End of explanation
nch0 = 512
obsfreq = 1780.
df = 1.5625
tint0 = 2.04780080448418e-05
#dm0 = 560.5
dm0 = 553.7  # equivalent to 560.5 in correct scaling
dmt_ao = dmdelay(dm0, obsfreq*1e-3, 1e4)/(24*3600)
bary_ao = 75.7555109 # seconds, from presto
mjd0 = int(57648) # subtract this for convenience. must be an integer mjd
#mjdfrac_ao = (37750 + 0.319370962786796)/(24*3600) - dmt_ao # start time from header
#mjdfrac_ao = (37750 + 0.336629511829426)/(24*3600) - dmt_ao # start time from header, second round
mjdfrac_ao = 57648.43692057008 - dmt_ao - mjd0 + bary_ao/(24*3600)
nint0 = 4096
dt_ao = tint0*nint0
Explanation: Define python names for Arecibo header info
End of explanation
tbin = 8
nint1 = nint0/tbin
tint0 = dt_ao/nint0
print('Orig time resolution {0}'.format(tint0))
tint1 = tint0*tbin
print('New time resolution {0}'.format(tint1))
data = fits.getdata(name, ext=4)
spec = data[0][9]  # spectrum read as (npol, nchan, nint) array
specI = spec[:2].mean(axis=0)  # average over first two to get stokes I **confirm**
specI = specI.reshape(nch0, nint1, tbin).mean(axis=2)  # use array slicing to rebin
# flag bad channels
flagch = range(0, 16) + [99, 100, 101, 286, 287, 310, 311] + range(320, 356) + range(nch0-110, nch0)
specI[flagch] = 0
specI = np.ma.masked_equal(specI, 0)
fbin = 2
nch1 = nch0/fbin
print('Orig freq resolution {0}'.format(df))
print('New freq resolution {0}'.format(df*fbin))
specAO = specI.reshape(nch1, fbin, nint1).mean(axis=1)  # use array slicing to rebin
meanspec = specAO.mean(axis=1)
specAO = specAO - meanspec[:,None]
specAO = (specAO-np.ma.mean(specAO))/np.ma.std(specAO)  # convert to snr per pixel
#specAO = specAO/5  # scale by gain relative to VLA?
print(specAO.shape)
tbary_ao = np.linspace(mjdfrac_ao, mjdfrac_ao+(dt_ao/(24*3600)), nint1)
freqs_ao = np.linspace(obsfreq - fbin*df*nch1, obsfreq, nch1)
extent_ao = (tbary_ao[0], tbary_ao[-1], freqs_ao[0], freqs_ao[-1])
fig = pl.figure(figsize=(12,8))
pl.imshow(specAO, interpolation='nearest', aspect='auto', extent=extent_ao)
pl.colorbar()
pl.figure(figsize=(8,8))
pl.subplot(211)
pl.plot(tbary_ao, specAO.mean(axis=0))
pl.ylabel('Amp (arb)')
pl.xlabel('Time (s)')
pl.subplot(212)
pl.plot(freqs_ao, meanspec)
pl.xlabel('Freq (MHz)')
pl.ylabel('Amp (arb)')
Explanation: Read dynamic spectrum and average down a bit
Define time/freq binning and barycentric time grid for Arecibo data
End of explanation
sp = specAO.reshape(nch1, nint1/16, 16).mean(axis=2)
print(sp.mean(axis=0).max()/sp.mean(axis=0)[:20].std())
tpk_ao = tbary_ao[np.where(specAO.mean(axis=0) == specAO.mean(axis=0).max())][0]
print('%.10f' % tpk_ao)
Explanation: Arecibo burst SNR
End of explanation
sigma = 4
peakbins = np.where(specAO.mean(axis=0) >= sigma*specAO.mean(axis=0)[:256].std())
peaktimes = tbary_ao[peakbins]
window = (peaktimes.max()-peaktimes.min())*24*3600
print('{0} sigma limit selects {1} ms of pulse ({2} bins)'.format(sigma, 
                                                                  window*1e3, 
                                                                  len(peaktimes)))
print('Window width is a bit wider than properly-dedispersed pulse, due to extra sweep')
Sint_sys = specAO.mean(axis=0)[peakbins].mean()
noise = specAO.mean(axis=0)[:peakbins[0][0]].std()/np.sqrt(len(peaktimes))
Sint = (3./np.sqrt(600e6*window*2))*Sint_sys/noise
print('Integrated flux density over {0} ms pulse: {1} mJy'.format(window, Sint*1e3))
Sspec_sys = specAO[:, peakbins[0]].mean(axis=1)
noise_spec = specAO[:, :peakbins[0][0]].std()/np.sqrt(len(peaktimes))
Sspec = (3./np.sqrt(3.125e6*window*2))*Sspec_sys/noise_spec
print('A (less good) integrated flux density over 2 ms pulse: {0} mJy'.format(int(Sspec.mean()*1e3)))
Explanation: crude integrated flux measurement
End of explanation
fig = pl.figure(figsize=(15,7))
ax = fig.add_subplot(111)
pl.plot(1e-3*freqs_ao[::-1], Sspec, 'k.')
pl.text(0.75, 0.88, '57648, Arecibo', horizontalalignment='left', fontsize=24,
        verticalalignment='center', transform=ax.transAxes)
pl.errorbar(1.600, 1.2*Sspec.max(), yerr=(3./np.sqrt(3.125e6*window*2)), fmt='k.', ecolor='k')
pl.ylim(-0.03, Sspec.max()*1.4)
pl.xlabel('Frequency (GHz)', fontsize=18)
pl.ylabel('Flux density (Jy)', fontsize=18)
xt = pl.setp(ax.get_xticklabels(), fontsize=18)
yt = pl.setp(ax.get_yticklabels(), fontsize=18)
ax.xaxis.set_tick_params(width=4, color='k')
ax.yaxis.set_tick_params(width=4, color='k')
fig.savefig('specAO_57648.pdf', format='pdf')
Explanation: Make nice Arecibo burst spectrum like the VLA ones
End of explanation
calstring = 2.0520    2.89698    0.00279
2.1800    *******    *******
2.3080    *******    *******
2.4360    3.53585    0.00377
2.5640    3.69554    0.00376
2.6920    3.85507    0.00423
2.8200    4.00438    0.00486
2.9480    4.11069    0.00562
3.0520    4.20375    0.00631
3.1800    4.29385    0.00662
3.3080    4.36557    0.00715
3.4360    4.43684    0.00786
3.5640    4.46937    0.00850
3.6920    4.52488    0.00860
3.8200    4.53571    0.00969
3.9480    4.54625    0.00859 
# parse flux scale
freq = []
flux = []
eflux = []
for line in calstring.split('\n'): 
    if '*' not in line:
        result = line.split()
        freq.append(float(result[0]))
        flux.append(float(result[1]))
        eflux.append(float(result[2]))
calfreq = np.array(freq)
calflux = np.array(flux)
print(calfreq, calflux)
def getscannum(sdmfile):
    sdm = rtpipe.parsesdm.getsdm(sdmfile)
    for scan in sdm.scans():
        try:
            print('Scan {0} binary data file: {1}'.format(scan.idx, scan.bdf.fname))
            bdfscan = int(scan.idx)
        except IOError:
            pass
    return bdfscan
def read_cut(sdmfile, scan, segment, dm=558., dt=1, gainfile=None, **kwargs):
    if not gainfile:
        gainfile = '.'.join(sdmfile.split('.')[:-1] + ['GN'])
    st = rtpipe.RT.set_pipeline(sdmfile, scan, dmarr=[dm], dtarr=[dt], flaglist=[('badap', 3., 0.2)], 
                                uvoversample=1.5, gainfile=gainfile, flagantsol=True, 
                                timesub='mean', logfile=False, savecands=False,
                                savenoise=False, **kwargs)
        
    data = rtpipe.RT.pipeline_reproduce(st, candloc=[segment,0,0,0,0], product='data')
    u, v, w = rtpipe.parsesdm.get_uvw_segment(st, segment)
    return st, data, u, v, w
def correctdata(st, data, u, v, w, corr='ph,dm', lm = (-3.835e-04,5.406e-04)):
     lm gives (ra, dec) = (5 31 58.703708986 33 8 52.5067634154)
    as quoted in Chatterjee et al (2017)
    
    data2 = data.copy()
    
    if 'ph' in corr:
        l1, m1 = lm
        rtlib.phaseshift_threaded(data2, st, l1, m1, u, v)
    if 'dm' in corr:
        rtlib.dedisperse_par(data2, st['freq'], st['inttime'], st['dmarr'][0], [0, st['nbl']])
    
    return data2
# get array2 for bin in array near value
def find_nearest(array, array2, value):
    idx = (np.abs(array-value)).argmin()
    return array2[idx]
def getscale(st):
    # get flux scaling at nearest frequency
    scale = []
    for i in range(len(st['freq'])):
        freq = st['freq'][i]
        scale.append(find_nearest(calfreq, calflux, freq))
#         print(i, st['freq'][i], scale)
    scale = np.array(scale, dtype='complex64')[None,None,:,None]
    return scale
def correct_all(st, data, u, v, w):
    scale = getscale(st)
    dataph = correctdata(st, data*scale, u, v, w, corr='ph')
    dataphdm = correctdata(st, dataph, u, v, w, corr='dm')
    return dataphdm
Explanation: VLA
Set calibration table and set useful functions for rtpipe
End of explanation
key = '57648'
read = {}
sdmfile = '16A-496_sb32698778_1_02h00m_000.57648.37452900463.cut/'
scannum = getscannum(sdmfile)
read[key] = read_cut(sdmfile, scannum, 7, npix_max=7400, chans=range(2,256))
Explanation: Read data with rtpipe and phase it to FRB 121102
End of explanation
st, data, u, v, w = read[key]
st['dmarr'] = [dm0]
scale = getscale(st)
dataph = correctdata(st, data*scale, u, v, w, corr='ph,dm')
intrange = (30, 60)
dint = intrange[1] - intrange[0]
specVLA = dataph[intrange[0]:intrange[1]].mean(axis=3).mean(axis=1).real
specVLA = (specVLA-specVLA.mean())/specVLA.std()  # in SNR units
print(specAO.shape, specVLA.shape)
Explanation: Select VLA data for comparison to Arecibo
End of explanation
dmt_vla = dmdelay(dm0, st['freq'][-1], 1e4)/(24*3600)
topo_vla = 75.7533100 # seconds, from presto
mjdfrac_vla = st['segmenttimes'][7][0] + intrange[0]*st['inttime']/(24*3600) - mjd0 - dmt_vla + topo_vla/(24*3600)
print('VLA dMJD: {0}'.format(mjdfrac_vla))
print('AO dMJD: {0}'.format(mjdfrac_ao))
print('Diff: {0} s'.format((mjdfrac_ao-mjdfrac_vla)*24*3600))
tbary_vla = np.linspace(mjdfrac_vla, mjdfrac_vla+(dint*st['inttime']/(24*3600)), dint)
freqs_vla = 1e3*st['freq']
extent_vla = (tbary_vla[0], tbary_vla[-1], freqs_vla[0], freqs_vla[-1])
# put them together
fig = pl.figure(figsize=(8,8))
fig.add_subplot(211)
pl.imshow(specVLA.transpose(), interpolation='nearest', origin='bottom', extent=extent_vla, aspect='auto')
fig.add_subplot(212)
pl.imshow(specAO, interpolation='nearest', aspect='auto', extent=extent_ao)
tpk_vla = tbary_vla[np.where(specVLA.mean(axis=1) == specVLA.mean(axis=1).max())][0]
print('%.10f' % tpk_vla)
Explanation: Calculate dm delay and topocentric correction for VLA
End of explanation
gap = 30
specregrid = np.zeros(shape=(nch1+st['nchan']+gap, len(tbary_ao)))
for idx in range(len(tbary_ao)):
    specregrid[254+gap:, idx] = specAO[:, idx]
idxs_vla = [np.argmin(np.abs(tbary_vla-tbary_ao[i])) for i in range(len(tbary_ao))]
for idx_ao in range(len(idxs_vla)):
    idx_vla = idxs_vla[idx_ao]
    specregrid[:254, idx_ao] += specVLA[idx_vla, ::-1]
Explanation: Regrid VLA and AO data to same fixed image grid
End of explanation
# assumes fixed relative gain between VLA and AO == 3
fig = pl.figure(figsize=(12,12))
ax = fig.add_subplot(211)
pl.imshow(specregrid, interpolation='nearest', aspect='equal', vmax=0.8*specregrid.max(), cmap='Greys')#, vmin=-0.8)
ax.fill_between(np.arange(0, len(tbary_ao)), 254*np.ones(len(tbary_ao)), 
                (254+gap)*np.ones(len(tbary_ao)), facecolor='k')
pl.xlabel('Time (ms)', fontsize=14)
pl.ylabel('Frequency (GHz)', fontsize=14)
xticks = np.arange(0, 600, 100)
pl.xticks(xticks, np.array(tbin*2e-2*xticks, dtype='int'))
ntot = nch1+len(st['freq'])+gap
yticks = np.arange(0, ntot, 80)
print(yticks)
pl.yticks(yticks)
pl.yticks(yticks, [st['freq'][-1], st['freq'][-80], st['freq'][-160], st['freq'][-240],
                   np.round(1e-3*freqs_ao[ntot-320], 3), 
                   np.round(1e-3*freqs_ao[ntot-400], 3), 
                   np.round(1e-3*freqs_ao[ntot-480], 3)])
pl.xlim(150, len(tbary_ao))
pl.ylim(ntot-50, 0)
xt = pl.setp(ax.get_xticklabels(), fontsize=14)
yt = pl.setp(ax.get_yticklabels(), fontsize=14)
ax.xaxis.set_tick_params(width=3, color='k')
ax.yaxis.set_tick_params(width=3, color='k')
ax.text(170, 330, "Arecibo", rotation=90, fontsize=20)
ax.text(170, 130, "VLA", rotation=90, fontsize=20)
fig.savefig('aovla_spec.pdf', format='pdf')
Explanation: Plot it
End of explanation
# some residual dm sweep?
print(24*3600*(tbary_ao[np.where(specAO[128:].mean(axis=0) == specAO[128:].mean(axis=0).max())][0] -
               tbary_ao[np.where(specAO[:128].mean(axis=0) == specAO[:128].mean(axis=0).max())][0]))
dmt_ao_delta = dmdelay(dm0, obsfreq*1e-3, st['freq'][-1]) - dmdelay(560, obsfreq*1e-3, st['freq'][-1])
dmt_ao_inter = dmdelay(dm0, obsfreq*1e-3, (obsfreq-500)*1e-3) - dmdelay(560, obsfreq*1e-3, (obsfreq-500)*1e-3)
print(dmt_ao_delta, dmt_ao_inter)
Explanation: Calculate residual sweep and compare to DM model
End of explanation
# 4 bursts with ao, vla, effelsberg coverage
# (ami-la covers all 4 vla bursts)
# 57643 (AO-C, *VLA-S)
# 57645 (AO-L, *VLA-S)
# 57648 (*AO-L, *VLA-S, Eff-C)
# 57649 (AO-L, *VLA-S, Eff-C)
# * shows detections
# limits assume 2 ms pulse width
#s43d = np.array([[3.0, 4.0], [0.508, 0.0036*5]])
#s43l = np.array([[4.0], [0.0036*5]])  # useless data!
s45d = np.array([[1.38, 3.0], [0.002*5, (5/2.)*0.064]])
s45l = np.array([[1.38], [0.002*5]])  # too high?
s48d = np.array([[1.4, 3.0, 4.85], [0.057, (5/2.)*0.111, 0.028*5]])  # fixed for 2 ms width
s48l = np.array([[4.85], [0.028*5]])
s49d = np.array([[1.42, 3.0, 4.9], [0.002*5, (5/2.)*0.167, 0.028*5]])
s49l = np.array([[1.42, 4.9], [0.002*5, 0.028*5]])
fig = pl.figure(figsize=(10,5))
ax = fig.add_subplot(111)
# overplot upper limit symbols
#ax.plot(s43l[0], s43l[1], 'cv', ms=10)
ax.plot(s45l[0], s45l[1], 'kv', ms=10)
ax.plot(s48l[0], s48l[1], 'kv', ms=10)
ax.plot(s49l[0], s49l[1], 'kv', ms=10)
# plot lines with points first
#ax.plot(s43d[0], s43d[1], 'c.-', ms=10, label='12 Sep 2016')
ax.plot(s45d[0], s45d[1], 'k.--', ms=10, label='57645')
ax.plot(s48d[0], s48d[1], 'k.-.', ms=10, label='57648')
ax.plot(s49d[0], s49d[1], 'k.-', ms=10, label='57649')
ax.set_xlabel('Frequency (GHz)', fontsize=14)
ax.set_ylabel('Integrated Flux density (Jy; 2 ms)', fontsize=14)
ax.legend(fontsize=14)
xt = pl.setp(ax.get_xticklabels(), fontsize=14)
yt = pl.setp(ax.get_yticklabels(), fontsize=14)
ax.xaxis.set_tick_params(width=3, color='k')
ax.yaxis.set_tick_params(width=3, color='k')
fig.savefig('multispec.pdf', format='pdf')
Explanation: Burst spectra with and without detections
Compare Effelsberg, VLA, Arecibo measured burst fluxes and limits
End of explanation
def speclim(s):
    freqs = s[0]
    fluxes = s[1]
    for i in range(len(freqs)-1):
        freqi = freqs[i]
        freqi1 = freqs[i+1]
        # correct for plotting offset
        if freqi <= 1.5:
            freqi = 1.4
        elif freqi1 > 4.2:
            freqi1 = 4.85
        print(freqi, freqi1)
        print(fluxes[i], fluxes[i+1])
        print(np.log10(fluxes[i]/fluxes[i+1])/np.log10(freqi/freqi1))
speclim(s45d)
speclim(s48d)
speclim(s49d)
Explanation: Calculate implied spectral index and limits
End of explanation
<END_TASK>