Coverage for dragonfly_sim/core/wall_distance.py: 90%

1""" 

2Distance to the wall, as a CG1 DerivedField with an exact mesh derivative. 

3 

4d at every CG1 node is the distance to the nearest wall facet: the exact 

5point-to-segment distance for the straight 2D wall facets, the distance to 

6the nearest wall vertex in 3D (an over-estimate between vertices). The 

7nearest facets are found with a KD-tree over the facet midpoints (2D) or 

8the wall vertices (3D). 

9 

10mesh_jacobian() is the analytic derivative of every nodal distance with 

11respect to the coordinates of the node itself and of the wall vertices of 

12its nearest facet. With P the nearest point on segment AB, t its parameter 

13and r = (x - P)/d, 

14 

15 dd/dx = r, dd/dA = -(1 - t) r, dd/dB = -t r 

16 

17(t is stationary at an interior foot point, so it does not move to first 

18order; at a clipped end, t = 0 or 1, the same formula holds). On the wall 

19itself, d = 0 and the node moves with the wall, so the row is zero. The 

20derivative is exact wherever the nearest facet does not switch, i.e. 

21almost everywhere. 

22""" 

23import numpy as np 

24import scipy.spatial 

25from petsc4py import PETSc 

26import dolfinx 

27 

28from dragonfly_sim.utils.derived_fields import DerivedField 

29 

30 

31def _nearest_segments(x, A, B, k=16): 

32 """Index of the nearest segment, foot parameter t and foot point, per point.""" 

33 tree = scipy.spatial.cKDTree(0.5 * (A + B)) 

34 k = min(k, len(A)) 

35 _, idx = tree.query(x, k=k) 

36 idx = idx.reshape(len(x), k) 

37 Ak, Bk = A[idx], B[idx] 

38 AB = Bk - Ak 

39 t = np.clip(np.einsum("nkj,nkj->nk", x[:, None, :] - Ak, AB) 

40 / np.maximum(np.einsum("nkj,nkj->nk", AB, AB), 1e-300), 0.0, 1.0) 

41 P = Ak + t[..., None] * AB 

42 dist = np.linalg.norm(x[:, None, :] - P, axis=2) 

43 j = np.argmin(dist, axis=1) 

44 rows = np.arange(len(x)) 

45 return idx[rows, j], t[rows, j], P[rows, j], dist[rows, j] 

46 

47 

48class WallDistance(DerivedField): 

49 

50 def __init__(self, mesh_obj, wall_tags): 

51 self.mesh_obj = mesh_obj 

52 self.wall_tags = list(wall_tags) 

53 msh = mesh_obj.mesh 

54 V = dolfinx.fem.functionspace(msh, ("Lagrange", 1)) 

55 Vx = mesh_obj.coordinate_space 

56 # The derivative maps d's dofs onto coordinate_space's: they must be 

57 # the same nodes in the same numbering. 

58 if not (np.array_equal(V.dofmap.list, Vx.dofmap.list) 

59 and V.dofmap.index_map.size_local == Vx.dofmap.index_map.size_local): 

60 raise RuntimeError("wall distance: CG1 dofmap differs from the coordinate space's") 

61 self.function = dolfinx.fem.Function(V, name="wall_distance") 

62 self.update_geometry() 

63 

64 def _wall_geometry(self): 

65 """Wall segment endpoints (2D) or vertices (3D) on every rank, with the 

66 global coordinate_space block index of each endpoint.""" 

67 msh = self.mesh_obj.mesh 

68 comm = msh.comm 

69 gdim = msh.geometry.dim 

70 fdim = msh.topology.dim - 1 

71 tags = self.mesh_obj.meshtags 

72 facets = tags.indices[np.isin(tags.values, self.wall_tags)] 

73 facets = facets[facets < msh.topology.index_map(fdim).size_local] 

74 Vx = self.mesh_obj.coordinate_space 

75 imap = Vx.dofmap.index_map 

76 # vertex -> coordinate_space block: through the cells' local numbering 

77 geom = dolfinx.mesh.entities_to_geometry(msh, fdim, facets, False) 

78 node_to_dof = self._geometry_node_to_dof() 

79 coords = msh.geometry.x[geom][:, :, :gdim] if len(facets) else np.zeros((0, 2, gdim)) 

80 gids = (imap.local_to_global(node_to_dof[geom.ravel()].astype(np.int32)).reshape(geom.shape) 

81 if len(facets) else np.zeros((0, 2), dtype=np.int64)) 

82 coords = np.concatenate([c for c in comm.allgather(coords) if c.size] or [np.zeros((0, 2, gdim))]) 

83 gids = np.concatenate([g for g in comm.allgather(gids) if g.size] or [np.zeros((0, 2), np.int64)]) 

84 if coords.shape[1] > 2 and msh.topology.dim == 2: 

85 raise NotImplementedError("wall distance for facets with {} geometry nodes".format(coords.shape[1])) 

86 return coords, gids 

87 

88 def _geometry_node_to_dof(self): 

89 """Local geometry node -> local coordinate_space block (both ghosted).""" 

90 msh = self.mesh_obj.mesh 

91 Vx = self.mesh_obj.coordinate_space 

92 g = msh.geometry.dofmaps[0] 

93 d = Vx.dofmap.list 

94 out = np.full(msh.geometry.x.shape[0], -1, dtype=np.int64) 

95 out[g.ravel()] = d.ravel() 

96 return out 

97 

98 def update_geometry(self): 

99 msh = self.mesh_obj.mesh 

100 gdim = msh.geometry.dim 

101 x = self.function.function_space.tabulate_dof_coordinates()[:, :gdim] 

102 segs, gids = self._wall_geometry() 

103 self._segs, self._gids = segs, gids 

104 if gdim == 3: 

105 verts = segs.reshape(-1, 3) 

106 vgids = gids.reshape(-1) 

107 tree = scipy.spatial.cKDTree(verts) 

108 dist, j = tree.query(x) 

109 self._nearest = ("vertex", j, vgids, verts) 

110 else: 

111 A, B = segs[:, 0, :], segs[:, -1, :] 

112 j, t, P, dist = _nearest_segments(x, A, B) 

113 self._nearest = ("segment", j, t, P) 

114 self._x = x 

115 self._dist = dist 

116 self.function.x.array[:] = dist 

117 self.function.x.scatter_forward() 

118 

119 def mesh_jacobian(self): 

120 msh = self.mesh_obj.mesh 

121 gdim = msh.geometry.dim 

122 V = self.function.function_space 

123 Vx = self.mesh_obj.coordinate_space 

124 imap = V.dofmap.index_map 

125 n_own = imap.size_local 

126 rows_g = imap.local_to_global(np.arange(n_own, dtype=np.int32)) 

127 D = PETSc.Mat().createAIJ( 

128 ((n_own, imap.size_global), (n_own * gdim, imap.size_global * gdim)), comm=msh.comm) 

129 D.setPreallocationNNZ((3 * gdim, 3 * gdim)) 

130 D.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, False) 

131 x, dist = self._x[:n_own], self._dist[:n_own] 

132 if self._nearest[0] == "segment": 

133 _, j, t, P = self._nearest 

134 j, t, P = j[:n_own], t[:n_own], P[:n_own] 

135 else: 

136 _, j, vgids, verts = self._nearest 

137 j = j[:n_own] 

138 for i in range(n_own): 

139 if dist[i] <= 1e-14: 

140 continue 

141 if self._nearest[0] == "segment": 

142 r = (x[i] - P[i]) / dist[i] 

143 ga, gb = self._gids[j[i], 0], self._gids[j[i], -1] 

144 entries = [(rows_g[i], r), (ga, -(1.0 - t[i]) * r), (gb, -t[i] * r)] 

145 else: 

146 r = (x[i] - verts[j[i]]) / dist[i] 

147 entries = [(rows_g[i], r), (vgids[j[i]], -r)] 

148 for block, val in entries: 

149 D.setValues(np.array([rows_g[i]], dtype=PETSc.IntType), 

150 np.array([block * gdim + k for k in range(gdim)], dtype=PETSc.IntType), 

151 val, addv=True) 

152 D.assemble() 

153 return D