#!/usr/bin/env python
# coding: utf-8

# # p16: Poisson equation in 2-D

# We solve the following Poisson problem
# 
# $$
# u_{xx} + u_{yy} = 10\sin(8x(y-1)), \qquad -1 < x,y < 1, \qquad u=0 \quad \mbox{on boundary}
# $$

# In[1]:


get_ipython().run_line_magic('config', "InlineBackend.figure_format='svg'")
from chebPy import cheb
from numpy import meshgrid,sin,dot,eye,kron,zeros,reshape,linspace
from matplotlib.pyplot import figure,subplot,plot,title,axis,xlabel,ylabel,spy
from matplotlib import cm
from scipy.linalg import solve
from scipy.interpolate import RegularGridInterpolator


# In[2]:


N = 24; D,x = cheb(N); y = x;
xx,yy = meshgrid(x[1:N],y[1:N],indexing='ij')
xx = reshape(xx,(N-1)**2,order='F')
yy = reshape(yy,(N-1)**2,order='F')
f = 10*sin(8*xx*(yy-1))
D2 = dot(D,D); D2 = D2[1:N,1:N]; I = eye(N-1)
L = kron(I,D2) + kron(D2,I)
# Plot sparsity pattern
figure(figsize=(6,6)), spy(L)
# Solve Lu=f
u = solve(L,f)
# Convert 1-d vectors to 2-d
uu = zeros((N+1,N+1)); uu[1:N,1:N] = reshape(u,(N-1,N-1),order='F')
[xx,yy] = meshgrid(x,y,indexing='ij')
value = uu[3*N//4,3*N//4]

# Interpolate to finer mesh just for visualization
f = RegularGridInterpolator((x,y),uu,method='cubic')
xxx = linspace(-1.0,1.0,50); X,Y = meshgrid(xxx,xxx,indexing='ij')
uuu = f((X,Y))
fig = figure(figsize=(6,6))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X,Y,uuu,rstride=1,cstride=1,cmap=cm.jet,edgecolor='black')
title("$u(2^{-1/2},2^{-1/2})$="+str(value))
xlabel("x"), ylabel("y");