cov.py 16.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
#
Tobias WEBER's avatar
Tobias WEBER committed
8
# for a good explanation of the covariance matrix method, see, e.g., (Arens 2015), pp. 795 and 1372.
Tobias WEBER's avatar
Tobias WEBER committed
9
#
10

Tobias WEBER's avatar
Tobias WEBER committed
11
import os
Tobias WEBER's avatar
Tobias WEBER committed
12
import tas
13

14
15
16
17
18
19
20
try:
	import numpy as np
	import numpy.linalg as la
except ImportError:
	print("Numpy could not be imported!")
	exit(-1)

21

Tobias WEBER's avatar
Tobias WEBER committed
22
23
24
25
26
27
try:
	np.set_printoptions(
		precision = 4,
		floatmode = "fixed")
except TypeError:
	print("Warning: Numpy print options could not be set.")
Tobias WEBER's avatar
Tobias WEBER committed
28

29

Tobias WEBER's avatar
Tobias WEBER committed
30
31
32
33
34
35
36
37
options = {
	"verbose" : True,		# console outputs
	"plot_results" : True,		# show plot window
	"plot_neutrons" : True,		# also plot neutron events
	"centre_on_Q" : False,		# centre plots on Q or zero?
	"ellipse_points" : 128,		# number of points to draw ellipses
	"symsize" : 32.,
	"dpi" : 600,
38

Tobias WEBER's avatar
Tobias WEBER committed
39
40
41
42
43
44
45
46
47
48
49
	# column indices in ki,kf files
	"ki_start_idx" : 0,		# start index of ki 3-vector
	"kf_start_idx" : 3,		# start index of kf 3-vector
	"wi_idx" : 9,			# start index of ki weight factor
	"wf_idx" : 10,			# start index of kf weight factor

	# column indices in Q,E files
	"Q_start_idx" : 0,
	"E_idx" : 3,
	"w_idx" : 4,
}
50

51

52
53
54
55
# constants
sig2hwhm = np.sqrt(2. * np.log(2.))
sig2fwhm = 2.*sig2hwhm

56

Tobias WEBER's avatar
Tobias WEBER committed
57
58
59
60
61
62
63
64
65
#
# 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
66
67
	theeps = 1e-4
	beloweps = lambda d: np.abs(d) <= theeps
Tobias WEBER's avatar
Tobias WEBER committed
68
69
70
71
72
73
74
	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
75
	return [Q, E, w]
Tobias WEBER's avatar
Tobias WEBER committed
76
77
78



79
80
81
#
# loads a list of neutron events in the [ ki_vec, kf_vec, pos_vec, wi, wf ] format
#
82
def load_events_kikf(filename):
83
	dat = np.loadtxt(filename)
Tobias WEBER's avatar
Tobias WEBER committed
84
85
86
87
	ki = dat[:, options["ki_start_idx"]:options["ki_start_idx"]+3]
	kf = dat[:, options["kf_start_idx"]:options["kf_start_idx"]+3]
	wi = dat[:, options["wi_idx"]]
	wf = dat[:, options["wf_idx"]]
88
89
90

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

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



97
98
99
100
101
#
# 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
102

Tobias WEBER's avatar
Tobias WEBER committed
103
104
105
	Q = dat[:, options["Q_start_idx"]:options["Q_start_idx"]+3]
	E = dat[:, options["E_idx"]]
	w = dat[:, options["w_idx"]]
106

Tobias WEBER's avatar
Tobias WEBER committed
107
	return filter_events(Q, E, w)
108
109
110



111
112
113
#
# calculates the covariance matrix of the (Q, E) 4-vectors
#
Tobias WEBER's avatar
Tobias WEBER committed
114
def calc_covar(Q, E, w, Qpara, Qperp):
115
116
117
	# make a [Q, E] 4-vector
	Q4 = np.insert(Q, 3, E, axis=1)

Tobias WEBER's avatar
Tobias WEBER committed
118
	# calculate the mean Q 4-vector
119
	Qmean = [ np.average(Q4[:,i], weights = w) for i in range(4) ]
Tobias WEBER's avatar
Tobias WEBER committed
120
	if options["verbose"]:
121
		print("Mean (Q, E) vector in lab system:\n%s\n" % Qmean)
122

Tobias WEBER's avatar
Tobias WEBER committed
123
	# get the weighted covariance matrix
Tobias WEBER's avatar
Tobias WEBER committed
124
125
	Qcov = np.cov(Q4, rowvar = False, aweights = np.abs(w), ddof = 0)
	if options["verbose"]:
126
		print("Covariance matrix in lab system:\n%s\n" % Qcov)
127

Tobias WEBER's avatar
Tobias WEBER committed
128
	# the resolution is the inverse of the covariance
129
	Qres = la.inv(Qcov)
Tobias WEBER's avatar
Tobias WEBER committed
130
	if options["verbose"]:
131
		print("Resolution matrix in lab system:\n%s\n" % Qres)
132
133
134


	# create a matrix to transform into the coordinate system with Q along x
Tobias WEBER's avatar
Tobias WEBER committed
135
136
137
138
139
140
141
142
143
	# 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)
144

Tobias WEBER's avatar
Tobias WEBER committed
145
	if options["verbose"]:
146
		print("Qpara = %s\nQperp = %s\nQup = %s\n" % (Qnorm, Qside, Qup))
Tobias WEBER's avatar
Tobias WEBER committed
147

Tobias WEBER's avatar
Tobias WEBER committed
148
	# trafo matrix
149
150
151
152
153
154
	T = np.transpose(np.array([
		np.insert(Qnorm, 3, 0),
		np.insert(Qside, 3, 0),
		np.insert(Qup, 3, 0),
		[0, 0, 0, 1] ]))

Tobias WEBER's avatar
Tobias WEBER committed
155
	if options["verbose"]:
156
		print("Transformation into (Qpara, Qperp, Qup, E) system:\n%s\n" % T)
157

Tobias WEBER's avatar
Tobias WEBER committed
158
	# transform mean Q vector
159
	Qmean_Q = np.dot(np.transpose(T), Qmean)
Tobias WEBER's avatar
Tobias WEBER committed
160
	if options["verbose"]:
161
		print("Mean (Q, E) vector in (Qpara, Qperp, Qup, E) system:\n%s\n" % Qmean_Q)
162

Tobias WEBER's avatar
Tobias WEBER committed
163
	# transform the covariance matrix
164
	Qcov_Q = np.dot(np.transpose(T), np.dot(Qcov, T))
Tobias WEBER's avatar
Tobias WEBER committed
165
	if options["verbose"]:
166
		print("Covariance matrix in (Qpara, Qperp, Qup, E) system:\n%s\n" % Qcov_Q)
167

Tobias WEBER's avatar
Tobias WEBER committed
168
	# the resolution is the inverse of the covariance
169
	Qres_Q = la.inv(Qcov_Q)
Tobias WEBER's avatar
Tobias WEBER committed
170
	if options["verbose"]:
171
		print("Resolution matrix in (Qpara, Qperp, Qup, E) system:\n%s\n" % Qres_Q)
172

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

176
	# transform all neutron events
177
	Q4_Q = np.array([])
Tobias WEBER's avatar
Tobias WEBER committed
178
	if options["plot_neutrons"]:
179
		Q4_Q = np.dot(Q4, T)
Tobias WEBER's avatar
Tobias WEBER committed
180
		if not options["centre_on_Q"]:
181
			Q4_Q -= Qmean_Q
182
183


184
	return [Qres_Q, Q4_Q, Qmean_Q]
185
186
187
188



#
189
# calculates the characteristics of a given ellipse by principal axis trafo
190
191
192
#
def descr_ellipse(quadric):
	[ evals, evecs ] = la.eig(quadric)
193
194
195
	#print("Evals: %s" % evals)

	fwhms = 1./np.sqrt(np.abs(evals)) * sig2fwhm
196
197
198
199
200

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

201
	return [fwhms, angles/np.pi*180., evecs]
202

203

204

205
#
206
# projects along one axis of the quadric
207
#
208
def proj_quad(_E, idx):
209
210
211
212
213
214
215
216
217
218
219
	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


220
#
221
# describes the ellipsoid by a principal axis trafo and by 2d cuts
222
223
#
def calc_ellipses(Qres_Q):
224
225
226
	# 4d ellipsoid
	[fwhms, angles, rot] = descr_ellipse(Qres_Q)

Tobias WEBER's avatar
Tobias WEBER committed
227
228
229
	axes_to_delete = [ [2, 1], [2, 0], [1,0], [3,2] ]
	slice_first = [ True, True, True, False ]
	results = []
230

Tobias WEBER's avatar
Tobias WEBER committed
231
232
233
234
235
	for ellidx in range(len(axes_to_delete)):
		# sliced 2d ellipse
		Qres = np.delete(np.delete(Qres_Q, axes_to_delete[ellidx][0], axis=0), axes_to_delete[ellidx][0], axis=1)
		Qres = np.delete(np.delete(Qres, axes_to_delete[ellidx][1], axis=0), axes_to_delete[ellidx][1], axis=1)
		[fwhms, angles, rot] = descr_ellipse(Qres)
236

Tobias WEBER's avatar
Tobias WEBER committed
237
238
239
240
241
242
243
244
		# projected 2d ellipse
		if slice_first[ellidx]:
			Qres_proj = np.delete(np.delete(Qres_Q, axes_to_delete[ellidx][0], axis=0), axes_to_delete[ellidx][0], axis=1)
			Qres_proj = proj_quad(Qres_proj, axes_to_delete[ellidx][1])
		else:
			Qres_proj = proj_quad(Qres_Q, axes_to_delete[ellidx][0])
			Qres_proj = np.delete(np.delete(Qres_proj, axes_to_delete[ellidx][1], axis=0), axes_to_delete[ellidx][1], axis=1)
		[fwhms_proj, angles_proj, rot_proj] = descr_ellipse(Qres_proj)
245

Tobias WEBER's avatar
Tobias WEBER committed
246
247
		results.append({ "fwhms" : fwhms, "angles" : angles, "rot" : rot,
			"fwhms_proj" : fwhms_proj, "angles_proj" : angles_proj, "rot_proj" : rot_proj })
248

249

Tobias WEBER's avatar
Tobias WEBER committed
250
251
252
	if options["verbose"]:
		print("4d resolution ellipsoid diagonal elements fwhm (coherent-elastic scattering) lengths:\n%s\n" \
			% (1./np.sqrt(np.abs(np.diag(Qres_Q))) * sig2fwhm))
253

254
		print("4d resolution ellipsoid principal axes fwhm: %s" % fwhms)
255

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

259
260
261
262
263
264
265
266
267
		print("Qx,E sliced ellipse fwhm and slope angle: %s, %.4f" % (results[0]["fwhms"], results[0]["angles"][0]))
		print("Qy,E sliced ellipse fwhm and slope angle: %s, %.4f" % (results[1]["fwhms"], results[1]["angles"][0]))
		print("Qz,E sliced ellipse fwhm and slope angle: %s, %.4f" % (results[2]["fwhms"], results[2]["angles"][0]))
		print("Qx,Qy sliced ellipse fwhm and slope angle: %s, %.4f" % (results[3]["fwhms"], results[3]["angles"][0]))
		print()
		print("Qx,E projected ellipse fwhm and slope angle: %s, %.4f" % (results[0]["fwhms_proj"], results[0]["angles_proj"][0]))
		print("Qy,E projected ellipse fwhm and slope angle: %s, %.4f" % (results[1]["fwhms_proj"], results[1]["angles_proj"][0]))
		print("Qz,E projected ellipse fwhm and slope angle: %s, %.4f" % (results[2]["fwhms_proj"], results[2]["angles_proj"][0]))
		print("Qx,Qy projected ellipse fwhm and slope angle: %s, %.4f" % (results[3]["fwhms_proj"], results[3]["angles_proj"][0]))
Tobias WEBER's avatar
Tobias WEBER committed
268

269

Tobias WEBER's avatar
Tobias WEBER committed
270
	return results
271

272
273
274
275
276
277



#
# shows the 2d ellipses
#
Tobias WEBER's avatar
Tobias WEBER committed
278
def plot_ellipses(file, Q4, w, Qmean, ellis, use_tex=False):
279
280
	try:
		import mpl_toolkits.mplot3d as mplot3d
Tobias WEBER's avatar
Tobias WEBER committed
281
		import matplotlib
282
283
284
285
		import matplotlib.pyplot as plot
	except ImportError:
		print("Matplotlib could not be imported!")
		exit(-1)
286

Tobias WEBER's avatar
Tobias WEBER committed
287
	matplotlib.rc("text", usetex=use_tex)
Tobias WEBER's avatar
Tobias WEBER committed
288
289
290
291
292

	thesymsize = options["symsize"] * w
	themarker = "."


293
294
295
296
	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
297
298
	# 2d plots
	fig = plot.figure()
299

Tobias WEBER's avatar
Tobias WEBER committed
300
301
302
	num_ellis = len(ellis)
	coord_axes = [[0,3], [1,3], [2,3], [0,1]]
	coord_names = ["Qpara (1/A)", "Qperp (1/A)", "Qup (1/A)", "E (meV)"]
303

Tobias WEBER's avatar
Tobias WEBER committed
304
	if use_tex:
Tobias WEBER's avatar
Tobias WEBER committed
305
306
307
		coord_names[0] = "$Q_{\parallel}$ (\AA$^{-1}$)"
		coord_names[1] = "$Q_{\perp}$ (\AA$^{-1}$)"
		coord_names[2] = "$Q_{up}$ (\AA$^{-1}$)"
Tobias WEBER's avatar
Tobias WEBER committed
308
309


Tobias WEBER's avatar
Tobias WEBER committed
310
311
312
313
	ellplots = []
	for ellidx in range(num_ellis):
		# centre plots on zero or mean Q vector ?
		QxE = np.array([[0], [0]])
314

Tobias WEBER's avatar
Tobias WEBER committed
315
316
		if options["centre_on_Q"]:
			QxE = np.array([[Qmean[coord_axes[ellidx][0]]], [Qmean[coord_axes[ellidx][0]]]])
317

318

Tobias WEBER's avatar
Tobias WEBER committed
319
		phi = np.linspace(0, 2.*np.pi, options["ellipse_points"])
320

Tobias WEBER's avatar
Tobias WEBER committed
321
322
323
		ell_QxE = ellfkt(ellis[ellidx]["fwhms"]*0.5, ellis[ellidx]["rot"], phi, QxE)
		ell_QxE_proj = ellfkt(ellis[ellidx]["fwhms_proj"]*0.5, ellis[ellidx]["rot_proj"], phi, QxE)
		ellplots.append({"sliced":ell_QxE, "proj":ell_QxE_proj})
324

Tobias WEBER's avatar
Tobias WEBER committed
325

Tobias WEBER's avatar
Tobias WEBER committed
326
327
328
329
330
331
332
		subplot_QxE = fig.add_subplot(221 + ellidx)
		subplot_QxE.set_xlabel(coord_names[coord_axes[ellidx][0]])
		subplot_QxE.set_ylabel(coord_names[coord_axes[ellidx][1]])
		if len(Q4.shape)==2 and len(Q4)>0 and len(Q4[0])==4:
			subplot_QxE.scatter(Q4[:, coord_axes[ellidx][0]], Q4[:, coord_axes[ellidx][1]], marker=themarker, s=thesymsize)
		subplot_QxE.plot(ell_QxE[0], ell_QxE[1], c="black", linestyle="dashed")
		subplot_QxE.plot(ell_QxE_proj[0], ell_QxE_proj[1], c="black", linestyle="solid")
333

334
	plot.tight_layout()
Tobias WEBER's avatar
Tobias WEBER committed
335

Tobias WEBER's avatar
Tobias WEBER committed
336

Tobias WEBER's avatar
Tobias WEBER committed
337
	# 3d plot
Tobias WEBER's avatar
Tobias WEBER committed
338
339
	fig3d = plot.figure()
	subplot3d = fig3d.add_subplot(111, projection="3d")
Tobias WEBER's avatar
Tobias WEBER committed
340

Tobias WEBER's avatar
Tobias WEBER committed
341
342
343
	subplot3d.set_xlabel(coord_names[0])
	subplot3d.set_ylabel(coord_names[1])
	subplot3d.set_zlabel(coord_names[3])
Tobias WEBER's avatar
Tobias WEBER committed
344

345
	subplot3d.scatter(Q4[:,0], Q4[:,1], Q4[:,3], marker=themarker, s=thesymsize)
Tobias WEBER's avatar
Tobias WEBER committed
346
	# xE
Tobias WEBER's avatar
Tobias WEBER committed
347
348
	subplot3d.plot(ellplots[0]["sliced"][0], ellplots[0]["sliced"][1], zs=0., zdir="y", c="black", linestyle="dashed")
	subplot3d.plot(ellplots[0]["proj"][0], ellplots[0]["proj"][1], zs=0., zdir="y", c="black", linestyle="solid")
Tobias WEBER's avatar
Tobias WEBER committed
349
	# yE
Tobias WEBER's avatar
Tobias WEBER committed
350
351
	subplot3d.plot(ellplots[1]["sliced"][0], ellplots[1]["sliced"][1], zs=0., zdir="x", c="black", linestyle="dashed")
	subplot3d.plot(ellplots[1]["proj"][0], ellplots[1]["proj"][1], zs=0., zdir="x", c="black", linestyle="solid")
Tobias WEBER's avatar
Tobias WEBER committed
352
	# xy
Tobias WEBER's avatar
Tobias WEBER committed
353
354
	subplot3d.plot(ellplots[3]["sliced"][0], ellplots[3]["sliced"][1], zs=0., zdir="z", c="black", linestyle="dashed")
	subplot3d.plot(ellplots[3]["proj"][0], ellplots[3]["proj"][1], zs=0., zdir="z", c="black", linestyle="solid")
Tobias WEBER's avatar
Tobias WEBER committed
355

Tobias WEBER's avatar
Tobias WEBER committed
356

Tobias WEBER's avatar
Tobias WEBER committed
357
	if file != "":
Tobias WEBER's avatar
Tobias WEBER committed
358
359
360
		splitext = os.path.splitext(file)
		file3d = splitext[0] + "_3d" + splitext[1]

Tobias WEBER's avatar
Tobias WEBER committed
361
		if options["verbose"]:
Tobias WEBER's avatar
Tobias WEBER committed
362
363
			print("Saving 2d plot to \"%s\"." % file)
			print("Saving 3d plot to \"%s\"." % file3d)
Tobias WEBER's avatar
Tobias WEBER committed
364
365
		fig.savefig(file, dpi=options["dpi"])
		fig3d.savefig(file3d, dpi=options["dpi"])
Tobias WEBER's avatar
Tobias WEBER committed
366

Tobias WEBER's avatar
Tobias WEBER committed
367
	if options["plot_results"]:
Tobias WEBER's avatar
Tobias WEBER committed
368
		plot.show()
369
370
371



372
373
374
375
376
377
378
379
380
381
382
383
384
#
# 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)



385
#
Tobias WEBER's avatar
Tobias WEBER committed
386
# entry point
387
#
Tobias WEBER's avatar
Tobias WEBER committed
388
389
def run_cov():
	print("This is a covariance matrix calculator using neutron events,\n\twritten by T. Weber <tweber@ill.fr>, 30 March 2019.\n")
390
391
	check_versions()

392
393
394
395
396
	try:
		import argparse as arg
	except ImportError:
		print("Argparse could not be imported!")
		exit(-1)
Tobias WEBER's avatar
Tobias WEBER committed
397
398
399
400

	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")
Tobias WEBER's avatar
Tobias WEBER committed
401
402
403
404
405
406
407
408
	args.add_argument("--ellipse", default=options["ellipse_points"], type=int, nargs="?", help="number of points to draw ellipses")
	args.add_argument("--ki", default=options["ki_start_idx"], type=int, nargs="?", help="index of ki vector's first column in kikf file")
	args.add_argument("--kf", default=options["kf_start_idx"], type=int, nargs="?", help="index of kf vector's first column in kikf file")
	args.add_argument("--wi", default=options["wi_idx"], type=int, nargs="?", help="index of ki weight factor column in kikf file")
	args.add_argument("--wf", default=options["wf_idx"], type=int, nargs="?", help="index of kf weight factor column in kikf file")
	args.add_argument("--w", default=options["w_idx"], type=int, nargs="?", help="index of neutron weight factor column in QE file")
	args.add_argument("--Q", default=options["Q_start_idx"], type=int, nargs="?", help="index of Q vector's first column in QE file")
	args.add_argument("--E", default=options["E_idx"], type=int, nargs="?", help="index of E column in QE file")
409
	args.add_argument("--QEfile", action="store_true", help="use the QE file type")
410
	args.add_argument("--centreonQ", action="store_true", help="centre plots on mean Q")
Tobias WEBER's avatar
Tobias WEBER committed
411
	args.add_argument("--noverbose", action="store_true", help="don't show console logs")
Tobias WEBER's avatar
Tobias WEBER committed
412
413
	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")
Tobias WEBER's avatar
Tobias WEBER committed
414
	args.add_argument("--symsize", default=options["symsize"], type=float, nargs="?", help="size of the symbols in plots")
Tobias WEBER's avatar
Tobias WEBER committed
415
416
417
418
419
420
	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")
421
422
423
424
425
426
	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)")
Tobias WEBER's avatar
Tobias WEBER committed
427
	args.add_argument("--dpi", default=options["dpi"], type=int, nargs="?", help="DPI of output plot file")
Tobias WEBER's avatar
Tobias WEBER committed
428
	argv = args.parse_args()
429

Tobias WEBER's avatar
Tobias WEBER committed
430
431
432
433
434
	options["verbose"] = (argv.noverbose==False)
	options["plot_results"] = (argv.noplot==False)
	options["plot_neutrons"] = (argv.noneutrons==False)
	options["centre_on_Q"] = argv.centreonQ
	options["dpi"] = argv.dpi
435

436
437
438
439
440
441
442
	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)

Tobias WEBER's avatar
Tobias WEBER committed
443
		if options["verbose"]:
444
445
446
			print("Crystal B matrix:\n%s\n" % B)


Tobias WEBER's avatar
Tobias WEBER committed
447
448
	infile = argv.file
	outfile = argv.save
Tobias WEBER's avatar
Tobias WEBER committed
449
450
451
452
453
454
	options["ellipse_points"] = argv.ellipse
	options["ki_start_idx"] = argv.ki
	options["kf_start_idx"] = argv.kf
	options["wi_idx"] = argv.wi
	options["wf_idx"] = argv.wf
	options["symsize"] = argv.symsize
Tobias WEBER's avatar
Tobias WEBER committed
455
456
	avec = [ argv.az, argv.ay, argv.az ]
	bvec = [ argv.bx, argv.by, argv.bz ]
Tobias WEBER's avatar
Tobias WEBER committed
457

458

459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
	try:
		# input file is in h k l E w format?
		if argv.QEfile:
			[Q, E, w] = load_events_QE(infile)
			# convert rlu to 1/A
			if len(B) != 0:
				Q = np.dot(Q, np.transpose(B))
		# input file is in the kix kiy kiz kfx kfy kfz wi wf format?
		else:
			[Q, E, w] = load_events_kikf(infile)
	except OSError:
		print("Could not load input file %s." % infile)
		exit(-1)
	except NameError:
		print("Error processing input file %s." % infile)
		exit(-1)

476

Tobias WEBER's avatar
Tobias WEBER committed
477
478
479
	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)
480

Tobias WEBER's avatar
Tobias WEBER committed
481
		# convert rlu to 1/A
482
483
484
		if len(B) != 0:
			Qpara = np.dot(B, Qpara)
			Qperp = np.dot(B, Qperp)
Tobias WEBER's avatar
Tobias WEBER committed
485
486
487
	else:
		Qpara = np.array([])
		Qperp = np.array([])
488

Tobias WEBER's avatar
Tobias WEBER committed
489
	[Qres, Q4, Qmean] = calc_covar(Q, E, w, Qpara, Qperp)
Tobias WEBER's avatar
Tobias WEBER committed
490
491
492
493
494
	calcedellis = calc_ellipses(Qres)

	if options["plot_results"] or outfile!="":
		plot_ellipses(outfile, Q4, w, Qmean, calcedellis)

Tobias WEBER's avatar
Tobias WEBER committed
495
496
497
498
499
500
501


#
# main
#
if __name__ == "__main__":
	run_cov()
Tobias WEBER's avatar
test    
Tobias WEBER committed
502
503
504
505
506
507
508
509
510
511
512

	# test
	#calc_covar([
	#	[1., 2., 3.],
	#	[0.8, 2.2, 3.1],
	#	[1.2, 2.5, 2.9],
	#	[0.9, 1.9, 3.7],
	#	[1.2, 1.9, 4.2]],
	#
	#	[4., 3.5, 4.1, 4.1, 3.9],
	#	[1., 1., 0.5, 0.5, 0.5], [], [])