import os
from glob import glob
import numpy as np
from matplotlib.mlab import csv2rec
from mpl_toolkits.mplot3d.axes3d import Axes3D
directory = '/Users/alex/Documents/PTV/test/res/'
list_ptv_is_files = glob(os.path.join(directory,'ptv_is.*'))
# print list_ptv_is_files
# reading the ptv_is.* files
frames = []
for counter, ptv_is_file in enumerate(list_ptv_is_files):
frame = csv2rec(ptv_is_file,skiprows=1,delimiter=' ',names=['p','n','x','y','z'])
frame = rec_append_fields(frame,['t','id'],[np.zeros_like(frame.x)+counter,np.zeros_like(frame.x)-999],dtypes=[np.int,np.int])
frames.append(frame)
# adding trajectory id = linking
id = 0
for i,f in enumerate(frames):
for j,l in enumerate(f):
if l.p == -1 and l.n != -2:
l.id = id
id += 1
elif l.p != -1:
l.id = frames[i-1].id[l.p]
for i,f in enumerate(frames):
ind = f.id == -999
frames[i] = f[~ind]
last_traj = max(frames[-1].id)
traj = [[] for k in range(last_traj+1)]
for f in frames:
for p in f:
traj[p.id].append(p)
fig = figure(figsize=(12,8))
ax = fig.add_subplot(1,1,1, projection='3d')
for t in traj:
x = [p.x for p in t]
y = [p.y for p in t]
z = [p.z for p in t]
ax.plot(x,y,z)