1
2
3 """
4 A tool set for reading and using molecular level populations.
5
6 Author: R. Lombaert and M. van de Sande
7
8 """
9
10 import os, collections
11 import numpy as np
12 from cc.tools.io import DataIO
13 from cc.tools.readers.Reader import Reader
14
15 import matplotlib.pyplot as p
16
17 from scipy.interpolate import interp1d
18 from scipy.interpolate import InterpolatedUnivariateSpline as spline1d
19
20
21
23
24 '''
25 The PopReader class.
26
27 Inherits from the Reader class.
28
29 Provides methods to manage level populations calculated with
30 radiative transfer such as through GASTRoNOoM and MCP/ALI. Provides basic
31 methods for reading level populations.
32
33 Subclasses provide the read and parse methods that are code-specific, in
34 case the basic methods do not apply:
35 - MlineReader which reads mline output and combines spectroscopic
36 molecular data and level populations.
37
38 Typically spectroscopic information is not available here but is given in
39 MolReader objects. Collision rates are given in CollisReader objects. For
40 MCP/ALI those are combined in LamdaReader.
41
42 '''
43
45
46 '''
47 Initialize an PopReader object. The filename and additional args/kwargs
48 are passed to the Reader parent class, where the dict is made and the
49 filename is stored.
50
51 Additional args/kwargs are used for the dict creation of the parent of
52 Reader.
53
54 Note that the filename can be given as None, or may not be an existing
55 file. In this case, nothing is read, and the user can set the impact
56 parameter and the level populations manually with setP and setPop.
57
58 @param fn: The filename of the level populations.
59 @type fn: str
60
61 '''
62
63 super(PopReader,self).__init__(fn,*args,**kwargs)
64
65
66 self['pars'] = {}
67 self['pop'] = {}
68
69
70 if self.__class__ == PopReader and os.path.isfile(self.fn):
71 self.read()
72
73
74
76
77 '''
78 Read the level populations as a function of impact parameter.
79
80 Each level is stored
81 as an index according to the column from which it was read.
82
83 '''
84
85
86
87 data = np.loadtxt(self.fn)
88
89
90 self['p'] = data[:,0]
91
92
93
94 self['pars']['ny'] = len(data[0])-1
95 for i in range(1,self['pars']['ny']+1):
96 self['pop'][i] = data[:,i]
97
98
99
101
102 '''
103 Set the number of levels. Usually the highest level index.
104
105 @param ny: The number of levels
106 @type ny: int
107
108 '''
109
110 self['pars']['ny'] = ny
111
112
113
115
116 '''
117 Instead of reading the impact parameters, set them explicitly.
118
119 @param p: the impact parameters (cm)
120 @type p: array
121
122 '''
123
124 self['p'] = p
125
126
127
129
130 '''
131 Instead of reading the populations, set them here per level.
132
133 @param index: The level index
134 @type index: int
135 @param n: The level populations as a function of impact parameter for
136 level with index
137 @type n: array
138
139 '''
140
141 self['pop'][index] = n
142
143
144
146
147 '''
148 Return the impact parameter grid.
149
150 @return: The impact parameter grid in cm
151 @rtype: array
152
153 '''
154
155 return self['p']
156
157
158
160
161 '''
162 Return the level populations for a set of level indices.
163
164 Note that this is the level index, not lower J. For CO, the J quantum
165 number would be index-1.
166
167 Note that the array approach is not used because that makes indexing
168 (0-based for python, 1-based for fortran) more confusing and prone to
169 mistakes. Hence, a dictionary explicity index-key approach is preferred.
170 This is internal only.
171
172 If an index is given as an iterable, an array is still returened
173 with shape = (len(index),len(self['p'])) for ease of use. This includes
174 when all level populations are requested, i.e. when index is None.
175
176 If one wants to use the dictionary itself, self['pop'] is of course
177 available.
178
179 @keyword index: The index of the level, if None all levels are returned
180 as a dict. Each level is then accessed through
181 pop[index].
182
183 (default: None)
184 @type index: int
185
186 @return: The level populations in the form of a 1d or 2d array,
187 depending if a single level or multiple levels are requested.
188 @rtype: array
189
190 '''
191
192 if index is None:
193 return np.array([self['pop'][i] for i in self.getLI()])
194 elif isinstance(index,collections.Iterable) \
195 and not isinstance(index,str):
196 return np.array([self['pop'][i] for i in index])
197 else:
198 return self['pop'][int(index)]
199
200
201
203
204 '''
205 Return the indices of the excitation levels included in the level pops.
206
207 @return: The level indices
208 @rtype: array
209
210 '''
211
212 return range(1,self['pars']['ny']+1)
213
214
215
216 - def setInterp(self,itype='spline',*args,**kwargs):
217
218 '''
219 Set the interpolator for the level populations.
220
221 Additional arguments can be passed to the interpolator object.
222
223 @keyword itype: The type of interpolator. Either spline or linear.
224
225 (default: 'spline)
226 @type itype: str
227
228 '''
229
230
231 itype = itype.lower()
232 if itype == 'linear':
233 interp = interp1d
234 else:
235 interp = spline1d
236
237 self['ipop'] = dict()
238
239
240 for i in self.getLI():
241 self['ipop'][i] = interp(x=self['p'],y=self['pop'][i],\
242 *args,**kwargs)
243
244
245
247
248 '''
249 Get the interpolator for a given level index.
250
251 @param index: The level index
252 @type index: int
253
254 @return: The level population interpolator as a function of impact
255 parameter in cm.
256 @rtype: interpolator
257
258 '''
259
260 return self['ipop'][index]
261
262
263
265
266 '''
267 Plot the level populations for all included levels.
268
269 @keyword fn: Filename of the plot, including path and extension. If None
270 the plot is shown in the python session.
271
272 (default: None)
273 @type fn: str
274
275 '''
276
277 p.clf()
278 fig1 = p.figure(1)
279 for ii in range(1,self['pars']['ny']+1):
280 p.loglog(self.getP(),self.getPop(ii))
281 p.xlabel('P (cm)', fontsize = 14)
282 p.ylabel('Level Populations', fontsize = 14)
283 p.ylim(ymin=1e-12)
284 p.xlim(xmin=self.getP()[0])
285 p.xlim(xmax=self.getP()[-1])
286 if not fn is None:
287 if not os.path.splitext(fn)[1]: fn += '.pdf'
288 p.savefig(fn,bbox_inches='tight')
289 else:
290 p.show()
291