Missing file for lappy.

1 files changed, 39 insertions, 0 deletions

diff --git a/petsc.py b/petsc.py
new file mode 100644
index 0000000..7d59bf8
--- /dev/null
+++ b/petsc.py
@@ -0,0 +1,39 @@
+from scipy.sparse import coo_matrix
+import struct
+import numpy as np
+
+def load(fname):
+    MATRIX_MAGIC = 1211216
+    VECTOR_MAGIC = 1211214
+    
+    with open(fname, "rb") as f:
+        magic = struct.unpack(">i", f.read(4))[0]
+
+        if magic == MATRIX_MAGIC:
+            m = struct.unpack(">i", f.read(4))[0]
+            n = struct.unpack(">i", f.read(4))[0]
+            nonzero = struct.unpack(">i", f.read(4))[0]
+            nonzero_row = struct.unpack(">%di" %(m), f.read(m*4))
+            ks = struct.unpack(">%di" %(nonzero), f.read(nonzero*4))
+            xs = np.array(struct.unpack(">%dd" %(nonzero), f.read(nonzero*8)))
+
+            ii = np.zeros(nonzero, dtype=int)
+            jj = np.zeros(nonzero, dtype=int)
+            a = 0
+            for i in range(0, m):
+                for j in range(0, nonzero_row[i]):
+                    ii[a] = i
+                    jj[a] = ks[a]
+                    a += 1
+
+            return coo_matrix((xs, (ii, jj)), shape=(m, n))
+
+        elif magic == VECTOR_MAGIC:
+            m = struct.unpack(">i", f.read(4))[0]
+            xs = np.array(struct.unpack(">%dd" %(m), f.read(m*8)))
+            return xs
+            
+        else:
+            sys.stderr.write("Incorrect PETSc magic number (read %d).\n" %(magic))
+            exit(1)
+