1from petsc4py import PETSc 2import unittest 3 4# -------------------------------------------------------------------- 5 6class BaseTestVec(object): 7 8 COMM = None 9 TYPE = None 10 11 def setUp(self): 12 v = PETSc.Vec() 13 v.create(self.COMM) 14 v.setSizes(100) 15 v.setType(self.TYPE) 16 self.vec = v 17 18 def tearDown(self): 19 self.vec.destroy() 20 self.vec = None 21 PETSc.garbage_cleanup() 22 23 def testDuplicate(self): 24 self.vec.set(1) 25 vec = self.vec.duplicate() 26 self.assertFalse(self.vec.equal(vec)) 27 self.assertEqual(self.vec.sizes, vec.sizes) 28 del vec 29 30 def testCopy(self): 31 self.vec.set(1) 32 vec = self.vec.duplicate() 33 self.vec.copy(vec) 34 self.assertTrue(self.vec.equal(vec)) 35 del vec 36 37 def testDot(self): 38 self.vec.set(1) 39 d = self.vec.dot(self.vec) 40 self.assertAlmostEqual(abs(d), self.vec.getSize()) 41 self.vec.dotBegin(self.vec) 42 d = self.vec.dotEnd(self.vec) 43 self.assertAlmostEqual(abs(d), self.vec.getSize()) 44 45 def testNorm(self): 46 from math import sqrt 47 self.vec.set(1) 48 n1 = self.vec.norm(PETSc.NormType.NORM_1) 49 n2 = self.vec.norm(PETSc.NormType.NORM_2) 50 ni = self.vec.norm(PETSc.NormType.NORM_INFINITY) 51 self.assertAlmostEqual(n1, self.vec.getSize()) 52 self.assertAlmostEqual(n2, sqrt(self.vec.getSize())) 53 self.assertAlmostEqual(n2, self.vec.norm()) 54 self.assertAlmostEqual(ni, 1.0) 55 self.vec.normBegin(PETSc.NormType.NORM_1) 56 nn1 = self.vec.normEnd(PETSc.NormType.NORM_1) 57 self.assertAlmostEqual(nn1, n1) 58 self.vec.normBegin() 59 nn2 = self.vec.normEnd() 60 self.assertAlmostEqual(nn2, n2) 61 self.vec.normBegin(PETSc.NormType.NORM_INFINITY) 62 nni = self.vec.normEnd(PETSc.NormType.NORM_INFINITY) 63 self.assertAlmostEqual(nni, ni) 64 65 def testNormalize(self): 66 from math import sqrt 67 self.vec.set(1) 68 n2 = self.vec.normalize() 69 self.assertAlmostEqual(n2, sqrt(self.vec.getSize())) 70 self.assertAlmostEqual(1, self.vec.norm()) 71 72 def testSumMinMax(self): 73 self.vec.set(1) 74 self.assertEqual(self.vec.sum(), self.vec.getSize()) 75 self.vec.set(-7) 76 self.assertEqual(self.vec.min()[1], -7) 77 self.vec.set(10) 78 self.assertEqual(self.vec.max()[1], 10) 79 80 def testSwap(self): 81 v1 = self.vec 82 v2 = v1.duplicate() 83 v1.set(1) 84 v2.set(2) 85 v1.swap(v2) 86 idx, _ = self.vec.getOwnershipRange() 87 self.assertEqual(v1[idx], 2) 88 self.assertEqual(v2[idx], 1) 89 90 def testBsize(self): 91 self.vec.setBlockSize(1) 92 self.assertEqual(self.vec.getBlockSize(), 1) 93 self.vec.setBlockSize(1) 94 95 def testGetSetVals(self): 96 start, end = self.vec.getOwnershipRange() 97 self.vec[start] = -7 98 self.vec[end-1] = -7 99 self.vec.assemble() 100 self.assertEqual(self.vec[start], -7) 101 self.assertEqual(self.vec[end-1], -7) 102 for i in range(start, end): self.vec[i] = i 103 self.vec.assemble() 104 values = [self.vec[i] for i in range(start, end)] 105 self.assertEqual(values, list(range(start, end))) 106 sz = self.vec.getSize() 107 self.assertEqual(self.vec.sum(), (sz-1)/2.0*sz) 108 109 def testGetSetValsBlocked(self): 110 return 111 lsize, gsize = self.vec.getSizes() 112 start, end = self.vec.getOwnershipRange() 113 bsizes = list(range(1, lsize+1)) 114 nblocks = list(range(1, lsize+1)) 115 compat = [(bs, nb) 116 for bs in bsizes if not (gsize%bs or lsize % bs) 117 for nb in nblocks if bs*nb <= lsize] 118 for bsize, nblock in compat: 119 self.vec.setBlockSize(bsize) 120 bindex = [start//bsize+i for i in range(nblock)] 121 bvalue = [float(i) for i in range(nblock*bsize)] 122 self.vec.setValuesBlocked(bindex, bvalue) 123 self.vec.assemble() 124 index = [start+i for i in range(nblock*bsize)] 125 value = self.vec.getValues(index) 126 self.assertEqual(bvalue, list(value)) 127 128 def testGetSetArray(self): 129 self.vec.set(1) 130 arr0 = self.vec.getArray().copy() 131 self.assertEqual(arr0.sum(), self.vec.getLocalSize()) 132 arr0 = self.vec.getArray().copy() 133 self.vec.setRandom() 134 arr1 = self.vec.getArray().copy() 135 self.vec.setArray(arr1) 136 arr1 = self.vec.getArray().copy() 137 arr2 = self.vec.getArray().copy() 138 self.assertTrue((arr1 == arr2).all()) 139 import numpy 140 refs = self.vec.getRefCount() 141 arr3 = numpy.asarray(self.vec) 142 self.assertEqual(self.vec.getRefCount(), refs+1) 143 self.assertTrue((arr1 == arr3).all()) 144 arr3[:] = 0 145 self.assertAlmostEqual(abs(self.vec.sum()), 0) 146 self.assertEqual(self.vec.max()[1], 0) 147 self.assertEqual(self.vec.min()[1], 0) 148 self.vec.set(1) 149 self.assertAlmostEqual(abs(arr3.sum()), self.vec.getLocalSize()) 150 self.assertEqual(arr3.min(), 1) 151 self.assertEqual(arr3.max(), 1) 152 del arr3 153 self.assertEqual(self.vec.getRefCount(), refs) 154 155 def testPlaceArray(self): 156 self.vec.set(1) 157 array = self.vec.getArray().copy() 158 self.vec.placeArray(array) 159 array[:] = 2 160 self.assertAlmostEqual(abs(self.vec.sum()), 2*self.vec.getSize()) 161 self.vec.resetArray() 162 self.assertAlmostEqual(abs(self.vec.sum()), self.vec.getSize()) 163 164 def testLocalVector(self): 165 rank = self.vec.getComm().Get_rank() 166 self.vec.getArray()[:] = rank + 1 167 ln = self.vec.getLocalSize() 168 lvec = self.vec.createLocalVector() 169 self.vec.getLocalVector(lvec) 170 self.assertEqual(abs(lvec.sum()), (rank+1)*ln) 171 self.vec.restoreLocalVector(lvec) 172 self.vec.getLocalVector(lvec,readonly=True) 173 self.assertEqual(abs(lvec.sum()), (rank+1)*ln) 174 self.vec.restoreLocalVector(lvec,readonly=True) 175 lvec.destroy() 176 177 def testSetOption(self): 178 opt1 = PETSc.Vec.Option.IGNORE_OFF_PROC_ENTRIES 179 opt2 = PETSc.Vec.Option.IGNORE_NEGATIVE_INDICES 180 for opt in [opt1, opt2]*2: 181 for flag in [True,False]*2: 182 self.vec.setOption(opt,flag) 183 184 def testGetSetItem(self): 185 v = self.vec 186 w = v.duplicate() 187 # 188 v[...] = 7 189 self.assertEqual(v.max()[1], 7) 190 self.assertEqual(v.min()[1], 7) 191 # 192 v.setRandom() 193 w[...] = v 194 self.assertTrue(w.equal(v)) 195 # 196 v.setRandom() 197 w[...] = v.getArray() 198 self.assertTrue(w.equal(v)) 199 # 200 s, e = v.getOwnershipRange() 201 v.setRandom() 202 w[s:e] = v.getArray().copy() 203 w.assemble() 204 self.assertTrue(w.equal(v)) 205 w1, v1 = w[s], v[s] 206 w2, v2 = w[e-1], v[e-1] 207 self.assertEqual(w1, v1) 208 self.assertEqual(w2, v2) 209 210 def testMAXPY(self): 211 y = self.vec 212 y.set(1) 213 x = [y.copy() for _ in range(3)] 214 a = [1]*len(x) 215 y.maxpy(a, x) 216 z = y.duplicate() 217 z.set(len(x)+1) 218 assert (y.equal(z)) 219 220 def testBinOp(self): 221 x = self.vec 222 x.set(1) 223 n = x.getSize() 224 y = 2 + 2*x + 1 - x*3 - 1 225 self.assertEqual(y.min()[1], 1) 226 self.assertEqual(y.max()[1], 1) 227 z = (4*x)/(2*y) 228 self.assertEqual(z.min()[1], 2) 229 self.assertEqual(z.max()[1], 2) 230 z = z/2 231 self.assertEqual(z.min()[1], 1) 232 self.assertEqual(z.max()[1], 1) 233 s = (+x) @ (-y) 234 self.assertEqual(s, -n) 235 # 236 M, N = n, 2*n 237 A = PETSc.Mat().createDense((M, N), comm=self.COMM) 238 A.setUp() 239 rs, re = A.getOwnershipRange() 240 cs, ce = A.getOwnershipRangeColumn() 241 a, b = 3, 5 242 for i in range(rs, re): 243 for j in range(N): 244 A[i, j] = a*i + b*j 245 A.assemble() 246 y = x @ A 247 self.assertEqual(y.getSize(), N) 248 for i in range(cs, ce): 249 self.assertEqual(y[i], a*M*(M-1)/2 + b*i*M) 250 y.set(1) 251 z = A @ y 252 self.assertEqual(z.getSize(), M) 253 for i in range(rs, re): 254 self.assertEqual(z[i], b*N*(N-1)/2 + a*i*N) 255 256# -------------------------------------------------------------------- 257 258class TestVecSeq(BaseTestVec, unittest.TestCase): 259 COMM = PETSc.COMM_SELF 260 TYPE = PETSc.Vec.Type.SEQ 261 262class TestVecMPI(BaseTestVec, unittest.TestCase): 263 COMM = PETSc.COMM_WORLD 264 TYPE = PETSc.Vec.Type.MPI 265 266class TestVecShared(BaseTestVec, unittest.TestCase): 267 if PETSc.COMM_WORLD.getSize() == 1: 268 TYPE = PETSc.Vec.Type.SHARED 269 else: 270 TYPE = PETSc.Vec.Type.MPI 271 COMM = PETSc.COMM_WORLD 272 273#class TestVecSieve(BaseTestVec, unittest.TestCase): 274# CLASS = PETSc.VecSieve 275# TARGS = ([],) 276 277#class TestVecGhost(BaseTestVec, unittest.TestCase): 278# CLASS = PETSc.VecGhost 279# TARGS = ([],) 280 281# -------------------------------------------------------------------- 282 283class TestVecWithArray(unittest.TestCase): 284 285 def testCreateSeq(self): 286 import numpy 287 a = numpy.zeros(5, dtype=PETSc.ScalarType) 288 289 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_SELF) 290 v2 = PETSc.Vec().createWithArray(a, size=5, comm=PETSc.COMM_SELF) 291 v3 = PETSc.Vec().createWithArray(a, size=3, comm=PETSc.COMM_SELF) 292 293 self.assertTrue(v1.size == 5) 294 self.assertTrue(v2.size == 5) 295 self.assertTrue(v3.size == 3) 296 297 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 298 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 299 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 300 301 def testCreateMPI(self): 302 import numpy 303 a = numpy.zeros(5, dtype=PETSc.ScalarType) 304 305 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_WORLD) 306 v2 = PETSc.Vec().createWithArray(a, size=(5,None), comm=PETSc.COMM_WORLD) 307 v3 = PETSc.Vec().createWithArray(a, size=(3,None), comm=PETSc.COMM_WORLD) 308 309 self.assertTrue(v1.local_size == 5) 310 self.assertTrue(v2.local_size == 5) 311 self.assertTrue(v3.local_size == 3) 312 313 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 314 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 315 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 316 317 def testSetMPIGhost(self): 318 import numpy 319 v = PETSc.Vec().create() 320 v.setType(PETSc.Vec.Type.MPI) 321 v.setSizes((5,None)) 322 ghosts = [i % v.size for i in range(v.owner_range[1],v.owner_range[1]+3)] 323 v.setMPIGhost(ghosts) 324 v.setArray(numpy.array(range(*v.owner_range),dtype=PETSc.ScalarType)) 325 v.ghostUpdate() 326 with v.localForm() as loc: 327 self.assertTrue((loc[0:v.local_size] == range(*v.owner_range)).all()) 328 self.assertTrue((loc[v.local_size:] == ghosts).all()) 329 330# -------------------------------------------------------------------- 331 332if __name__ == '__main__': 333 unittest.main() 334 335# -------------------------------------------------------------------- 336