from hllrkfluxlimit2d import hllrkfluxlimit2d as hll
import cPickle
import numpy as np
import pylab as py

py.ion()

u = cPickle.load(file("helmholtz2d-fast.dat","r"))
#u = cPickle.load(file("shock2d-vertical.dat","r"))
uc = u.copy()

gamma = 1.4
hll.init(gamma, 0.1, 0.1, 0.2, 1.0)
pressure = lambda u : (gamma - 1) * (u[:,:,3] - 0.5 * u[:,:,1]**2 / u[:,:,0])

fig = py.figure(figsize=(12,12))
title = fig.suptitle("t = %0.2f"%hll.t)
py.subplot(221)
py.title("Density")
density = py.imshow(uc[:,:,0])
py.subplot(222)
py.title("velocity_x")
momx = py.imshow(uc[:,:,1]/uc[:,:,0])
py.subplot(223)
py.title("velocity_y")
momy = py.imshow(uc[:,:,2]/uc[:,:,0])
py.subplot(224)
py.title("Pressure_x")
energy = py.imshow(pressure(u))
py.draw()
  xmax = 0.75 * u.shape[0]
  for x in range(u.shape[0]):
    for y in range(u.shape[1]):
      u[x,y,:] = np.array([1.0, vflow, amp*np.sin(np.pi*x/(u.shape[0]-1)), 0])
      if xmin < x < xmax:
        u[x,y,0] = drat
        u[x,y,1] = -drat*(vflow)
        u[x,y,2] = drat * amp * np.sin(np.pi*x/(u.shape[0]-1))
      u[x,y,3] = energy(gamma, P, u[x,y,:])
  return u

#uc = helmholtz2(np.zeros((1024,1024,4),dtype=np.float64))
uc = shock2ddiagonal(np.zeros((1024,1024,4)))

gamma = 1.4
hll.init(gamma, 2./uc.shape[0], 2./uc.shape[1], 0.8, 1.9)
pressure = lambda u : (gamma - 1) * (u[:,:,3] - 0.5 * (u[:,:,1]**2 + u[:,:,2]**2)/ u[:,:,0])

fig = py.figure(figsize=(12,12))
title = fig.suptitle("t = %0.2f"%hll.t)
py.subplot(221)
py.title("Density")
density = py.imshow(uc[:,:,0])
py.colorbar()
py.subplot(222)
py.title("velocity_x")
momx = py.imshow(uc[:,:,1]/uc[:,:,0])
py.colorbar()
py.subplot(223)
py.title("velocity_y")
momy = py.imshow(uc[:,:,2]/uc[:,:,0])