Skip to content

Instantly share code, notes, and snippets.

@lobodin
Created May 19, 2012 10:17
Show Gist options
  • Save lobodin/2730359 to your computer and use it in GitHub Desktop.
Save lobodin/2730359 to your computer and use it in GitHub Desktop.
Poisson equation solution using TDMA
import math
import matplotlib.pyplot as plt
def tdma(a, b, c, d):
'''
a[i] * x[i-1] + b[i] * x[i] + c[i] * x[i+1] = d[i]
'''
length = len(a)
for i in xrange(1, length):
tmp = a[i] / b[i-1]
b[i] = b[i] - tmp * c[i-1]
d[i] = d[i] - tmp * d[i-1]
x = a
x[-1] = d[-1] / b[-1]
for i in xrange(length - 2, -1, -1):
x[i] = (d[i] - c[i] * x[i+1]) / b[i]
return x
def poisson_equation(h, x0, xn, a0, b0, c0, an, bn, cn, function):
'''
Problem definition:
Laplace(u) = f(x), x0 <= x <= xn
a0 * du/dx + b0 * u(x) = c0, x = x0
an * du/dx + db * u(x) = cn, x = xn
Linear equations system:
(h * b0 - a0) * u0 + a0 * u1 = c0 * h;
u[i-1] - 2u[i] + u[i+1] = h^2 * f[i], i = 1..n-1;
-a0 * un + (an + h * bn) * u[n+1] = cn * h;
'''
n = int((xn - x0) / h)
a, b, c, d = [], [], [], []
for i in xrange(n + 1):
a.append(1)
b.append(-2)
c.append(1)
d.append(h * h * function(x0 + i * h))
a[n] = -an
b[0] = h * b0 - a0
b[n] = an + h * bn
c[0] = a0
d[0] = c0 * h
d[n] = cn * h
u = tdma(a, b, c, d)
return u
def test():
def example_function(x):
return math.sin(x)
def exact_solution(i, h, x0):
x = x0 + i * h
return -math.sin(x) - 0.0116 * x + 1.7082
def plot(u, exact_u):
plt.plot(range(len(u)), u)
plt.plot(range(len(exact_u)), exact_u)
plt.show()
h = 0.1 # step
x0 = 0.0 # left border
xn = 20.0 # right border
a0, b0, c0, an, db, cn = 0.7, 1.0, 1.0, 0.3, 2.0, 1.0 # coefficients of boundary values
u = poisson_equation(h, x0, xn, a0, b0, c0, an, db, cn, example_function)
exact_u = map(lambda i: exact_solution(i, h, x0), range(len(u)))
plot(u, exact_u)
test()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment