Add option to adjust spectral index column for beam

lsmtool/operations_lib.py

@@ -17,11 +17,10 @@
 # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
 import logging
-import sys
-import os
 
 
-def attenuate(beamMS, fluxes, RADeg, DecDeg, timeIndx=0.5):
+def attenuate(beamMS, fluxes, RADeg, DecDeg, spectralIndex=None, referenceFrequency=None,
+              timeIndx=0.5, invert=False):
     """
     Returns flux attenuated by primary beam.
 
@@ -35,20 +34,27 @@ def attenuate(beamMS, fluxes, RADeg, DecDeg, timeIndx=0.5):
         List of RA values in degrees
     DecDeg : list
         List of Dec values in degrees
+    spectralIndex : list, optional
+        List of spectral indices to adjust
+    referenceFrequency : list, optional
+        List of reference frequency of polynomial fit
     timeIndx : float (between 0 and 1), optional
         Time as fraction of that covered by the beamMS for which the beam is
         calculated
+    invert : bool, optional
 
     Returns
     -------
     attFluxes : numpy array
+        Attenuated fluxes
+    adjSpectralIndex : numpy array
+        Adjusted spectral indices. Returned only if spectralIndex is not None
 
     """
     import numpy as np
     import pyrap.tables as pt
     import lofar.stationresponse as lsr
 
-    log = logging.getLogger('LSMTool')
     t = pt.table(beamMS, ack=False)
     time = None
     ant1 = -1
@@ -64,11 +70,15 @@ def attenuate(beamMS, fluxes, RADeg, DecDeg, timeIndx=0.5):
     if timeIndx > 1.0:
         timeIndx = 1.0
     time = min(time) + (max(time) - min(time)) * timeIndx
-    log.debug('Applying beam attenuation using beam at time of {0}% point of '
-        'observation.'.format(int(timeIndx*100.0)))
+    sw = pt.table(beamMS+'::SPECTRAL_WINDOW', ack=False)
+    numchannels = sw.col('NUM_CHAN')[0]
+    startfreq = np.min(sw.col('CHAN_FREQ')[0])
+    channelwidth = sw.col('CHAN_WIDTH')[0][0]
+    sw.close()
 
     attFluxes = []
-    sr = lsr.stationresponse(beamMS, inverse=False, useElementResponse=False,
+    adjSpectralIndex = []
+    sr = lsr.stationresponse(beamMS, inverse=invert, useElementResponse=False,
                              useArrayFactor=True, useChanFreq=False)
     if type(fluxes) is not list:
         fluxes = list(fluxes)
@@ -77,17 +87,54 @@ def attenuate(beamMS, fluxes, RADeg, DecDeg, timeIndx=0.5):
     if type(DecDeg) is not list:
         DecDeg = list(DecDeg)
 
-    for flux, RA, Dec in zip(fluxes, RADeg, DecDeg):
-        # Use ant1, mid time, and mid channel to compute the beam
+    for i, (flux, RA, Dec) in enumerate(zip(fluxes, RADeg, DecDeg)):
+        # Use ant1, times, and n channel to compute the beam, where n is determined by
+        # the order of the spectral index polynomial
         sr.setDirection(RA*np.pi/180., Dec*np.pi/180.)
         beam = sr.evaluateStation(time, ant1)
-        r = abs(beam[int(len(beam)/2.)])
-        beam = ( r[0][0] + r[1][1] ) / 2.
-        attFluxes.append(flux * beam)
 
+        # Correct the source spectrum if needed
+        if spectralIndex is not None:
+            nspec = len(spectralIndex[i])
+            s = max(1, int(numchannels/nspec))
+            ch_indices = list(range(0, numchannels, s))
+            ch_indices.append(numchannels)
+        else:
+            ch_indices = [int(numchannels/2)]
+        freqs_new = []
+        fluxes_new = []
+        for ind in ch_indices:
+            beam = abs(sr.evaluateChannel(time, ant1, ind))
+            beam = beam[0][0]  # take XX only (XX and YY should be equal)
+            if spectralIndex is not None:
+                nu = (startfreq + ind*channelwidth) / referenceFrequency[i] - 1
+                freqs_new.append(nu)
+                fluxes_new.append(polynomial(flux, spectralIndex[i], nu) * beam)
+            else:
+                fluxes_new.append(flux * beam)
+        if spectralIndex is not None:
+            fit = np.polyfit(freqs_new, fluxes_new, len(spectralIndex[i]-1)).tolist()
+            fit.reverse()
+            attFluxes.append(fit[0])
+            adjSpectralIndex.append(fit[1:])
+        else:
+            attFluxes.append(fluxes_new[0])
+            adjSpectralIndex.append(None)
+
+    if spectralIndex is not None:
+        return np.array(attFluxes), np.array(adjSpectralIndex)
+    else:
         return np.array(attFluxes)
 
 
+def polynomial(stokesI, coeff, nu):
+    result = stokesI
+    for i in range(len(coeff)):
+        result += coeff[i] * nu**(i+1)
+
+    return result
+
+
 def radec2xy(RA, Dec, refRA=None, refDec=None, crdelt=None):
     """
     Returns x, y for input ra, dec.
@@ -322,4 +369,3 @@ def calculateSeparation(ra1, dec1, ra2, dec2):
     coord2 = SkyCoord(ra2, dec2, unit=(u.degree, u.degree), frame='fk5')
 
     return coord1.separation(coord2)
-

lsmtool/skymodel.py

@@ -1600,7 +1600,8 @@ class SkyModel(object):
             return dist.value
 
     def write(self, fileName=None, format='makesourcedb', clobber=False,
-              sortBy=None, lowToHigh=False, addHistory=True, applyBeam=False):
+              sortBy=None, lowToHigh=False, addHistory=True, applyBeam=False,
+              invertBeam=False, adjustSI=False):
         """
         Writes the sky model to a file.
 
@@ -1630,7 +1631,12 @@ class SkyModel(object):
             If True, the history of operations is written to the sky model
             header.
         applyBeam : bool, optional
-            If True, apparent fluxes will be written.
+            If True, fluxes will be adjusted for the beam before being written.
+        invertBeam : bool, optional
+            If True, the beam correction is inverted (i.e., from apparent sky to
+            true sky).
+        adjustSI : bool, optional
+            If True, adjust the spectral index column for the beam
 
         Examples
         --------
@@ -1650,6 +1656,7 @@ class SkyModel(object):
         """
         import os
         import numpy as np
+        from .operations_lib import attenuate
 
         if fileName is None:
             if self._fileName is None:
@@ -1666,12 +1673,22 @@ class SkyModel(object):
 
         table = self.table.copy()
 
-        # Apply beam attenuation to I column if desired
+        # Apply beam attenuation
         if applyBeam:
-            I_apparent = self.getColValues('I', applyBeam=True)
+            spectralIndex = self.getColValues('SpectralIndex')
+            referenceFrequency = self.getColValues('ReferenceFrequency')
+            I_orig = self.getColValues('I')
+            RADeg = self.getColValues('Ra')
+            DecDeg = self.getColValues('Dec')
+            I_adj, SI_adj = attenuate(self.beamMS, I_orig, RADeg, DecDeg,
+                                      timeIndx=self.beamTime, invert=invertBeam,
+                                      spectralIndex=spectralIndex,
+                                      referenceFrequency=referenceFrequency)
             units = self.table.columns['I'].unit
-            table['I'] = I_apparent
+            table['I'] = I_adj
             table.columns['I'].unit = units
+            if adjustSI:
+                table['SpectralIndex'] = SI_adj
 
         # Sort if desired. For 'factor' output, save the order of patches in the
         # table meta