Poisson problem with dirichlet b.c. as grad=0 at the unit square edges and a list of source nodes (useful for heat sensor data interpolation)

heat_equation_delta_sources.py

# Dirichlet bc eliminated
import numpy as np
import matplotlib.pyplot as plt
import scipy.sparse as sp
import scipy.sparse.linalg as spla

N = 64
sources = [
    (0.3, 0.3, 15),
    (0.7, 0.7, 25)
]

def at(i,j):
    return i*(N-1) + j

h = 1 / N
one = np.ones(N-1)
stamp = sp.dok_matrix(sp.spdiags([-one, 2*one, -one], [-1,0,1], N-1, N-1) / h**2)
stamp[0,1] = -stamp[0,0]
stamp[-1, -2] = -stamp[-1,-1]
Inm1 = sp.eye(N-1)
A = sp.dok_matrix(sp.kron(stamp, Inm1) + sp.kron(Inm1, stamp))
f = np.zeros((N-1)**2)

for x, y, z in sources:
    i = int(round(x*N))
    j = int(round(y*N))
    # sources can not lie directly adjecent to the boundary
    assert 0 < i < N and 0 < j < N
    # Dirichlet nodes - new stencil for the source node:
    (
        A[at(i,j),at(i-1,j-1)], A[at(i,j),at(i-1,j)], A[at(i,j),at(i-1,j+1)],
        A[at(i,j),at(i  ,j-1)], A[at(i,j),at(i  ,j)], A[at(i,j),at(i  ,j+1)],
        A[at(i,j),at(i+1,j-1)], A[at(i,j),at(i+1,j)], A[at(i,j),at(i+1,j+1)]) = (
        0, 0, 0,
        0, 1, 0,
        0, 0, 0
    )
    # s.t. we force u(x,y) to be f(x,y) as given
    f[at(i, j)] = z
    #for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
    #    A[at(i+di, j+dj), at(i, j)] = 0
    #    f[at(i+di, j+dj)] += z / h**2

u = spla.spsolve(sp.csr_matrix(A), f).reshape((N-1, N-1))
plt.imshow(u, vmin=0, vmax=30)
plt.colorbar()