# casa 4.0.1 - 
# assumes that model for solar system objects is set per-channel and includes main lines.
# i.e. no need to transfer gains between spw.
# if all spw cannot be used find and set refspwmap

# This starts after deriving a bandpass table and a phase correction
# table on a selection of baselines for the calibration sources

# It derives the flux density for the brightest point-like calibration
# source from a Solar system object on a restricted set of baselines,
# and then sets that in the ms and uses all baselines to derive the
# flux density of the phase-reference source.

# If you have already set models, clear these with delmod or clearcal
# but make sure you retain or reset the model for your primary e.g.
# solar system flux calibrator.

# Raw and fitted flux densities are plotted to the screen at the
# appropriate steps and printed to the *.fluxes.table.
##############################################################

# Set for your data:

thisvis = 'band9_calibration_hands-on.ms'
spw_without_atm='0:0~939;1230~3339,1:0~1339;2200~3639,2,3'
uvrange1='0~115'    
selectedants = ''
allcals = '0,1,3'    # Source numbers
fluxstandard = '1'
bp = '0'
ph = '3'
pointcals='0,3'
bandpasstable ='band9_calibration_hands-on.ms.bandpass_smooth20ch'
arefant='DV04'

thesteps = []
step_title = {0: 'Derive gain tables for all cal sources in restricted uv range',
              1: 'Find flux density of bp cal source wrt solar system object',
              2: 'Refine bp cal flux density using error-weighted fit',
              3: 'Set the flux density for bp cal source',
              4: 'Derive gain tables for both cal sources, all good baselines',
              5: 'Find flux density of phase ref source',
              6: 'Refine phase ref flux density using error-weighted fit',
              7: 'Set the flux density for phase ref source',
              8: 'Print values to file'}

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


#** WARNING - any old gain tables with the suffixes used are deleted
#** before creation here.

import numpy as np

mystep = 0

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
# 0) Derive gain tables for all cal sources in restricted uv range

  os.system('rm -rf *.ampli_short_30s_uvrange1')
  gaincal(vis = 'band9_calibration_hands-on.ms',
    caltable = thisvis+'.phase_short_30s_uvrange1',
    field = allcals,
    selectdata = T,
    antenna = selectedants,  
    uvrange = uvrange1,                   
    solint = '30s',   
    minsnr = 2.0,     
    minblperant = 2,  
    spw= spw_without_atm,
    refant = arefant,
    gaintype = 'G',  
    calmode = 'p',
    gaintable = bandpasstable)
  
  os.system('rm -rf *.ampli_short_inf')
  gaincal(vis = thisvis,
    caltable = thisvis+'.ampli_short_inf',
    field = allcals,
    selectdata = T,
    antenna = selectedants,  
    uvrange =  uvrange1,                 
    solint = 'inf',
    spw=spw_without_atm, 
    minblperant = 2, # 
    minsnr = 2.0,    # 
    refant = arefant,
    gaintype = 'T',
    calmode = 'ap',
    gaintable = [bandpasstable, thisvis+'.phase_short_30s_uvrange1'])

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

mystep = 1

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
# 1) Find flux density of brightest point-like source, usually BP cal,
#    wrt solar system object on a restricted uv range

  os.system('rm -rf *.ms.flux_short_inf')
  fluxfitbp = fluxscale(vis = thisvis,
      caltable = thisvis+'.ampli_short_inf',
      fluxtable=thisvis+'.flux_short_inf',
  #    refspwmap=fscale_spwmap,
      listfile='bp.fluxscale.list',
      reference = fluxstandard) 

# Print the fitted values to screen
  fluxfitbp

# The values for 3c279 aren't too bad but for J1625-254 there is a big
# jump between sidebands which is unlikely to be astrophysical.
# Extract the flux densities and uncertainties for 3c279.

  freqs=np.array(fluxfitbp['freq']) 
  bp_fluxes=np.array(fluxfitbp[bp]['fluxd'])
  bp_flerrs=np.array(fluxfitbp[bp]['fluxdErr'])

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

mystep = 2

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
# 2) We assume that it has a 1st order spectral index and make an
#    error- weighted fit to improve the accuracy

  logflerrs=bp_flerrs/bp_fluxes
  logfluxes=np.log(bp_fluxes)
  logfreqs=np.log(freqs)

  sumey2  =  sum(1./logflerrs**2)
  sumx2ey2=  sum((logfreqs**2)/(logflerrs**2))
  sumxey2 =  sum(logfreqs/(logflerrs**2))
  sumyey2 =  sum(logfluxes/(logflerrs**2))
  sumxyey2=  sum(logfluxes*logfreqs/(logflerrs**2))

  delta= sumey2*sumx2ey2 - sumxey2**2
  intercept= (1./delta)*(sumx2ey2*sumyey2 - sumxey2*sumxyey2)
  slope=(1./delta)*(sumey2*sumxyey2-sumxey2*sumyey2)
  bpslope=slope

  logfluxes_bpcorr=intercept+slope*logfreqs

  fluxes_bpcorr=e**logfluxes_bpcorr

# Print the fluxes to screen
  fluxes_bpcorr

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

mystep = 3

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
# 3) Set the flux densities for the bp cal

  for i in range(0,4):
    setjy(vis = thisvis,
      field = bp,
      spw = str(i),
      fluxdensity = [fluxes_bpcorr[i],0,0,0])

##################################################################
  
mystep = 4

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

# 4) Use all baselines to derive a gain table for the phase-ref and BP cal
#    and derive the phase-ref flux density from the BP cal 

  os.system('rm -rf *.ampli_all_30s_uvrange1')
  gaincal(vis = 'band9_calibration_hands-on.ms',
    caltable = thisvis+'.phase_all_30s_uvrange1',
    field = pointcals,
    selectdata = T,
    solint = '30s',   
    minsnr = 2.0,     
    minblperant = 2,  
    spw= spw_without_atm,
    refant = arefant,
    gaintype = 'G',  
    calmode = 'p',
    gaintable = bandpasstable)
  
  os.system('rm -rf *.ampli_all_inf')
  
  gaincal(vis = thisvis,
    caltable = thisvis+'.ampli_all_inf',
    field = pointcals,
    selectdata = T,
    solint = 'inf',
    spw=spw_without_atm, 
    minblperant = 2, 
    minsnr = 2.0,     
    refant = arefant,
    gaintype = 'T',
    calmode = 'ap',
    gaintable = [bandpasstable, thisvis+'.phase_all_30s_uvrange1'])
 
##################################################################

mystep = 5

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

#5  Find flux density of phase-ref with respect to BP cal,

  os.system('rm -rf *.ms.flux_all_inf')

  fluxfitph = fluxscale(vis = thisvis,
      caltable = thisvis+'.ampli_all_inf',
      fluxtable=thisvis+'.flux_all_inf',
  #    refspwmap=fscale_spwmap,
      listfile='ph.fluxscale.list',
      reference = bp) 

# Print the fitted values to screen
  fluxfitph

# Extract the flux densities and uncertainties for ph.

  freqs=np.array(fluxfitph['freq']) 
  ph_fluxes=np.array(fluxfitph[ph]['fluxd'])
  ph_flerrs=np.array(fluxfitph[ph]['fluxdErr'])

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

mystep = 6

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

# 6) We assume that it has a 1st order spectral index and make an
#    error- weighted fit to improve the accuracy

  logflerrs=ph_flerrs/ph_fluxes
  logfluxes=np.log(ph_fluxes)
  logfreqs=np.log(freqs)

  sumey2  =  sum(1./logflerrs**2)
  sumx2ey2=  sum((logfreqs**2)/(logflerrs**2))
  sumxey2 =  sum(logfreqs/(logflerrs**2))
  sumyey2 =  sum(logfluxes/(logflerrs**2))
  sumxyey2=  sum(logfluxes*logfreqs/(logflerrs**2))

  delta= sumey2*sumx2ey2 - sumxey2**2
  intercept= (1./delta)*(sumx2ey2*sumyey2 - sumxey2*sumxyey2)
  slope=(1./delta)*(sumey2*sumxyey2-sumxey2*sumyey2)
  phslope=slope

  logfluxes_phcorr=intercept+slope*logfreqs

  fluxes_phcorr=e**logfluxes_phcorr

# Print the fluxes to screen
  fluxes_phcorr

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

mystep = 7

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
# 7) Set the flux densities for the ph cal

  for i in range(0,4):
    setjy(vis = thisvis,
      field = ph,
      spw = str(i),
      fluxdensity = [fluxes_phcorr[i],0,0,0])

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

mystep = 8

if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
## 8) Print values to file

  os.system('rm -rf *fluxes.table')
  ftab = open(thisvis+'.fluxes.table', 'a')
  ghz=freqs/1e9
  print >> ftab, 'BP spectral index %4.1f' % (bpslope)
  print >> ftab, 'PH spectral index %4.1f' % (phslope)
  print >> ftab, '%7.3f %7.3f %7.3f %7.3f GHz' % (ghz[0], ghz[1], ghz[2], ghz[3])
  print >> ftab, 'BP orig %7.3f %7.3f %7.3f %7.3f Jy' % (bp_fluxes[0], bp_fluxes[1], bp_fluxes[2], bp_fluxes[3])
  print >> ftab, 'BP errs %7.3f %7.3f %7.3f %7.3f Jy' % (bp_flerrs[0], bp_flerrs[1], bp_flerrs[2], bp_flerrs[3])
  print >> ftab, 'BP fit  %7.3f %7.3f %7.3f %7.3f Jy' % (fluxes_bpcorr[0],fluxes_bpcorr [1], fluxes_bpcorr[2], fluxes_bpcorr[3])
  print >> ftab, 'PH orig %7.3f %7.3f %7.3f %7.3f Jy' % (ph_fluxes[0], ph_fluxes[1], ph_fluxes[2], ph_fluxes[3])
  print >> ftab, 'PH errs %7.3f %7.3f %7.3f %7.3f Jy' % (ph_flerrs[0], ph_flerrs[1], ph_flerrs[2], ph_flerrs[3])
  print >> ftab, 'PH fit  %7.3f %7.3f %7.3f %7.3f Jy' % (fluxes_phcorr[0],fluxes_phcorr [1], fluxes_phcorr[2], fluxes_phcorr[3])
  ftab.close()
  
#################################################################
# end.  amsr@jb.man.ac.uk Feb 2013