ArepoPostProcessing/annikaEllipsoid.py at master · lakew68/ArepoPostProcessing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
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
import numpy as np
import scipy.linalg as linalg

def calc_inertia(arr_in, axrat_in):
	"""calculate the modified moment of inertia tensor and get its
	eigenvalues and eigenvalues"""
	tensor=np.zeros([3,3])
	# given initial values for primary axes, compute ellipsoidal
	# radius of each particle
	rp2=(arr_in[:,0]**2 + arr_in[:,1]**2/axrat_in[0]**2 + arr_in[:,2]**2/axrat_in[1]**2)
	idxx = np.where(rp2!=0.)[0]
	# construct the moment of inerial tensor to be diagonalized:
	for i in range(3):
		for j in range(3): tensor[i,j]=(arr_in[idxx,i]*arr_in[idxx,j]/rp2[idxx]).sum();
	evecs=linalg.eig(tensor)
	# return a tuple with eigenvalues and eigenvectors;
	# eigenvalues=evecs[0], eigenvectors=evecs[1]
	return evecs

def axis(arr_in, rad, shell=False, axes_out=False, fix_volume=True, quiet=False):
	"""Compute axis ratios iteratively.

	WORK IN PROGRESS -- converting from gadget_profile.pro.
	Needs to be checked.

	May want to add capability to compute axis ratios at multiple
	different radii in one function call (w/ a loop?)

	See Dubinski & Carlberg (1991) for details

	INPUTS:
	arr_in: array of particle positions, assumed to be centered
	rad: radius to compute axis ratios.  If computing axis ratios in
	shells, rad should be a two-element list / array, with rad[0] =
	inner radius of shell and rad[1] = outer radius of shell.
	Otherwise, rad should be a real number equal to the radius within
	which to compute the axis ratios


	OPTIONAL INPUTS:
	shell=False: by default, compute cumulative axis ratio
	(i.e. for all particles within radius r).  If shell=True, compute
	axis ratio in an ellipsoidal shell instead.
	axes_out=False:  if True, also return principal axes (in ascending order)
	fix_volume=True: keep the volume of the ellipsoid constant while iterating.
	If false, keep the semi-major axis equal to the initial (spherical) search
	radius. This will result in a smaller effective volume.
	quiet=False: if set to true, suppress information that is printed
	on each iteration
	"""
	cnt=0
	# initial guess for principal axes:
	evs0=np.array([[1e0,0e0,0e0],[0e0, 1e0, 0e0], [0e0, 0e0, 1e0]])
	# initial guess for axis ratios:
	axes=np.array([1e0,1e0])
	avdiff=1e0
	rad=np.asarray(rad)
	while (avdiff > 0.01) & (cnt < 100):
		axtemp=axes.copy()
		# compute ellipsoidal radius of each particle:
		dist2=(arr_in[:,0]**2 + arr_in[:,1]**2/axes[0]**2 + arr_in[:,2]**2/axes[1]**2)**0.5
		if shell == True:
			# find locations of particles between rad[0] and rad[1]
			if not fix_volume: r_ell=rad;
			else: r_ell=(rad**3/axes[0]/axes[1])**(1e0/3e0);
			locs=((dist2 < r_ell[1]) & (dist2 > r_ell[0])).nonzero()[0];
		else:
			# find locations of particles with r < rad
			locs=(dist2 < rad).nonzero()
			#current version keeps volume same as original spherical
			# volume rather than keeping principal axis same as
			# initial radius
			# compute ellipsoidal radius = (a*b*c)^(1/3)
			if not fix_volume: r_ell=rad;
			else: r_ell=(rad**3/axes[0]/axes[1])**(1e0/3e0);
			# get eigenvectors and eigenvalues.  Note: the inertia tensor
			# is a real symmetric matrix, so the eigenvalues should be real.
			axrat=calc_inertia(arr_in[locs], axtemp)
			if abs(np.imag(axrat[0])).max() > 0.:
				#raise ValueError('Error: eigenvalues are not all real!')
				print 'Error: eigenvalues are not all real!'
				if axes_out: return [-1., -1.], [None, None, None];
				else: return [-1., -1.];
		evals=np.real(axrat[0])
		evecs=axrat[1]
		# get axis ratios from eigenvalues:
		axes=np.sqrt([evals[1]/evals[0], evals[2]/evals[0]])
		# rotate particles (and previous eigenvector array) into new
		# basis:
		arr_in=np.dot(arr_in, evecs)
		evs0=np.dot(evs0, evecs)
		# compute difference between current and previous iteration
		# for axis ratios and diagonal of eigenvector matrix
		avd0=abs((axtemp[0]-axes[0])/(axtemp[0]+axes[0]))
		avd1=abs((axtemp[1]-axes[1])/(axtemp[1]+axes[1]))
		# used to check how close most recent rotation is to zero
		# rotation (i.e. how close evecs is to the identity matrix)
		avd2=3e0-abs(evecs[0,0])-abs(evecs[1,1]) - abs(evecs[2,2])
		if quiet == False:
			# print axis ratios relative to major axis
			print 'axis ratios: ', (np.sort(evals/evals.max()))**0.5
			print 'deviations from previous iteration: ', avd0, avd1, avd2
			print 'number of particles in shell / sphere: ', np.size(locs)
		avdiff=max(avd0, avd1, avd2)
		cnt+=1

	evals=evals/evals.max()
	inds=evals.argsort()
	final_axrats=evals[inds[:2]]**0.5
	#print 'number of iterations: ', cnt
	#print 'number of particles in shell / sphere: ', np.size(locs), '\n'
	#print 'normalized eigenvalues (sorted in ascending order) ' + 'and corresponding unit eigenvectors are:'
	#print '%.*f' % (5, evals[inds[0]]), evs0[:,inds[0]]
	#print '%.*f' % (5, evals[inds[1]]), evs0[:,inds[1]]
	#print '%.*f' % (5, evals[inds[2]]), evs0[:,inds[2]]
	#print 'axis ratios: ', final_axrats
	if cnt == 100: print 'Failed to converge after 100 iterations';
	if axes_out: return final_axrats, [evs0[:,inds[0]], evs0[:,inds[1]], evs0[:,inds[2]]];
	else: return final_axrats;