Coverage for src/sdynpy/fem/sdynpy_visualization_reduction.py: 7%
166 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-18 21:32 +0000
1"""
2Functions to create visualization artifacts (e.g. tracelines) automatically such that a
3reduced geometry consisting of only a point cloud of nodes can be visualized.
5This module defines a Geometry object as well as all of the subcomponents of
6a geometry object: nodes, elements, tracelines and coordinate system. Geometry
7plotting is also handled in this module.
8"""
10"""
11Copyright 2022 National Technology & Engineering Solutions of Sandia,
12LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
13Government retains certain rights in this software.
15This program is free software: you can redistribute it and/or modify
16it under the terms of the GNU General Public License as published by
17the Free Software Foundation, either version 3 of the License, or
18(at your option) any later version.
20This program is distributed in the hope that it will be useful,
21but WITHOUT ANY WARRANTY; without even the implied warranty of
22MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23GNU General Public License for more details.
25You should have received a copy of the GNU General Public License
26along with this program. If not, see <https://www.gnu.org/licenses/>.
27"""
29import numpy as np
30from scipy.spatial import cKDTree
31from scipy.sparse import coo_matrix
32from scipy.sparse.csgraph import minimum_spanning_tree
34from ..core.sdynpy_geometry import Geometry, TracelineArray, traceline_array
37def build_node_id_to_index(node_ids):
38 node_ids = np.asarray(node_ids)
39 return {nid: i for i, nid in enumerate(node_ids)}
42def derive_node_block_sets(node_ids, elem_conn, elem_block_ids):
43 """
44 Derive, for each node index, the set of element blocks touching that node.
46 Parameters
47 ----------
48 node_ids : (n,) array of node IDs
49 elem_conn : sequence length m
50 Each entry is an iterable of node IDs for one element.
51 elem_block_ids : (m,) array-like
52 Block ID for each element.
54 Returns
55 -------
56 node_block_sets : list of sets, length n
57 node_block_sets[i] is the set of block IDs incident on node i.
58 """
59 id_to_index = build_node_id_to_index(node_ids)
60 n = len(node_ids)
62 node_block_sets = [set() for _ in range(n)]
64 for conn, blk in zip(elem_conn, elem_block_ids):
65 for nid in conn:
66 if nid in id_to_index:
67 i = id_to_index[nid]
68 node_block_sets[i].add(blk)
70 return node_block_sets
73def derive_node_adjacency_from_elements(node_ids, elem_conn):
74 """
75 Build a node adjacency graph from element connectivity.
77 Here, any two nodes appearing in the same element are marked adjacent.
78 This is a coarse notion of adjacency, but it is easy to compute and
79 works even for mixed element sizes.
81 Parameters
82 ----------
83 node_ids : (n,) array of node IDs
84 elem_conn : sequence length m
85 Each entry is an iterable of node IDs for one element.
87 Returns
88 -------
89 adjacency : list of sets, length n
90 adjacency[i] contains node indices adjacent to node i.
91 """
92 id_to_index = build_node_id_to_index(node_ids)
93 n = len(node_ids)
94 adjacency = [set() for _ in range(n)]
96 for conn in elem_conn:
97 idxs = [id_to_index[nid] for nid in conn if nid in id_to_index]
99 # connect all pairs within the element
100 # for visualization purposes this is acceptable as a coarse adjacency
101 for a_pos in range(len(idxs)):
102 i = idxs[a_pos]
103 for b_pos in range(a_pos + 1, len(idxs)):
104 j = idxs[b_pos]
105 adjacency[i].add(j)
106 adjacency[j].add(i)
108 return adjacency
111def block_overlap_penalty(block_set_i, block_set_j, penalty):
112 """
113 If block memberships overlap, no penalty.
114 Otherwise apply multiplicative penalty.
115 """
116 if block_set_i & block_set_j:
117 return 1.0
118 return 1.0 + penalty
121def build_visualization_graph_from_fe_data(
122 node_ids,
123 node_coords,
124 elem_conn,
125 elem_block_ids,
126 retained_node_ids=None,
127 k_candidates=12,
128 k_extra=2,
129 block_penalty=0.5,
130 max_length_factor=3.0,
131 use_node_adjacency=False,
132):
133 """
134 Build a sparse visualization graph for retained FE nodes.
136 Parameters
137 ----------
138 node_ids : (n,) array-like
139 All node IDs.
140 node_coords : (n,3) array-like
141 Coordinates for all nodes.
142 elem_conn : sequence length m
143 Each element is an iterable of node IDs.
144 elem_block_ids : (m,) array-like
145 Block ID per element.
146 retained_node_ids : array-like, optional
147 Subset of node IDs to include. If None, all nodes are retained.
148 k_candidates : int
149 Number of Euclidean nearest neighbors to consider.
150 k_extra : int
151 Number of extra local edges per node beyond the MST.
152 block_penalty : float
153 Penalty applied when nodes share no block membership.
154 max_length_factor : float
155 Local edges longer than this times median nearest-neighbor distance
156 are pruned. MST edges are always kept.
157 use_node_adjacency : bool
158 If True, require candidate edges to be between nodes that are adjacent
159 in the coarse element-derived node adjacency graph.
161 Returns
162 -------
163 retained_ids : (nr,) ndarray
164 Retained node IDs, in returned local ordering.
165 retained_coords : (nr,3) ndarray
166 Retained node coordinates.
167 edges : list of tuple(int, int)
168 Undirected edges in retained-node local indexing.
169 edge_lengths : ndarray
170 Euclidean edge lengths.
171 retained_block_sets : list of sets
172 Block-membership sets for retained nodes.
173 """
175 node_ids = np.asarray(node_ids)
176 node_coords = np.asarray(node_coords, dtype=float)
177 elem_block_ids = np.asarray(elem_block_ids)
179 if node_coords.ndim != 2 or node_coords.shape[1] != 3:
180 raise ValueError("node_coords must have shape (n, 3)")
182 if len(node_ids) != len(node_coords):
183 raise ValueError("node_ids and node_coords must have same length")
185 # Derive node block memberships on full node set
186 full_node_block_sets = derive_node_block_sets(node_ids, elem_conn, elem_block_ids)
188 # Optional full node adjacency
189 full_adjacency = None
190 if use_node_adjacency:
191 full_adjacency = derive_node_adjacency_from_elements(node_ids, elem_conn)
193 # Retained subset
194 if retained_node_ids is None:
195 retain_mask = np.ones(len(node_ids), dtype=bool)
196 else:
197 retained_node_ids = set(retained_node_ids)
198 retain_mask = np.array([nid in retained_node_ids for nid in node_ids], dtype=bool)
200 retained_ids = node_ids[retain_mask]
201 retained_coords = node_coords[retain_mask]
202 retained_block_sets = [full_node_block_sets[i] for i in np.where(retain_mask)[0]]
204 nr = len(retained_ids)
205 if nr < 2:
206 return retained_ids, retained_coords, [], np.array([]), retained_block_sets
208 # Map full-node index -> retained-node local index
209 full_to_retained = {}
210 retained_full_indices = np.where(retain_mask)[0]
211 for local_i, full_i in enumerate(retained_full_indices):
212 full_to_retained[full_i] = local_i
214 # Build retained adjacency if requested
215 retained_adjacency = None
216 if use_node_adjacency:
217 retained_adjacency = [set() for _ in range(nr)]
218 for full_i in retained_full_indices:
219 i_local = full_to_retained[full_i]
220 for full_j in full_adjacency[full_i]:
221 if full_j in full_to_retained:
222 j_local = full_to_retained[full_j]
223 retained_adjacency[i_local].add(j_local)
225 # Geometric neighbor search on retained nodes
226 tree = cKDTree(retained_coords)
227 k_query = min(k_candidates + 1, nr)
228 dists, nbrs = tree.query(retained_coords, k=k_query)
230 if k_query == 1:
231 dists = dists[:, None]
232 nbrs = nbrs[:, None]
234 nn_dists = []
235 for i in range(nr):
236 for jj in range(1, k_query):
237 j = int(nbrs[i, jj])
238 if i == j:
239 continue
240 if use_node_adjacency and j not in retained_adjacency[i]:
241 continue
242 nn_dists.append(dists[i, jj])
243 break
245 if not nn_dists:
246 return retained_ids, retained_coords, [], np.array([]), retained_block_sets
248 median_nn_dist = np.median(nn_dists)
249 max_allowed_length = max_length_factor * median_nn_dist
251 # Candidate edges
252 edge_data = {}
253 # edge_data[(i,j)] = (weighted_dist, euclidean_dist)
255 for i in range(nr):
256 for jj in range(1, k_query):
257 j = int(nbrs[i, jj])
258 if i == j:
259 continue
261 if use_node_adjacency and j not in retained_adjacency[i]:
262 continue
264 a, b = sorted((i, j))
265 euclid = np.linalg.norm(retained_coords[a] - retained_coords[b])
267 penalty_factor = block_overlap_penalty(
268 retained_block_sets[a], retained_block_sets[b], block_penalty
269 )
270 weighted = euclid * penalty_factor
272 key = (a, b)
273 if key not in edge_data or weighted < edge_data[key][0]:
274 edge_data[key] = (weighted, euclid)
276 if not edge_data:
277 return retained_ids, retained_coords, [], np.array([]), retained_block_sets
279 # Build weighted sparse graph for MST
280 rows, cols, data = [], [], []
281 for (i, j), (weighted, _) in edge_data.items():
282 rows.extend([i, j])
283 cols.extend([j, i])
284 data.extend([weighted, weighted])
286 graph = coo_matrix((data, (rows, cols)), shape=(nr, nr)).tocsr()
287 mst = minimum_spanning_tree(graph)
289 mst_edges = set()
290 mst_coo = mst.tocoo()
291 for i, j in zip(mst_coo.row, mst_coo.col):
292 a, b = sorted((int(i), int(j)))
293 mst_edges.add((a, b))
295 # Add extra local edges
296 extra_edges = set()
297 for i in range(nr):
298 local_candidates = []
300 for jj in range(1, k_query):
301 j = int(nbrs[i, jj])
302 if i == j:
303 continue
305 if use_node_adjacency and j not in retained_adjacency[i]:
306 continue
308 a, b = sorted((i, j))
309 euclid = np.linalg.norm(retained_coords[a] - retained_coords[b])
311 penalty_factor = block_overlap_penalty(
312 retained_block_sets[a], retained_block_sets[b], block_penalty
313 )
314 weighted = euclid * penalty_factor
316 local_candidates.append((weighted, euclid, a, b))
318 local_candidates.sort(key=lambda x: x[0])
320 added = 0
321 for weighted, euclid, a, b in local_candidates:
322 if euclid <= max_allowed_length:
323 extra_edges.add((a, b))
324 added += 1
325 if added >= k_extra:
326 break
328 all_edges = mst_edges | extra_edges
330 # Prune long edges, but keep MST
331 final_edges = set()
332 for a, b in all_edges:
333 euclid = np.linalg.norm(retained_coords[a] - retained_coords[b])
334 if (a, b) in mst_edges or euclid <= max_allowed_length:
335 final_edges.add((a, b))
337 final_edges = sorted(final_edges)
338 edge_lengths = np.array(
339 [np.linalg.norm(retained_coords[i] - retained_coords[j]) for i, j in final_edges]
340 )
342 return retained_ids, retained_coords, final_edges, edge_lengths, retained_block_sets
345def create_visualization_tracelines(
346 geometry: Geometry,
347 retained_node_ids=None,
348 k_candidates=12,
349 k_extra=2,
350 block_penalty=0.5,
351 max_length_factor=3.0,
352 use_node_adjacency=False,
353):
354 """Creates tracelines that can be used to improve visualization of a reduced geometry.
356 Parameters
357 ----------
358 geometry : Geometry
359 The geometry containing the original mesh with original elements and colors.
360 retained_node_ids : np.ndarray,
361 A numpy array consisting of the node IDs that will be retained.
362 k_candidates : int, optional
363 Number of nearest neighbors to consider, by default 12
364 k_extra : int, optional
365 Number of extra local edges per node beyond the MST, by default 2
366 block_penalty : float, optional
367 Penalty applied when nodes share no block membership, by default 0.5
368 max_length_factor : float, optional
369 Local edges longer than this times median nearest-neighbor distance
370 are pruned. MST edges are always kept. By default 3.0
371 use_node_adjacency : bool, optional
372 If True, require candidate edges to be between nodes that are adjacent
373 in the coarse element-derived node adjacency graph. By default, False.
375 Returns
376 -------
377 tracelines : TracelineArray
378 Returns a TracelineArray containing lines that can be used to visualize reduced geometry.
379 """
381 node_ids = geometry.node.id
382 coords = geometry.global_node_coordinate()
383 conn = geometry.element.connectivity
384 block_ids = geometry.element.color
386 ret_ids, ret_xyz, edges, lengths, block_sets = build_visualization_graph_from_fe_data(
387 node_ids=node_ids,
388 node_coords=coords,
389 elem_conn=conn,
390 elem_block_ids=block_ids,
391 retained_node_ids=retained_node_ids,
392 k_candidates=k_candidates,
393 k_extra=k_extra,
394 block_penalty=block_penalty,
395 max_length_factor=max_length_factor,
396 use_node_adjacency=use_node_adjacency,
397 )
398 connectivity = [ret_ids[np.array(edge)] for edge in edges]
399 colors = []
400 for edge in edges:
401 set_1 = block_sets[edge[0]]
402 set_2 = block_sets[edge[1]]
403 full_set = set_1 | set_2
404 # Pick the first item if there is more than 1
405 color, *_ = full_set
406 colors.append(color)
407 return traceline_array(
408 np.arange(len(connectivity)) + 1, color=colors, connectivity=connectivity
409 )