from math import ceil, floor, sqrt
from time import time
from numba import njit, prange
from numba.typed import List
import numpy as np
# from nvtx import annotate
from scipy.spatial import ConvexHull, Delaunay
from scipy.spatial._qhull import QhullError
from sphractal.constants import ATOMIC_RAD_DICT, METALLIC_RAD_DICT
[docs]def estDuration(func):
"""Return time taken to run a function."""
def wrap(*arg, **kwargs):
start = time()
result = func(*arg, **kwargs)
end = time()
duration = end - start
return result, duration
return wrap
[docs]@njit(fastmath=True, cache=True)
def getMinMaxXYZ(atomsXYZ):
"""Return the minimum and maximum values of each dimension of a set of coordinates."""
rowNum, colNum = atomsXYZ.shape
minXYZ, maxXYZ = np.empty(colNum), np.empty(colNum)
for col in range(colNum):
minVal = maxVal = atomsXYZ[0, col]
for row in range(1, rowNum):
val = atomsXYZ[row, col]
minVal, maxVal = min(minVal, val), max(maxVal, val)
minXYZ[col], maxXYZ[col] = minVal, maxVal
return max(maxXYZ - minXYZ), minXYZ, maxXYZ
# @annotate('readInp', color='cyan')
# noinspection GrazieInspection
[docs]def readInp(filePath, radType='atomic'):
"""Parse an xyz or a lmp file."""
radDict = ATOMIC_RAD_DICT if radType == 'atomic' else METALLIC_RAD_DICT
atomsEle, atomsRad, atomsXYZ = [], [], []
numLinesSkip = 9 if '.lmp' in filePath else 2
with open(filePath, 'r') as f:
for (i, line) in enumerate(f):
if i < numLinesSkip:
continue
eleXYZ = line.split()
ele = eleXYZ[-4]
atomsEle.append(ele)
atomsRad.append(radDict[ele])
atomX, atomY, atomZ = float(eleXYZ[-3]), float(eleXYZ[-2]), float(eleXYZ[-1])
atomsXYZ.append((atomX, atomY, atomZ))
atomsEle, atomsRad, atomsXYZ = np.array(atomsEle, dtype='U2'), np.array(atomsRad), np.array(atomsXYZ)
maxRange, minXYZ, maxXYZ = getMinMaxXYZ(atomsXYZ)
return atomsEle, atomsRad, atomsXYZ, maxRange, minXYZ, maxXYZ
[docs]@njit(fastmath=True, cache=True)
def allDirVecs():
"""Return a list of vectors corresponding to the standard 26 directions from a point."""
dirVecs = []
for x in range(-1, 2):
for y in range(-1, 2):
for z in range(-1, 2):
dirVecs.append((x, y, z))
return dirVecs
# @annotate('findNN', color='magenta')
[docs]@njit(fastmath=True, cache=True)
def findNN(atomsRad, atomsXYZ, minXYZ, maxXYZ, maxAtomRad, radMult=1.2, calcBL=False):
"""
Compute the nearest neighbour list and average bond length for each atom.
Parameters
----------
atomsRad : 1D ndarray of floats
Radius of each atom.
atomsXYZ : 2D ndarray of floats
Cartesian coordinates of each atom.
minXYZ : 1D ndarray of floats
Minimum values of each dimension in the Cartesian space.
maxXYZ : 1D ndarray of floats
Maximum values of each dimension in the Cartesian space.
maxAtomRad : Union[int,float]
Maximum value of atomic radius.
radMult : Union[int,float]
Multiplier to the atomic radii.
calcBL : bool, optional
Whether to compute the average bond length for each atom.
Returns
-------
atomsNeighIdxsPadded : 2D ndarray of ints
Neighbour atoms indices of each atom, padded with -1 for fixed number of columns.
atomsAvgBondLen : 1D ndarray of floats
Average bond lengths for each each.
"""
(minX, minY, minZ), (maxX, maxY, maxZ) = minXYZ, maxXYZ
atomsNeighIdxs = [[int(j) for j in range(0)] for _ in range(len(atomsRad))]
atomsAvgBondLen = np.zeros_like(atomsRad)
stepSize = ceil(maxAtomRad * 2 * radMult)
numX, numY, numZ = max(1, ceil((maxX-minX) / stepSize)), max(1, ceil((maxY-minY) / stepSize)), max(1, ceil((maxZ-minZ) / stepSize))
boxes = [[[[int(i) for i in range(0)] for _ in range(numZ)] for _ in range(numY)] for _ in range(numX)]
allDirections = allDirVecs()
for (i, atom1rad) in enumerate(atomsRad):
atom1X, atom1Y, atom1Z = atomsXYZ[i]
x, y, z = floor((atom1X-minX) / stepSize), floor((atom1Y-minY) / stepSize), floor((atom1Z-minZ) / stepSize)
for (dirX, dirY, dirZ) in allDirections:
if 0 <= x + dirX < numX and 0 <= y + dirY < numY and 0 <= z + dirZ < numZ:
for j in boxes[x + dirX][y + dirY][z + dirZ]:
atom2X, atom2Y, atom2Z = atomsXYZ[j]
atom2rad = atomsRad[j]
diffX, diffY, diffZ = abs(atom1X - atom2X), abs(atom1Y - atom2Y), abs(atom1Z - atom2Z)
sumOfSquares = diffX*diffX + diffY*diffY + diffZ*diffZ
if sumOfSquares < ((atom1rad+atom2rad)*radMult) ** 2:
atomsNeighIdxs[i].append(j)
atomsNeighIdxs[j].append(i)
if calcBL:
atomsAvgBondLen[i] += sqrt(sumOfSquares)
atomsAvgBondLen[j] += sqrt(sumOfSquares)
boxes[x][y][z].append(i)
# Turn list of lists into NumPy array padded with -1
maxNeighNum = max([len(atomNeighIdxs) for atomNeighIdxs in atomsNeighIdxs])
atomsNeighIdxsPadded = np.full((len(atomsNeighIdxs), maxNeighNum), -1)
for (i, atomNeighIdx) in enumerate(atomsNeighIdxs):
atomsNeighIdxsPadded[i][:len(atomNeighIdx)] = atomNeighIdx
if calcBL:
for atomID in range(len(atomsRad)):
atomsAvgBondLen[atomID] /= len(atomsNeighIdxs[atomID])
return atomsNeighIdxsPadded, atomsAvgBondLen
[docs]@njit(fastmath=True, cache=True)
def findTetras(tetraVtxsIdxs, r, alpha):
"""Return tetrahedrons with their circumsphere radii smaller than a specified alpha value."""
return tetraVtxsIdxs[r < alpha, :]
[docs]@njit(fastmath=True, cache=True)
def rmDupTris(tris):
"""Remove triangles that occur twice (internal triangles)."""
unqTris = []
for tri in tris:
if tris.count(tri) == 1:
unqTris.append(tri)
return unqTris
# @annotate('alphaShape', color='cyan')
[docs]def alphaShape(atomsXYZ, tetraVtxsIdxs, alpha):
"""
Return points that form the surface using alpha shape algorithm.
Algorithm modified from https://stackoverflow.com/questions/26303878/alpha-shapes-in-3d
Radius of the sphere fitting inside the tetrahedral < alpha (http://mathworld.wolfram.com/Circumsphere.html)
"""
# Find radius of the circumsphere
tetraVtxs = atomsXYZ[tetraVtxsIdxs]
norm2 = np.sum(tetraVtxs ** 2, axis=2)[:, :, np.newaxis]
ones = np.ones((tetraVtxs.shape[0], tetraVtxs.shape[1], 1))
a = np.linalg.det(np.concatenate((tetraVtxs, ones), axis=2))
Dx = np.linalg.det(np.concatenate((norm2, tetraVtxs[:, :, [1, 2]], ones), axis=2))
Dy = np.linalg.det(np.concatenate((norm2, tetraVtxs[:, :, [0, 2]], ones), axis=2))
Dz = np.linalg.det(np.concatenate((norm2, tetraVtxs[:, :, [0, 1]], ones), axis=2))
c = np.linalg.det(np.concatenate((norm2, tetraVtxs), axis=2))
with np.errstate(divide='ignore', invalid='ignore'):
r = np.sqrt(Dx**2 + Dy**2 + Dz**2 - 4*a*c) / (2*np.abs(a))
# Find tetrahedrons and triangles
tetras = findTetras(tetraVtxsIdxs, r, alpha)
triComb = np.array([(0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)])
tris = tetras[:, triComb].reshape(-1, 3)
tris = rmDupTris(List(tuple(i) for i in tris))
return np.unique(np.concatenate(tris))
# @annotate('findSurf', color='yellow')
[docs]def findSurf(atomsXYZ, atomsNeighIdxs, option='alphaShape', alpha=3.0, bulkCN=12):
"""
Return the indices of surface atoms.
Parameters
----------
atomsXYZ : 2D ndarray of floats
Cartesian coordinates of each atom.
atomsNeighIdxs : 2D ndarray of ints
Neighbour atoms indices of each atom.
option : {'alphaShape', 'convexHull', 'numNeigh'}, optional
Algorithm to identify the spheres on the surface.
'convexHull' tends to identify less surface atoms; 'numNeigh' tends to identify more surface atoms.
'alphaShape' is a generalisation of 'convexHull'.
alpha : Union[int, float], optional
'alpha' for the alpha shape algorithm, only used if 'option' is 'alphaShape'.
bulkCN : int, optional
Minimum number of neighbouring atoms for a non-surface atom.
Returns
-------
atomsSurfIdxs : 1D ndarray of ints
Indices of surface atoms.
Notes
-----
Other alternatives include:
- https://dl.acm.org/doi/abs/10.1145/2073304.2073339
- https://onlinelibrary.wiley.com/doi/pdf/10.1002/jcc.25384
- https://www.jstage.jst.go.jp/article/tmrsj/45/4/45_115/_article
"""
atomsSurfIdxs = np.zeros(len(atomsXYZ), dtype=np.bool_)
if option == 'convexHull':
for atomIdx in np.array(ConvexHull(atomsXYZ).vertices):
atomsSurfIdxs[atomIdx] = True
elif option == 'numNeigh':
for (atomIdx, atomNeighIdxs) in enumerate(atomsNeighIdxs):
if len(atomNeighIdxs[atomNeighIdxs > -1]) < bulkCN:
atomsSurfIdxs[atomIdx] = True
elif option == 'alphaShape':
try:
tetraVtxsIdxs = Delaunay(atomsXYZ).simplices
for atomIdx in alphaShape(atomsXYZ, tetraVtxsIdxs, alpha):
atomsSurfIdxs[atomIdx] = True
except QhullError:
atomsSurfIdxs = np.full(len(atomsXYZ), True)
atomsSurfIdxs = np.where(atomsSurfIdxs)[0]
return atomsSurfIdxs
[docs]@njit(fastmath=True, cache=True)
def findAtomNeighs(neighIdxs, atomsSurfIdxs):
"""Divide a list of atomic neighbours into two based on whether they lie on the surface."""
surfNeighIdxs, bulkNeighIdxs = [], []
for idx in neighIdxs:
if idx in atomsSurfIdxs:
surfNeighIdxs.append(idx)
else:
bulkNeighIdxs.append(idx)
return np.array(surfNeighIdxs, dtype=np.uint32), np.array(bulkNeighIdxs, dtype=np.uint32)
[docs]@njit(fastmath=True, cache=True) # parallel=True slows things down, incompatible with multiprocessing.Pool()
def calcDist(p1, p2):
"""Return Euclidean distance between two points."""
d = 0
for i in prange(3):
d += (p2[i] - p1[i]) ** 2
return d ** 0.5
[docs]@njit(fastmath=True, cache=True)
def getOrdSurfNeighCombs(surfNeighDists, surfNeighIdxs):
"""Compute ID pairs of a set of neighbouring atoms, ordered by distance from a given point."""
ordSurfNeighIdxs = surfNeighIdxs[np.argsort(surfNeighDists)]
surfNeighIdxPairs = np.stack(np.triu_indices(len(ordSurfNeighIdxs), k=1), axis=-1)
return [ordSurfNeighIdxs[idxPair] for idxPair in surfNeighIdxPairs]
[docs]@njit(fastmath=True, cache=True)
def closestSurfAtoms(pointXYZ, surfNeighIdxs, atomsXYZ, atomsNeighIdxs):
"""Return two closest surface atom pairs from a given point, that are neighbour of each other."""
if len(surfNeighIdxs) < 2:
return np.array([np.nan]), np.array([np.nan])
surfNeighXYZs = atomsXYZ[surfNeighIdxs]
surfNeighDists = np.array([calcDist(pointXYZ, surfNeighXYZ) for surfNeighXYZ in surfNeighXYZs])
for idxPair in getOrdSurfNeighCombs(surfNeighDists, surfNeighIdxs):
if idxPair[0] in atomsNeighIdxs[idxPair[1]]:
return atomsXYZ[idxPair[0]], atomsXYZ[idxPair[1]]
return np.array([np.nan]), np.array([np.nan])
[docs]@njit(fastmath=True, cache=True)
def oppositeInnerAtoms(pointXYZ, atom1XYZ, atomNeighIdxs,
atomsSurfIdxs, atomsXYZ, atomsNeighIdxs):
"""Return whether a point lies opposite of the average coordinates of neighbouring inner atoms of a given atom."""
surfNeighIdxs, bulkNeighIdxs = findAtomNeighs(atomNeighIdxs, atomsSurfIdxs)
atom2XYZ, atom3XYZ = closestSurfAtoms(pointXYZ, surfNeighIdxs, atomsXYZ, atomsNeighIdxs)
innerNeighXYZs = atomsXYZ[bulkNeighIdxs] if len(bulkNeighIdxs) > 0 else atomsXYZ[surfNeighIdxs]
avgInnerAtomXYZ = np.array([innerNeighXYZs[:, i].mean() for i in range(3)])
normal = avgInnerAtomXYZ - atom1XYZ if len(atom3XYZ) == 1 else np.cross(atom2XYZ - atom1XYZ, atom3XYZ - atom1XYZ)
# If the given point is on the opposite side of the surface plane, the sign of products of the dot products will be negative
return np.dot(normal, avgInnerAtomXYZ - atom1XYZ) * np.dot(normal, pointXYZ - atom1XYZ) < 0