source: trunk/python/selector.py @ 2875

Last change on this file since 2875 was 2875, 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: sel = sd.selector()

sel.set_msselection_field()
print sel._gettaql()

Put in Release Notes: No

Module(s): Module Names change impacts.

Description: Describe your changes here...

Fixed behavior when empty string is passed as an argument.
Also, documentation is added.

File size: 14.4 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:
226            query = 'SELECT FROM $1 WHERE ' + ' || '.join(query_list)
227            self._settaql(query)
[2870]228
229    def generate_query(self, selection_list):
230        for s in selection_list:
[2873]231            if s.isdigit() or re.match('^[<>]=?[0-9]*$', s) or \
232                    re.match('^[0-9]+~[0-9]+$', s):
[2870]233                #print '"%s" is ID selection using < or <='%(s)
[2873]234                a = FieldIdRegexGenerator(s)
235                yield '(%s)'%(a.get_regex())
[2875]236            elif len(s) > 0:
[2873]237                #print '"%s" is UNIX style pattern match'%(s)
238                yield '(SRCNAME == pattern(\'%s\'))'%(s)
[2870]239       
[948]240    def get_scans(self):
241        return list(self._getscans())
242    def get_cycles(self):
243        return list(self._getcycles())
244    def get_beams(self):
245        return list(self._getbeams())
246    def get_ifs(self):
247        return list(self._getifs())
248    def get_pols(self):
249        return list(self._getpols())
250    def get_poltypes(self):
251        return list(self._getpoltypes())
252    def get_order(self):
253        return list(self._getorder())
[1819]254    def get_types(self):
255        return list(self._gettypes())
[1930]256    def get_rows(self):
257        return list(self._getrows())
[1337]258    def get_query(self):
259        prefix = "SELECT FROM $1 WHERE "
260        return self._gettaql().replace(prefix, "")
[1542]261
[948]262    def get_name(self):
263        print "NYI"
264        s = self._gettaql()
[1337]265        return
266    def __str__(self):
267        out = ""
268        d = {"SCANNO": self.get_scans(),
269             "CYCLENO": self.get_cycles(),
270             "BEAMNO": self.get_beams(),
271             "IFNO": self.get_ifs(),
272             "Pol Type": self.get_poltypes(),
273             "POLNO": self.get_pols(),
274             "QUERY": self.get_query(),
[1930]275             "SRCTYPE": self.get_types(),
276             "ROWS": self.get_rows(),
[1337]277             "Sort Order": self.get_order()
278             }
279        for k,v in d.iteritems():
280            if v:
281                out += "%s: %s\n" % (k, v)
282        if len(out):
283            return out[:-1]
284        else:
285            return out
[2340]286
[948]287    def __add__(self, other):
288        """
289        Merge two selections.
290        """
[1596]291        if self.is_empty():
[2340]292            return selector(other)
[1596]293        elif other.is_empty():
[2340]294            return selector(self)
[948]295        union = selector()
296        gets = [[self._getscans(), other._getscans(), union._setscans],
297                [self._getcycles(), other._getcycles(),union._setcycles],
298                [self._getbeams(), other._getbeams(), union._setbeams],
299                [self._getifs(), other._getifs(), union._setifs],
300                [self._getpols(), other._getpols(), union._setpols]]
301        for v in gets:
302            vec = list(v[0]+v[1])
303            vec.sort()
304            v[2](unique(vec))
[1349]305        q = other.get_query()
306        qs = self.get_query()
307        if len(q) and len(qs):
308            union.set_query(qs +" AND " + q)
309        else:
310            if len(q):
311                union.set_query(q)
312            elif len(qs):
313                union.set_query(qs)
[948]314        return union
[2873]315
316class FieldIdRegexGenerator(object):
317    def __init__(self, pattern):
318        if pattern.isdigit():
319            self.regex = 'FIELDNAME == regex(\'.+__%s$\')'%(pattern)
320        else:
321            self.regex = None
322            ineq = None
323            if pattern.find('<') >= 0:
324                ineq = '<'
325                s = pattern.strip().lstrip(ineq).lstrip('=')
326                if not s.isdigit():
327                    raise RuntimeError('Invalid syntax: %s'%(pattern))
328                self.id = int(s) + (-1 if pattern.find('=') < 0 else 0)
329                self.template = string.Template('FIELDNAME == regex(\'.+__${reg}$\')')
330            elif pattern.find('>') >= 0:
331                ineq = '>'
332                s = pattern.strip().lstrip(ineq).lstrip('=')
333                if not s.isdigit():
334                    raise RuntimeError('Invalid syntax: %s'%(pattern))
335                self.id = int(s) + (-1 if pattern.find('=') >= 0 else 0)
336                self.template = string.Template('FIELDNAME == regex(\'.+__[0-9]+$\') && FIELDNAME != regex(\'.+__${reg}$\')')
337            elif pattern.find('~') >= 0:
338                s = map(string.strip, pattern.split('~'))
339                if len(s) == 2 and s[0].isdigit() and s[1].isdigit():
340                    self.id = [int(s[0])-1,int(s[1])]
341                else:
342                    raise RuntimeError('Invalid syntax: %s'%(pattern))
343                self.template = string.Template('FIELDNAME == regex(\'.+__${reg}$\') && FIELDNAME != regex(\'.+__${optreg}$\')')
344            else:
345                raise RuntimeError('Invalid syntax: %s'%(pattern))
346            #print 'self.id=',self.id
347
348    def get_regex(self):
349        if self.regex is not None:
350            # 'X'
351            return self.regex
352        elif isinstance(self.id, list):
353            # 'X~Y'
354            return self.template.safe_substitute(reg=self.__compile(self.id[1]),
355                                                 optreg=self.__compile(self.id[0]))
356        else:
357            # '<(=)X' or '>(=)X'
358            return self.template.safe_substitute(reg=self.__compile(self.id))
359
360    def __compile(self, idx):
361        pattern = ''
362        if idx >= 0:
363            if idx < 10:
364                num_digits = 1
365            else:
366                num_digits = int(math.log10(idx)) + 1
367            numerics = []
368            modulo = 10
369            divider = 1
370            for i in xrange(num_digits):
371                numerics.append(int((idx % modulo) / divider))
372                modulo *= 10
373                divider *= 10
374            #print 'numerics=',numerics
375            if num_digits == 1:
376                if numerics[0] == 0:
377                    pattern = '0'
378                else:
379                    pattern = '[0-%s]'%(numerics[0])
380            elif num_digits == 2:
381                pattern_list = ['[0-9]']
382                pattern_list.append('[1-%s][0-9]'%(numerics[1]-1))
383                if numerics[0] == 0:
384                    pattern_list.append('%s%s'%(numerics[1],numerics[0]))
385                else:
386                    pattern_list.append('%s[0-%s]'%(numerics[1],numerics[0]))
387                pattern = '(%s)'%('|'.join(pattern_list))
388            elif num_digits == 3:
389                pattern_list = ['[0-9]','[1-9][0-9]']
390                if numerics[2] == 2:
391                    pattern_list.append('1[0-9][0-9]')
392                elif numerics[2] > 2:
393                    pattern_list.append('[1-%s][0-9][0-9]'%(numerics[2]-1))
394                if numerics[1] == 0:
395                    if numerics[0] == 0:
396                        pattern_list.append('%s00'%(numerics[2]))
397                    else:
398                        pattern_list.append('%s0[0-%s]'%(numerics[2],numerics[0]))
399                else:
400                    pattern_list.append('%s[0-%s][0-9]'%(numerics[2],numerics[1]-1))
401                    if numerics[0] == 0:
402                        pattern_list.append('%s%s%s'%(numerics[2],numerics[1],numerics[0]))
403                    else:
404                        pattern_list.append('%s%s[0-%s]'%(numerics[2],numerics[1],numerics[0]))
405                pattern = '(%s)'%('|'.join(pattern_list))
406            else:
407                raise RuntimeError('ID > 999 is not supported')
408                pattern = ''
409        else:
410            raise RuntimeError('ID must be >= 0')
411        return pattern
Note: See TracBrowser for help on using the repository browser.