xref: /petsc/src/binding/petsc4py/demo/legacy/perftest/driver.py (revision 55a74a43bb44613d95e937906bec3b8c3581b432) 
1*55a74a43SLisandro Dalcin#!/usr/bin/env python
2*55a74a43SLisandro Dalcinimport sys, petsc4py
3*55a74a43SLisandro Dalcinpetsc4py.init(sys.argv)
4*55a74a43SLisandro Dalcinfrom petsc4py import PETSc
5*55a74a43SLisandro Dalcinimport numpy as np
6*55a74a43SLisandro Dalcin
7*55a74a43SLisandro Dalcintry:
8*55a74a43SLisandro Dalcin    from matplotlib import pylab
9*55a74a43SLisandro Dalcinexcept ImportError:
10*55a74a43SLisandro Dalcin    pylab = None
11*55a74a43SLisandro Dalcin
12*55a74a43SLisandro Dalcin# this user class is an application
13*55a74a43SLisandro Dalcin# context for the nonlinear problem
14*55a74a43SLisandro Dalcin# at hand; it contains some parameters
15*55a74a43SLisandro Dalcin# and knows how to compute residuals
16*55a74a43SLisandro Dalcin
17*55a74a43SLisandro Dalcinclass AppCtx:
18*55a74a43SLisandro Dalcin
19*55a74a43SLisandro Dalcin    def __init__(self, nx, ny, nz):
20*55a74a43SLisandro Dalcin        self.n = np.array([nx, ny, nz], dtype='i')
21*55a74a43SLisandro Dalcin        self.h = np.array([1.0/(n-1) for n in self.n], dtype='d')
22*55a74a43SLisandro Dalcin        from App import formFunction
23*55a74a43SLisandro Dalcin        from App import formInitial
24*55a74a43SLisandro Dalcin        self._formFunction = formFunction
25*55a74a43SLisandro Dalcin        self._formInitial = formInitial
26*55a74a43SLisandro Dalcin
27*55a74a43SLisandro Dalcin    def formInitial(self, t, X):
28*55a74a43SLisandro Dalcin        xx = X.getArray(readonly=0).reshape(self.n, order='f')
29*55a74a43SLisandro Dalcin        self._formInitial(self.h, t, xx)
30*55a74a43SLisandro Dalcin
31*55a74a43SLisandro Dalcin    def formFunction(self, ts, t, X, Xdot, F):
32*55a74a43SLisandro Dalcin        n = self.n
33*55a74a43SLisandro Dalcin        h = self.h
34*55a74a43SLisandro Dalcin        x = X.getArray(readonly=1).reshape(n, order='f')
35*55a74a43SLisandro Dalcin        xdot = Xdot.getArray(readonly=1).reshape(n, order='f')
36*55a74a43SLisandro Dalcin        f = F[...].reshape(n, order='f')
37*55a74a43SLisandro Dalcin        self._formFunction(h, t, x, xdot, f)
38*55a74a43SLisandro Dalcin
39*55a74a43SLisandro Dalcin    def plot(self, t, x):
40*55a74a43SLisandro Dalcin        nx, ny, nz = self.n
41*55a74a43SLisandro Dalcin        from numpy import mgrid
42*55a74a43SLisandro Dalcin        #
43*55a74a43SLisandro Dalcin        U = x.getArray(readonly=1).reshape(nx,ny,nz, order='f')
44*55a74a43SLisandro Dalcin        #
45*55a74a43SLisandro Dalcin        X, Y =  mgrid[0:1:1j*nx,0:1:1j*ny]
46*55a74a43SLisandro Dalcin        Z = U[:,:,nz//2]
47*55a74a43SLisandro Dalcin        pylab.figure(0)
48*55a74a43SLisandro Dalcin        pylab.contourf(X,Y,Z)
49*55a74a43SLisandro Dalcin        pylab.colorbar()
50*55a74a43SLisandro Dalcin        pylab.plot(X.ravel(),Y.ravel(),'.k')
51*55a74a43SLisandro Dalcin        pylab.title('z=0.50')
52*55a74a43SLisandro Dalcin        pylab.xlabel('x')
53*55a74a43SLisandro Dalcin        pylab.ylabel('y')
54*55a74a43SLisandro Dalcin        pylab.axis('equal')
55*55a74a43SLisandro Dalcin        #
56*55a74a43SLisandro Dalcin        X, Y =  mgrid[0:1:1j*nx,0:1:1j*nz]
57*55a74a43SLisandro Dalcin        Z = U[:,ny//4,:]
58*55a74a43SLisandro Dalcin        pylab.figure(1)
59*55a74a43SLisandro Dalcin        pylab.contourf(X,Y,Z)
60*55a74a43SLisandro Dalcin        pylab.colorbar()
61*55a74a43SLisandro Dalcin        pylab.plot(X.ravel(),Y.ravel(),'.k')
62*55a74a43SLisandro Dalcin        pylab.title('y=0.25')
63*55a74a43SLisandro Dalcin        pylab.xlabel('x')
64*55a74a43SLisandro Dalcin        pylab.ylabel('z')
65*55a74a43SLisandro Dalcin        pylab.axis('equal')
66*55a74a43SLisandro Dalcin        #
67*55a74a43SLisandro Dalcin        X, Y =  mgrid[0:1:1j*ny,0:1:1j*nz]
68*55a74a43SLisandro Dalcin        Z = U[nx//2,:,:]
69*55a74a43SLisandro Dalcin        pylab.figure(2)
70*55a74a43SLisandro Dalcin        pylab.contourf(X,Y,Z)
71*55a74a43SLisandro Dalcin        pylab.colorbar()
72*55a74a43SLisandro Dalcin        pylab.plot(X.ravel(),Y.ravel(),'.k')
73*55a74a43SLisandro Dalcin        pylab.title('x=0.50')
74*55a74a43SLisandro Dalcin        pylab.xlabel('y')
75*55a74a43SLisandro Dalcin        pylab.ylabel('z')
76*55a74a43SLisandro Dalcin        pylab.axis('equal')
77*55a74a43SLisandro Dalcin
78*55a74a43SLisandro Dalcin
79*55a74a43SLisandro Dalcindef run_test(nx,ny,nz,samples,plot=False):
80*55a74a43SLisandro Dalcin    ts  = PETSc.TS().create()
81*55a74a43SLisandro Dalcin    ts.setType('theta')
82*55a74a43SLisandro Dalcin    ts.setTheta(1.0)
83*55a74a43SLisandro Dalcin    ts.setTimeStep(0.01)
84*55a74a43SLisandro Dalcin    ts.setTime(0.0)
85*55a74a43SLisandro Dalcin    ts.setMaxTime(1.0)
86*55a74a43SLisandro Dalcin    ts.setMaxSteps(10)
87*55a74a43SLisandro Dalcin    eft = PETSc.TS.ExactFinalTime.STEPOVER
88*55a74a43SLisandro Dalcin    ts.setExactFinalTime(eft)
89*55a74a43SLisandro Dalcin
90*55a74a43SLisandro Dalcin    x = PETSc.Vec().createSeq(nx*ny*nz)
91*55a74a43SLisandro Dalcin    ts.setSolution(x)
92*55a74a43SLisandro Dalcin    app = AppCtx(nx, ny, nz)
93*55a74a43SLisandro Dalcin    f = PETSc.Vec().createSeq(nx*ny*nz)
94*55a74a43SLisandro Dalcin    ts.setIFunction(app.formFunction, f)
95*55a74a43SLisandro Dalcin    ts.snes.setUseMF(1)
96*55a74a43SLisandro Dalcin    ts.snes.ksp.setType('cg')
97*55a74a43SLisandro Dalcin    ts.setFromOptions()
98*55a74a43SLisandro Dalcin    ts.setUp()
99*55a74a43SLisandro Dalcin
100*55a74a43SLisandro Dalcin    wt = 1e300
101*55a74a43SLisandro Dalcin    for i in range(samples):
102*55a74a43SLisandro Dalcin        app.formInitial(0, x)
103*55a74a43SLisandro Dalcin        t1 = PETSc.Log.getTime()
104*55a74a43SLisandro Dalcin        ts.solve(x)
105*55a74a43SLisandro Dalcin        t2 = PETSc.Log.getTime()
106*55a74a43SLisandro Dalcin        wt = min(wt,t2-t1)
107*55a74a43SLisandro Dalcin
108*55a74a43SLisandro Dalcin    if plot and pylab:
109*55a74a43SLisandro Dalcin        app.plot(ts.time, x)
110*55a74a43SLisandro Dalcin
111*55a74a43SLisandro Dalcin    return wt
112*55a74a43SLisandro Dalcin
113*55a74a43SLisandro DalcinOptDB = PETSc.Options()
114*55a74a43SLisandro Dalcin
115*55a74a43SLisandro Dalcin
116*55a74a43SLisandro Dalcinstart = OptDB.getInt('start', 12)
117*55a74a43SLisandro Dalcinstep = OptDB.getInt('step', 4)
118*55a74a43SLisandro Dalcinstop = OptDB.getInt('stop', start)
119*55a74a43SLisandro Dalcinsamples = OptDB.getInt('samples', 1)
120*55a74a43SLisandro Dalcin
121*55a74a43SLisandro Dalcinplot = OptDB.getBool('plot', False)
122*55a74a43SLisandro Dalcinif plot and not pylab:
123*55a74a43SLisandro Dalcin    PETSc.Sys.Print("matplotlib not available")
124*55a74a43SLisandro Dalcin
125*55a74a43SLisandro Dalcinfor n in range(start, stop+step, step):
126*55a74a43SLisandro Dalcin    nx = ny = nz = n+1
127*55a74a43SLisandro Dalcin    wt = run_test(nx,ny,nz,samples,plot)
128*55a74a43SLisandro Dalcin    PETSc.Sys.Print("Grid %3d x %3d x %3d -> %f  seconds (%2d samples)"
129*55a74a43SLisandro Dalcin                    % (nx,ny,nz,wt,samples))
130*55a74a43SLisandro Dalcin
131*55a74a43SLisandro Dalcinif plot and pylab:
132*55a74a43SLisandro Dalcin    pylab.show()
133