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