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,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 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
