cov.py 17.8 KB
Newer Older
1
#!/usr/bin/env python
2
3
4
5
#
# calculate covariance from neutron events
# @author Tobias Weber <tweber@ill.fr>
# @date 30-mar-2019
Tobias WEBER's avatar
Tobias WEBER committed
6
# @license GPLv3, see 'LICENSE' file
7
8
#

Tobias WEBER's avatar
Tobias WEBER committed
9
import os
10
11
import numpy as np
import numpy.linalg as la
12
import tascalc as tas
13

14

Tobias WEBER's avatar
Tobias WEBER committed
15
16
17
18
np.set_printoptions(
	floatmode = "fixed", 
	precision = 4)

19
20
verbose = True			# console outputs
plot_results = True		# show plot window
21
plot_neutrons = True	# also plot neutron events
22
centre_on_Q = False		# centre plots on Q or zero?
23
ellipse_points = 128	# number of points to draw ellipses
24
symsize = 32.
25
dpi = 600
26

27
# column indices in kf,kf file
28
29
ki_start_idx = 0	# start index of ki 3-vector
kf_start_idx = 3	# start index of kf 3-vector
30
31
wi_idx = 9			# start index of ki weight factor
wf_idx = 10			# start index of kf weight factor
32

33
34
35
36
# column indices in Q,E file
Q_start_idx = 0
E_idx = 3
w_idx = 4
37

38

39
40
41
42
# constants
sig2hwhm = np.sqrt(2. * np.log(2.))
sig2fwhm = 2.*sig2hwhm

43

Tobias WEBER's avatar
Tobias WEBER committed
44
45
46
47
48
49
50
51
52
#
# normalises events and filters out too low probabilities
#
def filter_events(Q, E, w):
	# normalise intensity/probability
	w /= np.max(w)


	# filter out too low probabilities
53
54
	theeps = 1e-4
	beloweps = lambda d: np.abs(d) <= theeps
Tobias WEBER's avatar
Tobias WEBER committed
55
56
57
58
59
60
61
	nonzero_idx = [i for i in range(len(w)) if not beloweps(w[i])]

	Q = Q[nonzero_idx]
	E = E[nonzero_idx]
	w = w[nonzero_idx]


Tobias WEBER's avatar
Tobias WEBER committed
62
	return [Q, E, w]
Tobias WEBER's avatar
Tobias WEBER committed
63
64
65



66
67
68
#
# loads a list of neutron events in the [ ki_vec, kf_vec, pos_vec, wi, wf ] format
#
69
def load_events_kikf(filename):
70
	dat = np.loadtxt(filename)
71
72
73
74
	ki = dat[:, ki_start_idx:ki_start_idx+3]
	kf = dat[:, kf_start_idx:kf_start_idx+3]
	wi = dat[:, wi_idx]
	wf = dat[:, wf_idx]
75
76
77

	w = wi * wf
	Q = ki - kf
78
	E = tas.k2_to_E * (np.multiply(ki, ki).sum(1) - np.multiply(kf, kf).sum(1))
79

Tobias WEBER's avatar
Tobias WEBER committed
80
	return filter_events(Q, E, w)
81
82
83



84
85
86
87
88
#
# loads a list of neutron events in the [ h, k, l, E, w ] format
#
def load_events_QE(filename):
	dat = np.loadtxt(filename)
Tobias WEBER's avatar
Tobias WEBER committed
89

90
91
92
93
	Q = dat[:, Q_start_idx:Q_start_idx+3]
	E = dat[:, E_idx]
	w = dat[:, w_idx]

Tobias WEBER's avatar
Tobias WEBER committed
94
	return filter_events(Q, E, w)
95
96
97



98
99
100
#
# calculates the covariance matrix of the (Q, E) 4-vectors
#
Tobias WEBER's avatar
Tobias WEBER committed
101
def calc_covar(Q, E, w, Qpara, Qperp):
102
103
104
	# make a [Q, E] 4-vector
	Q4 = np.insert(Q, 3, E, axis=1)

Tobias WEBER's avatar
Tobias WEBER committed
105
	# calculate the mean Q 4-vector
106
	Qmean = [ np.average(Q4[:,i], weights = w) for i in range(4) ]
107
108
	if verbose:
		print("Mean (Q, E) vector in lab system:\n%s\n" % Qmean)
109

Tobias WEBER's avatar
Tobias WEBER committed
110
111
	# get the weighted covariance matrix
	Qcov = np.cov(Q4, rowvar = False, aweights = w, ddof = 0)
112
113
	if verbose:
		print("Covariance matrix in lab system:\n%s\n" % Qcov)
114

Tobias WEBER's avatar
Tobias WEBER committed
115
	# the resolution is the inverse of the covariance
116
	Qres = la.inv(Qcov)
117
118
	if verbose:
		print("Resolution matrix in lab system:\n%s\n" % Qres)
119
120
121


	# create a matrix to transform into the coordinate system with Q along x
Tobias WEBER's avatar
Tobias WEBER committed
122
123
124
125
126
127
128
129
130
	# choose given coordinate system
	if len(Qpara) == 3 and len(Qperp) == 3:
		Qnorm = Qpara / la.norm(Qpara)
		Qside = Qperp / la.norm(Qperp)
		Qup = np.cross(Qnorm, Qside)
	else:
		Qnorm = Qmean[0:3] / la.norm(Qmean[0:3])
		Qup = np.array([0, 1, 0])
		Qside = np.cross(Qup, Qnorm)
131

132
133
	if verbose:
		print("Qpara = %s\nQperp = %s\nQup = %s\n" % (Qnorm, Qside, Qup))
Tobias WEBER's avatar
Tobias WEBER committed
134

Tobias WEBER's avatar
Tobias WEBER committed
135
	# trafo matrix
136
137
138
139
140
141
	T = np.transpose(np.array([
		np.insert(Qnorm, 3, 0),
		np.insert(Qside, 3, 0),
		np.insert(Qup, 3, 0),
		[0, 0, 0, 1] ]))

142
143
	if verbose:
		print("Transformation into (Qpara, Qperp, Qup, E) system:\n%s\n" % T)
144

Tobias WEBER's avatar
Tobias WEBER committed
145
	# transform mean Q vector
146
	Qmean_Q = np.dot(np.transpose(T), Qmean)
147
148
	if verbose:
		print("Mean (Q, E) vector in (Qpara, Qperp, Qup, E) system:\n%s\n" % Qmean_Q)
149

Tobias WEBER's avatar
Tobias WEBER committed
150
	# transform the covariance matrix
151
	Qcov_Q = np.dot(np.transpose(T), np.dot(Qcov, T))
152
153
	if verbose:
		print("Covariance matrix in (Qpara, Qperp, Qup, E) system:\n%s\n" % Qcov_Q)
154

Tobias WEBER's avatar
Tobias WEBER committed
155
	# the resolution is the inverse of the covariance
156
	Qres_Q = la.inv(Qcov_Q)
157
158
	if verbose:
		print("Resolution matrix in (Qpara, Qperp, Qup, E) system:\n%s\n" % Qres_Q)
159

160
	#[ evals, evecs ] = la.eig(Qcov_Q)
161
	#print("Ellipsoid fwhm radii:\n%s\n" % (np.sqrt(np.abs(evals)) * sig2fwhm))
162

163
	# transform all neutron events
164
	Q4_Q = np.array([])
Tobias WEBER's avatar
Tobias WEBER committed
165
	if plot_neutrons:
166
167
168
		Q4_Q = np.dot(Q4, T)
		if not centre_on_Q:
			Q4_Q -= Qmean_Q
169
170


171
	return [Qres_Q, Q4_Q, Qmean_Q]
172
173
174
175



#
176
# calculates the characteristics of a given ellipse by principal axis trafo
177
178
179
#
def descr_ellipse(quadric):
	[ evals, evecs ] = la.eig(quadric)
180
181
182
	#print("Evals: %s" % evals)

	fwhms = 1./np.sqrt(np.abs(evals)) * sig2fwhm
183
184
185
186
187

	angles = np.array([])
	if len(quadric) == 2:
		angles = np.array([ np.arctan2(evecs[1][0], evecs[0][0]) ])

188
	return [fwhms, angles/np.pi*180., evecs]
189

190

191

192
#
193
# projects along one axis of the quadric
194
#
195
def proj_quad(_E, idx):
196
197
198
199
200
201
202
203
204
205
206
	E = np.delete(np.delete(_E, idx, axis=0), idx, axis=1)
	if np.abs(_E[idx, idx]) < 1e-8:
		return E

	v = (_E[idx,:] + _E[:,idx]) * 0.5
	vv = np.outer(v, v) / _E[idx, idx]
	vv = np.delete(np.delete(vv, idx, axis=0), idx, axis=1)

	return E - vv


207
#
208
# describes the ellipsoid by a principal axis trafo and by 2d cuts
209
210
#
def calc_ellipses(Qres_Q):
211
212
213
	# 4d ellipsoid
	[fwhms, angles, rot] = descr_ellipse(Qres_Q)

214
215
	if verbose:
		print("4d resolution ellipsoid diagonal elements fwhm (coherent-elastic scattering) lengths:\n%s\n" \
Tobias WEBER's avatar
Tobias WEBER committed
216
			% (1./np.sqrt(np.abs(np.diag(Qres_Q))) * sig2fwhm))
217
218

		print("4d resolution ellipsoid principal axes fwhm lengths:\n%s\n" % fwhms)
219

220
221
222
		Qres_proj = proj_quad(Qres_Q, 2)
		Qres_proj = proj_quad(Qres_proj, 1)
		Qres_proj = proj_quad(Qres_proj, 0)
Tobias WEBER's avatar
Tobias WEBER committed
223
		print("Incoherent-elastic fwhm width: %.4f meV\n" % (1./np.sqrt(np.abs(Qres_proj[0,0])) * sig2fwhm))
224

225
226
227
228

	# 2d sliced ellipses
	Qres_QxE = np.delete(np.delete(Qres_Q, 2, axis=0), 2, axis=1)
	Qres_QxE = np.delete(np.delete(Qres_QxE, 1, axis=0), 1, axis=1)
229
230
	[fwhms_QxE, angles_QxE, rot_QxE] = descr_ellipse(Qres_QxE)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
231
		print("2d Qx,E sliced ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QxE, angles_QxE[0]))
232
233
234

	Qres_QyE = np.delete(np.delete(Qres_Q, 2, axis=0), 2, axis=1)
	Qres_QyE = np.delete(np.delete(Qres_QyE, 0, axis=0), 0, axis=1)
235
236
	[fwhms_QyE, angles_QyE, rot_QyE] = descr_ellipse(Qres_QyE)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
237
		print("2d Qy,E sliced ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QyE, angles_QyE[0]))
238
239
240

	Qres_QzE = np.delete(np.delete(Qres_Q, 1, axis=0), 1, axis=1)
	Qres_QzE = np.delete(np.delete(Qres_QzE, 0, axis=0), 0, axis=1)
241
242
	[fwhms_QzE, angles_QzE, rot_QzE] = descr_ellipse(Qres_QzE)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
243
		print("2d Qz,E sliced ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QzE, angles_QzE[0]))
244
245
246
247
248

	Qres_QxQy = np.delete(np.delete(Qres_Q, 3, axis=0), 3, axis=1)
	Qres_QxQy = np.delete(np.delete(Qres_QxQy, 2, axis=0), 2, axis=1)
	[fwhms_QxQy, angles_QxQy, rot_QxQy] = descr_ellipse(Qres_QxQy)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
249
		print("2d Qx,Qy sliced ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QxQy, angles_QxQy[0]))
250
251
252
253


	# 2d projected ellipses
	Qres_QxE_proj = np.delete(np.delete(Qres_Q, 2, axis=0), 2, axis=1)
254
	Qres_QxE_proj = proj_quad(Qres_QxE_proj, 1)
255
256
	[fwhms_QxE_proj, angles_QxE_proj, rot_QxE_proj] = descr_ellipse(Qres_QxE_proj)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
257
		print("2d Qx,E projected ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QxE_proj, angles_QxE_proj[0]))
258
259

	Qres_QyE_proj = np.delete(np.delete(Qres_Q, 2, axis=0), 2, axis=1)
260
	Qres_QyE_proj = proj_quad(Qres_QyE_proj, 0)
261
262
	[fwhms_QyE_proj, angles_QyE_proj, rot_QyE_proj] = descr_ellipse(Qres_QyE_proj)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
263
		print("2d Qy,E projected ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QyE_proj, angles_QyE_proj[0]))
264

265
	Qres_QzE_proj = np.delete(np.delete(Qres_Q, 1, axis=0), 1, axis=1)
266
	Qres_QzE_proj = proj_quad(Qres_QzE_proj, 0)
267
268
	[fwhms_QzE_proj, angles_QzE_proj, rot_QzE_proj] = descr_ellipse(Qres_QzE_proj)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
269
		print("2d Qz,E projected ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QzE_proj, angles_QzE_proj[0]))
270

271
	Qres_QxQy_proj = proj_quad(Qres_Q, 3)
272
273
274
	Qres_QxQy_proj = np.delete(np.delete(Qres_QxQy_proj, 2, axis=0), 2, axis=1)
	[fwhms_QxQy_proj, angles_QxQy_proj, rot_QxQy_proj] = descr_ellipse(Qres_QxQy_proj)
	if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
275
		print("2d Qx,Qy projected ellipse fwhm lengths and slope angle:\n%s, %.4f\n" % (fwhms_QxQy_proj, angles_QxQy_proj[0]))
Tobias WEBER's avatar
Tobias WEBER committed
276

277
278
279
280
281
282


	return [fwhms_QxE, rot_QxE, fwhms_QyE, rot_QyE, \
		fwhms_QzE, rot_QzE, fwhms_QxQy, rot_QxQy, \
		fwhms_QxE_proj, rot_QxE_proj, fwhms_QyE_proj, rot_QyE_proj, \
		fwhms_QzE_proj, rot_QzE_proj, fwhms_QxQy_proj, rot_QxQy_proj]
283
284
285
286
287
288



#
# shows the 2d ellipses
#
289
290
291
292
293
294
def plot_ellipses(file, Q4, w, Qmean, \
	fwhms_QxE, rot_QxE, fwhms_QyE, rot_QyE, \
	fwhms_QzE, rot_QzE, fwhms_QxQy, rot_QxQy, \
	fwhms_QxE_proj, rot_QxE_proj, fwhms_QyE_proj, rot_QyE_proj, \
	fwhms_QzE_proj, rot_QzE_proj, fwhms_QxQy_proj, rot_QxQy_proj):

Tobias WEBER's avatar
Tobias WEBER committed
295
	import mpl_toolkits.mplot3d as mplot3d
296
297
	import matplotlib.pyplot as plot

298
299
300
301
	ellfkt = lambda rad, rot, phi, Qmean2d : \
		np.dot(rot, np.array([ rad[0]*np.cos(phi), rad[1]*np.sin(phi) ])) + Qmean2d


Tobias WEBER's avatar
Tobias WEBER committed
302
	# centre plots on zero or mean Q vector ?
303
304
305
306
307
308
309
310
311
312
313
	QxE = np.array([[0], [0]])
	QyE = np.array([[0], [0]])
	QzE = np.array([[0], [0]])
	QxQy = np.array([[0], [0]])

	if centre_on_Q:
		QxE = np.array([[Qmean[0]], [Qmean[3]]])
		QyE = np.array([[Qmean[1]], [Qmean[3]]])
		QzE = np.array([[Qmean[2]], [Qmean[3]]])
		QxQy = np.array([[Qmean[0]], [Qmean[1]]])

314
315

	phi = np.linspace(0, 2.*np.pi, ellipse_points)
316

317
318
319
320
	ell_QxE = ellfkt(fwhms_QxE*0.5, rot_QxE, phi, QxE)
	ell_QyE = ellfkt(fwhms_QyE*0.5, rot_QyE, phi, QyE)
	ell_QzE = ellfkt(fwhms_QzE*0.5, rot_QzE, phi, QzE)
	ell_QxQy = ellfkt(fwhms_QxQy*0.5, rot_QxQy, phi, QxQy)
321

322
323
324
325
	ell_QxE_proj = ellfkt(fwhms_QxE_proj*0.5, rot_QxE_proj, phi, QxE)
	ell_QyE_proj = ellfkt(fwhms_QyE_proj*0.5, rot_QyE_proj, phi, QyE)
	ell_QzE_proj = ellfkt(fwhms_QzE_proj*0.5, rot_QzE_proj, phi, QzE)
	ell_QxQy_proj = ellfkt(fwhms_QxQy_proj*0.5, rot_QxQy_proj, phi, QxQy)
326
327


Tobias WEBER's avatar
Tobias WEBER committed
328
	thesymsize = symsize * w
329
	themarker = "."
Tobias WEBER's avatar
Tobias WEBER committed
330

331

Tobias WEBER's avatar
Tobias WEBER committed
332
	# Qpara, E axis
333
334
335
	fig = plot.figure()
	subplot_QxE = fig.add_subplot(221)
	subplot_QxE.set_xlabel("Qpara (1/A)")
Tobias WEBER's avatar
Tobias WEBER committed
336
	subplot_QxE.set_ylabel("E (meV)")
337
	if len(Q4.shape)==2 and len(Q4)>0 and len(Q4[0])==4:
338
		subplot_QxE.scatter(Q4[:, 0], Q4[:, 3], marker=themarker, s=thesymsize)
339
	subplot_QxE.plot(ell_QxE[0], ell_QxE[1], c="black", linestyle="dashed")
340
	subplot_QxE.plot(ell_QxE_proj[0], ell_QxE_proj[1], c="black", linestyle="solid")
341

Tobias WEBER's avatar
Tobias WEBER committed
342
	# Qperp, E axis
343
344
345
	subplot_QyE = fig.add_subplot(222)
	subplot_QyE.set_xlabel("Qperp (1/A)")
	subplot_QyE.set_ylabel("E (meV)")
346
	if len(Q4.shape)==2 and len(Q4)>0 and len(Q4[0])==4:
347
		subplot_QyE.scatter(Q4[:, 1], Q4[:, 3], marker=themarker, s=thesymsize)
348
	subplot_QyE.plot(ell_QyE[0], ell_QyE[1], c="black", linestyle="dashed")
349
	subplot_QyE.plot(ell_QyE_proj[0], ell_QyE_proj[1], c="black", linestyle="solid")
350

Tobias WEBER's avatar
Tobias WEBER committed
351
	# Qup, E axis
352
353
354
	subplot_QzE = fig.add_subplot(223)
	subplot_QzE.set_xlabel("Qup (1/A)")
	subplot_QzE.set_ylabel("E (meV)")
355
	if len(Q4.shape)==2 and len(Q4)>0 and len(Q4[0])==4:
356
		subplot_QzE.scatter(Q4[:, 2], Q4[:, 3], marker=themarker, s=thesymsize)
357
	subplot_QzE.plot(ell_QzE[0], ell_QzE[1], c="black", linestyle="dashed")
358
	subplot_QzE.plot(ell_QzE_proj[0], ell_QzE_proj[1], c="black", linestyle="solid")
359

Tobias WEBER's avatar
Tobias WEBER committed
360
	# Qpara, Qperp axis
361
362
	subplot_QxQy = fig.add_subplot(224)
	subplot_QxQy.set_xlabel("Qpara (1/A)")
363
	subplot_QxQy.set_ylabel("Qperp (1/A)")
364
	if len(Q4.shape)==2 and len(Q4)>0 and len(Q4[0])==4:
365
		subplot_QxQy.scatter(Q4[:, 0], Q4[:, 1], marker=themarker, s=thesymsize)
366
	subplot_QxQy.plot(ell_QxQy[0], ell_QxQy[1], c="black", linestyle="dashed")
367
	subplot_QxQy.plot(ell_QxQy_proj[0], ell_QxQy_proj[1], c="black", linestyle="solid")
368
	plot.tight_layout()
Tobias WEBER's avatar
Tobias WEBER committed
369

Tobias WEBER's avatar
Tobias WEBER committed
370

Tobias WEBER's avatar
Tobias WEBER committed
371
	# 3d plot
Tobias WEBER's avatar
Tobias WEBER committed
372
373
	fig3d = plot.figure()
	subplot3d = fig3d.add_subplot(111, projection="3d")
Tobias WEBER's avatar
Tobias WEBER committed
374

Tobias WEBER's avatar
Tobias WEBER committed
375
376
377
378
	subplot3d.set_xlabel("Qpara (1/A)")
	subplot3d.set_ylabel("Qperp (1/A)")
	subplot3d.set_zlabel("E (meV)")

379
	subplot3d.scatter(Q4[:,0], Q4[:,1], Q4[:,3], marker=themarker, s=thesymsize)
Tobias WEBER's avatar
Tobias WEBER committed
380
381
	# xE
	subplot3d.plot(ell_QxE[0], ell_QxE[1], zs=0., zdir="y", c="black", linestyle="dashed")
382
	subplot3d.plot(ell_QxE_proj[0], ell_QxE_proj[1], zs=0., zdir="y", c="black", linestyle="solid")
Tobias WEBER's avatar
Tobias WEBER committed
383
384
	# yE
	subplot3d.plot(ell_QyE[0], ell_QyE[1], zs=0., zdir="x", c="black", linestyle="dashed")
385
	subplot3d.plot(ell_QyE_proj[0], ell_QyE_proj[1], zs=0., zdir="x", c="black", linestyle="solid")
Tobias WEBER's avatar
Tobias WEBER committed
386
387
	# xy
	subplot3d.plot(ell_QxQy[0], ell_QxQy[1], zs=0., zdir="z", c="black", linestyle="dashed")
388
	subplot3d.plot(ell_QxQy_proj[0], ell_QxQy_proj[1], zs=0., zdir="z", c="black", linestyle="solid")
Tobias WEBER's avatar
Tobias WEBER committed
389

Tobias WEBER's avatar
Tobias WEBER committed
390

Tobias WEBER's avatar
Tobias WEBER committed
391
	if file != "":
Tobias WEBER's avatar
Tobias WEBER committed
392
393
394
		splitext = os.path.splitext(file)
		file3d = splitext[0] + "_3d" + splitext[1]

Tobias WEBER's avatar
Tobias WEBER committed
395
		if verbose:
Tobias WEBER's avatar
Tobias WEBER committed
396
397
			print("Saving 2d plot to \"%s\"." % file)
			print("Saving 3d plot to \"%s\"." % file3d)
398
399
		fig.savefig(file, dpi=dpi)
		fig3d.savefig(file3d, dpi=dpi)
Tobias WEBER's avatar
Tobias WEBER committed
400
401
402

	if plot_results:
		plot.show()
403
404
405



406
407
408
409
410
411
412
413
414
415
416
417
418
#
# checks versions of needed packages
#
def check_versions():
	npver = np.version.version.split(".")
	if int(npver[0]) >= 2:
		return
	if int(npver[0]) < 1 or int(npver[1]) < 10:
		print("Numpy version >= 1.10 is required, but installed version is %s." % np.version.version)
		exit(-1)



419
420
421
#
# main
#
422
if __name__ == "__main__":
Tobias WEBER's avatar
Tobias WEBER committed
423
	print("This is a covariance matrix calculator using neutron events, written by T. Weber <tweber@ill.fr>, 30 March 2019.\n")
424
425
	check_versions()

Tobias WEBER's avatar
Tobias WEBER committed
426
427
428
429
430
431
	import argparse as arg

	args = arg.ArgumentParser(description="Calculates the covariance matrix of neutron scattering events.")
	args.add_argument("file", type=str, help="input file")
	args.add_argument("-s", "--save", default="", type=str, nargs="?", help="save plot to file")
	args.add_argument("--ellipse", default=ellipse_points, type=int, nargs="?", help="number of points to draw ellipses")
432
433
434
435
436
437
438
439
	args.add_argument("--ki", default=ki_start_idx, type=int, nargs="?", help="index of ki vector's first column in kikf file")
	args.add_argument("--kf", default=kf_start_idx, type=int, nargs="?", help="index of kf vector's first column in kikf file")
	args.add_argument("--wi", default=wi_idx, type=int, nargs="?", help="index of ki weight factor column in kikf file")
	args.add_argument("--wf", default=wf_idx, type=int, nargs="?", help="index of kf weight factor column in kikf file")
	args.add_argument("--w", default=w_idx, type=int, nargs="?", help="index of neutron weight factor column in QE file")
	args.add_argument("--Q", default=Q_start_idx, type=int, nargs="?", help="index of Q vector's first column in QE file")
	args.add_argument("--E", default=E_idx, type=int, nargs="?", help="index of E column in QE file")
	args.add_argument("--kikf", action="store_true", help="use the kikf file type")
440
	args.add_argument("--centreonQ", action="store_true", help="centre plots on mean Q")
Tobias WEBER's avatar
Tobias WEBER committed
441
	args.add_argument("--noverbose", action="store_true", help="don't show console logs")
Tobias WEBER's avatar
Tobias WEBER committed
442
443
444
445
446
447
448
449
450
	args.add_argument("--noplot", action="store_true", help="don't show any plot windows")
	args.add_argument("--noneutrons", action="store_true", help="don't show neutron events in plots")
	args.add_argument("--symsize", default=symsize, type=float, nargs="?", help="size of the symbols in plots")
	args.add_argument("--ax", default=None, type=float, nargs="?", help="x component of first orientation vector")
	args.add_argument("--ay", default=None, type=float, nargs="?", help="y component of first orientation vector")
	args.add_argument("--az", default=None, type=float, nargs="?", help="z component of first orientation vector")
	args.add_argument("--bx", default=None, type=float, nargs="?", help="x component of second orientation vector")
	args.add_argument("--by", default=None, type=float, nargs="?", help="y component of second orientation vector")
	args.add_argument("--bz", default=None, type=float, nargs="?", help="z component of second orientation vector")
451
452
453
454
455
456
	args.add_argument("--a", "--as", "-a", default=None, type=float, nargs="?", help="lattice constant a (only needed in case data is in rlu)")
	args.add_argument("--b", "--bs", "-b", default=None, type=float, nargs="?", help="lattice constant b (only needed in case data is in rlu)")
	args.add_argument("--c", "--cs", "-c", default=None, type=float, nargs="?", help="lattice constant c (only needed in case data is in rlu)")
	args.add_argument("--aa", "--alpha", default=90., type=float, nargs="?", help="lattice angle alpha (only needed in case data is in rlu)")
	args.add_argument("--bb", "--beta", default=90., type=float, nargs="?", help="lattice angle beta (only needed in case data is in rlu)")
	args.add_argument("--cc", "--gamma", default=90., type=float, nargs="?", help="lattice angle gamma (only needed in case data is in rlu)")
457
	args.add_argument("--dpi", default=dpi, type=int, nargs="?", help="DPI of output plot file")
Tobias WEBER's avatar
Tobias WEBER committed
458
	argv = args.parse_args()
459

Tobias WEBER's avatar
Tobias WEBER committed
460
461
462
	verbose = (argv.noverbose==False)
	plot_results = (argv.noplot==False)
	plot_neutrons = (argv.noneutrons==False)
463
	centre_on_Q = argv.centreonQ
464
	use_kikf_file = argv.kikf
465
	dpi = argv.dpi
466

467
468
469
470
471
472
473
474
475
476
477
	B = []
	if argv.a!=None and argv.b!=None and argv.c!=None and argv.aa!=None and argv.bb!=None and argv.cc!=None:
		lattice = np.array([ argv.a, argv.b, argv.c,  ])
		angles = np.array([ argv.aa, argv.bb, argv.cc ]) / 180.*np.pi

		B = tas.get_B(lattice, angles)

		if verbose:
			print("Crystal B matrix:\n%s\n" % B)


Tobias WEBER's avatar
Tobias WEBER committed
478
479
480
481
482
483
484
	infile = argv.file
	outfile = argv.save
	ellipse_points = argv.ellipse
	ki_start_idx = argv.ki
	kf_start_idx = argv.kf
	wi_idx = argv.wi
	wf_idx = argv.wf
Tobias WEBER's avatar
Tobias WEBER committed
485
	symsize = argv.symsize
Tobias WEBER's avatar
Tobias WEBER committed
486
487
	avec = [ argv.az, argv.ay, argv.az ]
	bvec = [ argv.bx, argv.by, argv.bz ]
Tobias WEBER's avatar
Tobias WEBER committed
488

489
490
491
492
493
	if use_kikf_file:
		[Q, E, w] = load_events_kikf(infile)
	else:
		[Q, E, w] = load_events_QE(infile)

494
495
		# convert rlu to 1/A 
		if len(B) != 0:
Tobias WEBER's avatar
Tobias WEBER committed
496
			Q = np.dot(Q, np.transpose(B))
497

Tobias WEBER's avatar
Tobias WEBER committed
498
499
500
	if avec[0]!=None and avec[1]!=None and avec[2]!=None and bvec[0]!=None and bvec[1]!=None and bvec[2]!=None:
		Qpara = np.array(avec)
		Qperp = np.array(bvec)
501
502
503
504
505

		# convert rlu to 1/A 
		if len(B) != 0:
			Qpara = np.dot(B, Qpara)
			Qperp = np.dot(B, Qperp)
Tobias WEBER's avatar
Tobias WEBER committed
506
507
508
	else:
		Qpara = np.array([])
		Qperp = np.array([])
509

Tobias WEBER's avatar
Tobias WEBER committed
510
	[Qres, Q4, Qmean] = calc_covar(Q, E, w, Qpara, Qperp)
511
512
513
514
515
	[fwhms_QxE, rot_QxE, fwhms_QyE, rot_QyE, \
	fwhms_QzE, rot_QzE, fwhms_QxQy, rot_QxQy, \
	fwhms_QxE_proj, rot_QxE_proj, fwhms_QyE_proj, rot_QyE_proj, \
	fwhms_QzE_proj, rot_QzE_proj, fwhms_QxQy_proj, rot_QxQy_proj] \
		= calc_ellipses(Qres)
516

Tobias WEBER's avatar
Tobias WEBER committed
517
	if plot_results or outfile!="":
518
519
520
521
522
		plot_ellipses(outfile, Q4, w, Qmean, \
			fwhms_QxE, rot_QxE, fwhms_QyE, rot_QyE, \
			fwhms_QzE, rot_QzE, fwhms_QxQy, rot_QxQy, \
			fwhms_QxE_proj, rot_QxE_proj, fwhms_QyE_proj, rot_QyE_proj, \
			fwhms_QzE_proj, rot_QzE_proj, fwhms_QxQy_proj, rot_QxQy_proj)