ofan666 · February 21, 2012 11:11 · Mar 9, 2012 · Mar 9, 2012 · Feb 25, 2012 · Feb 25, 2012
diff --git a/TDMAsolver.py b/TDMAsolver.py
@@ -1,6 +1,9 @@
-import numpy as np
+try:
+   import numpypy as np    # for compatibility with numpy in pypy
+except:
+   import numpy as np      # if using numpy in cpython
 
-## Tri Diagonal Natrix Algorithm(a.k.a Thomas algorithm) solver
+## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver
 def TDMAsolver(a, b, c, d):
     '''
     TDMA solver, a b c d can be NumPy array type or Python list type.

diff --git a/TDMAsolver.py b/TDMAsolver.py
@@ -7,18 +7,18 @@ def TDMAsolver(a, b, c, d):
     refer to http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
     '''
     nf = len(a)     # number of equations
-    ac, bc, cc, dc = map(np.copy, (a, b, c, d))     # copy the array
+    ac, bc, cc, dc = map(np.array, (a, b, c, d))     # copy the array
     for it in xrange(1, nf):
         mc = ac[it]/bc[it-1]
         bc[it] = bc[it] - mc*cc[it-1] 
         dc[it] = dc[it] - mc*dc[it-1]
 
-    xc = np.copy(ac)
-    xc[nf-1] = dc[nf-1]/bc[nf-1]
+    xc = ac
+    xc[-1] = dc[-1]/bc[-1]
 
     for il in xrange(nf-2, -1, -1):
         xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]
 
-    del ac, bc, cc, dc  # delete variables from memory
+    del bc, cc, dc  # delete variables from memory
 
     return xc
diff --git a/TDMAsolver.py b/TDMAsolver.py
@@ -13,7 +13,7 @@ def TDMAsolver(a, b, c, d):
         bc[it] = bc[it] - mc*cc[it-1] 
         dc[it] = dc[it] - mc*dc[it-1]
 
-    xc = np.arange(nf, dtype=np.float)
+    xc = np.copy(ac)
     xc[nf-1] = dc[nf-1]/bc[nf-1]
 
     for il in xrange(nf-2, -1, -1):

diff --git a/TDMAsolver.py b/TDMAsolver.py
@@ -13,7 +13,7 @@ def TDMAsolver(a, b, c, d):
         bc[it] = bc[it] - mc*cc[it-1] 
         dc[it] = dc[it] - mc*dc[it-1]
 
-    xc = np.arange(nf, dtype=np.float64)
+    xc = np.arange(nf, dtype=np.float)
     xc[nf-1] = dc[nf-1]/bc[nf-1]
 
     for il in xrange(nf-2, -1, -1):

diff --git a/TDMAsolver.py b/TDMAsolver.py
@@ -13,7 +13,7 @@ def TDMAsolver(a, b, c, d):
         bc[it] = bc[it] - mc*cc[it-1] 
         dc[it] = dc[it] - mc*dc[it-1]
 
-    xc = np.arange(nf)
+    xc = np.arange(nf, dtype=np.float64)
     xc[nf-1] = dc[nf-1]/bc[nf-1]
 
     for il in xrange(nf-2, -1, -1):

diff --git a/TDMAsolver.py b/TDMAsolver.py
@@ -0,0 +1,24 @@
+import numpy as np
+
+## Tri Diagonal Natrix Algorithm(a.k.a Thomas algorithm) solver
+def TDMAsolver(a, b, c, d):
+    '''
+    TDMA solver, a b c d can be NumPy array type or Python list type.
+    refer to http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
+    '''
+    nf = len(a)     # number of equations
+    ac, bc, cc, dc = map(np.copy, (a, b, c, d))     # copy the array
+    for it in xrange(1, nf):
+        mc = ac[it]/bc[it-1]
+        bc[it] = bc[it] - mc*cc[it-1] 
+        dc[it] = dc[it] - mc*dc[it-1]
+
+    xc = np.arange(nf)
+    xc[nf-1] = dc[nf-1]/bc[nf-1]
+
+    for il in xrange(nf-2, -1, -1):
+        xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]
+
+    del ac, bc, cc, dc  # delete variables from memory
+
+    return xc
No results found