Coverage for pySDC/implementations/datatype_classes/mesh.py: 100%
69 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 13:41 +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
11_NUMPY1 = np.lib.NumpyVersion(np.__version__) < '2.0.0'
14class mesh(np.ndarray):
15 """
16 Numpy-based datatype for serial or parallel meshes.
17 Can include a communicator and expects a dtype to allow complex data.
19 Attributes:
20 comm: MPI communicator or None
21 """
23 comm = None
24 xp = np
26 def __new__(cls, init, val=0.0, **kwargs):
27 """
28 Instantiates new datatype. This ensures that even when manipulating data, the result is still a mesh.
30 Args:
31 init: either another mesh or a tuple containing the dimensions, the communicator and the dtype
32 val: value to initialize
34 Returns:
35 obj of type mesh
37 """
38 if isinstance(init, mesh):
39 obj = np.ndarray.__new__(cls, shape=init.shape, dtype=init.dtype, **kwargs)
40 obj[:] = init[:]
41 # a fresh array of its own, so `__array_finalize__` has nothing to inherit from and the
42 # communicator has to be carried over by hand
43 obj.comm = init.comm
44 elif (
45 isinstance(init, tuple)
46 and (init[1] is None or isinstance(init[1], MPI.Intracomm))
47 and isinstance(init[2], np.dtype)
48 ):
49 obj = np.ndarray.__new__(cls, init[0], dtype=init[2], **kwargs)
50 obj.fill(val)
51 obj.comm = init[1]
52 else:
53 raise NotImplementedError(type(init))
54 return obj
56 def __array_wrap__(self, obj, context=None, return_scalar=False):
57 """
58 Hand a ufunc's result back as a mesh with this one's communicator, cf.
59 https://numpy.org/doc/stable/user/basics.subclassing.html#array-wrap-for-ufuncs-and-other-functions
61 There is no `__array_ufunc__` here: NumPy's own dispatch computes in place where asked, writes into an
62 explicit `out` and returns it, and only the result passes through this method. On the way in, a Python
63 override cost more than the arithmetic for all but large meshes.
65 Args:
66 obj (np.ndarray): the result as NumPy computed it
67 context: the ufunc and its arguments, unused
68 return_scalar (bool): whether NumPy wants a scalar, as from a reduction to a single value
70 Returns:
71 mesh or scalar: the result
72 """
73 if isinstance(obj, mesh): # written into a mesh passed as `out`: that one, with its own communicator
74 return obj
75 # NumPy 1 does not pass `return_scalar`, but there a ufunc turns a 0-d result into a scalar anyway
76 if return_scalar or (_NUMPY1 and obj.ndim == 0):
77 return obj[()]
78 # the result is a fresh array with no communicator to inherit through `__array_finalize__`
79 obj = obj.view(type(self))
80 obj.comm = self.comm
81 return obj
83 def __array_finalize__(self, obj):
84 """
85 Carry the communicator onto every array derived from this one.
87 Without this the communicator would live only on the instance it was given to, and a slice,
88 a sum or a copy would silently lose it -- so `__abs__` would stop reducing across ranks and
89 quietly return a local norm on a space-parallel run.
90 """
91 if obj is not None:
92 self.comm = getattr(obj, 'comm', None)
94 def __abs__(self):
95 """
96 Overloading the abs operator
98 Returns:
99 float: absolute maximum of all mesh values
100 """
101 # take absolute values of the mesh values
102 local_absval = float(np.max(np.ndarray.__abs__(self)))
104 if self.comm is not None:
105 if self.comm.size > 1:
106 global_absval = self.comm.allreduce(sendobj=local_absval, op=MPI.MAX)
107 else:
108 global_absval = local_absval
109 else:
110 global_absval = local_absval
112 return float(global_absval)
114 def isend(self, dest=None, tag=None, comm=None):
115 """
116 Routine for sending data forward in time (non-blocking)
118 Args:
119 dest (int): target rank
120 tag (int): communication tag
121 comm: communicator
123 Returns:
124 request handle
125 """
126 return comm.Issend(self[:], dest=dest, tag=tag)
128 def irecv(self, source=None, tag=None, comm=None):
129 """
130 Routine for receiving in time
132 Args:
133 source (int): source rank
134 tag (int): communication tag
135 comm: communicator
137 Returns:
138 None
139 """
140 return comm.Irecv(self[:], source=source, tag=tag)
142 def bcast(self, root=None, comm=None):
143 """
144 Routine for broadcasting values
146 Args:
147 root (int): process with value to broadcast
148 comm: communicator
150 Returns:
151 broadcasted values
152 """
153 comm.Bcast(self[:], root=root)
154 return self
157def _component_property(index, name, base):
158 """
159 Make a property that reads and writes the ``index``-th component of a multi-component mesh.
161 Args:
162 index (int): position of the component along the first axis
163 name (str): name of the component, used in error messages
164 base: the single-component datatype a component is a view of
166 Returns:
167 property: gives a view on the component, and writes into it on assignment
168 """
170 def getter(self):
171 if self.shape[0] != len(self.components):
172 raise AttributeError(f'Cannot access {name!r} in {type(self)!r} because the shape is unexpected.')
173 return self[index].view(base)
175 def setter(self, value):
176 getter(self)[:] = value
178 return property(getter, setter, doc=f'View on the {name!r} component of the mesh')
181class MultiComponentMeshMixin:
182 r"""
183 Generic mesh with multiple components.
185 Mix this into a mesh datatype to obtain the multi-component version of it, as ``MultiComponentMesh`` does for
186 ``mesh`` and ``CuPyMultiComponentMesh`` does for ``cupy_mesh``. To make a specific multi-component mesh, derive
187 from one of those and list the components as strings in the class attribute ``components``. An example:
189 .. code-block:: python
191 class imex_mesh(MultiComponentMesh):
192 components = ['impl', 'expl']
194 Instantiating such a mesh will expand the mesh along an added first dimension for each component and allow access
195 to the components with ``.``. Continuing the above example:
197 .. code-block:: python
199 init = ((100,), None, numpy.dtype('d'))
200 f = imex_mesh(init)
201 f.shape # (2, 100)
202 f.expl.shape # (100,)
204 The components are properties, generated when the subclass is created. Both ``f.expl[:] = ...`` and
205 ``f.expl = ...`` write into the mesh; the component is never replaced by an unrelated object. Because the
206 properties live on the class, you cannot name a component like something that is already an attribute of the
207 underlying mesh datatype -- doing so raises an ``AttributeError`` when the class is created rather than silently
208 shadowing the component.
210 There are a couple more things to keep in mind:
211 - Because a multi-component mesh is just an array with one more dimension, all components must have the same
212 shape.
213 - You can use the entire mesh like an array in operations that accept arrays, but make sure that you really want
214 to apply the same operation on all components if you do.
215 """
217 components = []
219 def __init_subclass__(cls, **kwargs):
220 """
221 Turn the names listed in ``components`` into properties giving a view on the corresponding slice of the mesh.
222 """
223 super().__init_subclass__(**kwargs)
225 # a single component is an instance of the datatype this multi-component mesh is built on
226 base = next(c for c in cls.__mro__ if not issubclass(c, MultiComponentMeshMixin))
228 for index, name in enumerate(cls.components):
229 if hasattr(base, name):
230 raise AttributeError(
231 f'Cannot use {name!r} as a component of {cls.__name__} because it is already an attribute of '
232 f'{base.__name__}!'
233 )
234 setattr(cls, name, _component_property(index, name, base))
236 def __new__(cls, init, *args, **kwargs):
237 if isinstance(init, tuple):
238 shape = (init[0],) if type(init[0]) is int else init[0]
239 obj = super().__new__(cls, ((len(cls.components), *shape), *init[1:]), *args, **kwargs)
240 else:
241 obj = super().__new__(cls, init, *args, **kwargs)
243 return obj
246class MultiComponentMesh(MultiComponentMeshMixin, mesh):
247 """Numpy-based mesh with multiple components, see ``MultiComponentMeshMixin``."""
250class imex_mesh(MultiComponentMesh):
251 """NumPy mesh with an implicit part ``impl`` and an explicit part ``expl``, for IMEX right-hand sides."""
253 components = ['impl', 'expl']
256class comp2_mesh(MultiComponentMesh):
257 """NumPy mesh with two parts ``comp1`` and ``comp2``, for multi-implicit right-hand sides."""
259 components = ['comp1', 'comp2']