# ALMA Data Reduction Script

# Calibration

thesteps = []
step_title = {0: 'Stokes dirty images of Polarization calibrator',
              1: 'Stokes clean images of Polarization calibrator',
              2: 'Stokes images of Phase calibrator (optional)',
              3: 'IQU Stokes clean images of the Target',
              4: 'Final image production'}


try:
  thesteps = mysteps
  print 'List of steps to be executed ...', mysteps
except:
  print 'global variable mysteps not set.'
if (thesteps==[]):
  thesteps = range(0,len(step_title))
  print 'Executing all steps: ', thesteps

# The Python variable 'mysteps' will control which steps
# are executed when you start the script using
#   execfile('scriptForImagingPol.py')
# e.g. setting
#   mysteps = [2,3,4]# before starting the script will make the script execute
# only steps 2, 3, and 4
# Setting mysteps = [] will make it execute all steps.

import re

if casadef.casa_version < '4.4.0' :
    sys.exit("Please use CASA version greater than or equal to 4.4.0 with this script")

# Imaging parameters

#>>> Modify the following parameters to match the project; for calibrators, use full names and not field numbers
calmsname = 'calibrated.ms'
targetmsname = 'calibrated_final.ms' # MS with target data split out
phasecal = 'J0701-4634'
polcal = 'J0522-3627'
check = 'J0718-4319'
target = 'L2_Pup'

refant = 'DV07'
outframe = 'LSRK'
#>>> Generally, you want 5-8 cells (i.e., pixels) across the narrowest
#>>> part of the beam. You can estimate the beam size using the following
#>>> equation: 206265.0/(longest baseline in wavelengths).  To determine
#>>> the longest baseline, use plotms with xaxis='uvwave' and
#>>> yaxis='amp'. Divide the estimated beam size by five to eight to get
#>>> your cell size. It's better to error on the side of slightly too
#>>> many cells per beam than too few. Once you have made an image,
#>>> please re-assess the cell size based on the beam of the image. You can
#>>> check your cell size using  au.pickCellSize('calibrated_final.ms'). Note
#>>> however, that this routine does not take into account the projection of
#>>> the baseline onto the sky, so the plotms method described above is overall
#>>> more accurate.

#>>> To determine the image size (i.e., the imsize parameter), first you
#>>> need to figure out whether the ms is a mosaic by either looking out
#>>> the output from listobs/vishead or checking the spatial setup in the OT. For
#>>> single fields, an imsize equal to the size of the primary beam is
#>>> usually sufficient. The ALMA 12m primary beam in arcsec scales as
#>>> 6300 / nu[GHz] and the ALMA 7m primary beam in arcsec scales as
#>>> 10608 / nu[GHz], where nu[GHz] is the sky frequency. However, if
#>>> there is significant point source and/or extended emission beyond
#>>> the edges of your initial images, you should increase the imsize to
#>>> incorporate more emission. For mosaics, you can get the imsize from
#>>> the spatial tab of the OT. The parameters "p length" and "q length"
#>>> specify the dimensions of the mosaic. If you're imaging a mosaic,
#>>> pad the imsize substantially to avoid artifacts. Note that the imsize
#>>> parameter is in PIXELS, not arcsec, so you will need to divide the image size
#>>> in arcsec by the pixel size to determine a value for imsize.

#>>> Note that you can check your image size using 
#>>> au.pickCellSize('calibrated_final.ms', imsize=True). This task
#>>> now works both mosaics and single fields, but has not been tested
#>>> extensively on mosaics. Please report any large issues to
#>>> Todd Hunter. Note that au.pickCellSize does not take
#>>> into account the projection of the baselines, so the plotms
#>>> method is more accurate.
cell='0.007arcsec'
imsize=[3072,3072]
mask='circle[[1536pix,1536pix],1000pix]'

### Stokes dirty images of the polarization calibrator
mystep = 0
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

#>>> au.genImageName(vis=calmsname, spw=[0,1,2,3], field=polcal, imtype='mfs', targettype='sci', stokes='IQUV', mous='', modtext='manual'); rename 'sci' to 'pol' (task does not offer targettype)
  polimage = 'J0522-3627_polleak.spw0_1_2_3.mfs.IQUV.manual' 

  os.system('rm -rf '+polimage+'_dirty*')
  tclean(vis=calmsname,
        field=polcal,
        imagename=polimage+'_dirty',
        cell=cell,
        imsize=imsize,
        outframe=outframe,
        stokes='IQUV',  
        specmode='mfs',
        deconvolver='clarkstokes',
        gridder='standard',
        interactive=False,
        mask=mask,
        niter=0)
  
### Stokes clean images of the polarization calibrator
mystep = 1
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -rf '+polimage+'.*')
  tclean(vis=calmsname,
        field=polcal,
        imagename=polimage,
        cell=cell,
        imsize=imsize,
        outframe=outframe,
        stokes='IQUV',
        specmode='mfs',
        deconvolver='clarkstokes',
        gridder='standard',
        interactive=False,
        mask=mask,
        niter=1000)

#>>>  for stokes in 'I','Q','U','V':
#>>>      print au.genImageName(vis=calmsname, spw=[0,1,2,3], field=polcal, imtype='mfs', targettype='sci', stokes=stokes, mous='', modtext='manual'); again rename the 'sci' to 'polleak'
  polimageI = 'J0522-3627_polleak.spw0_1_2_3.mfs.I.manual'
  polimageQ = 'J0522-3627_polleak.spw0_1_2_3.mfs.Q.manual'
  polimageU = 'J0522-3627_polleak.spw0_1_2_3.mfs.U.manual'
  polimageV = 'J0522-3627_polleak.spw0_1_2_3.mfs.V.manual'

  os.system('rm -rf '+polimageI+'.image')
  immath(imagename=polimage+'.image',outfile=polimageI+'.image',expr='IM0',stokes='I')
  os.system('rm -rf '+polimageQ+'.image')
  immath(imagename=polimage+'.image',outfile=polimageQ+'.image',expr='IM0',stokes='Q')
  os.system('rm -rf '+polimageU+'.image')
  immath(imagename=polimage+'.image',outfile=polimageU+'.image',expr='IM0',stokes='U')
  os.system('rm -rf '+polimageV+'.image')
  immath(imagename=polimage+'.image',outfile=polimageV+'.image',expr='IM0',stokes='V')

  os.system('rm -rf '+polimageI+'.pb')
  immath(imagename=polimage+'.pb',outfile=polimageI+'.pb',expr='IM0',stokes='I')
  os.system('rm -rf '+polimageQ+'.pb')
  immath(imagename=polimage+'.pb',outfile=polimageQ+'.pb',expr='IM0',stokes='Q')
  os.system('rm -rf '+polimageU+'.pb')
  immath(imagename=polimage+'.pb',outfile=polimageU+'.pb',expr='IM0',stokes='U')
  os.system('rm -rf '+polimageV+'.pb')
  immath(imagename=polimage+'.pb',outfile=polimageV+'.pb',expr='IM0',stokes='V')

  resI=imfit(imagename = polimageI+'.image', box = '1036,1036,2036,2036')
  resQ=imfit(imagename = polimageQ+'.image', box = '1036,1036,2036,2036')
  resU=imfit(imagename = polimageU+'.image', box = '1036,1036,2036,2036')

 
# and then we extract the flux and error values for each Stokes
 
  fluxI=resI['results']['component0']['flux']['value'][0]
  errorI=resI['results']['component0']['flux']['error'][0]
  
  fluxQ=resQ['results']['component0']['flux']['value'][1]
  errorQ=resQ['results']['component0']['flux']['error'][1]
 
  fluxU=resU['results']['component0']['flux']['value'][2]
  errorU=resU['results']['component0']['flux']['error'][2]
 

#Now we use these values to compute polarization intensity, angle and ratio, and their errors:


  fluxPI  = sqrt( fluxQ**2 + fluxU**2 )
  errorPI = sqrt( (fluxQ*errorU)**2 + (fluxU*errorQ)**2 ) / fluxPI
 
  fluxPImjy = 1000*fluxPI
  errPImjy  = 1000*errorPI
 
  polRatio = fluxPI / fluxI
  errPRat  = polRatio * sqrt( (errorPI/fluxPI)**2 + (errorI/fluxI)**2 )
 
  polAngle = 0.5 * degrees( atan2(fluxU,fluxQ) )
  errPA    = 0.5 * degrees( errorPI / fluxPI )


  print 'Pol ratio of Polarization calibrator: ', polRatio
  print 'Pol angle of Polarization calibrator: ', polAngle

#Pol ratio of Polarization calibrator:  0.0307198496102
#Pol angle of Polarization calibrator:  -64.7798071747


### Stokes images of the phase calibrator
mystep = 2
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

#>>> au.genImageName(vis=calmsname, spw=[0,1,2,3], field=phasecal, imtype='mfs', targettype='ph', stokes='IQUV', mous='', modtext='manual')
  phaseimage = 'J0701-4634_ph.spw0_1_2_3.mfs.IQUV.manual' #au.genImageName(vis=calmsname, spw=[0,1,2,3], field=phasecal, imtype='mfs', targettype='ph', stokes='IQUV', mous='', modtext='manual')

  os.system('rm -rf '+phaseimage+'.*')
  tclean(vis=calmsname,
        field=phasecal,
        imagename=phaseimage,
        cell=cell,
        imsize=imsize,
        outframe=outframe,
        stokes='IQUV',
        specmode='mfs',
        deconvolver='clarkstokes',
        gridder='standard',
        interactive=False,
        mask=mask,
        niter=500)

#>>>  for stokes in 'I','Q','U','V':
#>>>      print au.genImageName(vis=calmsname, spw=[0,1,2,3], field=phasecal, imtype='mfs', targettype='ph', stokes=stokes, mous='', modtext='manual')
  phaseimageI = 'J0701-4634_ph.spw0_1_2_3.mfs.I.manual'
  phaseimageQ = 'J0701-4634_ph.spw0_1_2_3.mfs.Q.manual'
  phaseimageU = 'J0701-4634_ph.spw0_1_2_3.mfs.U.manual'
  phaseimageV = 'J0701-4634_ph.spw0_1_2_3.mfs.V.manual'

  os.system('rm -rf '+phaseimageI+'.image')
  immath(imagename=phaseimage+'.image',outfile=phaseimageI+'.image',expr='IM0',stokes='I')
  os.system('rm -rf '+phaseimageQ+'.image')
  immath(imagename=phaseimage+'.image',outfile=phaseimageQ+'.image',expr='IM0',stokes='Q')
  os.system('rm -rf '+phaseimageU+'.image')
  immath(imagename=phaseimage+'.image',outfile=phaseimageU+'.image',expr='IM0',stokes='U')
  os.system('rm -rf '+phaseimageV+'.image')
  immath(imagename=phaseimage+'.image',outfile=phaseimageV+'.image',expr='IM0',stokes='V')

  os.system('rm -rf '+phaseimageI+'.pb')
  immath(imagename=phaseimage+'.pb',outfile=phaseimageI+'.pb',expr='IM0',stokes='I')
  os.system('rm -rf '+phaseimageQ+'.pb')
  immath(imagename=phaseimage+'.pb',outfile=phaseimageQ+'.pb',expr='IM0',stokes='Q')
  os.system('rm -rf '+phaseimageU+'.pb')
  immath(imagename=phaseimage+'.pb',outfile=phaseimageU+'.pb',expr='IM0',stokes='U')
  os.system('rm -rf '+phaseimageV+'.pb')
  immath(imagename=phaseimage+'.pb',outfile=phaseimageV+'.pb',expr='IM0',stokes='V')

  resI=imfit(imagename = phaseimageI+'.image', box = '1036,1036,2036,2036')
  resQ=imfit(imagename = phaseimageQ+'.image', box = '1036,1036,2036,2036')
  resU=imfit(imagename = phaseimageU+'.image', box = '1036,1036,2036,2036')

 
# and then we extract the flux and error values for each Stokes
 
  fluxI=resI['results']['component0']['flux']['value'][0]
  errorI=resI['results']['component0']['flux']['error'][0]
 
  fluxQ=resQ['results']['component0']['flux']['value'][1]
  errorQ=resQ['results']['component0']['flux']['error'][1]
 
  fluxU=resU['results']['component0']['flux']['value'][2]
  errorU=resU['results']['component0']['flux']['error'][2]
 

#Now we use these values to compute polarization intensity, angle and ratio, and their errors:


  fluxPI  = sqrt( fluxQ**2 + fluxU**2 )
  errorPI = sqrt( (fluxQ*errorU)**2 + (fluxU*errorQ)**2 )/fluxPI
 
  fluxPImjy = 1000*fluxPI
  errPImjy  = 1000*errorPI
 
  polRatio = fluxPI / fluxI
  errPRat  = polRatio * sqrt((errorPI/fluxPI)**2+(errorI/fluxI)**2)
 
  polAngle = 0.5 * degrees(atan2(fluxU,fluxQ))
  errPA    = 0.5 * degrees(errorPI/fluxPI)

  print 'Pol ratio of Phase calibrator: ', polRatio
  print 'Pol angle of Phase calibrator: ', polAngle

#Pol ratio of Phase calibrator:  0.0262798045037
#Pol angle of Phase calibrator:  -85.3885061033

# Also image the check source 

#>>> au.genImageName(vis=calmsname, spw=[0,1,2,3], field=check, imtype='mfs', targettype='chk', stokes='IQUV', mous='', modtext='manual')
  checkimage = 'J0718-4319_chk.spw0_1_2_3.mfs.IQUV.manual'

  os.system('rm -rf '+checkimage+'.*')
  tclean(vis=calmsname,
        field=check,
        imagename=checkimage,
        cell=cell,
        imsize=imsize,
        specmode='mfs',
        deconvolver='clarkstokes',
        gridder='standard',
        outframe=outframe,
        stokes='IQUV',
        interactive=False,
        mask=mask,
        niter=1000)


### IQU Stokes clean images of the Target
mystep = 3
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

#>>> au.genImageName(vis=targetmsname, spw=[0,1,2,3], field=target, imtype='mfs', targettype='sci', stokes='IQUV', mous='', modtext='manual')
  targetimage = 'L2_Pup_sci.spw0_1_2_3.mfs.IQUV.manual'  

  os.system('rm -rf '+targetimage+'.*')
  tclean(vis=targetmsname,
        field=target,
        imagename=targetimage,
        cell=cell,
        imsize=imsize,
        specmode='mfs',
        deconvolver='clarkstokes',
        gridder='standard',
        outframe=outframe,
        stokes='IQUV',
        interactive=False,
        mask=mask,
        threshold='0.0005Jy',
        niter=100000)

#>>>  for stokes in 'I','Q','U','V':
#>>>      print au.genImageName(vis=targetmsname, spw=[0,1,2,3], field=target, imtype='mfs', targettype='sci', stokes=stokes, mous='', modtext='manual')
  targetimageI = 'L2_Pup_sci.spw0_1_2_3.mfs.I.manual'
  targetimageQ = 'L2_Pup_sci.spw0_1_2_3.mfs.Q.manual'
  targetimageU = 'L2_Pup_sci.spw0_1_2_3.mfs.U.manual'
  targetimageV = 'L2_Pup_sci.spw0_1_2_3.mfs.V.manual'

  os.system('rm -rf '+targetimageI+'.image')
  immath(imagename=targetimage+'.image',outfile=targetimageI+'.image',expr='IM0',stokes='I')
  os.system('rm -rf '+targetimageQ+'.image')
  immath(imagename=targetimage+'.image',outfile=targetimageQ+'.image',expr='IM0',stokes='Q')
  os.system('rm -rf '+targetimageU+'.image')
  immath(imagename=targetimage+'.image',outfile=targetimageU+'.image',expr='IM0',stokes='U')
  os.system('rm -rf '+targetimageV+'.image')
  immath(imagename=targetimage+'.image',outfile=targetimageV+'.image',expr='IM0',stokes='V')

  os.system('rm -rf '+targetimageI+'.pb')
  immath(imagename=targetimage+'.pb',outfile=targetimageI+'.pb',expr='IM0',stokes='I')
  os.system('rm -rf '+targetimageQ+'.pb')
  immath(imagename=targetimage+'.pb',outfile=targetimageQ+'.pb',expr='IM0',stokes='Q')
  os.system('rm -rf '+targetimageU+'.pb')
  immath(imagename=targetimage+'.pb',outfile=targetimageU+'.pb',expr='IM0',stokes='U')
  os.system('rm -rf '+targetimageV+'.pb')
  immath(imagename=targetimage+'.pb',outfile=targetimageV+'.pb',expr='IM0',stokes='V')

### Final image production
mystep = 4
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  ##############################################
  # Apply a primary beam correction

  import glob

  myimages = glob.glob("*.image")

  rmtables('*.pbcor')
  for image in myimages:
    pbimage = image.rsplit('.',1)[0]+'.pb'
    outfile = image.rsplit('.',1)[0]+'.pbcor'
    impbcor(imagename=image, pbimage=pbimage, outfile = outfile)

  ##############################################



  # Polarization images production

#>>> au.genImageName(vis=targetmsname, spw=[0,1,2,3], field=target, imtype='mfs', targettype='sci', stokes='P', mous='', modtext='pbcor')
  polIname = 'L2_Pup_sci.spw0_1_2_3.mfs.P.pbcor'

  os.system('rm -rf '+polIname)
  immath(outfile=polIname,
         mode='poli',
         imagename=[targetimageQ+'.pbcor',targetimageU+'.pbcor'],
         sigma='0.0Jy/beam')

#>>> au.genImageName(vis=targetmsname, spw=[0,1,2,3], field=target, imtype='mfs', targettype='sci', stokes='A', mous='', modtext='pbcor')
  polAname = 'L2_Pup_sci.spw0_1_2_3.mfs.A.pbcor'

  os.system('rm -rf '+polAname)
  immath(outfile=polAname,
         mode='pola',
         imagename=[targetimageQ+'.pbcor',targetimageU+'.pbcor'],
         polithresh='0.01mJy/beam')

  # produce FITS files

  import glob

  myimages = glob.glob("*.pbcor")
  for image in myimages:
      exportfits(imagename=image, fitsimage=image+'.fits',overwrite=True)

  myimages = glob.glob("*.pb")
  for image in myimages:
      exportfits(imagename=image, fitsimage=image+'.fits',overwrite=True)