Coverage for /home/casatest/venv/lib/python3.12/site-packages/casatasks/private/task_uvcontsub_old.py: 7%
241 statements
« prev ^ index » next coverage.py v7.10.4, created at 2025-08-21 07:43 +0000
1import os
2import re
3import shutil
4import time
5import tempfile
6import numpy as np
8from casatools import calibrater, ms, table
9from casatasks import casalog, virtualconcat
11from .mstools import write_history
12from .update_spw import *
13from .parallel.parallel_data_helper import ParallelDataHelper
14from .parallel.parallel_task_helper import ParallelTaskHelper
16mytb = table( )
17mycb = calibrater( )
18myms = ms()
19_ms = ms() # task also referenced the old global ms object
21def uvcontsub_old(vis, field, fitspw, excludechans, combine, solint, fitorder, spw, want_cont):
23 if ParallelDataHelper.isMMSAndNotServer(vis):
24 helper = ParallelTaskHelper('uvcontsub_old', locals())
25 helper._consolidateOutput = False
26 retVar = helper.go()
28 # Gather the list of continuum subtraction-SubMSs
29 cont_subMS_list = []
30 contsub_subMS_list = []
31 for subMS in retVar:
32 if retVar[subMS]:
33 cont_subMS_list.append(subMS + ".cont")
34 contsub_subMS_list.append(subMS + ".contsub")
36 if len(cont_subMS_list) <= 0:
37 raise ValueError("No continuum-subtracted sub-MSs for concatenation","SEVERE")
39 # We have to sort the list because otherwise it
40 # depends on the time the engines dispatches their sub-MSs
41 cont_subMS_list.sort()
42 contsub_subMS_list.sort()
44 # deal with the pointing table
45 auxfile = "uvcontsub_aux2_"+str(time.time())
46 pnrows = 0
47 try:
48 mytb.open(vis+'/POINTING')
49 pnrows = mytb.nrows()
50 mytb.close()
51 if(pnrows>0):
52 shutil.copytree(os.path.realpath(vis+'/POINTING'), auxfile)
53 except Exception as instance:
54 raise RuntimeError("Error handling POINTING table %s: %s" %
55 (vis+'/POINTING',str(instance)))
57 if want_cont:
58 try:
59 virtualconcat(concatvis=helper._arg['vis'] + ".cont",vis=cont_subMS_list,
60 copypointing=False)
61 except Exception as instance:
62 raise RuntimeError("Error concatenating continuum sub-MSs %s: %s" %
63 (str(cont_subMS_list),str(instance)))
64 try:
65 virtualconcat(concatvis=helper._arg['vis'] + ".contsub",vis=contsub_subMS_list,
66 copypointing=False)
67 except Exception as instance:
68 raise RuntimeError("Error concatenating continuum-subtracted sub-MSs %s: %s" %
69 (str(contsub_subMS_list),str(instance)))
71 # Remove continuum subtraction-SubMSs
72 if want_cont:
73 for subMS in cont_subMS_list:
74 if (os.system("rm -rf " + subMS ) != 0):
75 casalog.post("Problem removing continuum sub-MS %s into working directory" % (subMS),'SEVERE')
76 for subMS in contsub_subMS_list:
77 if (os.system("rm -rf " + subMS ) != 0):
78 casalog.post("Problem removing continuum-subtracted sub-MS %s into working directory" % (subMS),'SEVERE')
80 if(pnrows>0): # put back the POINTING table
81 casalog.post("Restoring the POINTING table ...",'NORMAL')
82 try:
83 if want_cont:
84 theptab = os.path.realpath(helper._arg['vis'] + ".cont/POINTING")
85 os.system('rm -rf '+theptab)
86 shutil.copytree(auxfile, theptab)
87 theptab = os.path.realpath(helper._arg['vis'] + ".contsub/POINTING")
88 os.system('rm -rf '+theptab)
89 os.system('mv '+auxfile+' '+theptab)
90 except Exception as instance:
91 raise RuntimeError("Error restoring pointing table from %s: %s" %
92 (auxfile,str(instance)))
94 # Restore origin (otherwise gcwrap shows virtualconcat)
95 casalog.origin('uvcontsub_old')
97 return
99 # Run normal code
100 try:
101 casalog.origin('uvcontsub_old')
103 casalog.post('This task is deprecated and will be removed in an upcoming release.',
104 'WARN')
106 # Get these checks done and out of the way.
107 # This one is redundant - it is already checked at the XML level.
108 if not os.path.isdir(vis):
109 raise ValueError('Visibility data set not found - please verify the name')
111 #
112 if excludechans:
113 locfitspw=_quantityRangesToChannels(vis,field,fitspw,excludechans)
114 casalog.post("Exclude channels in fitspw: spws:channels will be used in the fit are:%s" % locfitspw)
115 elif fitspw.count('Hz'):
116 locfitspw=_quantityRangesToChannels(vis,field,fitspw,excludechans)
117 else:
118 locfitspw=fitspw
119 # casalog.post("locfitspw=",locfitspw)
120 mytb.open(vis + '/SPECTRAL_WINDOW')
121 allspw = '0~' + str(mytb.nrows() - 1)
122 mytb.close()
124 # Verify fitspw and spw as plausible
125 if len(locfitspw)>0:
126 fitspwerr=subtract_spws(locfitspw,allspw)
127 if fitspwerr:
128 badfitspw=''
129 for ispw in range(len(fitspwerr)):
130 badfitspw+=str(fitspwerr.pop())
131 raise ValueError("fitspw contains non-existent spw(s): "+badfitspw)
132 if len(spw)>0:
133 spwerr=subtract_spws(spw,allspw)
134 if spwerr:
135 badspw=''
136 for ispw in range(len(spwerr)):
137 badspw+=str(spwerr.pop())
138 raise ValueError("spw contains non-existent spw(s): "+badspw)
140 if 'spw' not in combine:
141 #spwmfitspw = subtract_spws(spw, fitspw)
142 spwmfitspw = subtract_spws(spw, locfitspw)
143 if spwmfitspw == 'UNKNOWN':
144 #spwmfitspw = subtract_spws(allspw, fitspw)
145 spwmfitspw = subtract_spws(allspw, locfitspw)
146 if spwmfitspw:
147 raise ValueError("combine must include 'spw' when the fit is being applied to spws outside fitspw.")
149 # cb will put the continuum-subtracted data in CORRECTED_DATA, so
150 # protect vis by splitting out its relevant parts to a working copy.
151 csvis = vis.rstrip('/') + '.contsub'
153 # The working copy will need all of the channels in fitspw + spw, so we
154 # may or may not need to filter it down to spw at the end.
155 #myfitspw = fitspw
156 myfitspw = locfitspw
157 myspw = spw
159 # We only need the spws in the union of fitspw and spw, but keep all
160 # the channels or the weights (as used by calibrater) will be messed up.
161 #tempspw = re.sub(r':[^,]+', '', join_spws(fitspw, spw))
162 tempspw = re.sub(r':[^,]+', '', join_spws(locfitspw, spw))
163 if tempspw == allspw:
164 tempspw = ''
165 if tempspw:
166 # The split will reindex the spws. Update spw and fitspw.
167 # Do fitspw first because the spws in spw are supposed to be
168 # a subset of the ones in fitspw.
169 casalog.post('split is being run internally, and the selected spws')
170 casalog.post('will be renumbered to start from 0 in the output!')
172 # Now get myfitspw.
173 #myfitspw = update_spwchan(vis, tempspw, fitspw)
174 myfitspw = update_spwchan(vis, tempspw, locfitspw)
175 myspw = update_spwchan(vis, tempspw, spw)
177 final_csvis = csvis
178 workingdir = os.path.abspath(os.path.dirname(vis.rstrip('/')))
179 csvis = tempfile.mkdtemp(prefix=csvis.split('/')[-1], dir=workingdir)
181 # ms does not have a colnames method, so open vis with tb even though
182 # it is already open with myms. Note that both use nomodify=True,
183 # however, and no problem was revealed in testing.
184 mytb.open(vis, nomodify=True)
185 if 'CORRECTED_DATA' in mytb.colnames():
186 whichcol = 'CORRECTED_DATA'
187 else:
188 # DON'T remind the user that split before uvcontsub wastes time -
189 # scratch columns will eventually go away.
190 whichcol = 'DATA'
191 mytb.close()
193 casalog.post('Preparing to add scratch columns.')
194 myfield = field
195 if field == '':
196 myfield = '*'
197 if myfield != '*' and set(_ms.msseltoindex(vis,
198 field=myfield)['field']) == set(_ms.msseltoindex(vis,
199 field='*')['field']):
200 myfield = '*'
201 if whichcol != 'DATA' or tempspw != '' or myfield != '*':
202 casalog.post('splitting to ' + csvis + ' with spw="'
203 + tempspw + '"')
204 myms.open(vis, nomodify=True)
205 split_result = myms.split(csvis, spw=tempspw, field=myfield, whichcol=whichcol)
206 myms.close()
207 # jagonzal (CAS-4113): Take care of the trivial parallelization
208 if not split_result:
209 msg = "MSSelectionNullSelection: %s does not have data for spw=%s, field=%s, bailing out!" % (vis,tempspw,field)
210 shutil.rmtree(csvis)
211 raise RuntimeError(msg)
212 else:
213 # This takes almost 30s/GB. (lustre, 8/2011)
214 casalog.post('Copying ' + vis + ' to ' + csvis + ' with cp.')
215 shutil.copytree(vis, csvis, symlinks=True, dirs_exist_ok=True)
217 # It is less confusing if we write the history now that the "root" MS
218 # is made, but before cb adds its messages.
219 #
220 param_names = uvcontsub_old.__code__.co_varnames[:uvcontsub_old.__code__.co_argcount]
221 local_vars = locals( )
222 param_vals = [local_vars[p] for p in param_names]
224 write_history(myms, csvis, 'uvcontsub_old', param_names, param_vals,
225 casalog)
227 if (type(csvis) == str) and os.path.isdir(csvis):
228 mycb.setvi(old=True,quiet=False); # old VI, for now
229 mycb.open(csvis)
230 else:
231 raise RuntimeError('Visibility data set not found - please verify the name')
233 # select the data for continuum subtraction
234 mycb.reset()
235 mycb.selectvis(spw=myfitspw)
237 # Set up the solve
238 amuellertab = tempfile.mkdtemp(prefix='Temp_contsub.tab',
239 dir=workingdir)
241 mycb.setsolve(type='A', t=solint, table=amuellertab, combine=combine,
242 fitorder=fitorder)
244 # solve for the continuum
245 mycb.solve()
247 # subtract the continuum
248 mycb.selectvis(spw=myspw)
249 aspwmap=-1
250 # if we combined on spw in solve, fanout the result with spwmap=-999;
251 if 'spw' in combine:
252 aspwmap = -999
253 mycb.setapply(table=amuellertab, spwmap=aspwmap)
255 # Generate CORRECTED_DATA without continuum
256 mycb.correct()
258 if want_cont:
259 # Generate MODEL_DATA with only the continuum model
260 mycb.corrupt()
262 mycb.close()
264 # Delete the temporary caltable
265 shutil.rmtree(amuellertab)
267 # Move the line data from CORRECTED_DATA to DATA, and do any
268 # final filtering by spw.
269 myms.open(csvis)
270 # Using ^ in spw is untested here!
271 myms.split(final_csvis, spw=myspw, whichcol='corrected')
272 myms.close()
274 #casalog.post("\"want_cont\" = \"%s\"" % want_cont, 'WARN')
275 #casalog.post("\"csvis\" = \"%s\"" % csvis, 'WARN')
276 if want_cont:
277 myms.open(csvis)
278 # No channel averaging (== skipping for fitorder == 0) is done
279 # here, but it would be a reasonable option. The user can always
280 # do it by running split again.
281 myms.split(final_csvis[:-3], # .contsub -> .cont
282 whichcol='MODEL_DATA',
283 spw=myspw)
284 myms.close()
286 #casalog.post("rming \"%s\"" % csvis, 'WARN')
287 shutil.rmtree(csvis)
289 finally:
290 mycb.close() # Harmless if cb is closed.
293def _quantityRangesToChannels(vis,field,infitspw,excludechans):
294 """
295 private function to convert frequnecy (in the future, velocity
296 as well) ranges in fitspw to channel indexes.
297 For excludeechans=True, it will select channels outside given by infitspw
298 returns: list containing new channel ranges
299 """
300 mytb.open(vis+'/SPECTRAL_WINDOW')
301 nspw=mytb.nrows()
302 mytb.close()
304 fullspwids=str(list(range(nspw))).strip('[,]')
305 tql={'field':field,'spw':fullspwids}
306 myms.open(vis)
307 myms.msselect(tql,True)
308 allsels=myms.msselectedindices()
309 myms.reset()
310 # input fitspw selection
311 tql['spw']=infitspw
312 myms.msselect(tql,True)
313 usersels=myms.msselectedindices()['channel']
314 myms.close()
315 # sort the arrays so that chan ranges are
316 # in order
317 usersels=usersels[np.lexsort((usersels[:,1],usersels[:,0]))]
318 spwsels=''
319 spwid=-1
320 prevspwid=None
321 newchanlist=[]
322 nsels=len(usersels)
323 # casalog.post("Usersels=",usersels)
324 if excludechans:
325 for isel in range(nsels):
326 prevspwid = spwid
327 spwid=usersels[isel][0]
328 lochan=usersels[isel][1]
329 hichan=usersels[isel][2]
330 stp=usersels[isel][3]
331 maxchanid=allsels['channel'][spwid][2]
332 # find left and right side ranges of the selected range
333 if spwid != prevspwid:
334 # first line in the selected spw
335 if lochan > 0:
336 outloL=0
337 outhiL=lochan-1
338 outloR= (0 if hichan+1>=maxchanid else hichan+1)
339 if outloR:
340 if isel<nsels-1 and usersels[isel+1][0]==spwid:
341 outhiR=usersels[isel+1][1]-1
342 else:
343 outhiR=maxchanid
344 else:
345 outhiR=0 # higher end of the user selected range reaches maxchanid
346 # so no right hand side range
347 # casalog.post("outloL,outhiL,outloR,outhiR==", outloL,outhiL,outloR,outhiR)
348 else:
349 # no left hand side range
350 outloL=0
351 outhiL=0
352 outloR=hichan+1
353 if isel<nsels-1 and usersels[isel+1][0]==spwid:
354 outhiR=usersels[isel+1][1]-1
355 else:
356 outhiR=maxchanid
357 else:
358 #expect the left side range is already taken care of
359 outloL=0
360 outhiL=0
361 outloR=hichan+1
362 if outloR>=maxchanid:
363 #No more boundaries to consider
364 outloR=0
365 outhiR=0
366 else:
367 if isel<nsels-1 and usersels[isel+1][0]==spwid:
368 outhiR=min(usersels[isel+1][1]-1,maxchanid)
369 else:
370 outhiR=maxchanid
371 if outloR > outhiR:
372 outloR = 0
373 outhiR = 0
375 #
376 if (not(outloL == 0 and outhiL == 0)) and outloL <= outhiL:
377 newchanlist.append([spwid,outloL,outhiL,stp])
378 if (not(outloR == 0 and outhiR == 0)) and outloR <= outhiR:
379 newchanlist.append([spwid,outloR,outhiR,stp])
380 # casalog.post("newchanlist=",newchanlist)
381 else:
382 # excludechans=False
383 newchanlist=usersels
385 #return newchanlist
386 # create spw selection string from newchanlist
387 spwstr=''
388 for irange in range(len(newchanlist)):
389 chanrange=newchanlist[irange]
390 spwstr+=str(chanrange[0])+':'+str(chanrange[1])+'~'+str(chanrange[2])
391 if irange != len(newchanlist)-1:
392 spwstr+=','
393 return spwstr