diff --git a/.github/workflows/test_package.yml b/.github/workflows/test_package.yml index 7c6dd2a..f9d9405 100644 --- a/.github/workflows/test_package.yml +++ b/.github/workflows/test_package.yml @@ -37,8 +37,16 @@ jobs: - name: Run tests working-directory: tests run: python3 -m pytest -vxs + + - name: Run tests in parallel (np=2) + working-directory: tests + run: mpiexec -n 2 python3 -m pytest -v -m "not serial" + + - name: Run tests in parallel (np=3) + working-directory: tests + run: mpiexec -n 3 python3 -m pytest -v -m "not serial" - - name: Run tests in parallel + - name: Run tests in parallel (np=4) working-directory: tests run: mpiexec -n 4 python3 -m pytest -v -m "not serial" diff --git a/src/networks_fenicsx/assembly.py b/src/networks_fenicsx/assembly.py index 853b792..2466acd 100644 --- a/src/networks_fenicsx/assembly.py +++ b/src/networks_fenicsx/assembly.py @@ -235,7 +235,11 @@ def compute_forms( dx_global = ufl.Measure("dx", domain=network_mesh.mesh) - tangent = self._network_mesh.tangent + J = ufl.Jacobian(network_mesh.mesh) + t = J[:, 0] + t /= ufl.sqrt(ufl.inner(t, t)) + + tangent = self._network_mesh.orientation * t for i, (submesh, entity_map, facet_marker) in enumerate( zip( network_mesh.submeshes, diff --git a/src/networks_fenicsx/mesh.py b/src/networks_fenicsx/mesh.py index 377201b..f77683c 100644 --- a/src/networks_fenicsx/mesh.py +++ b/src/networks_fenicsx/mesh.py @@ -9,7 +9,7 @@ """ import inspect -from typing import Any, Callable, Iterable +from typing import Callable, Iterable from mpi4py import MPI @@ -75,7 +75,8 @@ class NetworkMesh: _submesh_facet_markers: list[mesh.MeshTags] _edge_meshes: list[mesh.Mesh] _edge_entity_maps: list[mesh.EntityMap] - _tangent: fem.Function + _orientation: fem.Function + _boundary_values: npt.NDArray[np.int32] _lm_mesh: mesh.Mesh | None _lm_map: mesh.EntityMap | None @@ -175,7 +176,6 @@ def _build_mesh( assert isinstance(graph, nx.DiGraph), f"Directional graph not present of {graph_rank}" geom_dim = len(graph.nodes[1]["pos"]) edge_coloring = color_graph(graph, color_strategy) - num_edge_colors = len(set(edge_coloring.values())) cells_array = np.asarray([[u, v] for u, v in graph.edges()]) number_of_nodes = graph.number_of_nodes() @@ -196,14 +196,16 @@ def _build_mesh( bifurcation_in_offsets, len(bifurcation_in_color) + len(in_edges) ) bifurcation_in_color = np.append( - bifurcation_in_color, [edge_coloring[edge] for edge in in_edges] + bifurcation_in_color, + np.asarray([edge_coloring[edge] for edge in in_edges], dtype=np.int32), ) out_edges = graph.out_edges(bifurcation) bifurcation_out_offsets = np.append( bifurcation_out_offsets, len(bifurcation_out_color) + len(out_edges) ) bifurcation_out_color = np.append( - bifurcation_out_color, [edge_coloring[edge] for edge in out_edges] + bifurcation_out_color, + np.asarray([edge_coloring[edge] for edge in out_edges], dtype=np.int32), ) # Map boundary_values to inlet and outlet data from graph @@ -260,8 +262,11 @@ def _build_mesh( bifurcation_out_color, bifurcation_out_offsets ) + def _original_order(a, b): + """Order on edge entities to detect order reversals during mesh generation.""" + return a < b + # Generate mesh - tangents: npt.NDArray[np.float64] = np.empty((0, self._geom_dim), dtype=np.float64) if comm.rank == graph_rank: assert isinstance(graph, nx.DiGraph), ( f"No directional graph present on rank {comm.rank}" @@ -271,14 +276,12 @@ def _build_mesh( mesh_nodes = vertex_coords.copy() cells = [] cell_markers = [] - local_tangents = [] if len(line_weights) == 0: for segment in cells_array: cells.append(np.array([segment[0], segment[1]], dtype=np.int64)) cell_markers.append(edge_coloring[(segment[0], segment[1])]) start = vertex_coords[segment[0]] end = vertex_coords[segment[1]] - local_tangents.append(_compute_tangent(start, end)) else: for segment in cells_array: start_coord_pos = mesh_nodes.shape[0] @@ -286,9 +289,6 @@ def _build_mesh( end = vertex_coords[segment[1]] internal_line_coords = start * (1 - line_weights) + end * line_weights - tangent = _compute_tangent(start, end) - local_tangents.append(np.tile(tangent, (internal_line_coords.shape[0] + 1, 1))) - mesh_nodes = np.vstack((mesh_nodes, internal_line_coords)) cells.append(np.array([segment[0], start_coord_pos], dtype=np.int64)) segment_connectivity = ( @@ -314,14 +314,20 @@ def _build_mesh( dtype=np.int32, ) ) + cells_ = np.vstack(cells).astype(np.int64) cell_markers_ = np.array(cell_markers, dtype=np.int32) - tangents = np.vstack(local_tangents).astype(np.float64) + + orientations = np.full_like(cell_markers_, 1.0, dtype=np.float64) + orientations[~_original_order(cells_[:, 0], cells_[:, 1])] = -1.0 + + assert cell_markers_.shape == orientations.shape else: cells_ = np.empty((0, 2), dtype=np.int64) mesh_nodes = np.empty((0, self._geom_dim), dtype=np.float64) cell_markers_ = np.empty((0,), dtype=np.int32) - tangents = tangents[:, : self._geom_dim].copy() + orientations = np.empty(0, dtype=np.float64) + sig = inspect.signature(mesh.create_cell_partitioner) part_kwargs = {} if "max_facet_to_cell_links" in list(sig.parameters.keys()): @@ -342,49 +348,13 @@ def _build_mesh( ) self._msh = graph_mesh + tdim = self.mesh.topology.dim + + # Compute subdomain tags. local_entities, local_values = _io.distribute_entity_data( - self.mesh, - self.mesh.topology.dim, - cells_, - cell_markers_, + self.mesh, tdim, cells_, cell_markers_ ) - # Distribute tangents - tangent_space = fem.functionspace(graph_mesh, ("DG", 0, (self._geom_dim,))) - self._tangent = fem.Function(tangent_space) - - # Which cells from the input mesh (rank 0) correspond to the local cells on this rank - original_cell_indices = self.mesh.topology.original_cell_index[ - : self.mesh.topology.index_map(self.mesh.topology.dim).size_local - ] - - # First compute how much data we are getting from each rank - send_count = original_cell_indices.size - recv_counts = comm.gather(send_count, root=graph_rank) - if comm.rank == graph_rank: - assert isinstance(recv_counts, Iterable) - assert sum(recv_counts) == cells_.shape[0] - - # Send to rank 0 which tangents that we need - recv_buffer: None | list[npt.NDArray[np.int64] | list[Any] | MPI.Datatype] = None - if comm.rank == graph_rank: - assert recv_counts is not None - recv_buffer = [np.empty(cells_.shape[0], dtype=np.int64), recv_counts, MPI.INT64_T] - - comm.Gatherv(sendbuf=original_cell_indices, recvbuf=recv_buffer, root=graph_rank) - send_t_buffer = None - if comm.rank == graph_rank: - assert recv_buffer is not None - recv_data = recv_buffer[0] - assert isinstance(recv_data, np.ndarray) - send_t_buffer = [ - tangents[recv_data].flatten(), - np.array(recv_counts) * self._geom_dim, - ] - - comm.Scatterv(sendbuf=send_t_buffer, recvbuf=self.tangent.x.array, root=graph_rank) - self.tangent.x.scatter_forward() - self._subdomains = mesh.meshtags_from_entities( self.mesh, self.mesh.topology.dim, @@ -392,6 +362,43 @@ def _build_mesh( local_values, ) + # Distribute orientations as meshtag. + local_cells, local_orientations = _io.distribute_entity_data( + self.mesh, tdim, cells_, orientations + ) + + meshtag_orientation = mesh.meshtags_from_entities( + self.mesh, tdim, _graph.adjacencylist(local_cells), local_orientations + ) + + # Assemble orientations function + orientation_space = fem.functionspace(graph_mesh, ("DG", 0)) + self._orientation = fem.Function(orientation_space) + self._orientation.x.array[meshtag_orientation.indices] = meshtag_orientation.values + + # Correct orientations for possible reorder + e_idx = np.arange(self.mesh.topology.index_map(tdim).size_local, dtype=np.int32) + e_geo = mesh.entities_to_geometry(self.mesh, tdim, e_idx) + + global_input = self.mesh.geometry.input_global_indices + in_order = _original_order(global_input[e_geo[:, 0]], global_input[e_geo[:, 1]]) + + # Four cases that might arise (per edge), with naming i := input_in_order, n := in_order: + # 1) i & n: + # orientation = 1 (input=1, no action needed) + # 2) i & !n: + # orientation = -1 (input=1, need sign flip) + # 3) !i & n: + # orientation = -1 (input=-1, no action needed) + # 4) !i & !n + # orientation = 1 (input=-1, need sign flip) + # + # so we need to sign flip when + # (i & !n) | (!i & !n) = !n + self.orientation.x.array[: e_idx.size][~in_order] *= -1 + + self.orientation.x.scatter_forward() + self._in_marker = 3 * number_of_nodes self._out_marker = 5 * number_of_nodes if comm.rank == graph_rank: @@ -493,9 +500,9 @@ def entity_maps(self): return self._edge_entity_maps @property - def tangent(self): + def orientation(self): """Return DG-0 field containing the tangent vector of the graph.""" - return self._tangent + return self._orientation @property def bifurcation_values(self) -> npt.NDArray[np.int32]: @@ -529,23 +536,3 @@ def in_marker(self) -> int: @property def out_marker(self) -> int: return self._out_marker - - -def _compute_tangent( - start: npt.NDArray[np.float64], end: npt.NDArray[np.float64] -) -> npt.NDArray[np.float64]: - """Compute the tangent vector from start to end points. - - Tangent is oriented according to positive y-axis. - If perpendicular to y-axis, align with x-axis. - - Args: - start: Starting point coordinates. - end: Ending point coordinates. - - Returns: - Normalized tangent vector from start to end. - """ - tangent = end - start - tangent /= np.linalg.norm(tangent, axis=-1) - return tangent diff --git a/tests/test_orientation.py b/tests/test_orientation.py new file mode 100644 index 0000000..599926a --- /dev/null +++ b/tests/test_orientation.py @@ -0,0 +1,58 @@ +import networkx as nx +import numpy as np +import pytest + +import dolfinx +import ufl +from networks_fenicsx.mesh import NetworkMesh + + +def linear_graph(n: int, dim: int = 2, ordered=lambda _: True) -> nx.DiGraph: + G = nx.DiGraph() + G.add_nodes_from(range(n)) + + for i in range(n - 1): + if ordered(i): + G.add_edge(i, i + 1) + else: + G.add_edge(i + 1, i) + + for i in range(n): + pos = np.zeros(dim) + pos[0] = i / (n - 1) + G.nodes[i]["pos"] = pos + + return G + + +@pytest.mark.parametrize("n", [30]) +@pytest.mark.parametrize("order", ["in", "reverse", "alternating"]) +@pytest.mark.parametrize("N", [1, 4, 8]) +def test_orientation(n: int, order: str, N: int) -> None: + if order == "in": + ordered = lambda _: True + elif order == "reverse": + ordered = lambda _: False + elif order == "alternating": + ordered = lambda k: k % 2 + else: + raise RuntimeError() + + G = linear_graph(n, ordered=ordered) + + network_mesh = NetworkMesh(G, N=N) + + J = ufl.Jacobian(network_mesh.mesh) + t = J[:, 0] + t /= ufl.sqrt(ufl.inner(t, t)) + f = dolfinx.fem.form(ufl.inner(ufl.as_vector((1, 0)), t) * network_mesh.orientation * ufl.dx) + + val = network_mesh.comm.allreduce(dolfinx.fem.assemble_scalar(f)) + + if order == "in": + assert np.isclose(val, 1.0) + elif order == "reverse": + assert np.isclose(val, -1.0) + else: + edge_count = n - 1 + assert np.isclose(val, edge_count % 2 * -1 / edge_count)