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