Coverage for python/src/dolfinx_mpc/spider.py: 100%
71 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-07 09:58 +0000
1# Copyright (C) 2026 Jørgen S. Dokken
2#
3# This file is part of DOLFINX_MPC
4#
5# SPDX-License-Identifier: MIT
6"""Spider meshes: the point meshes holding the bodies of spiders (RBE2, RBE3).
8A spider has a *body*, a point of a point mesh, and *feet*, the dofs it is tied to. The bodies of
9the spiders of a problem are the points of a point mesh made by :func:`create_spider_mesh`. A
10spider is named by its input index, the original cell index of its point, and its coordinate is
11the coordinate of its dofs in the space on the point mesh: moving the point mesh, with
12:func:`move`, moves the spider.
13"""
15from __future__ import annotations
17import typing
19from mpi4py import MPI
21import dolfinx.fem as _fem
22import dolfinx.mesh as _mesh
23import numpy as np
24import numpy.typing as npt
25import ufl
26from dolfinx import default_real_type
27from dolfinx.common import local_range
29__all__ = ["create_spider_mesh", "create_spider_pair", "move", "spider_values"]
32def create_spider_mesh(comm: MPI.Comm, points: npt.ArrayLike, dtype: npt.DTypeLike | None = None) -> _mesh.Mesh:
33 """Create the point mesh holding the bodies of spiders.
35 The points are given on the first process. Every other process passes either no points or the
36 same points. Spider `k` is the point in row `k`, its input index; coinciding points are
37 distinct spiders, for instance two spiders joined by a spring. The points are spread over the
38 processes as the post office splits an index range: of `M` points, process `r` owns those in
39 :func:`dolfinx.common.local_range(r, M, size) <dolfinx.common.local_range>`. Spider `k` of any
40 two spider meshes of `M` points is therefore on the same process.
42 Args:
43 comm: The communicator of the meshes the spiders join
44 points: The points, shape `(num_points, 3)`, on the first process. On the others, no
45 points or the same points.
46 dtype: The coordinate type of the mesh. Defaults to the type of the first process's
47 points, or the default real type if they are not floating point.
49 Returns:
50 A mesh of points, from :func:`dolfinx.mesh.create_point_mesh`.
52 Raises:
53 ValueError: If a process passes points other than those of the first process. Raised on
54 every process.
56 Note:
57 Collective.
58 """
59 points = np.asarray(points)
60 if dtype is None:
61 own = points.dtype if np.issubdtype(points.dtype, np.floating) else np.dtype(default_real_type)
62 dtype = comm.bcast(own if comm.rank == 0 else None, root=0)
63 dtype = np.dtype(dtype)
64 local = points.astype(dtype).reshape(-1, 3)
65 reference = comm.bcast(local if comm.rank == 0 else None, root=0)
67 # The same points up to rounding in the coordinate type
68 eps = 100 * np.finfo(dtype).eps
69 scale = max(float(np.abs(reference).max(initial=0.0)), 1.0)
70 same = len(local) == 0 or (
71 local.shape == reference.shape and np.allclose(local, reference, rtol=eps, atol=eps * scale)
72 )
73 if not comm.allreduce(same, op=MPI.LAND):
74 raise ValueError("Every process must pass no points or the points of the first process")
75 start, end = local_range(comm.rank, len(reference), comm.size)
76 return _mesh.create_point_mesh(comm, reference[start:end])
79def _local_input_indices(spiders: _mesh.Mesh) -> npt.NDArray[np.int64]:
80 """The input index of each point of a spider mesh local to the process."""
81 imap = spiders.topology.index_map(0)
82 return np.asarray(spiders.topology.original_cell_index[: imap.size_local + imap.num_ghosts], dtype=np.int64)
85def create_spider_pair(spiders_A: _mesh.Mesh, spiders_B: _mesh.Mesh) -> _mesh.EntityMap:
86 r"""Relate spider `k` of one spider mesh to spider `k` of another, for every `k`.
88 The pairs are, for instance, the two ends of springs: order the points of the two meshes so that
89 row `k` of each holds the two ends of spring `k`. A form coupling the spaces on the two meshes,
90 such as a spring :math:`\int K (w_A - w_B) \cdot (v_A - v_B)`, is integrated over `spiders_A`
91 with the returned map in `entity_maps`. Each mesh keeps its own geometry, so the two spiders of
92 a pair may start at the same point and move apart.
94 Both meshes must have the same number of spiders. :func:`create_spider_mesh` then puts spider
95 `k` of both on the same process.
97 Args:
98 spiders_A: The spider mesh integrated over
99 spiders_B: The other spider mesh
101 Returns:
102 The entity map with `spiders_A` as topology and `spiders_B` as sub-topology.
104 Raises:
105 ValueError: If the meshes have different numbers of spiders, or spider `k` of the two is on
106 different processes, as for a point mesh not made by :func:`create_spider_mesh`. Raised
107 on every process.
109 Note:
110 Collective.
111 """
112 num_A = spiders_A.topology.index_map(0).size_global
113 num_B = spiders_B.topology.index_map(0).size_global
114 if num_A != num_B:
115 raise ValueError(f"The spider meshes must have the same number of spiders, not {num_A} and {num_B}")
117 # The local point of spiders_A with the input index of each local point of spiders_B
118 local_A = {k: i for i, k in enumerate(_local_input_indices(spiders_A).tolist())}
119 b_to_a = np.array([local_A.get(k, -1) for k in _local_input_indices(spiders_B).tolist()], dtype=np.int32)
120 if not spiders_A.comm.allreduce(bool((b_to_a >= 0).all()), op=MPI.LAND):
121 raise ValueError("Spider k of the two meshes is on different processes. Create both with create_spider_mesh.")
122 return _mesh.entity_map(spiders_A.topology, spiders_B.topology, 0, b_to_a)
125def move(
126 mesh: _mesh.Mesh,
127 u: _fem.Function | ufl.core.expr.Expr | typing.Callable[[npt.NDArray[np.floating]], npt.NDArray[np.inexact]],
128) -> None:
129 """Move the geometry nodes of a mesh by a displacement.
131 The displacement is interpolated into the space of the mesh's coordinate element, whose dofs
132 in each cell are the cell's geometry nodes, and added node by node. No relation between the
133 dofs of `u` and the geometry nodes is assumed. As in `scifem.move`.
135 Args:
136 mesh: The mesh to move
137 u: The displacement: a function on `mesh`, an expression, or a callable of the
138 coordinates, shape `(3, num_points)`. A function with more components than the
139 geometric dimension, such as one on a spider space with rotations, moves the mesh by
140 its first components, the translation. A complex displacement moves it by its real part.
141 """
142 gdim = mesh.geometry.dim
143 V_x = _fem.functionspace(mesh, mesh.ufl_domain().ufl_coordinate_element())
144 if isinstance(u, _fem.Function):
145 dtype = u.x.array.dtype
146 if int(np.prod(u.ufl_shape)) > gdim:
147 u = ufl.as_vector([u[i] for i in range(gdim)])
148 else:
149 dtype = mesh.geometry.x.dtype
150 if isinstance(u, ufl.core.expr.Expr) and not isinstance(u, _fem.Function):
151 dtype = np.result_type(dtype, np.dtype(mesh.geometry.x.dtype))
152 u = _fem.Expression(u, V_x.element.interpolation_points, dtype=dtype)
153 u_x = _fem.Function(V_x, dtype=dtype)
154 u_x.interpolate(u)
156 # Each node once: a node in several cells has the same value in each
157 nodes = mesh.geometry.dofmaps[0].reshape(-1)
158 displacement = np.zeros((mesh.geometry.x.shape[0], gdim), dtype=mesh.geometry.x.dtype)
159 displacement[nodes] = u_x.x.array.reshape(-1, gdim)[V_x.dofmap.list.reshape(-1)].real
160 mesh.geometry.x[:, :gdim] += displacement
163def spider_values(u: _fem.Function, spider: int) -> np.ndarray:
164 """The values of a function on a spider mesh at one spider, on every process.
166 A process only holds a spider's values if it owns the spider's dofs or has feet tied to it, so
167 the owner broadcasts them.
169 Args:
170 u: A function on a space of a mesh from :func:`create_spider_mesh`
171 spider: The input index of the spider
173 Returns:
174 The values at the spider, one per component of the space.
176 Note:
177 Collective.
178 """
179 V = u.function_space
180 comm = V.mesh.comm
181 topology = V.mesh.topology
182 num_cells = topology.index_map(topology.dim).size_local
183 cell = np.flatnonzero(np.asarray(topology.original_cell_index[:num_cells]) == spider)
184 owner = comm.allreduce(comm.rank if len(cell) > 0 else -1, op=MPI.MAX)
185 if owner < 0:
186 raise IndexError(f"The spider mesh has no point with input index {spider}")
187 values = None
188 if comm.rank == owner:
189 bs = V.dofmap.index_map_bs
190 dof = V.dofmap.list[cell[0], 0]
191 values = u.x.array[bs * dof : bs * (dof + 1)].copy()
192 return comm.bcast(values, root=owner)