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 testSetOption(self): 162 opt1 = PETSc.Vec.Option.IGNORE_OFF_PROC_ENTRIES 163 opt2 = PETSc.Vec.Option.IGNORE_NEGATIVE_INDICES 164 for opt in [opt1, opt2]*2: 165 for flag in [True,False]*2: 166 self.vec.setOption(opt,flag) 167 168 def testGetSetItem(self): 169 v = self.vec 170 w = v.duplicate() 171 # 172 v[...] = 7 173 self.assertEqual(v.max()[1], 7) 174 self.assertEqual(v.min()[1], 7) 175 # 176 v.setRandom() 177 w[...] = v 178 self.assertTrue(w.equal(v)) 179 # 180 v.setRandom() 181 w[...] = v.getArray() 182 self.assertTrue(w.equal(v)) 183 # 184 s, e = v.getOwnershipRange() 185 v.setRandom() 186 w[s:e] = v.getArray().copy() 187 self.assertTrue(w.equal(v)) 188 w1, v1 = w[s], v[s] 189 w2, v2 = w[e-1], v[e-1] 190 self.assertEqual(w1, v1) 191 self.assertEqual(w2, v2) 192 193 def testMAXPY(self): 194 y = self.vec 195 y.set(1) 196 x = [y.copy() for _ in range(3)] 197 a = [1]*len(x) 198 y.maxpy(a, x) 199 z = y.duplicate() 200 z.set(len(x)+1) 201 assert (y.equal(z)) 202 203 204# -------------------------------------------------------------------- 205 206class TestVecSeq(BaseTestVec, unittest.TestCase): 207 COMM = PETSc.COMM_SELF 208 TYPE = PETSc.Vec.Type.SEQ 209 210class TestVecMPI(BaseTestVec, unittest.TestCase): 211 COMM = PETSc.COMM_WORLD 212 TYPE = PETSc.Vec.Type.MPI 213 214class TestVecShared(BaseTestVec, unittest.TestCase): 215 if PETSc.COMM_WORLD.getSize() == 1: 216 TYPE = PETSc.Vec.Type.SHARED 217 else: 218 TYPE = PETSc.Vec.Type.MPI 219 COMM = PETSc.COMM_WORLD 220 221#class TestVecSieve(BaseTestVec, unittest.TestCase): 222# CLASS = PETSc.VecSieve 223# TARGS = ([],) 224 225#class TestVecGhost(BaseTestVec, unittest.TestCase): 226# CLASS = PETSc.VecGhost 227# TARGS = ([],) 228 229# -------------------------------------------------------------------- 230 231class TestVecWithArray(unittest.TestCase): 232 233 def testCreateSeq(self): 234 import numpy 235 a = numpy.zeros(5, dtype=PETSc.ScalarType) 236 237 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_SELF) 238 v2 = PETSc.Vec().createWithArray(a, size=5, comm=PETSc.COMM_SELF) 239 v3 = PETSc.Vec().createWithArray(a, size=3, comm=PETSc.COMM_SELF) 240 241 self.assertTrue(v1.size == 5) 242 self.assertTrue(v2.size == 5) 243 self.assertTrue(v3.size == 3) 244 245 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 246 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 247 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 248 249 def testCreateMPI(self): 250 import numpy 251 a = numpy.zeros(5, dtype=PETSc.ScalarType) 252 253 v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_WORLD) 254 v2 = PETSc.Vec().createWithArray(a, size=(5,None), comm=PETSc.COMM_WORLD) 255 v3 = PETSc.Vec().createWithArray(a, size=(3,None), comm=PETSc.COMM_WORLD) 256 257 self.assertTrue(v1.local_size == 5) 258 self.assertTrue(v2.local_size == 5) 259 self.assertTrue(v3.local_size == 3) 260 261 a1 = v1.getDict()['__array__']; self.assertTrue(a is a1) 262 a2 = v2.getDict()['__array__']; self.assertTrue(a is a2) 263 a3 = v3.getDict()['__array__']; self.assertTrue(a is a2) 264 265 def testSetMPIGhost(self): 266 import numpy 267 v = PETSc.Vec().create() 268 v.setType(PETSc.Vec.Type.MPI) 269 v.setSizes((5,None)) 270 ghosts = [i % v.size for i in range(v.owner_range[1],v.owner_range[1]+3)] 271 v.setMPIGhost(ghosts) 272 v.setArray(numpy.array(range(*v.owner_range))) 273 v.ghostUpdate() 274 with v.localForm() as loc: 275 self.assertTrue((loc[0:v.local_size] == range(*v.owner_range)).all()) 276 self.assertTrue((loc[v.local_size:] == ghosts).all()) 277 278# -------------------------------------------------------------------- 279 280if __name__ == '__main__': 281 unittest.main() 282 283# -------------------------------------------------------------------- 284