Various QOL updates for networks

.github/workflows/build_docs.yml

@@ -26,7 +26,7 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v6
+        uses: actions/checkout@v7
 
       - name: Install dependencies
         run: python3 -m pip install ".[docs]"
@@ -35,7 +35,7 @@ jobs:
         run: jupyter book build --keep-going .
 
       - name: Upload artifact
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           path: ${{ env.PUBLISH_DIR }}
           name: ${{ matrix.label }}-docs

.github/workflows/check_formatting.yml

@@ -24,10 +24,10 @@ jobs:
 
     steps:
       # This action sets the current path to the root of your github repo
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v7
         with:
           python-version: ${{ matrix.python-version }}
 

.github/workflows/deploy.yml

@@ -29,22 +29,22 @@ jobs:
     steps:
       - name: Download docs artifact
         # docs artifact is uploaded by build-docs job
-        uses: actions/download-artifact@v7
+        uses: actions/download-artifact@v8
         with:
           name: stable-docs
           path: "./public"
 
       - name: Upload artifact
-        uses: actions/upload-pages-artifact@v4
+        uses: actions/upload-pages-artifact@v5
         with:
           path: "./public"
 
       - name: Checkout
-        uses: actions/checkout@v6
+        uses: actions/checkout@v7
 
       - name: Setup Pages
-        uses: actions/configure-pages@v5
+        uses: actions/configure-pages@v6
 
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5

.github/workflows/pypi.yml

@@ -10,7 +10,7 @@ jobs:
       DEB_PYTHON_INSTALL_LAYOUT: deb_system
 
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
 
       - name: Install build dependencies
         run: |
@@ -21,7 +21,7 @@ jobs:
       - name: Build SDist and wheel
         run: python3 -m build --no-isolation --sdist
 
-      - uses: actions/upload-artifact@v6
+      - uses: actions/upload-artifact@v7
         with:
           path: dist/*
 
@@ -37,7 +37,7 @@ jobs:
       id-token: write
 
     steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
         with:
           name: artifact
           path: dist

.github/workflows/test_package.yml

@@ -14,14 +14,16 @@ jobs:
     container: ghcr.io/fenics/dolfinx/dolfinx:${{ matrix.label }}
     env:
       DEB_PYTHON_INSTALL_LAYOUT: deb_system
+      PRTE_MCA_rmaps_default_mapping_policy: ":oversubscribe"
+
 
     strategy:
       fail-fast: false
       matrix:
         label: ["stable", "nightly"]
 
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
 
       - name: Update pip
         run: python3 -m pip install --upgrade pip
@@ -48,4 +50,4 @@ jobs:
 
       - name: Run demos
         working-directory: demos
-        run: pytest . -sv
+        run: python3 -m pytest . -sv

src/networks_fenicsx/assembly.py

@@ -12,6 +12,7 @@
 import numpy.typing as npt
 
 import basix
+import dolfinx.fem.petsc as _petsc_fem
 import dolfinx.la.petsc as _petsc_la
 import ufl
 from dolfinx import common, fem
@@ -328,12 +329,12 @@ def network(self) -> NetworkMesh:
     @common.timed("nxfx:HydraulicNetworkAssembler:assemble")
     def assemble(
         self,
-        A: PETSc.Mat | None = None,  # type: ignore[name-defined]
-        b: PETSc.Mat | None = None,  # type: ignore[name-defined]
+        A: PETSc.Mat | None = None,
+        b: PETSc.Vec | None = None,
         assemble_lhs: bool = True,
         assemble_rhs: bool = True,
         kind: str | typing.Sequence[typing.Sequence[str]] | None = None,
-    ) -> tuple[PETSc.Mat, PETSc.Vec]:  # type: ignore[name-defined]
+    ) -> tuple[PETSc.Mat, PETSc.Vec]:
         """Assemble system matrix and rhs vector.
 
         Note:
@@ -351,20 +352,21 @@ def assemble(
         """
         if assemble_lhs:
             if A is None:
-                A = fem.petsc.create_matrix([[aij for aij in ai] for ai in self._a], kind=kind)
-            A = fem.petsc.assemble_matrix(A, self._a, bcs=[])  # type: ignore
+                A = _petsc_fem.create_matrix([[aij for aij in ai] for ai in self._a], kind=kind)
+            A = _petsc_fem.assemble_matrix(A, self._a, bcs=[])  # type: ignore
             A.assemble()
             kind = "nest" if A.getType() == PETSc.Mat.Type.NEST else kind  # type: ignore[attr-defined]
         if assemble_rhs:
             if b is None:
                 assert isinstance(kind, str) or kind is None
-                b = fem.petsc.create_vector(fem.extract_function_spaces(self._L), kind=kind)
-            b = fem.petsc.assemble_vector(b, self._L)  # type: ignore
+                b = _petsc_fem.create_vector(fem.extract_function_spaces(self._L), kind=kind)
+            b = _petsc_fem.assemble_vector(b, self._L)  # type: ignore
             _petsc_la._ghost_update(
                 b,
-                insert_mode=PETSc.InsertMode.ADD_VALUES,  # type: ignore[attr-defined]
-                scatter_mode=PETSc.ScatterMode.REVERSE,  # type: ignore[attr-defined]
+                insert_mode=PETSc.InsertMode.ADD_VALUES,  # type: ignore[arg-type]
+                scatter_mode=PETSc.ScatterMode.REVERSE,  # type: ignore[arg-type]
             )
+        assert A is not None and b is not None
         return (A, b)
 
     @property

src/networks_fenicsx/mesh.py

@@ -328,23 +328,27 @@ def _original_order(a, b):
             cell_markers_ = np.empty((0,), dtype=np.int32)
             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()):
-            part_kwargs["max_facet_to_cell_links"] = np.max(max_connections)
-        partitioner = mesh.create_cell_partitioner(mesh.GhostMode.shared_facet, **part_kwargs)
-        sig = inspect.signature(mesh.create_mesh)
-        kwargs = {}
-        if "max_facet_to_cell_links" in list(sig.parameters.keys()):
-            kwargs["max_facet_to_cell_links"] = np.max(max_connections)
+        max_facet_to_cell_links = np.max(max_connections)
+
+        if hasattr(mesh, "create_cell_partitioner"):
+            sig = inspect.signature(mesh.create_cell_partitioner)
+            if "max_facet_to_cell_links" in list(sig.parameters.keys()):
+                part = mesh.create_cell_partitioner(
+                    mesh.GhostMode.shared_facet,
+                    max_facet_to_cell_links=max_facet_to_cell_links,
+                )
+            else:
+                part = mesh.create_cell_partitioner(mode=mesh.GhostMode.shared_facet)  # type: ignore
+        else:
+            part = _graph.partitioner()
 
         graph_mesh = mesh.create_mesh(
             comm,
             x=mesh_nodes,
             cells=cells_,
             e=ufl.Mesh(basix.ufl.element("Lagrange", "interval", 1, shape=(self._geom_dim,))),
-            partitioner=partitioner,
-            **kwargs,
+            partitioner=part,
+            max_facet_to_cell_links=max_facet_to_cell_links,
         )
         self._msh = graph_mesh
 
@@ -516,13 +520,13 @@ def in_edges(self, bifurcation_idx: int) -> npt.NDArray[np.int32]:
         """Return the list of in-edge colors for a given bifurcation node.
         Index is is the index of the bifurcation in {py:meth}`self.bifurcation_values`."""
         assert bifurcation_idx < len(self.bifurcation_values)
-        return self._bifurcation_in_color.links(np.int32(bifurcation_idx))
+        return self._bifurcation_in_color.links(int(bifurcation_idx))
 
     def out_edges(self, bifurcation_idx: int) -> npt.NDArray[np.int32]:
         """Return the list of out-edge colors for a given bifurcation node.
         Index is is the index of the bifurcation in {py:meth}`self.bifurcation_values`."""
         assert bifurcation_idx < len(self.bifurcation_values)
-        return self._bifurcation_out_color.links(np.int32(bifurcation_idx))
+        return self._bifurcation_out_color.links(int(bifurcation_idx))
 
     @property
     def num_edge_colors(self) -> int:

src/networks_fenicsx/solver.py

@@ -38,7 +38,7 @@ def __init__(
     ):
         self._assembler = assembler
 
-        self._ksp = PETSc.KSP().create(self._assembler.network.comm)  # type: ignore[attr-defined]
+        self._ksp = PETSc.KSP().create(self._assembler.network.comm)  # type: ignore[arg-type]
 
         self._A = dolfinx.fem.petsc.create_matrix(self._assembler.bilinear_forms, kind=kind)
         kind = "nest" if self._A.getType() == "nest" else kind  # type: ignore[attr-defined]
@@ -66,25 +66,29 @@ def __init__(
         opts = PETSc.Options()  # type: ignore[attr-defined]
         opts.prefixPush(self.ksp.getOptionsPrefix())
         for key, value in petsc_options.items():
-            opts[key] = value
+            opts.setValue(key, value)
         self.ksp.setFromOptions()
         self._A.setFromOptions()
         self._b.setFromOptions()
         opts.prefixPop()
+        for key, value in petsc_options.items():
+            opts.delValue(f"{self.ksp.getOptionsPrefix()}{key}")
 
     @property
     def assembler(self) -> assembly.HydraulicNetworkAssembler:
         """The hydraulic network assembler."""
         return self._assembler
 
     @property
-    def A(self) -> PETSc.Mat:  # type: ignore[name-defined]
+    def A(self) -> PETSc.Mat:
         """System matrix."""
+        assert self._A is not None
         return self._A
 
     @property
-    def b(self) -> PETSc.Vec:  # type: ignore[name-defined]
+    def b(self) -> PETSc.Vec:
         """Right-hand side vector."""
+        assert self._b is not None
         return self._b
 
     def assemble(self, lhs: bool = True, rhs: bool = True):
@@ -101,13 +105,14 @@ def assemble(self, lhs: bool = True, rhs: bool = True):
         self.assembler.assemble(self._A, self._b, assemble_lhs=lhs, assemble_rhs=rhs)
 
     @property
-    def ksp(self) -> PETSc.KSP:  # type: ignore[name-defined]
+    def ksp(self) -> PETSc.KSP:
+        assert self._ksp is not None
         return self._ksp
 
     @dolfinx.common.timed("nxfx:Solver:solve")
     def solve(
-        self, functions: list[dolfinx.fem.Function] | None = None
-    ) -> list[dolfinx.fem.Function]:
+        self, functions: typing.Sequence[dolfinx.fem.Function] | None = None
+    ) -> typing.Sequence[dolfinx.fem.Function]:
         """Solve the linear system of equations and assign them to a set of corresponding
         DOLFINx functions.
 
@@ -123,15 +128,16 @@ def solve(
                 functions.append(dolfinx.fem.Function(Vi, name=f"flux_color_{i}"))
             functions.append(dolfinx.fem.Function(self.assembler.pressure_space, name="pressure"))
             functions.append(dolfinx.fem.Function(self.assembler.lm_space, name="global_flux"))
-
+        assert self._x is not None
         self.ksp.solve(self.b, self._x)
         dolfinx.la.petsc._ghost_update(
             self._x,
-            insert_mode=PETSc.InsertMode.INSERT,  # type: ignore[attr-defined]
-            scatter_mode=PETSc.ScatterMode.FORWARD,  # type: ignore[attr-defined]
+            insert_mode=PETSc.InsertMode.INSERT,  # type: ignore[arg-type]
+            scatter_mode=PETSc.ScatterMode.FORWARD,  # type: ignore[arg-type]
         )
-        assert isinstance(self._x, PETSc.Vec)  # type: ignore[attr-defined]
-        dolfinx.fem.petsc.assign(self._x, functions)
+        assert isinstance(self._x, PETSc.Vec)
+        assert isinstance(functions, typing.Sequence)
+        dolfinx.fem.petsc.assign(self._x, functions)  # type: ignore[arg-type]
         return functions
 
     def __del__(self):