added leanpipedt and whole new build instructions

Author: Casey Law

applycals2.py

@@ -0,0 +1,333 @@
+import numpy as n
+import os, glob
+import casautil
+
+tb = casautil.tools.table()
+
+class solutions():
+    """ Container for CASA caltable(s).
+    Provides tools for applying to data of shape (nints, nbl, nch, npol).
+    Initialize class based on input file(s) and selection criteria.
+    Optional flux scale gain file can be given. Should be gain file applied to source with setjy applied.
+    """
+
+    def __init__(self, gainfile, flagants=True):
+        """ Initialize with a table of CASA gain solutions. Can later add BP.
+        """
+
+        self.parsegain(gainfile)
+        self.flagants = flagants
+
+    def parsegain(self, gainfile):
+        """Takes .g1 CASA cal table and places values in numpy arrays.
+        """
+
+        tb.open(gainfile)
+        mjd = tb.getcol('TIME')/(24*3600)     # mjd days, as for telcal
+        spw = tb.getcol('SPECTRAL_WINDOW_ID')
+        gain = tb.getcol('CPARAM')    # dimensions of (npol, 1?, ntimes*nants)
+        snr = tb.getcol('SNR')
+        flagged = tb.getcol('FLAG')
+        tb.close()
+        tb.open(gainfile + '/ANTENNA')
+        antname = tb.getcol('NAME')   # ant number in order written to gain file
+        antnum = n.array([int(antname[i][2:]) for i in range(len(antname))])
+        tb.close()
+
+        # # need to find parent data MS to get some metadata
+        # mslist = glob.glob(gainfile[:-3] + '*.ms')
+        # try:
+        #     msfile = mslist[0]
+        #     print 'Found parent data MS %s' % msfile
+        # except IndexError:
+        #     print 'Could not find parent data MS for metadata...'
+
+        # tb.open(msfile + '/ANTENNA')
+        # antname = tb.getcol('NAME')      # one name per ant
+        # tb.close()
+        # tb.open(msfile + '/SPECTRAL_WINDOW')
+        # reffreq = 1e-6*(tb.getcol('REF_FREQUENCY')+tb.getcol('TOTAL_BANDWIDTH')/2)   # similar to telcal "skyfreq"
+        # specname = tb.getcol('NAME')
+        # tb.close()
+        # tb.open(msfile + '/SOURCE')
+        # source = [name for name in tb.getcol('NAME') if 'J' in name][0]          # should return single cal name **hack**
+        # tb.close()
+        # nsol = len(gain[0,0])
+
+        # ifid0R = specname[0][7] + '-' + specname[0][8]       # one value
+        # ifid0L = specname[0][9] + '-' + specname[0][10]       # one value
+        # ifid1R = specname[1][7] + '-' + specname[1][8]       # one value
+        # ifid1L = specname[1][9] + '-' + specname[1][10]       # one value
+
+        # # paste R,L end to end, so first loop over time, then spw, then pol
+        # mjd = n.concatenate( (time, time), axis=0)
+        # ifid = [ifid0R]*(nsol/2) + [ifid1R]*(nsol/2) + [ifid0L]*(nsol/2) + [ifid1L]*(nsol/2)   # first quarter is spw0,pol0, then spw1,pol0, ...
+        # skyfreq = n.concatenate( (reffreq[0]*n.ones(nsol/2), reffreq[1]*n.ones(nsol/2), reffreq[0]*n.ones(nsol/2), reffreq[1]*n.ones(nsol/2)), axis=0)
+        # gain = n.concatenate( (gain[0,0],gain[1,0]), axis=0)
+        # amp = n.abs(gain)
+        # phase = n.degrees(n.angle(gain))
+        # source = [source]*nsol*2
+        # flagged = n.concatenate( (flag[0,0],flag[1,0]), axis=0)
+                   
+        nants = len(n.unique(antnum))
+        nspw = len(n.unique(spw))
+        self.spwlist = n.unique(spw)
+        npol = len(gain)
+
+        # merge times less than 2s
+        nsol = 0
+        newmjd = [n.unique(mjd)[0]]
+        skip = []
+        for i in range(1, len(n.unique(mjd))):
+            if 24*3600*(n.unique(mjd)[i] - n.unique(mjd)[i-1]) < 30.:
+                skip.append(n.unique(mjd)[i])
+                continue
+            else:
+                newmjd.append(n.unique(mjd)[i])
+        
+        self.uniquemjd = n.array(newmjd)
+        nsol = len(self.uniquemjd)
+
+        print 'Parsed gain table solutions for %d solutions (skipping %d), %d ants, %d spw, and %d pols' % (nsol, len(skip), nants, nspw, npol)
+        print 'Unique solution times', self.uniquemjd
+
+        self.gain = n.zeros( (nsol, nants, nspw, npol), dtype='complex' )
+        flags = n.zeros( (nsol, nants, nspw, npol), dtype='complex' )
+        for sol in range(nsol):
+            for ant in range(nants):
+                for spw in range(nspw):
+                    for pol in range(npol):
+                        self.gain[sol, ant, spw, pol] = gain[pol,0,spw*nsol*nants+sol*nants+ant]
+                        flags[sol, ant, spw, pol] = flagged[pol,0,spw*nsol*nants+sol*nants+ant]
+        self.gain = n.ma.masked_array(self.gain, flags)
+
+#        gain = n.concatenate( (n.concatenate( (gain[0,0,:nants*nsol].reshape(nsol,nants,1,1), gain[1,0,:nants*nsol].reshape(nsol,nants,1,1)), axis=3), n.concatenate( (gain[0,0,nants*nsol:].reshape(nsol,nants,1,1), gain[1,0,nants*nsol:].reshape(nsol,nants,1,1)), axis=3)), axis=2)
+#        flagged = n.concatenate( (n.concatenate( (flagged[0,0,:nants*nsol].reshape(nsol,nants,1,1), flagged[1,0,:nants*nsol].reshape(nsol,nants,1,1)), axis=3), n.concatenate( (flagged[0,0,nants*nsol:].reshape(nsol,nants,1,1), flagged[1,0,nants*nsol:].reshape(nsol,nants,1,1)), axis=3)), axis=2)
+#        self.gain = n.ma.masked_array(gain, flagged == True)        
+
+        self.mjd = n.array(mjd); self.antnum = antnum
+
+        # make another version of ants array
+#        self.antnum = n.concatenate( (antnum, antnum), axis=0)
+#        self.amp = n.array(amp); self.phase = n.array(phase)
+#        self.antname = n.concatenate( (antname[antnum], antname[antnum]), axis=0)
+#        self.complete = n.arange(len(self.mjd))
+
+        # for consistency with telcal
+        #self.ifid = n.array(ifid); self.skyfreq = n.array(skyfreq); self.source = n.array(source)
+
+    def parsebp(self, bpfile, debug=False):
+        """ Takes bp CASA cal table and places values in numpy arrays.
+        Assumes two or fewer spw. :\
+        Assumes one bp solution per file.
+        """
+
+        # bandpass. taking liberally from Corder et al's analysisutilities
+        ([polyMode, polyType, nPolyAmp, nPolyPhase, scaleFactor, nRows, nSpws, nUniqueTimesBP, uniqueTimesBP,
+          nPolarizations, frequencyLimits, increments, frequenciesGHz, polynomialPhase,
+          polynomialAmplitude, timesBP, antennasBP, cal_desc_idBP, spwBP]) = openBpolyFile(bpfile, debug)
+
+        # index iterates over antennas, then times/sources (solution sets). each index has 2x npoly, which are 2 pols
+        polynomialAmplitude = n.array(polynomialAmplitude); polynomialPhase = n.array(polynomialPhase)
+        polynomialAmplitude[:,0] = 0.; polynomialAmplitude[:,nPolyAmp] = 0.
+        polynomialPhase[:,0] = 0.; polynomialPhase[nPolyPhase] = 0.
+        ampSolR, ampSolL = calcChebyshev(polynomialAmplitude, frequencyLimits, n.array(frequenciesGHz)*1e+9)
+        phaseSolR, phaseSolL = calcChebyshev(polynomialPhase, frequencyLimits, n.array(frequenciesGHz)*1e+9)
+
+        nants = len(n.unique(antennasBP))
+        self.bptimes = n.array(timesBP)
+        ptsperspec = 1000
+        npol = 2
+        print 'Parsed bp solutions for %d solutions, %d ants, %d spw, and %d pols' % (nUniqueTimesBP, nants, nSpws, nPolarizations)
+        self.bandpass = n.zeros( (nants, nSpws*ptsperspec, npol), dtype='complex')
+        for spw in range(nSpws):
+            ampSolR[spw*nants:(spw+1)*nants] += 1 - ampSolR[spw*nants:(spw+1)*nants].mean()     # renormalize mean over ants (per spw) == 1
+            ampSolL[spw*nants:(spw+1)*nants] += 1 - ampSolL[spw*nants:(spw+1)*nants].mean()
+            for ant in range(nants):
+                self.bandpass[ant, spw*ptsperspec:(spw+1)*ptsperspec, 0] = ampSolR[ant+spw*nants] * n.exp(1j*phaseSolR[ant+spw*nants])
+                self.bandpass[ant, spw*ptsperspec:(spw+1)*ptsperspec, 1] = ampSolL[ant+spw*nants] * n.exp(1j*phaseSolL[ant+spw*nants])
+
+        self.bpfreq = n.zeros( (nSpws*ptsperspec) )
+        for spw in range(nSpws):
+            self.bpfreq[spw*ptsperspec:(spw+1)*ptsperspec] = 1e9 * frequenciesGHz[nants*spw]
+
+#        bpSolR0 = ampSolR[:nants] * n.exp(1j*phaseSolR[:nants])
+#        bpSolR1 = ampSolR[nants:] * n.exp(1j*phaseSolR[nants:])
+#        bpSolL0 = ampSolL[:nants] * n.exp(1j*phaseSolL[:nants])
+#        bpSolL1 = ampSolL[nants:] * n.exp(1j*phaseSolL[nants:])
+
+        # structure close to tpipe data structure (nant, freq, pol). note that freq is oversampled to 1000 bins.
+#        self.bandpass = n.concatenate( (n.concatenate( (bpSolR0[:,:,None], bpSolR1[:,:,None]), axis=1), n.concatenate( (bpSolL0[:,:,None], bpSolL1[:,:,None]), axis=1)), axis=2)
+#        self.bpfreq = 1e9*n.concatenate( (frequenciesGHz[0], frequenciesGHz[nants]), axis=0)    # freq values at bp bins
+#        print 'Parsed bp table solutions for %d solutions, %d ants, %d spw, and %d pols' % (nUniqueTimesBP, nants, nSpws, nPolarizations)
+
+    def setselection(self, time, freqs, spws=[0,1], pols=[0,1], verbose=0):
+        """ Set select parameter that defines time, spw, and pol solutions to apply.
+        time defines the time to find solutions near in mjd.
+        freqs defines frequencies to select bandpass solution
+        spws is list of min/max indices to be used (e.g., [0,1])
+        pols is index of polarizations.
+        pols/spws not yet implemented beyond 2sb, 2pol.
+        """
+
+        # spw and pols selection not yet implemented beyond 2/2
+        self.spws = spws
+        self.pols = pols
+
+        # select by smallest time distance for source
+        mjddist = n.abs(time - self.uniquemjd)
+        closestgain = n.where(mjddist == mjddist.min())[0][0]
+
+        print 'Using gain solution at MJD %.5f, separated by %d min ' % (self.uniquemjd[closestgain], mjddist[closestgain]*24*60)
+        self.gain = self.gain[closestgain,:,spws[0]:spws[1]+1,pols[0]:pols[1]+1]
+
+        if hasattr(self, 'bandpass'):
+            bins = [n.where(n.min(n.abs(self.bpfreq-selfreq)) == n.abs(self.bpfreq-selfreq))[0][0] for selfreq in freqs]
+            self.bandpass = self.bandpass[:,bins,pols[0]:pols[1]+1]
+            self.freqs = freqs
+            if verbose:
+                print 'Using solution at BP bins: ', bins
+
+    def calc_flag(self, sig=3.0):
+        """ Calculates antennas to flag, based on bad gain and bp solutions.
+        """
+ 
+        if len(self.gain.shape) == 4:
+            gamp = n.abs(self.gain).mean(axis=0)   # mean gain amp for each ant over time
+        elif len(self.gain.shape) == 3:
+            gamp = n.abs(self.gain)   # gain amp for selected time
+
+#        badgain = n.where(gamp < gamp.mean() - sig*gamp.std())
+        badgain = n.where( (gamp < n.median(gamp) - sig*gamp.std()) | gamp.mask)
+        print 'Flagging low/bad gains for ant/spw/pol:', self.antnum[badgain[0]], badgain[1], badgain[2]
+
+        badants = badgain
+        return badants
+
+    def apply(self, data, blarr):
+        """ Applies calibration solution to data array. Assumes structure of (nint, nbl, nch, npol).
+        blarr is array of size 2xnbl that gives pairs of antennas in each baseline (a la tpipe.blarr).
+        """
+
+        ant1ind = [n.where(ant1 == n.unique(blarr))[0][0] for (ant1,ant2) in blarr]
+        ant2ind = [n.where(ant2 == n.unique(blarr))[0][0] for (ant1,ant2) in blarr]
+
+        # flag bad ants
+        if self.flagants:
+            badants = self.calc_flag()
+        else:
+            badants = n.array([[]])
+
+        # apply gain correction
+        if hasattr(self, 'bandpass'):
+            corr = n.ones_like(data)
+            flag = n.ones_like(data).astype('int')
+            chans_uncal = range(len(self.freqs))
+            for spw in range(len(self.gain[0])):
+                chsize = n.round(self.bpfreq[1]-self.bpfreq[0], 0)
+                ww = n.where( (self.freqs >= self.bpfreq[spw*1000]) & (self.freqs <= self.bpfreq[(spw+1)*1000-1]+chsize) )[0]
+                if len(ww) == 0:
+                    print 'Gain solution frequencies not found in data for spw %d.' % (self.spws[spw])
+                firstch = ww[0]
+                lastch = ww[-1]+1
+                for ch in ww:
+                    chans_uncal.remove(ch)
+                print 'Combining gain sol from spw=%d with BW chans from %d-%d' % (self.spws[spw], firstch, lastch)
+                for badant in n.transpose(badants):
+                    if badant[1] == spw:
+                        badbl = n.where((badant[0] == n.array(ant1ind)) | (badant[0] == n.array(ant2ind)))[0]
+                        flag[:, badbl, firstch:lastch, badant[2]] = 0
+
+                corr1 = self.gain[ant1ind, spw, :][None, :, None, :] * self.bandpass[ant1ind, firstch:lastch, :][None, :, :, :]
+                corr2 = (self.gain[ant2ind, spw, :][None, :, None, :] * self.bandpass[ant2ind, firstch:lastch, :][None, :, :, :]).conj()
+
+                corr[:, :, firstch:lastch, :] = corr1 * corr2
+            if len(chans_uncal):
+                print 'Setting data without bp solution to zero for chan range %d-%d.' % (chans_uncal[0], chans_uncal[-1])
+                flag[:, :, chans_uncal,:] = 0
+            data[:] *= flag/corr
+        else:
+            for spw in range(len(self.gain[0,0])):
+                pass
+
+def openBpolyFile(caltable, debug=False):
+#    mytb = au.createCasaTool(tbtool)    # from analysisutilities by corder
+    tb.open(caltable)
+    desc = tb.getdesc()
+    if ('POLY_MODE' in desc):
+        polyMode = tb.getcol('POLY_MODE')
+        polyType = tb.getcol('POLY_TYPE')
+        scaleFactor = tb.getcol('SCALE_FACTOR')
+        antenna1 = tb.getcol('ANTENNA1')
+        times = tb.getcol('TIME')
+        cal_desc_id = tb.getcol('CAL_DESC_ID')
+        nRows = len(polyType)
+        for pType in polyType:
+            if (pType != 'CHEBYSHEV'):
+                print "I do not recognized polynomial type = %s" % (pType)
+                return
+        # Here we assume that all spws have been solved with the same mode
+        uniqueTimesBP = n.unique(tb.getcol('TIME'))
+        nUniqueTimesBP = len(uniqueTimesBP)
+        if (nUniqueTimesBP >= 2):
+            if debug:
+                print "Multiple BP sols found with times differing by %s seconds. Using first." % (str(uniqueTimesBP-uniqueTimesBP[0]))
+            nUniqueTimesBP = 1
+            uniqueTimesBP = uniqueTimesBP[0]
+        mystring = ''
+        nPolyAmp = tb.getcol('N_POLY_AMP')
+        nPolyPhase = tb.getcol('N_POLY_PHASE')
+        frequencyLimits = tb.getcol('VALID_DOMAIN')
+        increments = 0.001*(frequencyLimits[1,:]-frequencyLimits[0,:])
+        frequenciesGHz = []
+        for i in range(len(frequencyLimits[0])):
+           freqs = (1e-9)*n.arange(frequencyLimits[0,i],frequencyLimits[1,i],increments[i])       # **for some reason this is nch-1 long?**
+           frequenciesGHz.append(freqs)
+        polynomialAmplitude = []
+        polynomialPhase = []
+        for i in range(len(polyMode)):
+            polynomialAmplitude.append([1])
+            polynomialPhase.append([0])
+            if (polyMode[i] == 'A&P' or polyMode[i] == 'A'):
+                polynomialAmplitude[i]  = tb.getcell('POLY_COEFF_AMP',i)[0][0][0]
+            if (polyMode[i] == 'A&P' or polyMode[i] == 'P'):
+                polynomialPhase[i] = tb.getcell('POLY_COEFF_PHASE',i)[0][0][0]
+  
+        tb.close()
+        tb.open(caltable+'/CAL_DESC')
+        nSpws = len(tb.getcol('NUM_SPW'))
+        spws = tb.getcol('SPECTRAL_WINDOW_ID')
+        spwBP = []
+        for c in cal_desc_id:
+            spwBP.append(spws[0][c])
+        tb.close()
+        nPolarizations = len(polynomialAmplitude[0]) / nPolyAmp[0]
+        if debug:
+            mystring += '%.3f, ' % (uniqueTimesBP/(24*3600))
+            print 'BP solution has unique time(s) %s and %d pols' % (mystring, nPolarizations)
+        
+        # This value is overridden by the new function doPolarizations in ValueMapping.
+        # print "Inferring %d polarizations from size of polynomial array" % (nPolarizations)
+        return([polyMode, polyType, nPolyAmp, nPolyPhase, scaleFactor, nRows, nSpws, nUniqueTimesBP,
+                uniqueTimesBP, nPolarizations, frequencyLimits, increments, frequenciesGHz,
+                polynomialPhase, polynomialAmplitude, times, antenna1, cal_desc_id, spwBP])
+    else:
+        tb.close()
+        return([])
+   # end of openBpolyFile()
+
+def calcChebyshev(coeffs, validDomain, freqs):
+    """
+    Given a set of coefficients,
+    this method evaluates a Chebyshev approximation.
+    Used for CASA bandpass reading.
+    input coeffs and freqs are numpy arrays
+    """
+
+    domain = (validDomain[1] - validDomain[0])[0]
+    bins = -1 + 2* n.array([ (freqs[i]-validDomain[0,i])/domain for i in range(len(freqs))])
+    ncoeffs = len(coeffs[0])/2
+    rr = n.array([n.polynomial.chebval(bins[i], coeffs[i,:ncoeffs]) for i in range(len(coeffs))])
+    ll = n.array([n.polynomial.chebval(bins[i], coeffs[i,ncoeffs:]) for i in range(len(coeffs))])
+
+    return rr,ll