#!/usr/bin/env python3
"""Python module of utility functions for findTDE calculations."""
from math import gcd
from fractions import Fraction
import numpy as np
# math functions
[docs]
def unit_vector(vector):
# https://stackoverflow.com/a/13849249
""" Returns the unit vector of the vector. """
return vector / np.linalg.norm(vector)
[docs]
def angle_between(v1, v2):
# https://stackoverflow.com/a/13849249
""" Returns the angle in radians between vectors 'v1' and 'v2'::
>>> angle_between((1, 0, 0), (0, 1, 0))
1.5707963267948966
>>> angle_between((1, 0, 0), (1, 0, 0))
0.0
>>> angle_between((1, 0, 0), (-1, 0, 0))
3.141592653589793
"""
v1_u = unit_vector(v1)
v2_u = unit_vector(v2)
return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))*(180/np.pi)
[docs]
def cart2sph(x, y, z):
dxy = np.sqrt(x**2 + y**2)
r = np.sqrt(dxy**2 + z**2)
theta = np.arctan2(y, x)
phi = np.arctan2(dxy, z)
theta, phi = np.rad2deg([theta, phi])
return r, theta % 360, phi
[docs]
def sph2cart(theta, phi, r=0.5):
theta, phi = np.deg2rad([theta, phi])
z = r * np.cos(phi)
rsinphi = r * np.sin(phi)
x = rsinphi * np.cos(theta)
y = rsinphi * np.sin(theta)
return x, y, z
# crystallography functions
# https://ssd.phys.strath.ac.uk/resources/crystallography/crystallographic-direction-calculator/
[docs]
def lattice_iconv(indices, ind_type='UVTW'):
"""
Function converting between rectangular and hexagonal lattice indices. Given a
tuple of indices and the initial index type as a keyword argument. Index type
is either [uvw] for rectangular or [UVTW] for hexagonal. Returns a tuple of
the opposite type.
"""
print(indices)
if ind_type == 'UVTW':
U, V, T, W = indices
u = 2*U + V
v = 2*V + U
w = W
n = gcd(u, gcd(v, w))
u /= n
v /= n
w /= n
new_indices = (int(u), int(v), int(w))
elif ind_type == 'uvw':
u, v, w = indices
U = Fraction(((2*u) - v), 3)
V = Fraction(((2*v) - u), 3)
T = -1*(U + V)
W = w
denom_gcd = gcd(U.denominator, gcd(V.denominator, T.denominator))
U *= denom_gcd
V *= denom_gcd
T *= denom_gcd
W *= denom_gcd
new_indices = (int(U), int(V), int(T), int(W))
return new_indices
[docs]
def lat2sph(uvw, ai):
"""
Function converting from lattice directions to spherical coordinates. Given two 2D
arrays, one for lattice directions and one for lattice vectors, returns a 2D array
of the spherical coordinates corresponding to the lattice directions.
"""
xyz, rpt = np.zeros(uvw.shape), np.zeros(uvw.shape)
for i in range(uvw.shape[0]):
xyz[i, :] = uvw[i, :]@ai
rpt[i, :] = np.around(cart2sph(xyz[i, 0], xyz[i, 1], xyz[i, 2]), decimals=2)
rpt[i, 0] = 1.00
return rpt
[docs]
def sph2lat(rpt, ai):
"""
Function converting from lattice directions to spherical coordinates. Given two 2D
arrays, one for lattice directions and one for lattice vectors, returns a 2D array
of the spherical coordinates corresponding to the lattice directions.
"""
xyz, uvw = np.zeros(rpt.shape), np.zeros(rpt.shape)
for i in range(rpt.shape[0]):
xyz[i, :] = sph2cart(rpt[i, 1], rpt[i, 2], r=rpt[i, 0])
uvw[i, :] = xyz[i, :]@np.linalg.inv(ai)
n = gcd(round(uvw[i, 0]), gcd(round(uvw[i, 1]), round(uvw[i, 2])))
uvw[i, :] /= n
return uvw
[docs]
def scale_C3z_symmetry(rpt):
"""
Function to scale all given spherical coordinates to a single zone based on threefold
rotation symmetry about the c-axis. Given a 2D array of spherical coordinates, returns
another 2D array of spherical coordinates with polar angle scaled by 120 deg increments
to the desired range.
"""
for i in range(rpt.shape[0]):
if rpt[i, 1] >= 30. and rpt[i, 1] <= 150.:
pass
elif rpt[i, 1] >= 0. and rpt[i, 1] < 30.:
rpt[i, 1] += 120.
elif rpt[i, 1] > 150. and rpt[i, 1] <= 270.:
rpt[i, 1] -= 120.
elif rpt[i, 1] > 270. and rpt[i, 1] <= 360.:
rpt[i, 1] -= 240.
else:
raise Exception('Angle outside of 0-360 deg')
return rpt
[docs]
def scale_C6z_symmetry(rpt):
"""
Function to scale all given spherical coordinates to a single zone based on sixfold
rotation symmetry about the c-axis. Given a 2D array of spherical coordinates, returns
another 2D array of spherical coordinates with polar angle scaled by 60 deg increments
to the desired range.
"""
rpt_C3z = scale_C3z_symmetry(rpt)
for i in range(rpt_C3z.shape[0]):
if rpt_C3z[i, 1] >= 30. and rpt_C3z[i, 1] <= 90.:
pass
elif rpt_C3z[i, 1] > 90. and rpt_C3z[i, 1] <= 150.:
rpt[i, 1] = 180. - rpt[i, 1]
else:
raise Exception('Angle outside of 30-150 deg')
return rpt
# interpolation functions
[docs]
def idw(samples, tx, ty, P=5):
"""
Function to compute a single IDW interpolation of sample data.
Change P value for different interpolation (P>2 is recommended).
"""
def dist(a, b):
return ((a[0]-b[0])**2 + (a[1]-b[1])**2)**(1./2.)
num = 0.
den = 0.
for i in range(0, len(samples)):
d = (dist(samples[i], [tx, ty]))**P
if(d < 1.e-5):
return samples[i][2]
w = 1/d
num += w*samples[i][2]
den += w
return num/den
[docs]
def idw_heatmap(inputdata, RES=360, P=5):
"""
Perform full IDW interpolation on a (nsamples x 3) array (x, y, f(x, y))
and produce data able to be used in a heatmap. From Victor.
Change P value for different interpolation (P>2 is recommended).
RES is the resolution of the final image.
"""
# from Victor, plot heatmap
# Setup z as a function of interpolated x, y
minx, maxx = min([d[0] for d in inputdata]), max([d[0] for d in inputdata])
miny, maxy = min([d[1] for d in inputdata]), max([d[1] for d in inputdata])
minz, maxz = min([d[2] for d in inputdata]), max([d[2] for d in inputdata]) # useful to scale 0 < z < 1
dx, dy = (maxx - minx)/(RES-1), (maxy - miny)/(RES-1)
xs = [minx + i*dx for i in range(0, RES)]
ys = [miny + i*dy for i in range(0, RES)]
zs = [[None for i in range(0, RES)] for j in range(0, RES)]
for i in range(0, RES):
for j in range(0, RES):
zs[i][j] = idw(samples=inputdata, tx=xs[i], ty=ys[j], P=P)
# print(xs, ys, zs)
return (xs, ys, np.transpose(zs))