Skip to content

Instantly share code, notes, and snippets.

@1f0
Created March 4, 2021 03:09
Show Gist options
  • Save 1f0/753826db3fbc54014603f80dcf4845e7 to your computer and use it in GitHub Desktop.
Save 1f0/753826db3fbc54014603f80dcf4845e7 to your computer and use it in GitHub Desktop.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
n = 25
q = n * n
A = np.eye(q)
for i in range(1, n-1):
for j in range(1, n-1):
idx = n * i + j
neigh = [n*(i-1)+j, n*(i+1)+j, n*i+j-1, n*i+j+1]
for k in neigh:
A[idx][k] = -1 / len(neigh)
b = np.zeros(q)
b[n*(n//2) + n//2] = 1
z=np.linalg.solve(A, b)
zmat = z.reshape(n, n)
# plt.imshow(zmat)
# plt.colorbar()
# plt.show()
import mpl_toolkits.mplot3d
x = np.arange(n)
y = np.arange(n)
x, y = np.meshgrid(x, y)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf=ax.plot_surface(x, y, zmat, cmap=cm.coolwarm)
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment