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