1import unittest 2from petsc4py import PETSc 3from sys import getrefcount 4 5# -------------------------------------------------------------------- 6 7class MyODE: 8 """ 9 du/dt + u**2 = 0; 10 u0,u1,u2 = 1,2,3 11 """ 12 def __init__(self): 13 self.rhsfunction_calls = 0 14 self.rhsjacobian_calls = 0 15 self.ifunction_calls = 0 16 self.ijacobian_calls = 0 17 self.presolve_calls = 0 18 self.update_calls = 0 19 self.postsolve_calls = 0 20 self.monitor_calls = 0 21 22 def rhsfunction(self,ts,t,u,F): 23 # print ('MyODE.rhsfunction()') 24 self.rhsfunction_calls += 1 25 f = -(u * u) 26 f.copy(F) 27 28 def rhsjacobian(self,ts,t,u,J,P): 29 # print ('MyODE.rhsjacobian()') 30 self.rhsjacobian_calls += 1 31 P.zeroEntries() 32 diag = -2 * u 33 P.setDiagonal(diag) 34 P.assemble() 35 if J != P: J.assemble() 36 return True # same_nz 37 38 def ifunction(self,ts,t,u,du,F): 39 # print ('MyODE.ifunction()') 40 self.ifunction_calls += 1 41 f = du + u * u 42 f.copy(F) 43 44 def ijacobian(self,ts,t,u,du,a,J,P): 45 # print ('MyODE.ijacobian()') 46 self.ijacobian_calls += 1 47 P.zeroEntries() 48 diag = a + 2 * u 49 P.setDiagonal(diag) 50 P.assemble() 51 if J != P: J.assemble() 52 return True # same_nz 53 54 def monitor(self, ts, s, t, u): 55 self.monitor_calls += 1 56 dt = ts.time_step 57 ut = ts.vec_sol.norm() 58 #prn = PETSc.Sys.Print 59 #prn('TS: step %2d, T:%f, dT:%f, u:%f' % (s,t,dt,ut)) 60 61 62class BaseTestTSNonlinear(object): 63 64 TYPE = None 65 66 def setUp(self): 67 self.ts = PETSc.TS().create(PETSc.COMM_SELF) 68 eft = PETSc.TS.ExactFinalTime.STEPOVER 69 self.ts.setExactFinalTime(eft) 70 ptype = PETSc.TS.ProblemType.NONLINEAR 71 self.ts.setProblemType(ptype) 72 self.ts.setType(self.TYPE) 73 if PETSc.ScalarType().dtype.char in 'fF': 74 snes = self.ts.getSNES() 75 snes.setTolerances(rtol=1e-6) 76 77 def tearDown(self): 78 self.ts = None 79 PETSc.garbage_cleanup() 80 81 82class BaseTestTSNonlinearRHS(BaseTestTSNonlinear): 83 84 def testSolveRHS(self): 85 ts = self.ts 86 dct = self.ts.getDict() 87 self.assertTrue(dct is not None) 88 self.assertTrue(type(dct) is dict) 89 90 ode = MyODE() 91 J = PETSc.Mat().create(ts.comm) 92 J.setSizes(3); 93 J.setFromOptions() 94 J.setUp() 95 u, f = J.createVecs() 96 97 ts.setAppCtx(ode) 98 ts.setRHSFunction(ode.rhsfunction, f) 99 ts.setRHSJacobian(ode.rhsjacobian, J, J) 100 ts.setMonitor(ode.monitor) 101 102 ts.snes.ksp.pc.setType('none') 103 104 T0, dT, nT = 0.00, 0.1, 10 105 T = T0 + nT*dT 106 ts.setTime(T0) 107 ts.setTimeStep(dT) 108 ts.setMaxTime(T) 109 ts.setMaxSteps(nT) 110 ts.setFromOptions() 111 u[0], u[1], u[2] = 1, 2, 3 112 ts.solve(u) 113 114 self.assertTrue(ode.rhsfunction_calls > 0) 115 self.assertTrue(ode.rhsjacobian_calls > 0) 116 117 dct = self.ts.getDict() 118 self.assertTrue('__appctx__' in dct) 119 self.assertTrue('__rhsfunction__' in dct) 120 self.assertTrue('__rhsjacobian__' in dct) 121 self.assertTrue('__monitor__' in dct) 122 123 n = ode.monitor_calls 124 ts.monitor(ts.step_number, ts.time) 125 self.assertEqual(ode.monitor_calls, n+1) 126 n = ode.monitor_calls 127 ts.monitorCancel() 128 ts.monitor(ts.step_number, ts.time) 129 self.assertEqual(ode.monitor_calls, n) 130 131 def testFDColorRHS(self): 132 ts = self.ts 133 ode = MyODE() 134 J = PETSc.Mat().create(ts.comm) 135 J.setSizes(5); J.setType('aij') 136 J.setPreallocationNNZ(nnz=1) 137 u, f = J.createVecs() 138 139 ts.setAppCtx(ode) 140 ts.setRHSFunction(ode.rhsfunction, f) 141 ts.setRHSJacobian(ode.rhsjacobian, J, J) 142 ts.setMonitor(ode.monitor) 143 144 T0, dT, nT = 0.00, 0.1, 10 145 T = T0 + nT*dT 146 ts.setTime(T0) 147 ts.setTimeStep(dT) 148 ts.setMaxTime(T) 149 ts.setMaxSteps(nT) 150 ts.setFromOptions() 151 u[0], u[1], u[2] = 1, 2, 3 152 153 ts.setSolution(u) 154 ode.rhsjacobian(ts,0,u,J,J) 155 ts.setUp() 156 ts.snes.setUseFD(True) 157 ts.solve(u) 158 159 def testResetAndSolveRHS(self): 160 self.ts.reset() 161 self.ts.setStepNumber(0) 162 self.testSolveRHS() 163 self.ts.reset() 164 self.ts.setStepNumber(0) 165 self.testSolveRHS() 166 self.ts.reset() 167 168class BaseTestTSNonlinearI(BaseTestTSNonlinear): 169 170 def testSolveI(self): 171 ts = self.ts 172 dct = self.ts.getDict() 173 self.assertTrue(dct is not None) 174 self.assertTrue(type(dct) is dict) 175 176 ode = MyODE() 177 J = PETSc.Mat().create(ts.comm) 178 J.setSizes(3); 179 J.setFromOptions() 180 J.setUp() 181 u, f = J.createVecs() 182 183 ts.setAppCtx(ode) 184 ts.setIFunction(ode.ifunction, f) 185 ts.setIJacobian(ode.ijacobian, J, J) 186 ts.setMonitor(ode.monitor) 187 188 ts.snes.ksp.pc.setType('none') 189 190 T0, dT, nT = 0.00, 0.1, 10 191 T = T0 + nT*dT 192 ts.setTime(T0) 193 ts.setTimeStep(dT) 194 ts.setMaxTime(T) 195 ts.setMaxSteps(nT) 196 ts.setFromOptions() 197 u[0], u[1], u[2] = 1, 2, 3 198 ts.solve(u) 199 200 self.assertTrue(ode.ifunction_calls > 0) 201 self.assertTrue(ode.ijacobian_calls > 0) 202 203 dct = self.ts.getDict() 204 self.assertTrue('__appctx__' in dct) 205 self.assertTrue('__ifunction__' in dct) 206 self.assertTrue('__ijacobian__' in dct) 207 self.assertTrue('__monitor__' in dct) 208 209 n = ode.monitor_calls 210 ts.monitor(ts.step_number, ts.time) 211 self.assertEqual(ode.monitor_calls, n+1) 212 n = ode.monitor_calls 213 ts.monitorCancel() 214 ts.monitor(ts.step_number, ts.time) 215 self.assertEqual(ode.monitor_calls, n) 216 217 def testFDColorI(self): 218 ts = self.ts 219 ode = MyODE() 220 J = PETSc.Mat().create(ts.comm) 221 J.setSizes(5); J.setType('aij') 222 J.setPreallocationNNZ(nnz=1) 223 J.setFromOptions() 224 u, f = J.createVecs() 225 226 ts.setAppCtx(ode) 227 ts.setIFunction(ode.ifunction, f) 228 ts.setIJacobian(ode.ijacobian, J, J) 229 ts.setMonitor(ode.monitor) 230 231 T0, dT, nT = 0.00, 0.1, 10 232 T = T0 + nT*dT 233 ts.setTime(T0) 234 ts.setTimeStep(dT) 235 ts.setMaxTime(T) 236 ts.setMaxSteps(nT) 237 ts.setFromOptions() 238 u[0], u[1], u[2] = 1, 2, 3 239 240 ts.setSolution(u) 241 ode.ijacobian(ts,0,u,0*u,1,J,J) 242 ts.setUp() 243 ts.snes.setUseFD(True) 244 ts.solve(u) 245 246 def testResetAndSolveI(self): 247 self.ts.reset() 248 self.ts.setStepNumber(0) 249 self.testSolveI() 250 self.ts.reset() 251 self.ts.setStepNumber(0) 252 self.testSolveI() 253 self.ts.reset() 254 255class TestTSBeuler(BaseTestTSNonlinearRHS,BaseTestTSNonlinearI, 256 unittest.TestCase): 257 TYPE = PETSc.TS.Type.BEULER 258 259class TestTSCN(BaseTestTSNonlinearRHS,BaseTestTSNonlinearI, 260 unittest.TestCase): 261 TYPE = PETSc.TS.Type.CN 262 263class TestTSTheta(BaseTestTSNonlinearRHS, BaseTestTSNonlinearI, 264 unittest.TestCase): 265 TYPE = PETSc.TS.Type.THETA 266 267class TestTSAlpha(BaseTestTSNonlinearRHS, BaseTestTSNonlinearI, 268 unittest.TestCase): 269 TYPE = PETSc.TS.Type.ALPHA 270 271# -------------------------------------------------------------------- 272 273if __name__ == '__main__': 274 unittest.main() 275 276# -------------------------------------------------------------------- 277