#!/usr/bin/env python

#
# $Id: vort.py,v 1.3 2018/10/01 14:44:11 dintrans Exp $
# Animate the vorticity snapshots with the velocity field superimposed
#

import matplotlib
matplotlib.use('TKAgg')

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

vort, t = pc.read_slices(field='oo2', proc=0)
print "vort=", vort.shape
ux, t = pc.read_slices(field='uu1', proc=0)
uz, t = pc.read_slices(field='uu3', proc=0)

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)
vort=vort.reshape(nt, dim.nz, dim.nx)
ux=ux.reshape(nt, dim.nz, dim.nx)
uz=uz.reshape(nt, dim.nz, dim.nx)

qs1=N.random.random_integers(0,dim.nx-1, 100)
qs2=N.random.random_integers(0,dim.nz-1, 100)
xx,zz=P.meshgrid(grid.x, grid.z)
frame=par.xyz0[0],par.xyz1[0],par.xyz0[2],par.xyz1[2]

P.ion()
for n in range(nt):
    f=vort[n,...]
    a=ux[n, qs2, qs1]**2+uz[n, qs2, qs1]**2 ; norm=N.sqrt(a.max())
    ux[n, qs2, qs1]/=norm ; uz[n, qs2, qs1]/=norm
    if n==0:
        im=P.imshow(f, extent=frame, origin='lower', aspect='auto')
        vel=P.quiver(xx[qs2, qs1], zz[qs2, qs1], ux[n, qs2, qs1],
        uz[n, qs2, qs1], scale=7)
        st=P.figtext(0.8,0.2,'t=%.1f'%t[n], color='k')
    else:
        im.set_data(f)        
        vel.set_UVC(ux[n, qs2, qs1], uz[n, qs2, qs1])
        st.set_text('t=%.1f'%t[n])
    P.draw()
P.xlabel('x')
P.ylabel('z')

P.ioff()
P.show()