casa  $Rev:20696$
 All Classes Namespaces Files Functions Variables
task_uvcontsub2.py
Go to the documentation of this file.
00001 import os
00002 import re
00003 import shutil
00004 import tempfile
00005 
00006 # shutil.copytree is useless with directories created by tempfile
00007 # (or any directories that already exist).
00008 from distutils.dir_util import copy_tree
00009 
00010 from taskinit import *
00011 from update_spw import *
00012 
00013 mycb, myms, mytb = gentools(['cb', 'ms', 'tb'])
00014 
00015 def uvcontsub2(vis, field, fitspw, combine, solint, fitorder, spw, want_cont):
00016     try:
00017         casalog.origin('uvcontsub2')
00018         casalog.post('uvcontsub is now the same as uvcontsub2 and should meet your needs.', 'INFO')
00019         casalog.post('In the future, the uvcontsub2 task may be retired or used for experimentation.', 'INFO')
00020 
00021         # Get these checks done and out of the way.
00022         # This one is redundant - it is already checked at the XML level.
00023         if not os.path.isdir(vis):
00024             raise Exception, 'Visibility data set not found - please verify the name'
00025 
00026         
00027         mytb.open(vis + '/SPECTRAL_WINDOW')
00028         allspw = '0~' + str(mytb.nrows() - 1)
00029         mytb.close()
00030         if 'spw' not in combine:
00031             spwmfitspw = subtract_spws(spw, fitspw)
00032             if spwmfitspw == 'UNKNOWN':
00033                 spwmfitspw = subtract_spws(allspw, fitspw)
00034             if spwmfitspw:
00035                 raise Exception, "combine must include 'spw' when the fit is being applied to spws outside fitspw."
00036         
00037         # cb will put the continuum-subtracted data in CORRECTED_DATA, so
00038         # protect vis by splitting out its relevant parts to a working copy.
00039         csvis = vis.rstrip('/') + '.contsub'
00040         
00041         # The working copy will need all of the channels in fitspw + spw, so we
00042         # may or may not need to filter it down to spw at the end.
00043         myfitspw = fitspw
00044         myspw = spw
00045 
00046         # We only need the spws in the union of fitspw and spw, but keep all
00047         # the channels or the weights (as used by calibrater) will be messed up.
00048         tempspw = re.sub(r':[^,]+', '', join_spws(fitspw, spw))
00049         if tempspw == allspw:
00050             tempspw = ''
00051         if tempspw:
00052             # The split will reindex the spws.  Update spw and fitspw.
00053             # Do fitspw first because the spws in spw are supposed to be
00054             # a subset of the ones in fitspw.
00055             casalog.post('split is being run internally, and the selected spws')
00056             casalog.post('will be renumbered to start from 0 in the output!')
00057 
00058             # Now get myfitspw.
00059             myfitspw = update_spwchan(vis, tempspw, fitspw)
00060             myspw = update_spwchan(vis, tempspw, spw)
00061 
00062         final_csvis = csvis
00063         workingdir = os.path.abspath(os.path.dirname(vis.rstrip('/')))
00064         csvis = tempfile.mkdtemp(prefix=csvis.split('/')[-1], dir=workingdir)
00065 
00066         # ms does not have a colnames method, so open vis with tb even though
00067         # it is already open with myms.  Note that both use nomodify=True,
00068         # however, and no problem was revealed in testing.
00069         mytb.open(vis, nomodify=True)
00070         if 'CORRECTED_DATA' in mytb.colnames():
00071             whichcol = 'CORRECTED_DATA'
00072         else:
00073             # DON'T remind the user that split before uvcontsub2 wastes time -
00074             # scratch columns will eventually go away.
00075             whichcol = 'DATA'
00076         mytb.close()
00077 
00078         casalog.post('Preparing to add scratch columns.')
00079         myfield = field
00080         if field == '':
00081             myfield = '*'
00082         if myfield != '*' and set(ms.msseltoindex(vis,
00083                                field=myfield)['field']) == set(ms.msseltoindex(vis,
00084                                                                  field='*')['field']):
00085             myfield = '*'
00086         if whichcol != 'DATA' or tempspw != '' or myfield != '*':
00087             casalog.post('splitting to ' + csvis + ' with spw="'
00088                          + tempspw + '"')
00089             myms.open(vis, nomodify=True)
00090             myms.split(csvis, spw=tempspw, field=myfield, whichcol=whichcol)
00091             myms.close()
00092         else:
00093             # This takes almost 30s/GB.  (lustre, 8/2011)
00094             casalog.post('Copying ' + vis + ' to ' + csvis + ' with cp.')
00095             copy_tree(vis, csvis)
00096 
00097         # It is less confusing if we write the history now that the "root" MS
00098         # is made, but before cb adds its messages.
00099         #
00100         # Not a dict, because we want to maintain the order.
00101         param_names = uvcontsub2.func_code.co_varnames[:uvcontsub2.func_code.co_argcount]
00102         param_vals = [eval(p) for p in param_names]
00103             
00104         write_history(myms, csvis, 'uvcontsub2', param_names, param_vals,
00105                       casalog)
00106 
00107         if (type(csvis) == str) and os.path.isdir(csvis):
00108             mycb.open(csvis)
00109         else:
00110             raise Exception, 'Visibility data set not found - please verify the name'
00111 
00112         ## for now, forbid requests for fitorder>0 
00113         #if (fitorder>0):
00114         #    raise Exception, "Sorry, uvcontsub2 currently only supports fitorder=0."
00115         
00116         # select the data for continuum subtraction
00117         mycb.reset()
00118         mycb.selectvis(spw=myfitspw)
00119 
00120         # Set up the solve
00121         amuellertab = tempfile.mkdtemp(prefix='Temp_contsub.tab',
00122                                        dir=workingdir)
00123 
00124         mycb.setsolve(type='A', t=solint, table=amuellertab, combine=combine,
00125                       fitorder=fitorder)
00126 
00127         # solve for the continuum
00128         mycb.solve()
00129 
00130         # subtract the continuum
00131         mycb.selectvis(spw=myspw)
00132         aspwmap=-1
00133         # if we combined on spw in solve, fanout the result with spwmap=-999;
00134         if 'spw' in combine:
00135             aspwmap = -999
00136         mycb.setapply(table=amuellertab, spwmap=aspwmap)
00137 
00138         # Generate CORRECTED_DATA without continuum
00139         mycb.correct()
00140 
00141         if want_cont:
00142             # Generate MODEL_DATA with only the continuum model
00143             mycb.corrupt()
00144 
00145         mycb.close()
00146 
00147         # Delete the temporary caltable
00148         shutil.rmtree(amuellertab)
00149 
00150         # Move the line data from CORRECTED_DATA to DATA, and do any
00151         # final filtering by spw.
00152         myms.open(csvis)
00153         # Using ^ in spw is untested here!
00154         myms.split(final_csvis, spw=myspw, whichcol='corrected')
00155         myms.close()
00156 
00157         #casalog.post("\"want_cont\" = \"%s\"" % want_cont, 'WARN')
00158         #casalog.post("\"csvis\" = \"%s\"" % csvis, 'WARN')
00159         if want_cont:
00160             myms.open(csvis)
00161             # No channel averaging (== skipping for fitorder == 0) is done
00162             # here, but it would be a reasonable option.  The user can always
00163             # do it by running split again.
00164             myms.split(final_csvis[:-3],             # .contsub -> .cont
00165                        whichcol='MODEL_DATA',
00166                        spw=myspw)
00167             myms.close()
00168 
00169         #casalog.post("rming \"%s\"" % csvis, 'WARN')
00170         shutil.rmtree(csvis)
00171 
00172     except Exception, instance:
00173         casalog.post('Error in uvcontsub2: ' + str(instance), 'SEVERE')
00174         mycb.close()                        # Harmless if cb is closed.
00175         raise Exception