xref: /petsc/lib/petsc/bin/PetscBinaryIO_tests.py (revision dc97da86830806a16cfaacfab740baa2e54c8135) 
1import unittest
2from petsc4py import PETSc
3import os
4from PetscBinaryIO import *
5
6class TestPetscBinaryIO(unittest.TestCase):
7    def setUp(self):
8        try:
9            os.remove('test.dat')
10        except OSError:
11            pass
12        try:
13            os.remove('test.dat.info')
14        except OSError:
15            pass
16
17    def test_VecRead(self):
18        """Test reading a Vec"""
19        array = np.array([1.1, 2.2, 3.3])
20        vec = PETSc.Vec().createSeq(3)
21        vec[...] = array
22        viewer = PETSc.Viewer().createBinary('test.dat', PETSc.Viewer.Mode.W)
23        vec.view(viewer)
24        viewer.destroy()
25        vec.destroy()
26
27        result, = PetscBinaryIO().readBinaryFile('test.dat')
28        self.assertTrue(np.allclose(array, result))
29
30    def test_VecWrite(self):
31        """Test writing a Vec"""
32        array = np.array([1.1, 2.2, 3.3])
33        PetscBinaryIO().writeBinaryFile('test.dat', [array.view(Vec),])
34
35        vec = PETSc.Vec().createSeq(3)
36        vec.set(0.)
37        viewer = PETSc.Viewer().createBinary('test.dat', PETSc.Viewer.Mode.R)
38        vec.load(viewer)
39        viewer.destroy()
40
41        self.assertTrue(np.allclose(array, vec[...]))
42        vec.destroy()
43
44    def test_ISRead(self):
45        """Test reading an IS"""
46        indices = np.array([3,4,5])
47        anis = PETSc.IS().createGeneral(list(indices))
48        viewer = PETSc.Viewer().createBinary('test.dat', PETSc.Viewer.Mode.W)
49        anis.view(viewer)
50        viewer.destroy()
51        anis.destroy()
52
53        result, = PetscBinaryIO().readBinaryFile('test.dat')
54        self.assertTrue((indices == result).all())
55
56    def test_MatRead(self):
57        """Test reading a Mat"""
58        mat = PETSc.Mat().createAIJ(2)
59        mat[0,0] = 1.1
60        mat[0,1] = 2.1
61        mat[1,1] = 3.1
62        mat.assemble()
63
64        vals = np.array([1.1,2.1,3.1])
65        counts = np.array([0,2,3])
66        cols = np.array([0,1,1])
67
68        viewer = PETSc.Viewer().createBinary('test.dat', PETSc.Viewer.Mode.W)
69        mat.view(viewer)
70        viewer.destroy()
71        mat.destroy()
72
73        result, = PetscBinaryIO().readBinaryFile('test.dat')
74        self.assertTrue(np.allclose(vals, result[1][2]))
75        self.assertTrue((counts == result[1][0]).all())
76        self.assertTrue((cols == result[1][1]).all())
77        self.assertTrue((2,2) == result[0])
78
79    def test_MatWrite(self):
80        """Test writing a Mat"""
81        vals = np.array([1.1,2.1,3.1])
82        counts = np.array([0,2,3])
83        cols = np.array([0,1,1])
84        mat = MatSparse(((2,2),(counts,cols,vals)))
85
86        dense = np.array([1.1,2.1,0.0,3.1])
87
88        PetscBinaryIO().writeBinaryFile('test.dat', [mat,])
89
90        mat = PETSc.Mat().createAIJ(2)
91        viewer = PETSc.Viewer().createBinary('test.dat', PETSc.Viewer.Mode.R)
92        mat.load(viewer)
93        viewer.destroy()
94
95        self.assertTrue(np.allclose(dense, mat[:,:].ravel()))
96        mat.destroy()
97
98
99if __name__ == '__main__':
100    unittest.main()
101    try:
102        os.remove('test.dat')
103    except OSError:
104        pass
105    try:
106        os.remove('test.dat.info')
107    except OSError:
108        pass
109