#!/usr/bin/env python

#
# $Id: pvid.py,v 1.1 2009/04/29 13:25:49 dintrans Exp $
# Animate the magnetic field snapshots with the velocity field superimposed
#

import numpy as N
import pylab as P
import pencil as pc

dat, t = pc.read_slices(field='ss', proc=0, extension='xy')
ux, t = pc.read_slices(field='uu1', proc=0, extension='xy')
uy, t = pc.read_slices(field='uu2', proc=0, extension='xy')

dim=pc.read_dim()
par=pc.read_param(quiet=True)
par2=pc.read_param(quiet=True, param2=True)
grid=pc.read_grid(param=par, quiet=True, trim=True)
nt=len(t)
dat=dat.reshape(nt, dim.ny, dim.nx)
ux=ux.reshape(nt, dim.ny, dim.nx)
uy=uy.reshape(nt, dim.ny, dim.nx)

qs1=N.random.random_integers(0,dim.nx-1, 1000)
qs2=N.random.random_integers(0,dim.ny-1, 1000)
xx,yy=P.meshgrid(grid.x, grid.y)
frame=par.xyz0[0],par.xyz1[0],par.xyz0[1],par.xyz1[1]

P.ion()
im=P.imshow(dat[0,...], extent=frame, origin='lower', aspect='auto')
a=ux[0, qs2, qs1]**2+uy[0, qs2, qs1]**2
norm=N.sqrt(a.max())
ux[0, qs2, qs1]=ux[0, qs2, qs1]/norm
uy[0, qs2, qs1]=uy[0, qs2, qs1]/norm
vel=P.quiver(xx[qs2, qs1], yy[qs2, qs1], ux[0, qs2, qs1],
uy[0, qs2, qs1], scale=7)
st=P.figtext(0.8,0.15,'t=%.1f'%t[0], color='w')
P.xlabel('x')
P.ylabel('y')
P.title('Kelvin-Helmholtz instability')

for i in range(1, nt):
  im.set_data(dat[i, ...])
  a=ux[i, qs2, qs1]**2+uy[i, qs2, qs1]**2
  norm=N.sqrt(a.max())
  ux[i, qs2, qs1]=ux[i, qs2, qs1]/norm
  uy[i, qs2, qs1]=uy[i, qs2, qs1]/norm
  vel.set_UVC(ux[i, qs2, qs1], uy[i, qs2, qs1])
  st.set_text('t=%.1f'%t[i])

  P.draw()

P.show()