Coverage for /wheeldirectory/casa-6.7.0-12-py3.10.el8/lib/py/lib/python3.10/site-packages/casatasks/private/update_spw.py: 67%
319 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-11-01 07:19 +0000
1#! /usr/bin/env python
2# The above is for running the doctest, not normal use.
4"""
5A set of functions for manipulating spw:chan selection strings.
7If this is run from a shell (i.e. not in casapy), doctest will be used to run
8several unit tests from the doc strings, including the one below:
10Example:
11>>> from update_spw import update_spw
12>>> update_spw('0~2,5', None)[0]
13'0~2,3'
14>>> update_spw('0~2,5', None)[1]['5'] # doctest warning! dicts don't always print out in the same order!
15'3'
16"""
18import copy
19import os
21from casatools import ms
22from casatools import table as tbtool
23_ms = ms( )
25def update_spw(spw, spwmap=None):
26 """
27 Given an spw:chan selection string, return what it should be after the spws
28 have been remapped (i.e. by split), and a map from input to output spws
29 (spwmap). It does not change spw OR the *channels* part of the output spw
30 string! (See update_spwchan)
32 If given, spwmap will be used as a dictionary from (string) input spw to
33 (string) output spws. Otherwise it will be freshly calculated. Supplying
34 spwmap doesn't just save work: it is also necessary for chaining
35 update_spw() calls when the first selection includes more spws than the
36 subsequent one(s). HOWEVER, if given, spwmap must have slots for all the
37 spws that will appear in the output MS, i.e. it can't be grown once made.
39 Examples:
40 >>> from update_spw import update_spw
41 >>> myfitspw, spws = update_spw('0~3,5;6:1~7;11~13', None)
42 >>> myfitspw
43 '0~3,4;5:1~7;11~13'
44 >>> myspw = update_spw('1,5;6:8~10', spws)[0]
45 >>> myspw # not '0,1,2:8~10'
46 '1,4;5:8~10'
47 >>> update_spw('0~3,5;6:1~7;11~13,7~9:0~3,11,7~8:6~8', None)[0]
48 '0~3,4;5:1~7;11~13,6~8:0~3,9,6~7:6~8'
50 # Let's say we want updates of both fitspw and spw, but fitspw and spw
51 # are disjoint (in spws).
52 >>> fitspw = '1~10:5~122,15~22:5~122'
53 >>> spw = '6~14'
55 # Initialize spwmap with the union of them.
56 >>> spwmap = update_spw(join_spws(fitspw, spw), None)[1]
58 >>> myfitspw = update_spw(fitspw, spwmap)[0]
59 >>> myfitspw
60 '0~9:5~122,14~21:5~122'
61 >>> myspw = update_spw(spw, spwmap)[0]
62 >>> myspw
63 '5~13'
64 >>> myspw = update_spw('0,1,3;5~8:20~30;44~50^2', None)[0]
65 >>> myspw
66 '0,1,2;3~6:20~30;44~50^2'
67 """
68 # Blank is valid. Blank is good.
69 if not spw:
70 return '', {}
72 # A list of [spw, chan] pairs. The chan parts will not be changed.
73 spwchans = []
75 make_spwmap = False
76 if not spwmap:
77 spwmap = {}
78 make_spwmap = True
79 spws = set([])
81 # Because ; means different things when it separates spws and channel
82 # ranges, I can't think of a better way to construct spwchans than an
83 # explicit state machine. (But $spws_alone =~ s/:[^,]+//g;)
84 inspw = True # until a : is encountered.
85 spwgrp = ''
86 chagrp = ''
88 def store_spwchan(sstr, cstr):
89 spwchans.append([sstr, cstr])
90 if make_spwmap:
91 for sgrp in sstr.split(';'):
92 if sgrp.find('~') > -1:
93 start, end = map(int, sgrp.split('~'))
94 spws.update(range(start, end + 1))
95 else:
96 spws.add(int(sgrp))
98 for c in spw:
99 if c == ',': # Start new [spw, chan] pair.
100 # Store old one.
101 store_spwchan(spwgrp, chagrp)
103 # Initialize new one.
104 spwgrp = ''
105 chagrp = ''
106 inspw = True
107 elif c == ':':
108 inspw = False
109 elif inspw:
110 spwgrp += c
111 else:
112 chagrp += c
113 # Store final [spw, chan] pair.
114 store_spwchan(spwgrp, chagrp)
116 # casalog.post("spwchans ={}".format(spwchans))
117 # casalog.post("spws ={}".format(spws))
119 # Update spw (+ fitspw)
120 if make_spwmap:
121 i = 0
122 for s in sorted(spws):
123 spwmap[str(s)] = str(i)
124 i += 1
125 outstr = ''
126 for sc in spwchans:
127 sgrps = sc[0].split(';')
128 for sind in range(len(sgrps)):
129 sgrp = sgrps[sind]
130 if sgrp.find('~') > -1:
131 start, end = sgrp.split('~')
132 sgrps[sind] = spwmap[start] + '~' + spwmap[end]
133 else:
134 sgrps[sind] = spwmap[sgrp]
135 outstr += ';'.join(sgrps)
136 if sc[1]:
137 outstr += ':' + sc[1]
138 outstr += ','
140 return outstr.rstrip(','), spwmap # discard final comma.
142def spwchan_to_ranges(vis, spw):
143 """
144 Returns the spw:chan selection string spw as a dict of channel selection
145 ranges for vis, keyed by spectral window ID.
147 The ranges are stored as tuples of (start channel,
148 end channel (inclusive!),
149 step).
151 Note that '' returns an empty set! Use '*' to select everything!
153 Example:
154 >>> from update_spw import spwchan_to_ranges
155 >>> selranges = spwchan_to_ranges('uid___A002_X1acc4e_X1e7.ms', '7:10~20^2;40~55')
156 ValueError: spwchan_to_ranges() does not support multiple channel ranges per spw.
157 >>> selranges = spwchan_to_ranges('uid___A002_X1acc4e_X1e7.ms', '0~1:1~3,5;7:10~20^2')
158 >>> selranges
159 {0: (1, 3, 1), 1: (1, 3, 1), 5: (10, 20, 2), 7: (10, 20, 2)}
160 """
161 selarr = _ms.msseltoindex(vis, spw=spw)['channel']
162 nspw = selarr.shape[0]
163 selranges = {}
164 for s in range(nspw):
165 if selarr[s][0] in selranges:
166 raise ValueError('spwchan_to_ranges() does not support multiple channel ranges per spw.')
167 selranges[selarr[s][0]] = tuple(selarr[s][1:])
168 return selranges
170def spwchan_to_sets(vis, spw):
171 """
172 Returns the spw:chan selection string spw as a dict of sets of selected
173 channels for vis, keyed by spectral window ID.
175 Note that '' returns an empty set! Use '*' to select everything!
177 Example (16.ms has spws 0 and 1 with 16 chans each):
178 >>> from update_spw import spwchan_to_sets
179 >>> vis = casa['dirs']['data'] + '/regression/unittest/split/unordered_polspw.ms'
180 >>> spwchan_to_sets(vis, '0:0')
181 {0: set([0])}
182 >>> selsets = spwchan_to_sets(vis, '1:1~3;5~9^2,9') # 9 is a bogus spw.
183 >>> selsets
184 {1: [1, 2, 3, 5, 7, 9]}
185 >>> spwchan_to_sets(vis, '1:1~3;5~9^2,8')
186 {1: set([1, 2, 3, 5, 7, 9]), 8: set([0])}
187 >>> spwchan_to_sets(vis, '')
188 {}
189 """
190 if not spw: # _ms.msseltoindex(vis, spw='')['channel'] returns a
191 return {} # different kind of empty array. Skip it.
193 # Currently distinguishing whether or not vis is a valid MS from whether it
194 # just doesn't have all the channels in spw is a bit crude. Sanjay is
195 # working on adding some flexibility to _ms.msseltoindex.
196 if not os.path.isdir(vis):
197 raise ValueError(str(vis) + ' is not a valid MS.')
199 sets = {}
200 try:
201 scharr = _ms.msseltoindex(vis, spw=spw)['channel']
202 for scr in scharr:
203 if not scr[0] in sets:
204 sets[scr[0]] = set([])
206 # scr[2] is the last selected channel. Bump it up for range().
207 scr[2] += 1
208 sets[scr[0]].update(range(*scr[1:]))
209 except:
210 # spw includes channels that aren't in vis, so it needs to be trimmed
211 # down to make _ms.msseltoindex happy.
212 allrec = _ms.msseltoindex(vis, spw='*')
213 # casalog.post("Trimming {}".format(spw))
214 spwd = spw_to_dict(spw, {}, False)
215 for s in spwd:
216 if s in allrec['spw']:
217 endchan = allrec['channel'][s, 2]
218 if not s in sets:
219 sets[s] = set([])
220 if spwd[s] == '':
221 # We need to get the spw's # of channels without using
222 # _ms.msseltoindex.
223 mytb = tbtool()
224 mytb.open(vis + '/SPECTRAL_WINDOW')
225 spwd[s] = range(mytb.getcell('NUM_CHAN', s))
226 mytb.close()
227 sets[s].update([c for c in spwd[s] if c <= endchan])
228 return sets
230def set_to_chanstr(chanset, totnchan=None):
231 """
232 Essentially the reverse of expand_tilde. Given a set or list of integers
233 chanset, returns the corresponding string form. It will not use non-unity
234 steps (^) if multiple ranges (;) are necessary, but it will use ^ if it
235 helps to eliminate any ;s.
237 totnchan: the total number of channels for the input spectral window, used
238 to abbreviate the return string.
240 It returns '' for the empty set and '*' if
242 Examples:
243 >>> from update_spw import set_to_chanstr
244 >>> set_to_chanstr(set([0, 1, 2, 4, 5, 6, 7, 9, 11, 13]))
245 '0~2;4~7;9;11;13'
246 >>> set_to_chanstr(set([7, 9, 11, 13]))
247 '7~13^2'
248 >>> set_to_chanstr(set([7, 9]))
249 '7~9^2'
250 >>> set_to_chanstr([0, 1, 2])
251 '0~2'
252 >>> set_to_chanstr([0, 1, 2], 3)
253 '*'
254 >>> set_to_chanstr([0, 1, 2, 6], 3)
255 '*'
256 >>> set_to_chanstr([0, 1, 2, 6])
257 '0~2;6'
258 >>> set_to_chanstr([1, 2, 4, 5, 6, 7, 8, 9, 10, 11], 12)
259 '1~2;4~11'
260 """
261 if totnchan:
262 mylist = [c for c in chanset if c < totnchan]
263 else:
264 mylist = list(chanset)
266 if totnchan == len(mylist):
267 return '*'
269 mylist.sort()
271 retstr = ''
272 if len(mylist) > 1:
273 # Check whether the same step can be used throughout.
274 step = mylist[1] - mylist[0]
275 samestep = True
276 for i in range(2, len(mylist)):
277 if mylist[i] - mylist[i - 1] != step:
278 samestep = False
279 break
280 if samestep:
281 retstr = str(mylist[0]) + '~' + str(mylist[-1])
282 if step > 1:
283 retstr += '^' + str(step)
284 else:
285 sc = mylist[0]
286 oldc = sc
287 retstr = str(sc)
288 nc = len(mylist)
289 for i in range(1, nc):
290 cc = mylist[i]
291 if (cc > oldc + 1) or (i == nc - 1):
292 if (i == nc - 1) and (cc == oldc + 1):
293 retstr += '~' + str(cc)
294 else:
295 if oldc != sc:
296 retstr += '~' + str(oldc)
297 retstr += ';' + str(cc)
298 sc = cc
299 oldc = cc
300 elif len(mylist) > 0:
301 retstr = str(mylist[0])
302 return retstr
304def sets_to_spwchan(spwsets, nchans={}):
305 """
306 Returns a spw:chan selection string for a dict of sets of selected
307 channels keyed by spectral window ID.
309 nchans is a dict of the total number of channels keyed by spw, used to
310 abbreviate the return string.
312 Examples:
313 >>> from update_spw import sets_to_spwchan
314 >>> # Use nchans to get '1' instead of '1:0~3'.
315 >>> sets_to_spwchan({1: [0, 1, 2, 3]}, {1: 4})
316 '1'
317 >>> sets_to_spwchan({1: set([1, 2, 3, 5, 7, 9]), 8: set([0])})
318 '1:1~3;5;7;9,8:0'
319 >>> sets_to_spwchan({0: set([4, 5, 6]), 1: [4, 5, 6], 2: [4, 5, 6]})
320 '0~2:4~6'
321 >>> sets_to_spwchan({0: [4], 1: [4], 3: [0, 1], 4: [0, 1], 7: [0, 1]}, {3: 2, 4: 2, 7: 2})
322 '0~1:4,3~4,7'
323 """
324 # Make a list of spws for each channel selection.
325 csd = {}
326 for s in spwsets:
327 # Convert the set of channels to a string.
328 if spwsets[s]:
329 cstr = set_to_chanstr(spwsets[s], nchans.get(s))
331 if cstr:
332 if not cstr in csd:
333 csd[cstr] = []
334 csd[cstr].append(s)
336 # Now convert those spw lists into strings, inverting as we go so the final
337 # string can be sorted by spw:
338 scd = {}
339 while csd:
340 cstr, slist = csd.popitem()
341 slist.sort()
342 startspw = slist[0]
343 oldspw = startspw
344 sstr = str(startspw)
345 nselspw = len(slist)
346 for sind in range(1, nselspw):
347 currspw = slist[sind]
348 if (currspw > oldspw + 1) or (sind == nselspw - 1):
349 if currspw > oldspw + 1:
350 if oldspw != startspw:
351 sstr += '~' + str(oldspw)
352 sstr += ';' + str(currspw)
353 startspw = currspw
354 else: # The range has come to an end on the last spw.
355 sstr += '~' + str(currspw)
356 oldspw = currspw
357 scd[sstr] = cstr
358 spwgrps = sorted(scd.keys())
360 # Finally stitch together the final string.
361 scstr = ''
362 for sstr in spwgrps:
363 scstr += sstr
364 if scd[sstr] != '*':
365 scstr += ':' + scd[sstr]
366 scstr += ','
367 return scstr.rstrip(',')
369def update_spwchan(vis, sch0, sch1, truncate=False, widths={}):
370 """
371 Given an spw:chan selection string sch1, return what it must be changed to
372 to get the same result if used with the output of split(vis, spw=sch0).
374 '' is taken to mean '*' in the input but NOT the output! For the output
375 '' means sch0 and sch1 do not intersect.
377 truncate: If True and sch0 only partially overlaps sch1, return the update
378 of the intersection.
379 If (False and sch0 does not cover sch1), OR
380 there is no intersection, raises a ValueError.
382 widths is a dictionary of averaging widths (default 1) for each spw.
384 Examples:
385 >>> from update_spw import update_spwchan
386 >>> newspw = update_spwchan('anything.ms', 'anything', 'anything')
387 >>> newspw
388 '*'
389 >>> vis = casa['dirs']['data'] + '/regression/unittest/split/unordered_polspw.ms'
390 >>> update_spwchan(vis, '0~1:1~3,5;7:10~20^2', '0~1:2~3,5;7:12~18^2')
391 '0~1:1~2,2~3:1~4'
392 >>> update_spwchan(vis, '7', '3')
393 ValueError: '3' is not a subset of '7'.
394 >>> update_spwchan(vis, '7:10~20^2', '7:12~18^3')
395 ValueError: '7:12~18^3' is not a subset of '7:10~20^2'.
396 >>> update_spwchan(vis, '7:10~20^2', '7:12~18^3', truncate=True)
397 '0:1~4^3'
398 >>> update_spwchan(vis, '7:10~20^2', '7:12~18^3', truncate=True, widths={7: 2})
399 '0:0~2^2'
400 """
401 # Convert '' to 'select everything'.
402 if not sch0:
403 sch0 = '*'
404 if not sch1:
405 sch1 = '*'
407 # Short circuits
408 if sch1 == '*':
409 return '*'
410 elif sch1 in (sch0, '*'):
411 return '*'
413 sch0sets = spwchan_to_sets(vis, sch0)
414 sch1sets = spwchan_to_sets(vis, sch1)
416 outsets = {}
417 outspw = 0
418 s0spws = sorted(sch0sets.keys())
419 s1spws = sorted(sch1sets.keys())
420 ns0spw = len(s0spws)
421 nchans = {}
422 for s in s1spws:
423 if s in s0spws:
424 s0 = sch0sets[s]
425 s1 = sch1sets[s]
427 # Check for and handle (throw or dispose) channels in sch1 that aren't in
428 # sch0.
429 if s1.difference(s0):
430 if truncate:
431 s1.intersection_update(s0)
432 if not s1:
433 raise ValueError("'%s' does not overlap '%s'." % (sch1, sch0))
434 else:
435 raise ValueError("'%s' is not a subset of '%s'." % (sch1, sch0))
437 # Adapt s1 for a post-s0 world.
438 s0list = sorted(list(s0))
439 s1list = sorted(list(s1))
440 outchan = 0
441 nc0 = len(s0list)
442 for s1ind in range(len(s1list)):
443 while (outchan < nc0) and (s0list[outchan] < s1list[s1ind]):
444 outchan += 1
445 if outchan == nc0: # Shouldn't happen
446 outchan -= 1
447 s1list[s1ind] = outchan // widths.get(s, 1)
449 # Determine outspw.
450 while (outspw < ns0spw) and (s0spws[outspw] < s):
451 outspw += 1
452 if outspw == ns0spw: # Shouldn't happen
453 outspw -= 1
455 outsets[outspw] = set(s1list)
457 # Get the number of channels per spw that are selected by s0.
458 nchans[outspw] = len(s0)
459 elif not truncate:
460 raise ValueError(str(s) + ' is not a selected spw of ' + sch0)
462 return sets_to_spwchan(outsets, nchans)
464def expand_tilde(tstr, conv_multiranges=False):
465 """
466 Expands a string like '8~11' to [8, 9, 10, 11].
467 Returns '*' if tstr is ''!
469 conv_multiranges: If True, '*' will be returned if tstr contains ';'.
470 (split can't yet handle multiple channel ranges per spw.)
472 Examples:
473 >>> from update_spw import expand_tilde
474 >>> expand_tilde('8~11')
475 [8, 9, 10, 11]
476 >>> expand_tilde(None)
477 '*'
478 >>> expand_tilde('3~7^2;9~11')
479 [3, 5, 7, 9, 10, 11]
480 >>> expand_tilde('3~7^2;9~11', True)
481 '*'
482 """
483 tstr = str(tstr) # Allows bare ints.
484 if (not tstr) or (conv_multiranges and tstr.find(';') > -1):
485 return '*'
487 tstr = tstr.replace("'", '') # Dequote
488 tstr = tstr.replace('"', '')
490 numset = set([])
492 for numrang in tstr.split(';'):
493 step = 1
494 try:
495 if numrang.find('~') > -1:
496 if numrang.find('^') > -1:
497 numrang, step = numrang.split('^')
498 step = int(step)
499 start, end = map(int, numrang.split('~'))
500 else:
501 start = int(numrang)
502 end = start
503 except:
504 raise ValueError('numrang = ' + numrang + ', tstr = ' + tstr + ', conv_multiranges = ' + str(conv_multiranges))
505 numset.update(range(start, end + 1, step))
506 return sorted(list(numset))
508def spw_to_dict(spw, spwdict={}, conv_multiranges=True):
509 """
510 Expand an spw:chan string to {s0: [s0chans], s1: [s1chans, ...], ...}
511 where s0, s1, ... are integers for _each_ selected spw, and s0chans is a
512 set of selected chans (as integers) for s0. '' instead of a channel set
513 means that all of the channels are selected.
515 The spw:chan dict is unioned with spwdict.
517 Returning an empty dict means everything should be selected (i.e. spw = '').
518 (split can't yet handle multiple channel ranges per spw.)
520 conv_multiranges: If True, any spw with > 1 channel range selected will
521 have ALL of its channels selected.
522 (split can't yet handle multiple channel ranges per spw.)
524 Examples:
525 >>> from update_spw import spw_to_dict
526 >>> spw_to_dict('', {})
527 {}
528 >>> spw_to_dict('6~8:2~5', {})[6]
529 set([2, 3, 4, 5])
530 >>> spw_to_dict('6~8:2~5', {})[8]
531 set([2, 3, 4, 5])
532 >>> spw_to_dict('6~8:2~5', {6: ''})[6]
533 ''
534 >>> spw_to_dict('6~8:2~5', {6: '', 7: set([1, 7])})[7]
535 set([1, 2, 3, 4, 5, 7])
536 >>> spw_to_dict('7', {6: '', 7: set([1, 7])})[7]
537 ''
538 >>> spw_to_dict('7:123~127;233~267', {6: '', 7: set([1, 7])})[7] # Multiple chan ranges
539 ''
540 >>> spw_to_dict('5,7:123~127;233~267', {6: '', 7: set([1, 7])})[5]
541 ''
542 >>> spw_to_dict('5:3~5,7:123~127;233~267', {6: '', 7: set([1, 7])})[5]
543 set([3, 4, 5])
544 """
545 if not spw:
546 return {}
548 myspwdict = copy.deepcopy(spwdict)
550 # Because ; means different things when it separates spws and channel
551 # ranges, I can't think of a better way to construct myspwdict than an
552 # explicit state machine. (But $spws_alone =~ s/:[^,]+//g;)
553 inspw = True # Must start with an spw.
554 spwgrp = ''
555 chagrp = ''
557 def enter_ranges(spwg, chag):
558 spwrange = expand_tilde(spwg)
559 if spwrange == '*': # This shouldn't happen.
560 return {}
561 else:
562 charange = expand_tilde(chag, conv_multiranges)
563 for s in spwrange:
564 if charange == '*':
565 myspwdict[s] = ''
566 else:
567 if not s in myspwdict:
568 myspwdict[s] = set([])
569 if myspwdict[s] != '':
570 myspwdict[s].update(charange)
572 for c in spw:
573 if c == ',' or (inspw and c == ';'): # Start new [spw, chan] pair.
574 # Store old one.
575 enter_ranges(spwgrp, chagrp)
577 # Initialize new one.
578 spwgrp = ''
579 chagrp = ''
580 inspw = True
581 elif c == ':':
582 inspw = False
583 elif inspw:
584 spwgrp += c
585 else:
586 chagrp += c
588 # Store final [spw, chan] pair.
589 enter_ranges(spwgrp, chagrp)
590 return myspwdict
592def join_spws(spw1, spw2, span_semicolon=True):
593 """
594 Returns the union of spw selection strings spw1 and spw2.
596 span_semicolon (default True): If True, for any spws
597 that have > 1 channel range, the entire spw will be selected.
599 Examples:
600 >>> from update_spw import join_spws
601 >>> join_spws('0~2:3~5,3:9~13', '')
602 ''
603 >>> join_spws('0~2:3~5,3:9~13', '1~3:4~7')
604 '0:3~5,1~2:3~7,3'
605 >>> join_spws('1~10:5~122,15~22:5~122', '1~10:5~122,15~22:5~122')
606 '1~10:5~122,15~22:5~122'
607 >>> join_spws('', '')
608 ''
609 >>> join_spws('1~10:5~122,15~22:5~122', '0~21')
610 '0~21,22:5~122'
611 """
612 if not spw1 or not spw2:
613 return ''
615 spwdict = spw_to_dict(spw1, {})
616 spwdict = spw_to_dict(spw2, spwdict)
618 res = ''
619 # Convert channel sets to strings
620 for s in spwdict:
621 cstr = ''
622 if isinstance(spwdict[s], set):
623 cstr = set_to_chanstr(spwdict[s])
624 if span_semicolon and ';' in cstr:
625 cstr = ''
626 spwdict[s] = cstr
628 # If consecutive spws have the same channel selection, merge them.
629 slist = list(spwdict.keys())
630 slist.sort()
631 res = str(slist[0])
632 laststart = 0
633 for i in range(1, len(slist)):
634 # If consecutive spws have the same channel list,
635 if slist[i] == slist[i - 1] + 1 and spwdict[slist[i]] == spwdict[slist[i - 1]]:
636 if slist[i] == slist[laststart] + 1:
637 res += '~' # Continue the spw range.
638 else: # Terminate it and start a new one.
639 if res[-1] == '~': # if start != end
640 res += str(slist[i - 1])
641 if spwdict[slist[i - 1]] != '': # Add channel range, if any.
642 res += ':' + spwdict[slist[i - 1]]
643 res += ',' + str(slist[i])
644 laststart = i
645 if res[-1] == '~': # Finish the last range if it is dangling.
646 res += str(slist[-1])
647 if spwdict[slist[-1]] != '': # Add channel range, if any.
648 res += ':' + spwdict[slist[-1]]
649 return res
651def intersect_spws(spw1, spw2):
652 """
653 Almost the opposite of join_spws(), this returns the list of spws that the
654 spw:chan selection strings spw1 and spw2 have in common. Unlike join_spws(),
655 channel ranges are ignored. '' in the input counts as 'select everything',
656 so the intersection of '' with anything is anything. If the intersection
657 really is everything, '' is returned instead of a set.
659 Examples:
660 >>> from update_spw import intersect_spws
661 >>> intersect_spws('0~2:3~5,3:9~13', '')
662 set([0, 1, 2, 3])
663 >>> intersect_spws('0~2:3~5,3:9~13', '0~2:7~9,5')
664 set([0, 1, 2])
665 >>> intersect_spws('0~2:3~5;10~13,3:9~13', '0~2:7~9,5')
666 set([0, 1, 2])
667 >>> intersect_spws('0~2:3~5,3:9~13', '10~12:7~9,5') # Empty set
668 set([])
669 >>> intersect_spws('', '') # Everything
670 ''
671 """
672 if spw1 == '':
673 if spw2 == '':
674 return '' # intersection('', '') = ''
675 else: # intersection('', spw2) = spw2
676 return set(spw_to_dict(spw2, {}).keys()) # Just the spws, no chan ranges
677 elif spw2 == '': # intersection('', spw1) = spw1
678 return set(spw_to_dict(spw1, {}).keys()) # Just the spws, no chan ranges
679 else:
680 spwset1 = set(spw_to_dict(spw1, {}).keys()) # spws are the keys, chan
681 spwset2 = set(spw_to_dict(spw2, {}).keys()) # ranges are the values.
682 return spwset1.intersection(spwset2)
684def subtract_spws(spw1, spw2):
685 """
686 Returns the set of spws of spw selection string spw1 that are not in spw2.
687 Like intersect_spws(), this intentionally ignores channel ranges. It
688 assumes that spw1 and spw2 refer to the same MS (this only matters for '').
689 subtract_spws('', '0~5') is a tough case: it is impossible to know whether
690 '' is equivalent to '0~5' without reading the MS's SPECTRAL_WINDOW
691 subtable, so it returns 'UNKNOWN'.
693 Examples:
694 >>> from update_spw import subtract_spws
695 >>> subtract_spws('0~2:3~5,3:9~13', '') # Anything - Everything
696 set([])
697 >>> subtract_spws('0~2:3~5,3:9~13', '0~2:7~9,5')
698 set([3])
699 >>> subtract_spws('', '0~2:7~9,5') # Everything - Something
700 'UNKNOWN'
701 >>> subtract_spws('0~2,3:9~13', '4~7:7') # Something - Something Else
702 set([0, 1, 2, 3])
703 >>> subtract_spws('', '') # Everything - Everything
704 set([])
705 """
706 if spw1 == '':
707 if spw2 == '':
708 return set([])
709 else:
710 return 'UNKNOWN'
711 elif spw2 == '':
712 return set([])
713 else:
714 spwset1 = set(spw_to_dict(spw1, {}).keys()) # spws are the keys, chan
715 spwset2 = set(spw_to_dict(spw2, {}).keys()) # ranges are the values.
716 return spwset1.difference(spwset2)
718if __name__ == '__main__':
719 import doctest, sys
720 doctest.testmod(verbose=True)