sho-lesson/sho/plot.py
2018-12-08 19:40:22 +01:00

35 lines
906 B
Python

import numpy as np
from matplotlib import cm
import itertools
def surface(ax, shape, f):
Z = np.zeros( shape )
for y in range(shape[0]):
for x in range(shape[1]):
Z[y][x] = f( (x/shape[0],y/shape[1]), 0.5 )
X = np.arange(0,shape[0],1)
Y = np.arange(0,shape[1],1)
X,Y = np.meshgrid(X,Y)
ax.plot_surface(X, Y, Z, cmap=cm.viridis)
def path(ax, shape, history):
def pairwise(iterable):
a, b = itertools.tee(iterable)
next(b, None)
return zip(a, b)
k=0
for i,j in pairwise(range(len(history)-1)):
xi = history[i][1][0]*shape[0]
yi = history[i][1][1]*shape[1]
zi = history[i][0]
xj = history[j][1][0]*shape[0]
yj = history[j][1][1]*shape[1]
zj = history[j][0]
x = [xi, xj]
y = [yi, yj]
z = [zi, zj]
ax.plot(x,y,z, color=cm.RdYlBu(k))
k+=1