Poisson equation in 2D

import numpy as np;
import matplotlib.pyplot as plt;
plt.rcParams.update({"text.usetex":True});
%config InlineBackend.figure_format = "svg"
import scipy.sparse as scs
from scipy.sparse.linalg import spsolve

Nx = Ny = 30; h = 1/(Nx-1);
A = np.zeros((Ny**2, Nx**2)); b = np.zeros((Nx**2,1)); u = np.zeros((Nx**2,1));

def get_b(i,j,h):
    x = i*h; y = j*h;
    r = 2*(1-6*x**2)*y**2*(1-y**2) + 2*(1-6*y**2)*x**2*(1-x**2)
    return r


for j in np.arange(0,Ny):
    for i in np.arange(0,Nx):
        l = i+j*Nx
        if i !=0 and j !=0 and i != Nx-1 and j != Ny-1:
            A[l,l] = 4; 
            A[l,l-1] = -1; A[l,l+1] = -1;
            A[l,l-Nx] = -1; A[l,l+Nx] = -1;
            b[l,0] = h**2*get_b(i,j,h);
        else: # We are at the boundary
            A[l,l] = 1;

# Convert A into sparse CSR
A = scs.csr_matrix(A);
b = scs.csr_matrix(b);
print(A.shape)
print(b.shape)

(900, 900)
(900, 1)

u = spsolve(A,b);

xi = yi = np.linspace(0,1,Nx); [Xi,Yi] = np.meshgrid(xi,yi);
U = np.reshape(u, (Ny, Nx));

xa = ya = np.linspace(0,1,50); [Xa, Ya] = np.meshgrid(xa, ya);
Ua = (Xa**2-Xa**4)*(Ya**4-Ya**2)

plt.contour(Xi,Yi,U);
plt.contour(Xa,Ya,Ua,cmap='jet')
ax = plt.gca(); ax.set_aspect(1)