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 221# -------------------------------------------------------------------- 222 223class TestVecSeq(BaseTestVec, unittest.TestCase): 224 COMM = PETSc.COMM_SELF 225 TYPE = PETSc.Vec.Type.SEQ 226 227class TestVecMPI(BaseTestVec, unittest.TestCase): 228 COMM = PETSc.COMM_WORLD 229 TYPE = PETSc.Vec.Type.MPI 230 231class TestVecShared(BaseTestVec, unittest.TestCase): 232 if PETSc.COMM_WORLD.getSize() == 1: 233 TYPE = PETSc.Vec.Type.SHARED 234 else: 235 TYPE = PETSc.Vec.Type.MPI 236 COMM = PETSc.COMM_WORLD 237 238#class TestVecSieve(BaseTestVec, unittest.TestCase): 239# CLASS = PETSc.VecSieve 240# TARGS = ([],) 241 242#class TestVecGhost(BaseTestVec, unittest.TestCase): 243# CLASS = PETSc.VecGhost 244# TARGS = ([],) 245 246# -------------------------------------------------------------------- 247 248class TestVecWithArray(unittest.TestCase): 249 250 def testCreateSeq(self): 251 import numpy 252 a = numpy.zeros(5, dtype=PETSc.ScalarType) 253 254 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_SELF) 255 v2 = PETSc.Vec().createWithArray(a, size=5, comm=PETSc.COMM_SELF) 256 v3 = PETSc.Vec().createWithArray(a, size=3, comm=PETSc.COMM_SELF) 257 258 self.assertTrue(v1.size == 5) 259 self.assertTrue(v2.size == 5) 260 self.assertTrue(v3.size == 3) 261 262 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 263 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 264 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 265 266 def testCreateMPI(self): 267 import numpy 268 a = numpy.zeros(5, dtype=PETSc.ScalarType) 269 270 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_WORLD) 271 v2 = PETSc.Vec().createWithArray(a, size=(5,None), comm=PETSc.COMM_WORLD) 272 v3 = PETSc.Vec().createWithArray(a, size=(3,None), comm=PETSc.COMM_WORLD) 273 274 self.assertTrue(v1.local_size == 5) 275 self.assertTrue(v2.local_size == 5) 276 self.assertTrue(v3.local_size == 3) 277 278 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 279 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 280 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 281 282 def testSetMPIGhost(self): 283 import numpy 284 v = PETSc.Vec().create() 285 v.setType(PETSc.Vec.Type.MPI) 286 v.setSizes((5,None)) 287 ghosts = [i % v.size for i in range(v.owner_range[1],v.owner_range[1]+3)] 288 v.setMPIGhost(ghosts) 289 v.setArray(numpy.array(range(*v.owner_range),dtype=PETSc.ScalarType)) 290 v.ghostUpdate() 291 with v.localForm() as loc: 292 self.assertTrue((loc[0:v.local_size] == range(*v.owner_range)).all()) 293 self.assertTrue((loc[v.local_size:] == ghosts).all()) 294 295# -------------------------------------------------------------------- 296 297if __name__ == '__main__': 298 unittest.main() 299 300# -------------------------------------------------------------------- 301