## Purpose: Calculate the stress of individual particles in a colloidal crystal, using only particle positions
##          From the "SALSA" method of Lin et al. :
##          "Measuring nonlinear stresses generated by defects in 3D colloidal crystals"
##          Nature Materials 15, 1172-1176 (2016)
##
##          Stress calculated with equation 1
##
##          Input: Particle coordinate data: x, y, frame
##          Output: Particle coordinate data + stress tensor components s_ij: x, y, frame, s_xx, s_xy/s_yx, s_yy
import numpy as np
import scipy.spatial as sps
from scipy.spatial import Voronoi, cKDTree, ConvexHull
import pylab as pyl

## Find the volumes of the voronoi cells
## Input: Point coordinates x, y
## Output: List with
def voronoi_volumes(points):
    v = Voronoi(points)
    vol = np.zeros(v.npoints)
    for i, reg_num in enumerate(v.point_region):
        indices = v.regions[reg_num]
        if -1 in indices:
            vol[i] = np.inf
        else:
            vol[i] = ConvexHull(v.vertices[indices]).volume
    return vol

delt = 0.2 # Cutoff distance from contact (see Lin et al.)
umpx = 0.1004 # Micrometre pixel^-1
pDiam = 2.85 # Particle diameter in micrometres

fileLoc = 'H:\\Data\\2020-09-18\\Act 0\\'
dircL = fileLoc + 'small_coords_tracked.dat'
sData = np.loadtxt(dircS, dtype = float) # x, y, frame

tMax = int(np.amax(sData[:,2]))

tCnt = 0
pStress = np.column_stack((sData, np.zeros((sData.shape[0],3)) # x, y, frame, s_xx, s_xy/s_yx, s_yy
for t in range(0,tMax+1):


    print("Frame: " + str(t))
    frameWhr = np.where(sData[:,2] == t)
    stData = sData[frameWhr]


    vol = voronoi_volumes(stData[:,:2])
    tree = cKDTree(stData[:,:2]) # Neighbours

    stress = []
    for x in range(vol.shape[0]): # Go over particles
        xC = stData[x,:2]
        xN = tree.query_ball_point(xC, (delt+pDiam)/umpx) # Neighbouring particles that cause stress
        xN.remove(x)

        sTen = np.zeros((2,2)) # tensor

        if len(xN) == 0: # Has no neighbours
            stress.append([0,0,0])
            continue

        nV = stData[xN,:2]
        xVec = nV - xC # Vectors between particle and neighbours

        # To unit vectors
        normV = np.sqrt(xVec[:,0]**2 + xVec[:,1]**2)
        xVec[:,0] /= normV
        xVec[:,1] /= normV

        # local structural anisotropy
        sTen[0,0] = np.sum(xVec[:,0]**2)/xVec.shape[0]
        sTen[1,1] = np.sum(xVec[:,1]**2)/xVec.shape[0]
        crsM = np.sum(xVec[:,0]*xVec[:,1])/xVec.shape[0]
        sTen[0,1], sTen[1,0] = crsM, crsM


        sTen *= (pDiam/delt)/vol[x]

        # Add stress components to overall array
        pStress[frameWhr[0][x],4] = sTes[0,0]
        pStress[frameWhr[0][x],5] = sTes[0,1]
        pStress[frameWhr[0][x],6] = sTes[1,1]

np.savetxt(fileLoc+"small_stress_components.dat", pStress, fmt = "%.4f\t%.4f\t%.4f\t%.4f\t%.4f\t%.4f\r")
