Coverage for pySDC/implementations/datatype_classes/mesh.py: 96%
57 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-09-09 14:59 +0000
1import numpy as np
3try:
4 # TODO : mpi4py cannot be imported before dolfin when using fenics mesh
5 # see https://github.com/Parallel-in-Time/pySDC/pull/285#discussion_r1145850590
6 # This should be dealt with at some point
7 from mpi4py import MPI
8except ImportError:
9 MPI = None
12class mesh(np.ndarray):
13 """
14 Numpy-based datatype for serial or parallel meshes.
15 Can include a communicator and expects a dtype to allow complex data.
17 Attributes:
18 comm: MPI communicator or None
19 """
21 comm = None
22 xp = np
24 def __new__(cls, init, val=0.0, **kwargs):
25 """
26 Instantiates new datatype. This ensures that even when manipulating data, the result is still a mesh.
28 Args:
29 init: either another mesh or a tuple containing the dimensions, the communicator and the dtype
30 val: value to initialize
32 Returns:
33 obj of type mesh
35 """
36 if isinstance(init, mesh):
37 obj = np.ndarray.__new__(cls, shape=init.shape, dtype=init.dtype, **kwargs)
38 obj[:] = init[:]
39 elif (
40 isinstance(init, tuple)
41 and (init[1] is None or isinstance(init[1], MPI.Intracomm))
42 and isinstance(init[2], np.dtype)
43 ):
44 obj = np.ndarray.__new__(cls, init[0], dtype=init[2], **kwargs)
45 obj.fill(val)
46 cls.comm = init[1]
47 else:
48 raise NotImplementedError(type(init))
49 return obj
51 def __array_ufunc__(self, ufunc, method, *inputs, out=None, **kwargs):
52 """
53 Overriding default ufunc, cf. https://numpy.org/doc/stable/user/basics.subclassing.html#array-ufunc-for-ufuncs
54 """
55 args = []
56 for _, input_ in enumerate(inputs):
57 if isinstance(input_, mesh):
58 args.append(input_.view(np.ndarray))
59 else:
60 args.append(input_)
62 results = super().__array_ufunc__(ufunc, method, *args, **kwargs).view(type(self))
63 return results
65 def __abs__(self):
66 """
67 Overloading the abs operator
69 Returns:
70 float: absolute maximum of all mesh values
71 """
72 # take absolute values of the mesh values
73 local_absval = float(np.amax(np.ndarray.__abs__(self)))
75 if self.comm is not None:
76 if self.comm.Get_size() > 1:
77 global_absval = 0.0
78 global_absval = max(self.comm.allreduce(sendobj=local_absval, op=MPI.MAX), global_absval)
79 else:
80 global_absval = local_absval
81 else:
82 global_absval = local_absval
84 return float(global_absval)
86 def isend(self, dest=None, tag=None, comm=None):
87 """
88 Routine for sending data forward in time (non-blocking)
90 Args:
91 dest (int): target rank
92 tag (int): communication tag
93 comm: communicator
95 Returns:
96 request handle
97 """
98 return comm.Issend(self[:], dest=dest, tag=tag)
100 def irecv(self, source=None, tag=None, comm=None):
101 """
102 Routine for receiving in time
104 Args:
105 source (int): source rank
106 tag (int): communication tag
107 comm: communicator
109 Returns:
110 None
111 """
112 return comm.Irecv(self[:], source=source, tag=tag)
114 def bcast(self, root=None, comm=None):
115 """
116 Routine for broadcasting values
118 Args:
119 root (int): process with value to broadcast
120 comm: communicator
122 Returns:
123 broadcasted values
124 """
125 comm.Bcast(self[:], root=root)
126 return self
129class MultiComponentMesh(mesh):
130 r"""
131 Generic mesh with multiple components.
133 To make a specific multi-component mesh, derive from this class and list the components as strings in the class
134 attribute ``components``. An example:
136 ```
137 class imex_mesh(MultiComponentMesh):
138 components = ['impl', 'expl']
139 ```
141 Instantiating such a mesh will expand the mesh along an added first dimension for each component and allow access
142 to the components with ``.``. Continuing the above example:
144 ```
145 init = ((100,), None, numpy.dtype('d'))
146 f = imex_mesh(init)
147 f.shape # (2, 100)
148 f.expl.shape # (100,)
149 ```
151 Note that the components are not attributes of the mesh: ``"expl" in dir(f)`` will return False! Rather, the
152 components are handled in ``__getattr__``. This function is called if an attribute is not found and returns a view
153 on to the component if appropriate. Importantly, this means that you cannot name a component like something that
154 is already an attribute of ``mesh`` or ``numpy.ndarray`` because this will not result in calls to ``__getattr__``.
156 There are a couple more things to keep in mind:
157 - Because a ``MultiComponentMesh`` is just a ``numpy.ndarray`` with one more dimension, all components must have
158 the same shape.
159 - You can use the entire ``MultiComponentMesh`` like a ``numpy.ndarray`` in operations that accept arrays, but make
160 sure that you really want to apply the same operation on all components if you do.
161 - If you omit the assignment operator ``[:]`` during assignment, you will not change the mesh at all. Omitting this
162 leads to all kinds of trouble throughout the code. But here you really cannot get away without.
163 """
165 components = []
167 def __new__(cls, init, *args, **kwargs):
168 if isinstance(init, tuple):
169 shape = (init[0],) if type(init[0]) is int else init[0]
170 obj = super().__new__(cls, ((len(cls.components), *shape), *init[1:]), *args, **kwargs)
171 else:
172 obj = super().__new__(cls, init, *args, **kwargs)
174 return obj
176 def __getattr__(self, name):
177 if name in self.components:
178 if self.shape[0] == len(self.components):
179 return self[self.components.index(name)].view(mesh)
180 else:
181 raise AttributeError(f'Cannot access {name!r} in {type(self)!r} because the shape is unexpected.')
182 else:
183 raise AttributeError(f"{type(self)!r} does not have attribute {name!r}!")
186class imex_mesh(MultiComponentMesh):
187 components = ['impl', 'expl']
190class comp2_mesh(MultiComponentMesh):
191 components = ['comp1', 'comp2']