Instantly share code, notes, and snippets.

# wiseodd/pde_numpy.py

Created January 8, 2017 05:48
Show Gist options
• Save wiseodd/c08d5a2b02b1957a16f886ab7044032d to your computer and use it in GitHub Desktop.
Implementation of Finite Difference solution of Laplace Equation in Numpy and Theano
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
 import matplotlib import matplotlib.pyplot as plt import numpy as np import time from mpl_toolkits.mplot3d import Axes3D from matplotlib import cm plt.ion() fig = plt.figure() ax = fig.add_subplot(111, projection='3d') ax.set_zlim(-1.01, 1.01) def draw_plot(x, y, U): ax.clear() ax.set_zlim(-1.01, 1.01) ax.plot_surface(x, y, U, rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=True) plt.pause(1e-5) # Create 21x21 mesh grid m = 21 mesh_range = np.arange(-1, 1, 2/(m-1)) x, y = np.meshgrid(mesh_range, mesh_range) # Initial condition U = np.exp(-5 * (x**2 + y**2)) draw_plot(x, y, U) n = list(range(1, m-1)) + [m-2] e = n s = [0] + list(range(0, m-2)) w = s def pde_step(U): """ PDE calculation at a single time step t """ return (U[n, :]+U[:, e]+U[s, :]+U[:, w])/4. k = 5 U_step = U for it in range(500): U_step = pde_step(U_step) # Every k steps, draw the graphics if it % k == 0: draw_plot(x, y, U_step) while True: plt.pause(1e-5)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
 import matplotlib import matplotlib.pyplot as plt import numpy as np import time from mpl_toolkits.mplot3d import Axes3D from matplotlib import cm import theano as th from theano import tensor as T plt.ion() fig = plt.figure() ax = fig.add_subplot(111, projection='3d') ax.set_zlim(-1.01, 1.01) def draw_plot(x, y, U): ax.clear() ax.set_zlim(-1.01, 1.01) ax.plot_surface(x, y, U, rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=True) plt.pause(1e-5) # Create 21x21 mesh grid m = 21 mesh_range = np.arange(-1, 1, 2/(m-1)) x_arr, y_arr = np.meshgrid(mesh_range, mesh_range) # Initialize variables x, y = th.shared(x_arr), th.shared(y_arr) U = T.exp(-5 * (x**2 + y**2)) draw_plot(x_arr, y_arr, U.eval()) n = list(range(1, m-1)) + [m-2] e = n s = [0] + list(range(0, m-2)) w = s def pde_step(U): """ PDE calculation at a single time step t """ return (U[n, :]+U[:, e]+U[s, :]+U[:, w])/4. k = 5 # Batch process the PDE calculation, calculate together k steps result, updates = th.scan(fn=pde_step, outputs_info=U, n_steps=k) final_result = result[-1] calc_pde = th.function(inputs=[U], outputs=final_result, updates=updates) U_step = U.eval() for it in range(100): # Every k steps, draw the graphics U_step = calc_pde(U_step) draw_plot(x_arr, y_arr, U_step) while True: plt.pause(1e-5)