source: trunk/python/selector.py@ 2950

Last change on this file since 2950 was 2881, checked in by Takeshi Nakazato, 11 years ago

New Development: No

JIRA Issue: Yes CAS-5858

Ready for Test: Yes

Interface Changes: No

What Interface Changed: Please list interface changes

Test Programs: List test programs

Put in Release Notes: No

Module(s): Module Names change impacts.

Description: Describe your changes here...

Bug fix on handling existing TaQL query string.
If original query is related to FIELDNAME or SRCNAME,
it is overwritten. Otherwise, the new query is appended
to the original one.


File size: 14.8 KB
RevLine 
[1875]1import re
[2874]2import math
[2870]3import string
[1875]4from asap._asap import selector as _selector, srctype
[1826]5from asap.utils import unique, _to_list
[948]6
[932]7class selector(_selector):
[948]8 """
9 A selection object to be applied to scantables to restrict the
10 scantables to specific rows.
11 """
[1930]12 fields = ["pols", "ifs", "beams", "scans", "cycles", "name", "query", "types", "rows"]
[932]13
[1576]14 def __init__(self, *args, **kw):
15 if len(args) == 1:
16 if isinstance(args[0], self.__class__) \
17 or isinstance(args[0], _selector):
18 _selector.__init__(self, args[0])
19 else:
20 raise TypeError("Argument can only be a selector object")
21 else:
22 _selector.__init__(self)
23 for k,v in kw.items():
24 if k in self.fields:
25 func = getattr(self, "set_%s" % k)
26 func(v)
27
[932]28 def reset(self):
29 """
30 Unset all selections.
31 """
32 self._reset()
33
[948]34 def is_empty(self):
35 """
36 Has anything been set?
37 """
38 return self._empty()
39
[932]40 def set_polarisations(self, pols=[]):
41 """
42 Set the polarisations to be selected in the scantable.
43 Parameters:
44 pols: a list of integers of 0-3, or strings, e.g ["I","Q"].
45 Default [] is no selection
46 Example:
47 sel = selector()
48 # These are equivalent if data is 'linear'
49 sel.set_polarisations(["XX","Re(XY)"])
50 sel.set_polarisations([0,2])
51 # reset the polarisation selection
52 sel.set_polarisations()
[948]53
[932]54 """
[954]55 vec = _to_list(pols, str) or _to_list(pols, int)
[1045]56 if isinstance(vec, list): # is an empty and/or valid vector
[932]57 if len(vec) and isinstance(vec[-1],str):
[954]58 self._setpolstrings(vec)
[932]59 return
60 self._setpols(vec)
61 else:
62 raise TypeError('Unknown pol type. Please use [0,1...] or ["XX","YY"...]')
[1576]63
[1349]64 # for the americans
[1542]65 set_polarizations = set_polarisations
66 # for the lazy
67 set_pols = set_polarisations
[932]68
69 def set_ifs(self, ifs=[]):
70 """
71 Set a sequence of IF numbers (0-based).
72 Parameters:
73 ifs: a list of integers. Default [] is to unset the selection.
74 """
75 vec = _to_list(ifs, int)
[1045]76 if isinstance(vec,list):
[932]77 self._setifs(vec)
78 else:
79 raise TypeError('Unknown IFno type. Use lists of integers.')
80
81 def set_scans(self, scans=[]):
82 """
83 Set a sequence of Scan numbers (0-based).
84 Parameters:
85 scans: a list of integers. Default [] is to unset the selection.
86 """
87 vec = _to_list(scans, int)
[1045]88 if isinstance(vec,list):
[932]89 self._setscans(vec)
90 else:
91 raise TypeError('Unknown Scan number type. Use lists of integers.')
92
93 def set_beams(self, beams=[]):
94 """
95 Set a sequence of Beam numbers (0-based).
96 Parameters:
97 beams: a list of integers. Default [] is to unset the selection.
98 """
99 vec = _to_list(beams, int)
[1045]100 if isinstance(vec,list):
[932]101 self._setbeams(vec)
102 else:
103 raise TypeError('Unknown Beam number type. Use lists of integers.')
104
105 def set_cycles(self, cycles=[]):
106 """
107 Set a sequence of IF numbers (0-based).
108 Parameters:
109 cycless: a list of integers. Default [] is to unset the selection.
110 """
111 vec = _to_list(cycles, int)
[1045]112 if isinstance(vec,list):
[932]113 self._setcycles(vec)
114 else:
115 raise TypeError('Unknown Cycle number type. Use lists of integers.')
116
[948]117
[932]118 def set_name(self, name):
119 """
120 Set a selection based on a name. This can be a unix pattern , e.g. "*_R"
121 Parameters:
122 name: a string containing a source name or pattern
123 Examples:
124 # select all reference scans which start with "Orion"
125 selection.set_name("Orion*_R")
126 """
127 if isinstance(name, str):
[952]128 self._setname(name)
[932]129 else:
130 raise TypeError('name must be a string')
[948]131
[932]132 def set_tsys(self, tsysmin=0.0, tsysmax=None):
133 """
134 Select by Tsys range.
135 Parameters:
136 tsysmin: the lower threshold. Default 0.0
137 tsysmax: the upper threshold. Default None.
138 Examples:
139 # select all spectra with Tsys <= 500.0
140 selection.set_tsys(tsysmax=500.0)
[948]141
[932]142 """
[963]143 taql = "SELECT FROM $1 WHERE TSYS[0] >= %f" % (tsysmin)
[932]144 if isinstance(tsysmax, float):
[963]145 taql = taql + " AND TSYS[0] <= %f" % ( tsysmax)
[932]146 self._settaql(taql)
147
148 def set_query(self, query):
149 """
150 Select by Column query. Power users only!
151 Example:
152 # select all off scans with integration times over 60 seconds.
[1875]153 selection.set_query("SRCTYPE == PSOFF AND INTERVAL > 60.0")
[932]154 """
[1875]155 rx = re.compile("((SRCTYPE *[!=][=] *)([a-zA-Z.]+))", re.I)
156 for r in rx.findall(query):
157 sval = None
158 stype = r[-1].lower()
159 if stype.find('srctype.') == -1:
160 stype = ".".join(["srctype", stype])
161 try:
162 sval = eval(stype)
163 sval = "%s%d" % (r[1], sval)
164 except:
165 continue
166 query = query.replace(r[0], sval)
[932]167 taql = "SELECT FROM $1 WHERE " + query
168 self._settaql(taql)
[948]169
170 def set_order(self, order):
171 """
172 Set the order the scantable should be sorted by.
173 Parameters:
[1045]174 order: The list of column names to sort by in order
[948]175 """
176 self._setorder(order)
177
[1819]178 def set_rows(self, rows=[]):
179 """
180 Set a sequence of row numbers (0-based). Power users Only!
181 NOTICE row numbers can be changed easily by sorting,
[1826]182 prior selection, etc.
[1819]183 Parameters:
184 rows: a list of integers. Default [] is to unset the selection.
185 """
[1826]186 vec = _to_list(rows, int)
[1819]187 if isinstance(vec,list):
188 self._setrows(vec)
189 else:
190 raise TypeError('Unknown row number type. Use lists of integers.')
191
192 def set_types(self, types=[]):
193 """
[1826]194 Set a sequence of source types.
[1819]195 Parameters:
196 types: a list of integers. Default [] is to unset the selection.
197 """
[1826]198 vec = _to_list(types, int)
[1819]199 if isinstance(vec,list):
200 self._settypes(vec)
201 else:
202 raise TypeError('Unknown row number type. Use lists of integers.')
203
[2870]204 def set_msselection_field(self, selection):
205 """
[2875]206 Set a field selection in msselection syntax. The msselection
207 suppports the following syntax:
208
209 pattern match:
210 - UNIX style pattern match for source name using '*'
211 (compatible with set_name)
212
213 field id selection:
214 - simple number in string ('0', '1', etc.)
215 - range specification using '~' ('0~1', etc.)
216 - range specification using '>' or '<' in combination
217 with '=' ('>=1', '<3', etc.)
218
219 comma separated multiple selection:
220 - selections can be combined by using ',' ('0,>1',
221 'mysource*,2~4', etc.)
[2870]222 """
223 selection_list = map(string.strip, selection.split(','))
224 query_list = list(self.generate_query(selection_list))
[2875]225 if len(query_list) > 0:
[2881]226 original_query = self.get_query()
227 if len(original_query) == 0 or re.match('.*(SRC|FIELD)NAME.*',original_query):
228 query = 'SELECT FROM $1 WHERE ' + ' || '.join(query_list)
229 else:
230 query = 'SELECT FROM $1 WHERE (' + original_query + ') && (' + ' || '.join(query_list) + ')'
[2875]231 self._settaql(query)
[2870]232
233 def generate_query(self, selection_list):
234 for s in selection_list:
[2873]235 if s.isdigit() or re.match('^[<>]=?[0-9]*$', s) or \
236 re.match('^[0-9]+~[0-9]+$', s):
[2870]237 #print '"%s" is ID selection using < or <='%(s)
[2873]238 a = FieldIdRegexGenerator(s)
239 yield '(%s)'%(a.get_regex())
[2875]240 elif len(s) > 0:
[2873]241 #print '"%s" is UNIX style pattern match'%(s)
242 yield '(SRCNAME == pattern(\'%s\'))'%(s)
[2870]243
[948]244 def get_scans(self):
245 return list(self._getscans())
246 def get_cycles(self):
247 return list(self._getcycles())
248 def get_beams(self):
249 return list(self._getbeams())
250 def get_ifs(self):
251 return list(self._getifs())
252 def get_pols(self):
253 return list(self._getpols())
254 def get_poltypes(self):
255 return list(self._getpoltypes())
256 def get_order(self):
257 return list(self._getorder())
[1819]258 def get_types(self):
259 return list(self._gettypes())
[1930]260 def get_rows(self):
261 return list(self._getrows())
[1337]262 def get_query(self):
263 prefix = "SELECT FROM $1 WHERE "
264 return self._gettaql().replace(prefix, "")
[1542]265
[948]266 def get_name(self):
267 print "NYI"
268 s = self._gettaql()
[1337]269 return
270 def __str__(self):
271 out = ""
272 d = {"SCANNO": self.get_scans(),
273 "CYCLENO": self.get_cycles(),
274 "BEAMNO": self.get_beams(),
275 "IFNO": self.get_ifs(),
276 "Pol Type": self.get_poltypes(),
277 "POLNO": self.get_pols(),
278 "QUERY": self.get_query(),
[1930]279 "SRCTYPE": self.get_types(),
280 "ROWS": self.get_rows(),
[1337]281 "Sort Order": self.get_order()
282 }
283 for k,v in d.iteritems():
284 if v:
285 out += "%s: %s\n" % (k, v)
286 if len(out):
287 return out[:-1]
288 else:
289 return out
[2340]290
[948]291 def __add__(self, other):
292 """
293 Merge two selections.
294 """
[1596]295 if self.is_empty():
[2340]296 return selector(other)
[1596]297 elif other.is_empty():
[2340]298 return selector(self)
[948]299 union = selector()
300 gets = [[self._getscans(), other._getscans(), union._setscans],
301 [self._getcycles(), other._getcycles(),union._setcycles],
302 [self._getbeams(), other._getbeams(), union._setbeams],
303 [self._getifs(), other._getifs(), union._setifs],
304 [self._getpols(), other._getpols(), union._setpols]]
305 for v in gets:
306 vec = list(v[0]+v[1])
307 vec.sort()
308 v[2](unique(vec))
[1349]309 q = other.get_query()
310 qs = self.get_query()
311 if len(q) and len(qs):
312 union.set_query(qs +" AND " + q)
313 else:
314 if len(q):
315 union.set_query(q)
316 elif len(qs):
317 union.set_query(qs)
[948]318 return union
[2873]319
320class FieldIdRegexGenerator(object):
321 def __init__(self, pattern):
322 if pattern.isdigit():
323 self.regex = 'FIELDNAME == regex(\'.+__%s$\')'%(pattern)
324 else:
325 self.regex = None
326 ineq = None
327 if pattern.find('<') >= 0:
328 ineq = '<'
329 s = pattern.strip().lstrip(ineq).lstrip('=')
330 if not s.isdigit():
331 raise RuntimeError('Invalid syntax: %s'%(pattern))
332 self.id = int(s) + (-1 if pattern.find('=') < 0 else 0)
333 self.template = string.Template('FIELDNAME == regex(\'.+__${reg}$\')')
334 elif pattern.find('>') >= 0:
335 ineq = '>'
336 s = pattern.strip().lstrip(ineq).lstrip('=')
337 if not s.isdigit():
338 raise RuntimeError('Invalid syntax: %s'%(pattern))
339 self.id = int(s) + (-1 if pattern.find('=') >= 0 else 0)
340 self.template = string.Template('FIELDNAME == regex(\'.+__[0-9]+$\') && FIELDNAME != regex(\'.+__${reg}$\')')
341 elif pattern.find('~') >= 0:
342 s = map(string.strip, pattern.split('~'))
343 if len(s) == 2 and s[0].isdigit() and s[1].isdigit():
[2881]344 id0 = int(s[0])
345 id1 = int(s[1])
346 if id0 == 0:
347 self.id = id1
348 self.template = string.Template('FIELDNAME == regex(\'.+__${reg}$\')')
349 else:
350 self.id = [id0-1,id1]
351 self.template = string.Template('FIELDNAME == regex(\'.+__${reg}$\') && FIELDNAME != regex(\'.+__${optreg}$\')')
[2873]352 else:
353 raise RuntimeError('Invalid syntax: %s'%(pattern))
354 else:
355 raise RuntimeError('Invalid syntax: %s'%(pattern))
356 #print 'self.id=',self.id
357
358 def get_regex(self):
359 if self.regex is not None:
360 # 'X'
361 return self.regex
362 elif isinstance(self.id, list):
363 # 'X~Y'
364 return self.template.safe_substitute(reg=self.__compile(self.id[1]),
365 optreg=self.__compile(self.id[0]))
366 else:
367 # '<(=)X' or '>(=)X'
368 return self.template.safe_substitute(reg=self.__compile(self.id))
369
370 def __compile(self, idx):
371 pattern = ''
372 if idx >= 0:
[2876]373 numerics = map(int,list(str(idx)))
374 #numerics.reverse()
375 num_digits = len(numerics)
[2873]376 #print 'numerics=',numerics
377 if num_digits == 1:
378 if numerics[0] == 0:
379 pattern = '0'
380 else:
381 pattern = '[0-%s]'%(numerics[0])
382 elif num_digits == 2:
[2876]383 pattern = '(%s)'%('|'.join(
384 list(self.__gen_two_digit_pattern(numerics))))
[2873]385 elif num_digits == 3:
[2876]386 pattern = '(%s)'%('|'.join(
387 list(self.__gen_three_digit_pattern(numerics))))
[2873]388 else:
389 raise RuntimeError('ID > 999 is not supported')
390 else:
391 raise RuntimeError('ID must be >= 0')
392 return pattern
[2876]393
394 def __gen_two_digit_pattern(self, numerics):
395 assert len(numerics) == 2
396 yield '[0-9]'
397 if numerics[0] == 2:
398 yield '1[0-9]'
399 elif numerics[0] > 2:
400 yield '[1-%s][0-9]'%(numerics[0]-1)
401 if numerics[1] == 0:
402 yield '%s%s'%(numerics[0],numerics[1])
403 else:
404 yield '%s[0-%s]'%(numerics[0],numerics[1])
405
406 def __gen_three_digit_pattern(self, numerics):
407 assert len(numerics) == 3
408 yield '[0-9]'
409 yield '[1-9][0-9]'
410 if numerics[0] == 2:
411 yield '1[0-9][0-9]'
412 elif numerics[0] > 2:
413 yield '[1-%s][0-9][0-9]'%(numerics[0]-1)
414 if numerics[1] == 0:
415 if numerics[2] == 0:
416 yield '%s00'%(numerics[0])
417 else:
418 yield '%s0[0-%s]'%(numerics[0],numerics[2])
419 else:
420 if numerics[1] > 1:
421 yield '%s[0-%s][0-9]'%(numerics[0],numerics[1]-1)
422 elif numerics[1] == 1:
423 yield '%s0[0-9]'%(numerics[0])
424 if numerics[0] == 0:
425 yield '%s%s%s'%(numerics[0],numerics[1],numerics[2])
426 else:
427 yield '%s%s[0-%s]'%(numerics[0],numerics[1],numerics[2])
Note: See TracBrowser for help on using the repository browser.