source: trunk/python/selector.py@ 2873

Last change on this file since 2873 was 2873, checked in by Takeshi Nakazato, 12 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: Yes/No

Module(s): Module Names change impacts.

Description: Describe your changes here...

Updated sd.selector.set_msselection_field so that it is able to handle
field id selection such as '<1', '<=1', '>1', '>=1', '1~2'.


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