14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494 | class DataPlot(object):
"""
Plot transition probabilities and collision strengths from different data sets
"""
def __init__(self, elem=None, spec=None, all_data=[], atom=None, n_tem_points=10000,
ref_tem=None, OmegaInterp='linear',NLevels=None):
"""
Parameters:
elem: atomic elem
spec: ionization stage in spectroscopic notation (I = 1, II = 2, etc.)
atom: e.g. 'O3'
all_data: dictionary of all_data to be compared (see above for format)
n_tem_points: number of points in the fit (default=100; increase if fit is not smooth)
ref_tem: array of temperature values to be signaled in the plots
OmegaInterp: interpolating function between Omega values ('Cheb' [default], 'Linear')
**Example:**
dataplot = pn.DataPlot('O', 3) # initializes the plot
dataplot.plotA() # transition probabilities plot
dataplot.plotRelA() # relative transition probabilities plot
dataplot.plotOmega() # collision strength plot
"""
colors = np.array(['r', 'g', 'b', 'm', 'c', 'y'])
self.calling = 'DataPlot'
if atom is not None:
self.atom = str.capitalize(atom)
self.elem = parseAtom(self.atom)[0]
self.spec = int(parseAtom(self.atom)[1])
else:
self.elem = str.capitalize(elem)
self.spec = int(spec)
self.atom = self.elem + str(self.spec)
old_data = pn.atomicData.getDataFile(self.atom)
# Check if matplotlib installed
if not pn.config.INSTALLED['plt']:
pn.log_.warn('Matplotlib not installed!', calling=self.calling)
# Separate omega and A data sets
atom_data = []
coll_data = []
if all_data == []:
all_data = pn.atomicData.getAllAvailableFiles(self.atom, mark_current=False)
i_colors = 0
for file_ in all_data:
ID = file_.split('_')[-1]
type_ = (file_.split('_')[2]).split('.')[0]
if type_ in ['atom', 'coll']:
data_list = type_ + '_data'
vars()[data_list].append({'ID': ID, 'file_': file_, 'type': type_, 'color': colors[i_colors % len(colors)]})
i_colors += 1
self.atom_data = atom_data
self.coll_data = coll_data
# Temperature values to be signaled by vertical lines in omega plots. Can be changed by defining ref_tem
if ref_tem is None:
self.ref_tem = np.log10([5000., 10000., 20000.])
else:
self.ref_tem = ref_tem
# If it's not standard effective collision strengths, we don't want it
coll_data_temp=[]
for item in self.coll_data:
coll_data_temp.append(item)
for data in coll_data_temp:
pn.atomicData.setDataFile(data['file_'])
atom = pn.Atom(self.elem, self.spec, OmegaInterp=OmegaInterp, NLevels=NLevels)
"""
try:
if 'O_UNIT' in atom.CollData.comments.keys():
self.coll_data.remove(data)
except:
pass
"""
# For each data set, an atom is built.
del(atom)
for data in self.atom_data + self.coll_data:
pn.atomicData.setDataFile(data['file_'])
atom = pn.Atom(self.elem, self.spec, OmegaInterp=OmegaInterp, NLevels=NLevels)
data['atom'] = atom
for data in old_data:
if data is not None:
pn.atomicData.setDataFile(data)
self.atom_rom = (self.elem + '_' + int_to_roman(int(self.spec))).lower()
self.n_tem_points = n_tem_points
self.atom_n_max = 0
for at in self.atom_data:
if at['atom'].atomFileType != 'chianti':
if at['atom'].atomNLevels > self.atom_n_max:
self.atom_n_max = at['atom'].atomNLevels
self.coll_n_max = 0
for at in self.coll_data:
if at['atom'].collFileType != 'chianti':
if at['atom'].collNLevels > self.coll_n_max:
self.coll_n_max = at['atom'].collNLevels
def plotA(self, save=False, figsize=(18, 12), fignum=None, NLevels=None):
"""
Plot the log of the A values of each data set
**Parameters:**
save: if True, saves the plot in a file
figsize: figure size (default: [18, 12])
fignum: figure Number
"""
if NLevels is None:
atom_n_max = self.atom_n_max
else:
atom_n_max = NLevels
if not pn.config.INSTALLED['plt']:
pn.log_.error('Matplotlib not installed!', calling=self.calling)
plt.figure(fignum, figsize=figsize)
plt.clf()
ax = plt.subplot(111)
x = np.arange(atom_n_max * (atom_n_max - 1) / 2.)
ticks = x
# Inventory of markers. Chosen to be distinguished even when overlapped
mark = ['<', (5, 1), '>', 'o', '|']
i_marker = 0
# Background colors to distinguish lower levels
bg_color = 'grbymc'
tick_label = []
for i in range(atom_n_max - 1):
# x0, x1 are start and end points of each level
x0 = i * (atom_n_max - 1) - i * (i - 1) / 2 - 0.5
width = atom_n_max - i - 1
x1 = x0 + width
plt.axvspan(x0, x1, facecolor=bg_color[i%6], alpha=0.05)
# The x axis must stretch the maximum range (although some data set might have a lower n_level)
for j in range(i + 1, atom_n_max):
tick_label.append('(' + str(j + 1) + ', ' + str(i + 1) + ')')
for data in self.atom_data:
n_levels = data['atom'].atomNLevels
Ay = []
A = data['atom'].getA()
color = data['color']
try:
for i in range(n_levels - 1):
for j in range(i + 1, n_levels):
if A[j, i] > 0:
Ay.append(np.log10(A[j, i]))
else:
Ay.append(np.NaN)
plt.scatter(x, Ay, marker=mark[i_marker], s=300., c=color, alpha=0.35, linewidths=1, label='%s' % (data['ID']))
ax.set_xticks(ticks)
except:
pn.log_.warn('Problem in plotting A', calling=self.calling + '.plotA')
i_marker += 1
ax.set_xticklabels(tick_label)
# Plot features
plt.xlabel('Transition')
plt.ylabel('Log A(j, i)')
plt.title('Transition probabilities for [%s %s]' % (self.elem, int_to_roman(int(self.spec))))
plt.legend(loc='lower right', markerscale=1., scatterpoints=1, borderpad=1, labelspacing=1)
plt.show()
if save:
plt.figure(figsize=[18, 12])
plt.savefig(self.atom_rom + '_' + data['ID'] + "-" + self.ref_data + "_A.pdf")
def plotAllA(self, save=False, figsize=(18, 12), fignum=None, NLevels=None):
"""
Plot the log of the A values of each data set
Parameters:
save: if True, saves the plot in a file
figsize: figure size (default: [18, 12])
fignum: figure Number
"""
if NLevels is None:
atom_n_max = self.atom_n_max
else:
atom_n_max = NLevels
if not pn.config.INSTALLED['plt']:
pn.log_.error('Matplotlib not installed!', calling=self.calling)
fig = plt.figure(fignum, figsize=figsize)
plt.clf()
max_y_ticks = 6
ticks = np.arange(len(self.atom_data)+1)[1:]
A = np.zeros([len(self.atom_data), atom_n_max, atom_n_max])
color=[]
for i, data in enumerate(self.atom_data):
A_tmp = data['atom'].getA()
if A_tmp.shape[0] > atom_n_max:
A[i, :] = A_tmp[0:atom_n_max, 0:atom_n_max]
else:
A[i, 0:A_tmp.shape[0], 0:A_tmp.shape[1]] = A_tmp
color.append(data['color'])
A[np.where((A<=0))] = np.NaN
lgA = np.log10(A)
x = ticks
color = np.array(color)
for j in range(2, atom_n_max+1):
for i in range(1, j):
y = lgA[:, j-1, i-1]
ax = plt.subplot(atom_n_max - 1, atom_n_max - 1, (atom_n_max-1)*(j-2) + i)
ax.set_xticks(x)
plt.xlim((min(x)-0.5, max(x)+0.5))
ax.yaxis.set_major_locator(MaxNLocator(max_y_ticks-1))
try:
plt.scatter(x, y, c=color, label='_nolegend_', s=40, edgecolor='None')
lbl = '({0} -> {1})'.format(j,i)
ax.text(0.95, 0.95, lbl, fontsize=10, color="#660066", transform=ax.transAxes, ha="right", va="top")
except:
pn.log_.warn('Problem with plotting a subplot {} {}'.format(i,j), calling=self.calling)
if (j==atom_n_max) & (i==1):
plt.xlabel('Reference #', fontsize=8)
plt.ylabel('Log(A$_{ji}$)', fontsize=8)
plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0.12)
plt.gca().get_yaxis().get_major_formatter().set_useOffset(False)
ax.set_xticks(x)
#fig.text(.5, .05, "i (lower level)", fontsize=12, ha='center', color="#660066", va='top')
#fig.text(.05, .5, "j (upper level)", ha='left', va='center', fontsize=12, color="#660066", rotation=90)
title = "Available transition probabilities for {0} {1}".format(self.elem, int_to_roman(int(self.spec)))
fig.text(.5, .95, title, color="#191970", fontsize=14, ha='center')
for i, data in enumerate(self.atom_data):
x_txt = .95
y_txt = .85 - .04*i
fig.text(x_txt, y_txt, "Ref. %d: %s" %(i+1, data['ID']), color=color[i], ha = 'right')
plt.tight_layout(pad=0.1)
if save:
plt.savefig(self.atom_rom + "_all_As.pdf")
plt.show()
def plotRelA(self, ref_data=None, save=False, figsize=None, fignum=None, NLevels=None, fig=None):
"""
Plot the relative difference of the A of each data set with respect to the reference one
Parameters:
ref_data: reference data set for comparing transition probabilities (default=first data ID)
save: if True, save the plot in a file (default: False)
figsize: figure size (default: [18, 12])
fignum: figure number
"""
if NLevels is None:
atom_n_max = self.atom_n_max
else:
atom_n_max = NLevels
if not pn.config.INSTALLED['plt']:
pn.log_.error('Matplotlib not installed!', calling=self.calling)
# Commented out because it produced an extra, empty frame - VL 30 Jul 2015
#if fig is None:
# fig = plt.figure(fignum, figsize=figsize)
#plt.clf()
ticks = range(atom_n_max + 1)[1:]
# The As of the reference data set are stored
if ref_data is None:
ref_data = self.atom_data[0]['ID']
for data in self.atom_data:
if (data['ID'] is ref_data):
ref_A = data['atom'].getA()
# Only non-zero values are taken into account to prevent dividing by zero
nonzero_ref_A_indexes = np.nonzero(ref_A)
for data in self.atom_data:
if (data['ID'] is not ref_data):
A = data['atom'].getA()
try:
# The data to be compared might have fewer levels than the reference one. Only the ones in common are considered.
tmp_indexes = np.asarray(nonzero_ref_A_indexes)
up_indexes = tmp_indexes[0][tmp_indexes[0] < data['atom'].atomNLevels]
lo_indexes = tmp_indexes[1][tmp_indexes[0] < data['atom'].atomNLevels]
# Indexes for which the reference data set is not zero
ref_indexes = (up_indexes, lo_indexes)
# Ratio of A/A_ref for non-zero A_ref values
A_ratio = A[ref_indexes] / ref_A[ref_indexes]
nonzero_A_ratio_indexes = np.nonzero(A_ratio)
rel_A = np.log10(A_ratio[nonzero_A_ratio_indexes])
# Plotting starts
#fig = plt.figure()
if fig is None:
fig = plt.figure(fignum, figsize=figsize)
plt.clf()
ax = plt.subplot(111)
fig.subplots_adjust(top=0.9)
# Physical levels = array levels + 1
x = np.asarray(lo_indexes) + 1
y = np.asarray(up_indexes) + 1
# Size is proportional to difference between As
# Numerical value of size adjusted empirically, no particular meaning
size = np.multiply(10000., np.abs(rel_A))
# A different color is assigned to data points depending on whether they are smaller or larger
# than the reference data.
# The ratio is not compared to exactly one to include only truly different values
index_grt = np.where(rel_A > 0.0000)
color_grt = np.array(['#FF8C00']).repeat(x[index_grt].size)
index_sml = np.where(rel_A <= 0.0)
color_sml = np.array(['#FF1480']).repeat(x[index_sml].size)
# The range is adjusted to include only values for which there is a difference
xmin = np.max([np.min(x[index_sml]), np.min(x[index_grt])])
xmax = np.min([np.max(x[index_sml]), np.max(x[index_sml])])
ax.set_xticks(ticks)
ax.set_xlim((xmin - 0.25, xmax + 0.25)[:])
# The x-axis is plotted on the top to mimic the data array in the fits file
for tick in ax.xaxis.get_major_ticks():
tick.label1On = False
tick.label2On = True
ymin = np.min([np.min(y[index_sml]), np.min(y[index_grt])])
ymax = np.max([np.max(y[index_sml]), np.max(y[index_sml])])
ax.set_yticks(ticks)
# The y axis increases downward to mimic the data array in the fits file
ax.set_ylim((ymin - 0.25, ymax + 0.25)[::-1])
# Fake points, just to generate the labels (I don't know how to generate a reasonable label
# for data points with markers of different size)
plt.scatter(x[index_grt][0], y[index_grt][0], marker='o', s=.0002, c=color_grt[0],
label='Positive (max = %.2f)' % (np.max(rel_A)))
plt.scatter(x[index_grt][0], y[index_grt][0], marker='o', s=.0002, c=color_sml[0],
label='Negative (min = %.2f)' % (np.min(rel_A)))
# True data; size reflects divergence from reference data, color whether positive or negative
plt.scatter(x[index_grt], y[index_grt], marker='o', s=size[index_grt], c=np.asarray(color_grt), alpha=1.0)
plt.scatter(x[index_sml], y[index_sml], marker='o', s=size[index_sml], c=np.asarray(color_sml), alpha=1.0)
# X-axis label
ax.text(.5, 1.05, '$i$ (lower level)', transform=ax.transAxes, ha='center', va='bottom')
#plt.xlabel('$i$ (lower level)')
# Y-axis label
plt.ylabel('$j$ (upper level)')
# Legend
plt.legend(loc='upper right', markerscale=1000., scatterpoints=1, borderpad=1,
labelspacing=1, prop=dict(size=12), title=u'Log (%s / %s)' % (data['ID'], ref_data))
# Plot title
ax.text(.5, -.10, "Relative difference between $A$s data for [%s %s]"
% (self.elem, int_to_roman(int(self.spec))),
transform=ax.transAxes, ha='center', va='bottom', color='#191970', size=12)
except:
pn.log_.warn('Problem in plotting relA', calling='DatasetPlot')
if save:
plt.savefig(self.atom_rom + '_' + data['ID'] + "-" + ref_data + "_relA.pdf")
plt.show()
def plotOmega(self, save=False, figsize=(18, 12), fignum=1, scan_orders=None, NLevels=None,
fig=None):
"""
Plot the tabulated collision strengths of each data set and the fit that is performed by PyNeb
Parameters:
save: Boolean. Determine if the plot is automatically saved in a file (default: False)
figsize: List. figure size in inches (default: [18, 12])
fignum: Figure Number DEPRECATED!!!
scan_orders: = None or (min_order, max_order) or (min_order, -1) to go until the max. DEPRECATED!!!
fig: DEPRECATED!!!
"""
if NLevels is None:
coll_n_max = self.coll_n_max
else:
coll_n_max = NLevels
# Plotting range somewhat larger than actual data range (less tight)
if not pn.config.INSTALLED['plt']:
pn.log_.error('Matplotlib not installed!', calling=self.calling)
if fig is None:
#fig = plt.figure(fignum, figsize=figsize)
fig, axes = plt.subplots(coll_n_max-1, coll_n_max-1, figsize=figsize)
plt.autoscale(tight=False)
# I need two
first = True
first_done = False
legend_text = []
legend_lines = []
axes_plotaxis = np.zeros_like(axes, dtype=bool)
# style_dic = ['-', '--', '-.', ':'] * 10
for data in self.coll_data:
# tem_points = tabulated temperature points (different for each data set)
tem_points = self.tem_in_K(data['atom'].tem_units, data['atom'].getTemArray())
tem_min = min(tem_points)
tem_max = max(tem_points)
# tem_funct = array of temperature values for which the fit is evaluated
tem_funct = np.linspace(tem_min, tem_max, self.n_tem_points)
x_dots = np.log10(tem_points)
x_lines = np.log10(tem_funct)
if 'COEFF' in data['atom'].CollData.comments.keys():
coeff = float(data['atom'].CollData.comments['COEFF'])
else:
coeff = 1.0
# Loops over all levels
for i in range(1, coll_n_max+1):
for j in range(i + 1, coll_n_max+1):
# N levels require an N-1 x N-1 array of plots
# The subplots are arranged in rows and columns according to the upper and lower levels
ax = axes[j-2, i-1]
axes_plotaxis[j-2, i-1] = True
if (i <= data['atom'].collNLevels) and (j <= data['atom'].collNLevels):
y_dots = data['atom'].getOmega(tem_points, j, i)
try:
if 'O_UNIT' not in data['atom'].CollData.comments:
y_dots = coeff * data['atom'].getOmegaArray(j, i)
if y_dots.sum() > 0.0:
ax.plot(x_dots, y_dots, color=data['color'],
marker='*', linestyle='None', label='_nolegend_', markersize=10)
y_dots = data['atom'].getOmega(tem_points, j, i)
if y_dots.sum() > 0.0:
ax.plot(x_dots, y_dots, color=data['color'],
marker='o', linestyle='None', label='_nolegend_')
y_lines = data['atom'].getOmega(tem_funct, j, i)
if y_lines.sum() > 0.:
ax.plot(x_lines, y_lines,
color=data['color'], label='_nolegend_')
except:
pn.log_.warn('Problem with plotting a data set', calling=self.calling)
# Draws vertical lines at selected temperature values (only once for each subplot, hence "first")
if (first and (data['atom'].collNLevels == coll_n_max)):
for tem in self.ref_tem:
ax.axvline(tem, c='blue', alpha=0.4, ls=':')
lbl = "$\Omega$" + "(" + str(j) + "," + str(i) + ")"
ax.text(0.95, 0.95, lbl, transform=ax.transAxes, ha="right", va="top")
first_done = True
# Maximum number of ticks to avoid overcrowding
ax.xaxis.set_major_locator(MaxNLocator(4))
ax.yaxis.set_major_locator(MaxNLocator(3))
# The line type and identifier of each data set are stored for the legend
legend_lines.append(mpl.lines.Line2D(tem_funct, tem_funct, color=data['color'], linestyle='-'))
legend_text.append(data['ID'])
if first_done:
first = False
# The label of the reference temperature values are only plotted above the last plot
for tem in self.ref_tem:
ax.text(tem, ax.get_ylim()[1], ' %0.0f K' % (10 ** tem), ha="left", va="bottom",
rotation=65, fontsize=12, color='#FF0000', alpha=0.35)
# Axis labels, title and legend
fig.text(.5, .05, "Log($T_e/K$)", fontsize=16, ha='center', va='top')
fig.text(.05, .5, "$\Omega$", va='center', fontsize=20, rotation=90)
fig.text(.5, .95, "[%s %s] collision strengths" % (self.elem, int_to_roman(int(self.spec))),
color="#191970", fontsize=16, ha='center')
# plt.legend(legend_lines, legend_text, loc='upper right', borderpad=1,
# labelspacing=1, bbox_to_anchor=(1, 1 * coll_n_max))
axes[0, coll_n_max-2].legend(legend_lines, legend_text, loc='upper right', borderpad=1,
labelspacing=1)
for ax, ax_plt in zip(axes.ravel(), axes_plotaxis.ravel()):
if not ax_plt:
ax.set_axis_off()
#plt.tight_layout()
if save:
fig.savefig(self.atom_rom + "_CS.pdf")
fig.show()
def tem_in_K(self, tem_units, tem):
"""
Convert the temperature from the unit of the fits file into K
Parameters:
tem_units: 'log(K)' or 'K/1000'
tem: temperature
"""
if (tem_units == "log(K)"):
return np.power(10., tem)
elif (tem_units == "K/10000"):
return np.multiply(tem, 1.e4)
else: #T in Kelvin in the fits file
return tem
|