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.assertEqual(self.vec[start], -7) 100 self.assertEqual(self.vec[end-1], -7) 101 for i in range(start, end): self.vec[i] = i 102 values = [self.vec[i] for i in range(start, end)] 103 self.assertEqual(values, list(range(start, end))) 104 sz = self.vec.getSize() 105 self.assertEqual(self.vec.sum(), (sz-1)/2.0*sz) 106 107 def testGetSetValsBlocked(self): 108 return 109 lsize, gsize = self.vec.getSizes() 110 start, end = self.vec.getOwnershipRange() 111 bsizes = list(range(1, lsize+1)) 112 nblocks = list(range(1, lsize+1)) 113 compat = [(bs, nb) 114 for bs in bsizes if not (gsize%bs or lsize % bs) 115 for nb in nblocks if bs*nb <= lsize] 116 for bsize, nblock in compat: 117 self.vec.setBlockSize(bsize) 118 bindex = [start//bsize+i for i in range(nblock)] 119 bvalue = [float(i) for i in range(nblock*bsize)] 120 self.vec.setValuesBlocked(bindex, bvalue) 121 self.vec.assemble() 122 index = [start+i for i in range(nblock*bsize)] 123 value = self.vec.getValues(index) 124 self.assertEqual(bvalue, list(value)) 125 126 def testGetSetArray(self): 127 self.vec.set(1) 128 arr0 = self.vec.getArray().copy() 129 self.assertEqual(arr0.sum(), self.vec.getLocalSize()) 130 arr0 = self.vec.getArray().copy() 131 self.vec.setRandom() 132 arr1 = self.vec.getArray().copy() 133 self.vec.setArray(arr1) 134 arr1 = self.vec.getArray().copy() 135 arr2 = self.vec.getArray().copy() 136 self.assertTrue((arr1 == arr2).all()) 137 import numpy 138 refs = self.vec.getRefCount() 139 arr3 = numpy.asarray(self.vec) 140 self.assertEqual(self.vec.getRefCount(), refs+1) 141 self.assertTrue((arr1 == arr3).all()) 142 arr3[:] = 0 143 self.assertAlmostEqual(abs(self.vec.sum()), 0) 144 self.assertEqual(self.vec.max()[1], 0) 145 self.assertEqual(self.vec.min()[1], 0) 146 self.vec.set(1) 147 self.assertAlmostEqual(abs(arr3.sum()), self.vec.getLocalSize()) 148 self.assertEqual(arr3.min(), 1) 149 self.assertEqual(arr3.max(), 1) 150 del arr3 151 self.assertEqual(self.vec.getRefCount(), refs) 152 153 def testPlaceArray(self): 154 self.vec.set(1) 155 array = self.vec.getArray().copy() 156 self.vec.placeArray(array) 157 array[:] = 2 158 self.assertAlmostEqual(abs(self.vec.sum()), 2*self.vec.getSize()) 159 self.vec.resetArray() 160 self.assertAlmostEqual(abs(self.vec.sum()), self.vec.getSize()) 161 162 def testLocalVector(self): 163 rank = self.vec.getComm().Get_rank() 164 self.vec.getArray()[:] = rank + 1 165 ln = self.vec.getLocalSize() 166 lvec = self.vec.createLocalVector() 167 self.vec.getLocalVector(lvec) 168 self.assertEqual(abs(lvec.sum()), (rank+1)*ln) 169 self.vec.restoreLocalVector(lvec) 170 self.vec.getLocalVector(lvec,readonly=True) 171 self.assertEqual(abs(lvec.sum()), (rank+1)*ln) 172 self.vec.restoreLocalVector(lvec,readonly=True) 173 lvec.destroy() 174 175 def testSetOption(self): 176 opt1 = PETSc.Vec.Option.IGNORE_OFF_PROC_ENTRIES 177 opt2 = PETSc.Vec.Option.IGNORE_NEGATIVE_INDICES 178 for opt in [opt1, opt2]*2: 179 for flag in [True,False]*2: 180 self.vec.setOption(opt,flag) 181 182 def testGetSetItem(self): 183 v = self.vec 184 w = v.duplicate() 185 # 186 v[...] = 7 187 self.assertEqual(v.max()[1], 7) 188 self.assertEqual(v.min()[1], 7) 189 # 190 v.setRandom() 191 w[...] = v 192 self.assertTrue(w.equal(v)) 193 # 194 v.setRandom() 195 w[...] = v.getArray() 196 self.assertTrue(w.equal(v)) 197 # 198 s, e = v.getOwnershipRange() 199 v.setRandom() 200 w[s:e] = v.getArray().copy() 201 self.assertTrue(w.equal(v)) 202 w1, v1 = w[s], v[s] 203 w2, v2 = w[e-1], v[e-1] 204 self.assertEqual(w1, v1) 205 self.assertEqual(w2, v2) 206 207 def testMAXPY(self): 208 y = self.vec 209 y.set(1) 210 x = [y.copy() for _ in range(3)] 211 a = [1]*len(x) 212 y.maxpy(a, x) 213 z = y.duplicate() 214 z.set(len(x)+1) 215 assert (y.equal(z)) 216 217 218# -------------------------------------------------------------------- 219 220class TestVecSeq(BaseTestVec, unittest.TestCase): 221 COMM = PETSc.COMM_SELF 222 TYPE = PETSc.Vec.Type.SEQ 223 224class TestVecMPI(BaseTestVec, unittest.TestCase): 225 COMM = PETSc.COMM_WORLD 226 TYPE = PETSc.Vec.Type.MPI 227 228class TestVecShared(BaseTestVec, unittest.TestCase): 229 if PETSc.COMM_WORLD.getSize() == 1: 230 TYPE = PETSc.Vec.Type.SHARED 231 else: 232 TYPE = PETSc.Vec.Type.MPI 233 COMM = PETSc.COMM_WORLD 234 235#class TestVecSieve(BaseTestVec, unittest.TestCase): 236# CLASS = PETSc.VecSieve 237# TARGS = ([],) 238 239#class TestVecGhost(BaseTestVec, unittest.TestCase): 240# CLASS = PETSc.VecGhost 241# TARGS = ([],) 242 243# -------------------------------------------------------------------- 244 245class TestVecWithArray(unittest.TestCase): 246 247 def testCreateSeq(self): 248 import numpy 249 a = numpy.zeros(5, dtype=PETSc.ScalarType) 250 251 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_SELF) 252 v2 = PETSc.Vec().createWithArray(a, size=5, comm=PETSc.COMM_SELF) 253 v3 = PETSc.Vec().createWithArray(a, size=3, comm=PETSc.COMM_SELF) 254 255 self.assertTrue(v1.size == 5) 256 self.assertTrue(v2.size == 5) 257 self.assertTrue(v3.size == 3) 258 259 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 260 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 261 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 262 263 def testCreateMPI(self): 264 import numpy 265 a = numpy.zeros(5, dtype=PETSc.ScalarType) 266 267 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_WORLD) 268 v2 = PETSc.Vec().createWithArray(a, size=(5,None), comm=PETSc.COMM_WORLD) 269 v3 = PETSc.Vec().createWithArray(a, size=(3,None), comm=PETSc.COMM_WORLD) 270 271 self.assertTrue(v1.local_size == 5) 272 self.assertTrue(v2.local_size == 5) 273 self.assertTrue(v3.local_size == 3) 274 275 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 276 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 277 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 278 279 def testSetMPIGhost(self): 280 import numpy 281 v = PETSc.Vec().create() 282 v.setType(PETSc.Vec.Type.MPI) 283 v.setSizes((5,None)) 284 ghosts = [i % v.size for i in range(v.owner_range[1],v.owner_range[1]+3)] 285 v.setMPIGhost(ghosts) 286 v.setArray(numpy.array(range(*v.owner_range),dtype=PETSc.ScalarType)) 287 v.ghostUpdate() 288 with v.localForm() as loc: 289 self.assertTrue((loc[0:v.local_size] == range(*v.owner_range)).all()) 290 self.assertTrue((loc[v.local_size:] == ghosts).all()) 291 292# -------------------------------------------------------------------- 293 294if __name__ == '__main__': 295 unittest.main() 296 297# -------------------------------------------------------------------- 298