summaryrefslogtreecommitdiff
path: root/petsc.py
blob: 7d59bf8fd58570101a14a9de9fd391819d9b822f (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
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)