Coverage for pySDC/implementations/datatype_classes/petsc_vec.py: 95%
40 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-20 14:51 +0000
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-20 14:51 +0000
1from petsc4py import PETSc
3from pySDC.core.errors import DataError
6class petsc_vec(PETSc.Vec):
7 __array_priority__ = 1000 # otherwise rmul with float64 does not work (don't ask, won't tell)
9 def __new__(cls, init=None, val=0.0):
10 if isinstance(init, petsc_vec) or isinstance(init, PETSc.Vec):
11 obj = PETSc.Vec.__new__(cls)
12 init.copy(obj)
13 elif isinstance(init, PETSc.DMDA):
14 tmp = init.createGlobalVector()
15 obj = petsc_vec(tmp)
16 objarr = init.getVecArray(obj)
17 objarr[:] = val
18 else:
19 obj = PETSc.Vec.__new__(cls)
20 return obj
22 def __abs__(self):
23 """
24 Overloading the abs operator
26 Returns:
27 float: absolute maximum of all vec values
28 """
29 # take absolute values of the mesh values (INF = 3)
30 return self.norm(3)
32 def isend(self, dest=None, tag=None, comm=None):
33 """
34 Routine for sending data forward in time (non-blocking)
36 Args:
37 dest (int): target rank
38 tag (int): communication tag
39 comm: communicator
41 Returns:
42 request handle
43 """
44 return comm.Issend(self.getArray(), dest=dest, tag=tag)
46 def irecv(self, source=None, tag=None, comm=None):
47 """
48 Routine for receiving in time
50 Args:
51 source (int): source rank
52 tag (int): communication tag
53 comm: communicator
55 Returns:
56 None
57 """
58 return comm.Irecv(self.getArray(), source=source, tag=tag)
60 def bcast(self, root=None, comm=None):
61 """
62 Routine for broadcasting values
64 Args:
65 root (int): process with value to broadcast
66 comm: communicator
68 Returns:
69 broadcasted values
70 """
71 comm.Bcast(self.getArray(), root=root)
72 return self
75class petsc_vec_imex(object):
76 """
77 RHS data type for Vec with implicit and explicit components
79 This data type can be used to have RHS with 2 components (here implicit and explicit)
81 Attributes:
82 impl (petsc_vec): implicit part
83 expl (petsc_vec): explicit part
84 """
86 def __init__(self, init, val=0.0):
87 """
88 Initialization routine
90 Args:
91 init: can either be a tuple (one int per dimension) or a number (if only one dimension is requested)
92 or another imex_mesh object
93 val (float): an initial number (default: 0.0)
94 Raises:
95 DataError: if init is none of the types above
96 """
98 if isinstance(init, type(self)):
99 self.impl = petsc_vec(init.impl)
100 self.expl = petsc_vec(init.expl)
101 elif isinstance(init, PETSc.DMDA):
102 self.impl = petsc_vec(init, val=val)
103 self.expl = petsc_vec(init, val=val)
104 # something is wrong, if none of the ones above hit
105 else:
106 raise DataError('something went wrong during %s initialization' % type(self))
109class petsc_vec_comp2(object):
110 """
111 RHS data type for Vec with implicit and explicit components
113 This data type can be used to have RHS with 2 components (here implicit and explicit)
115 Attributes:
116 impl (petsc_vec): implicit part
117 expl (petsc_vec): explicit part
118 """
120 def __init__(self, init, val=0.0):
121 """
122 Initialization routine
124 Args:
125 init: can either be a tuple (one int per dimension) or a number (if only one dimension is requested)
126 or another imex_mesh object
127 val (float): an initial number (default: 0.0)
128 Raises:
129 DataError: if init is none of the types above
130 """
132 if isinstance(init, type(self)):
133 self.comp1 = petsc_vec(init.comp1)
134 self.comp2 = petsc_vec(init.comp2)
135 elif isinstance(init, PETSc.DMDA):
136 self.comp1 = petsc_vec(init, val=val)
137 self.comp2 = petsc_vec(init, val=val)
138 # something is wrong, if none of the ones above hit
139 else:
140 raise DataError('something went wrong during %s initialization' % type(self))