poisson_2d_dst.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import dstn, idstn


N = 256  # Number of points
L = 45.0  # Simulation size
S = np.linspace(-L/2.0, L/2.0, N)
X, Y = np.meshgrid(S, S)
DX = S[1] - S[0]  # Spatial step size
f = np.arange(0, N)
f[0] = 1e-60
P = np.pi*(f + 1)/L # Momenta in 1D
PX, PY = np.meshgrid(P, P) # Momenta in the x and y directions

def gaussian(a, sigma, x, y):
    return a*np.exp(-0.5*((X/L - x)**2 + (Y/L - y)**2)/sigma**2)


params = [[1.0, 0.01, 0.0, 0.0],
          [1.0, 0.01, -0.15, -0.15],
          [1.0, 0.01, 0.15, 0.15]]
rho = sum([gaussian(a, sigma, x, y) for a, sigma, x, y in params])


rho_p = dstn(rho)
phi_p = rho_p*(-1.0/(PX**2 + PY**2))
phi = idstn(phi_p)
plt.imshow(np.abs(phi))
plt.title('Solution to Poisson\'s Equation Using DST')
plt.show()
plt.close()

poisson_2d_fft.py

"""
https://en.wikipedia.org/wiki/Poisson%27s_equation
"""
import numpy as np
import matplotlib.pyplot as plt


N = 256  # Number of points
L = 45.0  # Simulation size
S = np.linspace(-L/2.0, L/2.0, N)
X, Y = np.meshgrid(S, S)
DX = S[1] - S[0]  # Spatial step size
f = np.fft.fftfreq(N)
f[0] = 1e-9
P = 2.0*np.pi*f*N/L # Momenta in 1D
PX, PY = np.meshgrid(P, P) # Momenta in the x and y directions

def gaussian(a, sigma, x, y):
    return a*np.exp(-0.5*((X/L - x)**2 + (Y/L - y)**2)/sigma**2)


params = [[1.0, 0.01, 0.0, 0.0],
          [1.0, 0.01, -0.15, -0.15],
          [1.0, 0.01, 0.15, 0.15]]
rho = sum([gaussian(a, sigma, x, y) for a, sigma, x, y in params])

rho_p = np.fft.fftn(rho)
phi_p = rho_p*(-1.0/(PX**2 + PY**2))
phi = np.fft.ifftn(phi_p)
plt.imshow(np.abs(phi) - np.amin(np.abs(phi)))
plt.title('Solution to Poisson\'s Equation Using FFT')
plt.show()
plt.close()

poisson_2d_implicit_method.py

import numpy as np
from scipy import sparse
from scipy.sparse import linalg
import matplotlib.pyplot as plt


# Constants
N = 256  # Number of points
L = 45.0  # Simulation size
S = np.linspace(-L/2.0, L/2.0, N)
X, Y = np.meshgrid(S, S)
DX = S[1] - S[0]  # Spatial step size

diag = (1.0/DX**2)*np.ones([N])
diags = np.array([diag, -2.0*diag, diag])
laplacian_1d = sparse.spdiags(diags, np.array([-1.0, 0.0, 1.0]), N, N)
LAPLACIAN = sparse.kronsum(laplacian_1d, laplacian_1d)


def gaussian(a, sigma, x, y):
    return a*np.exp(-0.5*((X/L - x)**2 + (Y/L - y)**2)/sigma**2)


params = [[1.0, 0.01, 0.0, 0.0],
          [1.0, 0.01, -0.15, -0.15],
          [1.0, 0.01, 0.15, 0.15]]
rho = sum([gaussian(a, sigma, x, y) for a, sigma, x, y in params])


phi = linalg.gcrotmk(LAPLACIAN, rho.reshape([N*N]))[0]

plt.imshow(np.abs(phi.reshape([N, N])))
plt.title('Solution to Poisson\'s Equation Using Implicit Methods')
plt.show()
plt.close()