from clawpack import pyclaw
from clawpack import riemann

claw = pyclaw.Controller()
claw.tfinal = 0.6
claw.num_output_times = 40

riemann_solver = riemann.euler_4wave_2D
claw.solver = pyclaw.ClawSolver2D(riemann_solver)
claw.solver.all_bcs = pyclaw.BC.extrap

grid_size = (300,300)
domain = pyclaw.Domain( (0.,0.), (1.,1.), grid_size)

claw.solution = pyclaw.Solution(claw.solver.num_eqn,domain)
gam = 1.4
claw.solution.problem_data['gamma']  = gam

# Set initial data
q = claw.solution.q
xx,yy = domain.grid.p_centers
l = xx<0.5; r = xx>=0.5; b = yy<0.5; t = yy>=0.5
q[0,...] = 2.*l*t + 1.*l*b + 1.*r*t + 3.*r*b
q[1,...] = 0.75*t - 0.75*b
q[2,...] = 0.5*l  - 0.5*r
q[3,...] = 0.5*q[0,...]*(q[1,...]**2+q[2,...]**2) + 1./(gam-1.)

claw.keep_copy = True       # Keep solution data in memory for plotting
claw.output_format = None   # Don't write solution data to file

status = claw.run()

%pylab inline

frame = claw.frames[40]
density = frame.q[0,:,:]
(vx,vy) = np.gradient(density)
vs = np.sqrt(vx**2 + vy**2)
x, y = frame.state.grid.c_centers    

plt.pcolormesh(x, y, vs, cmap='RdBu')
plt.axis('image')

from matplotlib import animation
import matplotlib.pyplot as plt
from clawpack.visclaw.JSAnimation import IPython_display
import numpy as np

fig = plt.figure(figsize=[4,4])

frame = claw.frames[0]
density = frame.q[0,:,:]
(vx,vy) = np.gradient(density)
vs = np.sqrt(vx**2 + vy**2)

x, y = frame.state.grid.c_centers    

# This essentially does a pcolor plot, but it returns the appropriate object
# for use in animation.  See http://matplotlib.org/examples/pylab_examples/pcolor_demo.html.
# Note that it's necessary to transpose the data array because of the way imshow works.
im = plt.imshow(vs.T, cmap='Greys', vmin=vs.min(), vmax=vs.max()/20,
           extent=[x.min(), x.max(), y.min(), y.max()],
           interpolation='nearest', origin='lower')

def fplot(frame_number):
    frame = claw.frames[frame_number]
    density = frame.q[0,:,:]
    (vx,vy) = np.gradient(density)
    vs = np.sqrt(vx**2 + vy**2)
    im.set_data(vs.T)
    return im,

animation.FuncAnimation(fig, fplot, frames=len(claw.frames), interval=20)

import plotly
un='jackp'
k='11m2qbzob9'
py = plotly.plotly(username_or_email=un, key=k)

frame = claw.frames[40]
density = frame.q[0,:,:]
(vx,vy) = np.gradient(density)
vs = np.sqrt(vx**2 + vy**2)

# Note that it's necessary to transpose the data array because plotly's heatmap uses imshow.
py.iplot({'z': vs.T, 'type':'heatmap'})

def setplot(plotdata):    
    plotfigure = plotdata.new_plotfigure(name='Density', figno=0)
    plotaxes = plotfigure.new_plotaxes()
    plotaxes.title = 'Density'
    plotaxes.scaled = True      # so aspect ratio is 1
    plotitem = plotaxes.new_plotitem(plot_type='2d_schlieren')
    plotitem.plot_var = 0    
    return plotdata

from clawpack.visclaw import data
from clawpack.visclaw import frametools
plotdata = data.ClawPlotData()
plotdata.setplot = claw.setplot
claw.plotdata = frametools.call_setplot(claw.setplot,plotdata)

frame = claw.load_frame(40)
claw.plot_frame(frame)

claw.plot()