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