def Y(self,trjs,N):
"""
y = Y(t,o)
"""
T,Z,O,P = rx.stack(trjs)
t = np.hstack(T)
o = np.vstack(O)
x,y,dx,dy,leg,TT,V,L,Fn,Bn = o.T
tp = t; yp = y; dyp = dy
t = np.linspace(t[0],t[-1],N,endpoint=False)
y = np.interp(t,tp,yp,left=np.nan,right=np.nan)
dy = np.interp(t,tp,dyp,left=np.nan,right=np.nan)
return t,np.array((y)).T
def tweak_grd(th0,dth0,ls='r',dashes=None,lw=4.,zorder=5,label=''):
st = time.clock()
x_ = np.asarray([0.,z0_,th0,dx_,0.,dth0])
x__,trjs = step(x_,p=p)
print '%0.2f sec; dx0 = %0.2f, dx1 = %0.2f' % (time.clock() - st,x_[3],x__[3])
T,X,O,P = rx.stack(trjs)
vars = O[0].keys()
t = np.hstack(T)
# observations until max compression
o = Struct(**dict([( v, np.vstack([oo[v] for oo,pp in zip(O,P) if (pp.j is not None) and (pp.j[0] < 0)]) ) for v in vars]))
z,th = o.q[:,[2,4]].T
if dashes is None:
axs['grd'].plot(th,z,ls,lw=lw,zorder=zorder,label=label)
else:
axs['grd'].plot(th,z,ls,lw=lw,zorder=zorder,label=label,dashes=dashes)
axs['grd'].arrow(th[-20],z[-20],np.diff(th[-20:])[0],np.diff(z[-20:])[0],
lw=lw,fc='k',head_width=.04,zorder=10)
def seq(x_,p=p,fn='',xlim=None,
props=dict(position=position,offv=550,text=text,pad_inches=0.,
light=0.,time=False,fn='',plots=['sch'],exts=exts,
alpha=0,offz=0,clf=True,ylim=(-0.1,1.45),fill=fill) ):
x__,trjs = step(x_,p=p)
props = props.copy()
T,X,O,P = rx.stack(trjs)
vars = O[0].keys()
o = Struct(**dict([( v, np.vstack([oo[v] for oo in O]) ) for v in vars]))
if len(trjs) == 7:
trjs.pop(5); trjs.pop(1)
n = len(trjs) / 2
offx = -trjs[n].p.foot[:,0].mean()
# enclose whole trajectory in figure
#props.update(xlim=.7*trjs[0].p.lamb[0]*np.asarray([-1.,+1.])+[trjs[0].x[0,0],trjs[-1].x[0,0]])
if xlim is None:
xlim = (.7*trjs[0].p.lamb[0]*np.asarray([-1.,+1.])
+[trjs[0].x[0,0],trjs[-1].x[0,0]]+offx)
props.update(xlim=xlim,offx=offx)
figs,axs = p.poly.sch(trjs[n].t[0],trjs[n].x[0],trjs[n].p,**props)
props.update(figs=figs,axs=axs,clf=False,alpha=1,offz=0,lw=6,time=True)
for k,trj in enumerate(trjs):
#props.update(alpha=float(k+1)/len(trjs))
if k+1 == len(trjs):
props.update(fn=di+fn)
props.update(offz=10*k,light=.9*(1.-float(k+1)/len(trjs)),offt=+1 - 2*(k == n))
p.poly.sch(trj.t[0]-trjs[n].t[0],trj.x[0],trj.p,**props)
#axs['sch'].plot(trj.x[:,0]+offx,trj.x[:,2],color=[.6,.6,.6],lw=6,zorder=1)
axs['sch'].plot(o.q[:,0]+offx,o.q[:,2],color=.3*np.ones(3),lw=4,zorder=10*(k+1))
axs['sch'].arrow(o.q[-1,0]+offx,o.q[-1,2],0.10,0.,head_width=0.06,lw=4,color=.3*np.ones(3),zorder=10*(k+1))
plt.draw()
return trjs,figs,axs,props
def plot(self,trjs,fn=None,dpi=None,pxw=1000,pxh=600,offv=0,figN=1,
plots=['pos','rot','eng'],
alpha=1.,col=('b','g','r'),lp='',lw=4,slab='',
clf=True,legend=True,fill=True,legloc='lower left'):
"""
Plot trajectory
Inputs:
trjs - list - Poly trajectory
Outputs:
(T,Z,O,P),(Te,Ze,Oe,Pe) - observations
"""
def fill(Te,Pe,ylim):
for te,pe in zip(Te,Pe):
if fill:
fc = 1 - .33*np.sum( pe.b ) / pe.n
ax.fill([te[0],te[0],te[1],te[1]],
[ylim[1],ylim[0],ylim[0],ylim[1]],
fc=np.ones(3)*fc,ec='none',zorder=-1)
#fc=np.ones(3)*fc,ec='k',zorder=-1)
#ax.text(.5*(te[0]+te[1]),.5*(ylim[0]+ylim[1]),np.sum( pe.b ),
# horizontalalignment='center',verticalalignment='center')
Te,Xe,Oe,Pe = rx.trans(trjs)
T,X,O,P = rx.stack(trjs)
vars = O[0].keys()
o = Struct(**dict([( v, np.vstack([oo[v] for oo in O]) ) for v in vars]))
legprop = {'size':12};
ms = 20.; mew = 2.; ma = 0.5
t = np.hstack(T)
xlim = (min(t),max(t))
fig = plt.figure(figN)
if clf:
fig.clf()
sps = len(plots); sp = 1;
if 'pos' in plots:
ax = plt.subplot(sps,1,sp); sp += 1; plt.grid('on')
H = (ax.plot(t,o.q[:,0],col[0]+lp,label='$x'+slab+'$',alpha=alpha),
ax.plot(t,o.q[:,1],col[1]+lp,label='$y'+slab+'$',alpha=alpha),
ax.plot(t,o.q[:,2],col[2]+lp,label='$z'+slab+'$',alpha=alpha))
[[hh.set_linewidth(lw) for hh in h] for h in H]
ax.plot(t,o.hip[:,:,0],'k:')
ax.plot(t,o.hip[:,:,1],'k:')
ax.plot(t,o.hip[:,:,2],'k:')
ax.plot(t,o.foot[:,:,0],'k--')
ax.plot(t,o.foot[:,:,1],'k--')
ax.plot(t,o.foot[:,:,2],'k--')
ylim = np.array(ax.get_ylim())#-0.25
fill(Te,Pe,ylim)
if clf:
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.set_ylabel('position')
if legend:
leg = ax.legend(ncol=6,loc=legloc,prop=legprop)
if 'rot' in plots:
ax = plt.subplot(sps,1,sp); sp += 1; plt.grid('on')
H = (ax.plot(t,o.q[:,3],col[0]+lp,label='$\\theta_x'+slab+'$',alpha=alpha),
ax.plot(t,o.q[:,4],col[1]+lp,label='$\\theta_y'+slab+'$',alpha=alpha),
ax.plot(t,o.q[:,5],col[2]+lp,label='$\\theta_z'+slab+'$',alpha=alpha))
[[hh.set_linewidth(lw) for hh in h] for h in H]
ylim = np.array(ax.get_ylim())#-0.25
fill(Te,Pe,ylim)
if clf:
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.set_ylabel('rotation')
if legend:
leg = ax.legend(ncol=6,loc=legloc,prop=legprop)
if 'eng' in plots:
ax = plt.subplot(sps,1,sp); sp += 1; plt.grid('on')
H = (ax.plot(t,o.E[:,0],col[0]+lp,label='$T'+slab+'$',alpha=alpha),
ax.plot(t,o.E[:,3],col[1]+lp,label='$V_k'+slab+'$',alpha=alpha),
ax.plot(t,o.E[:,2],col[2]+lp,label='$V_g'+slab+'$',alpha=alpha),
ax.plot(t,o.E[:,0]+o.E[:,1],'k'+lp,label='$L'+slab+'$',alpha=alpha))
[[hh.set_linewidth(lw) for hh in h] for h in H]
ylim = np.array(ax.get_ylim())#-0.25
fill(Te,Pe,ylim)
if clf:
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.set_ylabel('energy')
if legend:
leg = ax.legend(ncol=6,loc=legloc,prop=legprop)
ax.set_xlabel('time (t)')
if fn is not None:
plt.savefig(fn,bbox_inches='tight',pad_inches=0.1,dpi=dpi)
return o
def seq(x_,
p=p,
fn='',
xlim=None,
props=dict(position=position,
offv=550,
text=text,
pad_inches=0.,
light=0.,
time=False,
fn='',
plots=['sch'],
exts=exts,
alpha=0,
offz=0,
clf=True,
ylim=(-0.1, 1.45),
fill=fill)):
x__, trjs = step(x_, p=p)
props = props.copy()
T, X, O, P = rx.stack(trjs)
vars = O[0].keys()
o = Struct(**dict([(v, np.vstack([oo[v] for oo in O])) for v in vars]))
if len(trjs) == 7:
trjs.pop(5)
trjs.pop(1)
n = len(trjs) / 2
offx = -trjs[n].p.foot[:, 0].mean()
# enclose whole trajectory in figure
#props.update(xlim=.7*trjs[0].p.lamb[0]*np.asarray([-1.,+1.])+[trjs[0].x[0,0],trjs[-1].x[0,0]])
if xlim is None:
xlim = (.7 * trjs[0].p.lamb[0] * np.asarray([-1., +1.]) +
[trjs[0].x[0, 0], trjs[-1].x[0, 0]] + offx)
props.update(xlim=xlim, offx=offx)
figs, axs = p.poly.sch(trjs[n].t[0], trjs[n].x[0], trjs[n].p, **props)
props.update(figs=figs,
axs=axs,
clf=False,
alpha=1,
offz=0,
lw=6,
time=True)
for k, trj in enumerate(trjs):
#props.update(alpha=float(k+1)/len(trjs))
if k + 1 == len(trjs):
props.update(fn=di + fn)
props.update(offz=10 * k,
light=.9 * (1. - float(k + 1) / len(trjs)),
offt=+1 - 2 * (k == n))
p.poly.sch(trj.t[0] - trjs[n].t[0], trj.x[0], trj.p, **props)
#axs['sch'].plot(trj.x[:,0]+offx,trj.x[:,2],color=[.6,.6,.6],lw=6,zorder=1)
axs['sch'].plot(o.q[:, 0] + offx,
o.q[:, 2],
color=.3 * np.ones(3),
lw=4,
zorder=10 * (k + 1))
axs['sch'].arrow(o.q[-1, 0] + offx,
o.q[-1, 2],
0.10,
0.,
head_width=0.06,
lw=4,
color=.3 * np.ones(3),
zorder=10 * (k + 1))
plt.draw()
return trjs, figs, axs, props
def do_anim():
q=[1.,2.,10.,5.,2.,2.0,2*np.pi,np.pi,2.]
m,M,k,b,l0,a,om,ph,g = q
debug = False
hop = Hop()
p = hop.P(q,debug)
ph0 = np.pi/2.
t0 = ph0*np.pi/om
tf = (10*2 + ph0)*np.pi/om
n = np.infty
t = (t0,tf)
j = 1
z = np.array([0.2,1.75,0.,0.]) # nice transient
#j = 2
#z = np.array([1.0,0.0])
#z = np.array([1.6,0.6])
#z = np.array([1.4,0.6])
p.j = j
h = 1e-2
eps = 1e-12
Zeno = False
import time;
start = time.clock()
trjs = flow(hop, p, h, eps, (t0*om,tf*om), z)
#hop.plot(trjs[-10:])
hop.plot(trjs,fn='fig/hoptrj.eps')
xlim = plt.xlim(); ylim = plt.ylim()
plt.xlim(xlim); plt.ylim(ylim)
q0 = q[:]; p0 = hop.P(q0,debug); p0.j = 1; z0 = np.array([0.25,1.75,0.,0.])
hop.sch(3*np.pi/2,[0.25,2.15,0,0],p0,text=True,fn='fig/hopA.eps')
#hop.sch(3*np.pi/2,[0.25,2.15,0,0],p0,text=False,fn='fig/hopAplain.eps')
p0.j = 2
hop.sch(np.pi/4,[1.5,0],p0,text=True,fn='fig/hopG.eps')
#hop.sch(np.pi/4,[1.5,0],p0,text=False,fn='fig/hopGplain.eps')
#hop.sch(t0,z,p,text=True,fn='fig/hop.pdf')
p0.j = 1
print "%0.2f sec for sim & plot" % (time.clock() - start)
ph0 = 0#np.pi
t0 = ph0*np.pi/om
tf = (4*2 + ph0)*np.pi/om
n = np.infty
t = (t0,tf)
j = 2
z = np.array([1.952,1.908])
p.j = j
h = 1e-2
eps = 1e-12
Zeno = False
import time;
start = time.clock()
trjs = flow(hop, p, h, eps, (t0*om,tf*om), z)
hop.plot(trjs)
xlim = plt.xlim(); ylim = plt.ylim()
plt.xlim(xlim); plt.ylim(ylim)
lt = plt.plot([t0,t0],ylim,'y',lw=5.,alpha=0.75)
T,Z,O,P = rx.stack(trjs)
fps = 20.
fn = 0
pxh=400;
import os;
plt.ion()
for t,z,p in zip(T,Z,P):
for tt,zz in zip(t,z):
if tt >= t0 + fn/fps:
print '%4d : %0.2f , %0.2f' % (fn, t0+fn/fps, tt)
#plt.ioff()
plt.figure(1)
lt[0].set_xdata([tt,tt])
plt.draw()
#plt.savefig('vid/trj%04d.png'%fn);
plt.savefig('vid/trj%04d.png'%fn,dpi=pxh/4.);
plt.figure(2)
hop.sch(tt,zz,p)
#plt.ion()
plt.draw()
plt.savefig('vid/hop%04d.png'%fn,bbox_inches='tight',pad_inches=-0.1,dpi=pxh/5.8)
os.system('convert vid/hop%04d.png vid/trj%04d.png +append vid/hoptrj%04d.png'%(fn,fn,fn))
fn += 1
#break
os.system('rm vid/hop.mp4 vid/hoptrj.mp4')
os.system('ffmpeg -r %d -sameq -i vid/hop%%04d.png vid/hop.mp4' % fps)
os.system('ffmpeg -r %d -sameq -i vid/hoptrj%%04d.png vid/hoptrj.mp4' % fps)
def plot(self,trjs,fn='',dpi=None,pxh=None,pxw=None,
alpha=1.,lw=4,clf=True):
"""
Plot trajectory
Inputs:
trjs - list - Hop trajectory
Outputs:
(T,Z,O,P) - observations
"""
T,Z,O,P = rx.stack(trjs)
Te,Ze,Oe,Pe = rx.trans(trjs)
m,M,k,b,l0,a,om,ph,g = P[0].q
fsz = (8.,4.)
if pxh:
dpi = fsz(1)/pxh
elif pxw:
dpi = fsz(0)/pxw
fig = plt.figure(1,figsize=fsz)
#fig = plt.figure(1)
if clf:
fig.clf()
legprop = {'size':10};
sc = 1.
ms = 20.; mew = 2.; ma = 0.5
t = np.hstack(T)
o = np.vstack(O)
te = np.hstack(Te)
oe = np.vstack(Oe)
x,y,dx,dy,leg,TT,V,L,Fn,Bn = o.T
xe,ye,dxe,dye,lege,TTe,Ve,Le,Fne,Bne = oe.T
Vm = V.min()
V = V - Vm
E = TT+V
ax = plt.subplot(2,1,1); plt.grid('on')
H = (ax.plot(t,sc*x,'r',label='$x(t)$',alpha=alpha),
ax.plot(t,sc*y,'b',label='$y(t)$',alpha=alpha))
[[hh.set_linewidth(lw) for hh in h] for h in H]
#H = (ax.plot(t,0.5*np.sin(om*t+ph),label='$f(t) \\propto \\sin(\\omega t+\pi)$',color='k',zorder=-1,lw=2.,alpha=alpha))
ylim = np.array(ax.get_ylim())#-0.25
ylim = [-1.5,3.0]
for te,pe in zip(Te,Pe):
if pe.j == 2:
ax.fill([te[0],te[0],te[1],te[1]],
[ylim[1],ylim[0],ylim[0],ylim[1]],
fc=np.array([1.,1.,1.])*0.75,ec='none',zorder=-1)
ax.set_ylim(ylim)
ax.set_ylabel('height')
#leg = ax.legend(ncol=4,loc='lower center')
if clf:
leg = ax.legend(ncol=4,loc='lower left',prop=legprop)
if clf:
ax.set_xlim((t[0],t[-1]))
ax = plt.subplot(2,1,2); plt.grid('on')
H = (ax.plot(t,sc*dx,'r',label='$\\dot{x}(t)$',alpha=alpha),
ax.plot(t,sc*dy,'b',label='$\\dot{y}(t)$',alpha=alpha))
[[hh.set_linewidth(lw) for hh in h] for h in H]
#H = (ax.plot(t,0.5*np.sin(om*t+ph),label='$f(t) \\propto \\sin(\\omega t+\pi)$',color='k',zorder=-1,lw=2.,alpha=alpha))
ylim = np.array(ax.get_ylim())#-0.25
#ylim = [-1.5,3.0]
for te,pe in zip(Te,Pe):
if pe.j == 2:
ax.fill([te[0],te[0],te[1],te[1]],
[ylim[1],ylim[0],ylim[0],ylim[1]],
fc=np.array([1.,1.,1.])*0.75,ec='none',zorder=-1)
ax.set_ylim(ylim)
ax.set_xlabel('time (t)')
ax.set_ylabel('velocity')
if clf:
leg = ax.legend(ncol=4,loc='lower left',prop=legprop)
# nominal spring length
#ax.plot([t[0],t[-1]],[l0,l0],'k--',lw=2)
plt.subplot(2,1,1)
if clf:
ax.set_xlim((t[0],t[-1]))
if not fn == '':
plt.savefig(fn,bbox_inches='tight',pad_inches=0.1,dpi=dpi)
#plt.show()
return T,Z,O,P
def plot(self,trjs):
"""
Plot double pendulum trajectory
Inputs:
trjs - list - DP trajectory
Outputs:
(T,X,O,P) - observations
"""
T,X,O,P = stack(trjs)
Te,Xe,Oe,Pe = trans(trjs)
m1,m2,L1,L2,b1,b2,a,om,c,g = P[0].q
fig = plt.figure(1, figsize=(6,4))
fig.clf()
legprop = {'size':16};
sc = 180/np.pi
lw = 2.;
ms = 20.; mew = 2.; ma = 0.5
t = np.hstack(T)
o = np.vstack(O)
te = np.hstack(Te)
oe = np.vstack(Oe)
t1,t2,dt1,dt2,lam,TT,V,L,Fn,Bn = o.T
t1e,t2e,dt1e,dt2e,lame,TTe,Ve,Le,Fne,Bne = oe.T
Vm = V.min()
V = V - Vm
E = TT+V
ax = plt.subplot(1,1,1);
H = (ax.plot(t,sc*t1,'r',label='$\\theta_1$'),
ax.plot(t,sc*t2,'b',label='$\\theta_2$'))
[[hh.set_linewidth(lw) for hh in h] for h in H]
dy = np.diff(ax.get_ylim())
H = (ax.plot(t,0.25*dy*np.sin(om*t),label='$\\tau$',color='gray',
zorder=-1))
[h.set_linewidth(2*lw) for h in H]
H = (ax.plot(t,0.25*dy*lam/lam.max(),label='$\\lambda$',color='green',
zorder=-1))
[h.set_linewidth(2*lw) for h in H]
ylim = ax.get_ylim()
for te,pe in zip(Te,Pe):
if pe.j == 1:
ax.plot([te[-1],te[-1]],ylim,'k')
elif pe.j == 2:
ax.plot([te[-1],te[-1]],ylim,'k--')
M = (np.diff(1*(dt1 >= 0)) == -1).nonzero()
for m in M:
ax.plot([t[m],t[m]],ylim,'k:')
ax.set_ylim(ylim)
#H = (ax.plot(t,sc*dt1,'r-.',label='$\\dot{\\theta}_1$'),
# ax.plot(t,sc*dt2,'b--',label='$\\dot{\\theta}_2$'))
#[[hh.set_linewidth(lw) for hh in h] for h in H]
#H = (ax.plot(te,sc*t1e,'g.'),
# ax.plot(te,sc*t2e,'g.'))
#for h in H:
# for hh in h:
# hh.set_ms(ms); hh.set_mew(mew); hh.set_alpha(ma)
ax.set_xlim((t[0],t[-1]))
ax.set_xlabel('Time (sec)')
ax.set_ylabel('State')
#leg = ax.legend(prop=legprop,ncol=4,loc='lower center')
leg = ax.legend(ncol=4,loc='lower center')
if 0:
ax = plt.subplot(2,1,2)
lw = 1.; al=0.1
ms = 5.; mew = 2.; ma = 0.5
for jj,(t,o,te,oe) in enumerate(zip(T,O,Te,Oe)):
if jj == 0:
continue
t1,t2,dt1,dt2,lam,TT,V,L,Fn,Bn = o.T
t1e,t2e,dt1e,dt2e,lame,TTe,Ve,Le,Fne,Bne = oe.T
H = (ax.plot(t-t[0],sc*t1,'r.',label='$\\theta_1$',alpha=al),
ax.plot(t-t[0],sc*t2,'b.',label='$\\theta_2$',alpha=al))
[[hh.set_linewidth(lw) for hh in h] for h in H]
[[hh.set_ms(2.) for hh in h] for h in H]
H = (ax.plot(te-t[0],sc*t1e,'g.'),
ax.plot(te-t[0],sc*t2e,'g.'))
for h in H:
for hh in h:
hh.set_ms(ms); hh.set_mew(mew); hh.set_alpha(ma)
ax.set_xlabel('Time (sec)')
ax.set_ylabel('State')
#plt.show()
return T,X,O,P