animation_data.py

'''
Interpolate DPP movement data to fit
the eye generated by the CG-model initialization script.
'''

import os
import sys
import math
import json

import numpy as np
from scipy.spatial import cKDTree as KDTree
    
def to_spherical(x,y,z):
    phi = math.atan2(y, x)
    theta = math.acos(z/math.sqrt(x**2+y**2+z**2))
    return theta, phi

def plot(locations, points):
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    lmax = np.max(locations)
    pmax = np.max(points)
    
    ax.scatter(*[[p[i]/lmax for p in locations] for i in range(3)])
    ax.scatter(*[[p[i]/pmax for p in points] for i in range(3)])
    plt.show()

    
def main():
    
    movement_data = []
    
    if len(sys.argv) == 4:
        side = sys.argv[1]
        locations_savefn = sys.argv[2]
        rootdir = sys.argv[3]
        debug_plot = False
    else:
        side = 'right'
        locations_savefn = 'tmp/locations.npy'
        debug_plot = True
    
    # Get ommatidia locations from tmp file
    locations = np.load(locations_savefn)
    
    # GET MOVEMENT DATA FROM PUPIL ANALYSER EXPORT FILE
    fn = "data/3d_vectors_averaged_wtb_01_fullscan_wtb_02_fullscan_wtb_03_fullscan_wtb_04_fullscan_wtb_05_fullscan_2020-04-02 10:57:26.928333.json"
    with open(fn, 'r') as fp:
        data = json.load(fp)
    
    data = data[side]
    
    # 90degree rotation along x-axis to have same orientation as the 3D model
    data['points'] = [[x, -z, y] for x,y,z in data['points']]
    data['vectors'] = [[x, -z, y] for x,y,z in data['vectors']]
    
    spherical_points = [to_spherical(*p) for p in data['points']]
    kdtree = KDTree(spherical_points, workers=-1)
    
    # For each location in spherical coordinates,
    # find the closest point of movement data
    for location in locations:
        theta, phi = to_spherical(*location)
        
        dist, indx = kdtree.query([theta,phi])
        
        movement_data.append(data['vectors'][indx])
    
    savefn = os.path.join(os.path.dirname(locations_savefn), 'movement_data.npy')
    np.save(savefn, movement_data)
    
    if debug_plot:
        plot(locations, data['points'])
        
            
if __name__ == "__main__":
    main()