xref: /petsc/src/binding/petsc4py/demo/legacy/poisson3d/del2mat.py (revision 5a48edb989d3ea10d6aff6c0e26d581c18691deb)
1# file: del2mat.py
2
3from numpy import zeros
4from del2lib import del2apply
5
6class Del2Mat:
7
8    def __init__(self, n=1):
9        self.N = (n, n, n)
10        self.F = zeros([n+2]*3, order='f')
11
12    def create(self, A):
13        N = self.N
14        mat_size = A.getSize()
15        grid_eqs = N[0]*N[1]*N[2]
16        assert mat_size[0] == grid_eqs
17        assert mat_size[1] == grid_eqs
18
19    def mult(self, A, x, y):
20        "y <- A * x"
21        N, F = self.N, self.F
22        # get 3D arrays from vectos
23        xx = x.getArray(readonly=1).reshape(N, order='f')
24        yy = y.getArray(readonly=0).reshape(N, order='f')
25        # call Fortran subroutine
26        del2apply(F, xx, yy)
27
28    def multTranspose(self, A, x, y):
29        "y <- A' * x"
30        self.mult(x, y)
31
32    def getDiagonal(self, A, D):
33        "D[i] <- A[i,i]"
34        D[...] = 6.0
35