source: branches/polybatch/python/selector.py@ 2820

Last change on this file since 2820 was 1875, checked in by Malte Marquarding, 14 years ago

Added srctype enum handling
File size: 8.6 KB
Line 
1import re
2from asap._asap import selector as _selector, srctype
3from asap.utils import unique, _to_list
4
5class selector(_selector):
6 """
7 A selection object to be applied to scantables to restrict the
8 scantables to specific rows.
9 """
10 fields = ["pols", "ifs", "beams", "scans", "cycles", "name", "query"]
11
12 def __init__(self, *args, **kw):
13 if len(args) == 1:
14 if isinstance(args[0], self.__class__) \
15 or isinstance(args[0], _selector):
16 _selector.__init__(self, args[0])
17 else:
18 raise TypeError("Argument can only be a selector object")
19 else:
20 _selector.__init__(self)
21 for k,v in kw.items():
22 if k in self.fields:
23 func = getattr(self, "set_%s" % k)
24 func(v)
25
26 def reset(self):
27 """
28 Unset all selections.
29 """
30 self._reset()
31
32 def is_empty(self):
33 """
34 Has anything been set?
35 """
36 return self._empty()
37
38 def set_polarisations(self, pols=[]):
39 """
40 Set the polarisations to be selected in the scantable.
41 Parameters:
42 pols: a list of integers of 0-3, or strings, e.g ["I","Q"].
43 Default [] is no selection
44 Example:
45 sel = selector()
46 # These are equivalent if data is 'linear'
47 sel.set_polarisations(["XX","Re(XY)"])
48 sel.set_polarisations([0,2])
49 # reset the polarisation selection
50 sel.set_polarisations()
51
52 """
53 vec = _to_list(pols, str) or _to_list(pols, int)
54 if isinstance(vec, list): # is an empty and/or valid vector
55 if len(vec) and isinstance(vec[-1],str):
56 self._setpolstrings(vec)
57 return
58 self._setpols(vec)
59 else:
60 raise TypeError('Unknown pol type. Please use [0,1...] or ["XX","YY"...]')
61
62 # for the americans
63 set_polarizations = set_polarisations
64 # for the lazy
65 set_pols = set_polarisations
66
67 def set_ifs(self, ifs=[]):
68 """
69 Set a sequence of IF numbers (0-based).
70 Parameters:
71 ifs: a list of integers. Default [] is to unset the selection.
72 """
73 vec = _to_list(ifs, int)
74 if isinstance(vec,list):
75 self._setifs(vec)
76 else:
77 raise TypeError('Unknown IFno type. Use lists of integers.')
78
79 def set_scans(self, scans=[]):
80 """
81 Set a sequence of Scan numbers (0-based).
82 Parameters:
83 scans: a list of integers. Default [] is to unset the selection.
84 """
85 vec = _to_list(scans, int)
86 if isinstance(vec,list):
87 self._setscans(vec)
88 else:
89 raise TypeError('Unknown Scan number type. Use lists of integers.')
90
91 def set_beams(self, beams=[]):
92 """
93 Set a sequence of Beam numbers (0-based).
94 Parameters:
95 beams: a list of integers. Default [] is to unset the selection.
96 """
97 vec = _to_list(beams, int)
98 if isinstance(vec,list):
99 self._setbeams(vec)
100 else:
101 raise TypeError('Unknown Beam number type. Use lists of integers.')
102
103 def set_cycles(self, cycles=[]):
104 """
105 Set a sequence of IF numbers (0-based).
106 Parameters:
107 cycless: a list of integers. Default [] is to unset the selection.
108 """
109 vec = _to_list(cycles, int)
110 if isinstance(vec,list):
111 self._setcycles(vec)
112 else:
113 raise TypeError('Unknown Cycle number type. Use lists of integers.')
114
115
116 def set_name(self, name):
117 """
118 Set a selection based on a name. This can be a unix pattern , e.g. "*_R"
119 Parameters:
120 name: a string containing a source name or pattern
121 Examples:
122 # select all reference scans which start with "Orion"
123 selection.set_name("Orion*_R")
124 """
125 if isinstance(name, str):
126 self._setname(name)
127 else:
128 raise TypeError('name must be a string')
129
130 def set_tsys(self, tsysmin=0.0, tsysmax=None):
131 """
132 Select by Tsys range.
133 Parameters:
134 tsysmin: the lower threshold. Default 0.0
135 tsysmax: the upper threshold. Default None.
136 Examples:
137 # select all spectra with Tsys <= 500.0
138 selection.set_tsys(tsysmax=500.0)
139
140 """
141 taql = "SELECT FROM $1 WHERE TSYS[0] >= %f" % (tsysmin)
142 if isinstance(tsysmax, float):
143 taql = taql + " AND TSYS[0] <= %f" % ( tsysmax)
144 self._settaql(taql)
145
146 def set_query(self, query):
147 """
148 Select by Column query. Power users only!
149 Example:
150 # select all off scans with integration times over 60 seconds.
151 selection.set_query("SRCTYPE == PSOFF AND INTERVAL > 60.0")
152 """
153 rx = re.compile("((SRCTYPE *[!=][=] *)([a-zA-Z.]+))", re.I)
154 for r in rx.findall(query):
155 sval = None
156 stype = r[-1].lower()
157 if stype.find('srctype.') == -1:
158 stype = ".".join(["srctype", stype])
159 try:
160 sval = eval(stype)
161 sval = "%s%d" % (r[1], sval)
162 except:
163 continue
164 query = query.replace(r[0], sval)
165 taql = "SELECT FROM $1 WHERE " + query
166 self._settaql(taql)
167
168 def set_order(self, order):
169 """
170 Set the order the scantable should be sorted by.
171 Parameters:
172 order: The list of column names to sort by in order
173 """
174 self._setorder(order)
175
176 def set_rows(self, rows=[]):
177 """
178 Set a sequence of row numbers (0-based). Power users Only!
179 NOTICE row numbers can be changed easily by sorting,
180 prior selection, etc.
181 Parameters:
182 rows: a list of integers. Default [] is to unset the selection.
183 """
184 vec = _to_list(rows, int)
185 if isinstance(vec,list):
186 self._setrows(vec)
187 else:
188 raise TypeError('Unknown row number type. Use lists of integers.')
189
190 def set_types(self, types=[]):
191 """
192 Set a sequence of source types.
193 Parameters:
194 types: a list of integers. Default [] is to unset the selection.
195 """
196 vec = _to_list(types, int)
197 if isinstance(vec,list):
198 self._settypes(vec)
199 else:
200 raise TypeError('Unknown row number type. Use lists of integers.')
201
202 def get_scans(self):
203 return list(self._getscans())
204 def get_cycles(self):
205 return list(self._getcycles())
206 def get_beams(self):
207 return list(self._getbeams())
208 def get_ifs(self):
209 return list(self._getifs())
210 def get_pols(self):
211 return list(self._getpols())
212 def get_poltypes(self):
213 return list(self._getpoltypes())
214 def get_order(self):
215 return list(self._getorder())
216 def get_types(self):
217 return list(self._gettypes())
218 def get_query(self):
219 prefix = "SELECT FROM $1 WHERE "
220 return self._gettaql().replace(prefix, "")
221
222 def get_name(self):
223 print "NYI"
224 s = self._gettaql()
225 return
226 def __str__(self):
227 out = ""
228 d = {"SCANNO": self.get_scans(),
229 "CYCLENO": self.get_cycles(),
230 "BEAMNO": self.get_beams(),
231 "IFNO": self.get_ifs(),
232 "Pol Type": self.get_poltypes(),
233 "POLNO": self.get_pols(),
234 "QUERY": self.get_query(),
235 "Sort Order": self.get_order()
236 }
237 for k,v in d.iteritems():
238 if v:
239 out += "%s: %s\n" % (k, v)
240 if len(out):
241 return out[:-1]
242 else:
243 return out
244 def __add__(self, other):
245 """
246 Merge two selections.
247 """
248 if self.is_empty():
249 return other
250 elif other.is_empty():
251 return self
252 union = selector()
253 gets = [[self._getscans(), other._getscans(), union._setscans],
254 [self._getcycles(), other._getcycles(),union._setcycles],
255 [self._getbeams(), other._getbeams(), union._setbeams],
256 [self._getifs(), other._getifs(), union._setifs],
257 [self._getpols(), other._getpols(), union._setpols]]
258 for v in gets:
259 vec = list(v[0]+v[1])
260 vec.sort()
261 v[2](unique(vec))
262 q = other.get_query()
263 qs = self.get_query()
264 if len(q) and len(qs):
265 union.set_query(qs +" AND " + q)
266 else:
267 if len(q):
268 union.set_query(q)
269 elif len(qs):
270 union.set_query(qs)
271 return union
Note: See TracBrowser for help on using the repository browser.