%matplotlib inline
from matplotlib import pyplot as plt
from cavity_flow2d import cavity_flow2D

nx, ny = 128, 128

x = cavity_flow2D(nx, ny, 1, 1.0, 1.0)

Z = x[...].reshape(nx,ny,4)
fig, axs = plt.subplots(2,2, figsize=(10,8))
titles = ['u', 'v', 'vorticity', 'temperature']
for idx, ax in enumerate(axs.ravel()):
    cs = ax.contourf(Z[:,:,idx], 100)
    fig.colorbar(cs, ax=ax, shrink=0.9)
    ax.set_title(titles[idx])

from petsc4py import PETSc
PETSc.Options().setValue('-snes_monitor', 'cavity.conv')
PETSc.Options().setValue('-snes_view', 'ascii:cavity.conv::append')

PETSc.Options().setValue('-pc_type', 'lu')

nx, ny = 32, 32

x = cavity_flow2D(nx, ny, 1, 1.0, 1.0)
with file('cavity.conv', 'r') as f:
    print ''.join(f.readlines())

Z = x[...].reshape(nx,ny,4)
fig, axs = plt.subplots(2,2, figsize=(10,8))
titles = ['u', 'v', 'vorticity', 'temperature']
for idx, ax in enumerate(axs.ravel()):
    cs = ax.contourf(Z[:,:,idx], 100)
    fig.colorbar(cs, ax=ax, shrink=0.9)
    ax.set_title(titles[idx])

from IPython.parallel import Client
c = Client()
view = c[:]

%load_ext parallelmagic
view.activate()
view.block = True

%autopx

from cavity_flow2d import cavity_flow2D
nx, ny = 32, 32
x = cavity_flow2D(nx, ny, 1, 1.0, 1.0)

da = PETSc.DMDA().create([nx,ny],dof=4, stencil_width=1, stencil_type='star')
xn = da.createNaturalVec()
da.globalToNatural(x, xn)

from petsc4py import PETSc
sc, X = PETSc.Scatter.toZero(x)
sc.begin(xn, X, True, 'forward')
sc.end(xn, X, True, 'forward')
if not rank: Z = X[...].reshape(nx,ny,4)
else: Z = None

x.getLocalSize()

X.getLocalSize()

rank = PETSc.COMM_WORLD.Get_rank()

%autopx

Z = view['Z'][view['rank'].index(0)]
fig, axs = plt.subplots(2,2, figsize=(10,8))
titles = ['u', 'v', 'vorticity', 'temperature']
for idx, ax in enumerate(axs.ravel()):
    cs = ax.contourf(Z[:,:,idx], 100)
    fig.colorbar(cs, ax=ax, shrink=0.9)
    ax.set_title(titles[idx])