;==============================

PRO compute_zeta, dat, DO_PLOT = do_plot, _EXTRA = ex

;==============================
; A procedure to compute vertical vorticity.
;
; INPUT/OUTPUT
;    dat = DAT_DATA OBJECT
;
; KEYWORD INPUT
;    DO_PLOT = INT = If set, do a plot of zeta
;
; NOTES
; The field 'zeta' will be added to dat.
;
; All keywords accepted by quickplot_3d.pro are accepted here.
;
; CALL
; IDL> compute_zeta, dat 
;==========

nx = dat->query('nx')
ny = dat->query('ny')
nz = dat->query('nz')

dx = dat->query('dx')
dy = dat->query('dy')

dv_dx = FLTARR(nx,ny,nz)
du_dy = FLTARR(nx,ny,nz)
zeta  = FLTARR(nx,ny,nz)

v = dat->query('v')
u = dat->query('u')


;  dv/dx

FOR k = 0, nz - 1 DO BEGIN 
  FOR j = 0, ny - 1 DO BEGIN 
    FOR i = 1, nx - 2 DO BEGIN

    dv_dx[i,j,k] = ( 1./(2.*dx) ) * ( v[i+1,j,k] - v[i-1,j,k] )

    ENDFOR
    
    dv_dx[0,j,k] = (1./dx) * ( v[1,j,k] - v[0,j,k] )
    dv_dx[nx-1,j,k] = (1./dx) * ( v[nx-1,j,k] - v[nx-2,j,k] )
    
  ENDFOR  ; end j loop
ENDFOR    ; end k loop

;  du/dy

FOR k = 0, nz - 1 DO BEGIN 
  FOR i = 0, nx - 1 DO BEGIN 
    FOR j = 1, ny - 2 DO BEGIN

    du_dy[i,j,k] = ( 1./(2.*dy) ) * ( u[i,j+1,k] - u[i,j-1,k] )

    ENDFOR
    
    du_dy[i,0,k] = (1./dy) * ( u[i,1,k] - u[i,0,k] )
    du_dy[i,ny-1,k] = (1./dy) * ( u[i,ny-1,k] - u[i,ny-2,k] )
    
  ENDFOR  ; end i loop
ENDFOR    ; end k loop

;  vertical vorticity

FOR i = 0, nx - 1 DO BEGIN
  FOR j = 0, ny - 1 DO BEGIN
    FOR k = 0, nz - 1 DO BEGIN
    
    zeta[i,j,k] = dv_dx[i,j,k] - du_dy[i,j,k]
    
    ENDFOR
  ENDFOR
ENDFOR

dat->assign, 'zeta', zeta, /OVERRIDE

IF KEYWORD_SET(do_plot) THEN BEGIN
   quickplot_3d, dat, 'zeta', _EXTRA = ex
ENDIF

END