Update subspace management in tests and tutorials to follow best prac…

…tices as suggested upstream
multiphenics · Dec 28, 2023 · f4af8d0 · f4af8d0
1 parent 21ab832
commit f4af8d0
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 19 deletions.
diff --git a/tests/unit/fem/test_assembler_restriction.py b/tests/unit/fem/test_assembler_restriction.py
@@ -262,13 +262,14 @@ def _get_boundary_conditions(V: dolfinx.fem.FunctionSpace) -> typing.List[dolfin
         else:
             bc1 = list()
             for i in range(num_sub_elements):
-                bc1_fun = dolfinx.fem.Function(V.sub(i).collapse()[0])
+                Vi = V.sub(i)
+                bc1_fun = dolfinx.fem.Function(Vi.collapse()[0])
                 bc1_vector = bc1_fun.vector
                 with bc1_vector.localForm() as local_form:
                     local_form.set(i + 1. + offset)
                 bc1_vector.destroy()
-                bdofs = locate_boundary_dofs(V.sub(i), bc1_fun.function_space)
-                bc1.append(dolfinx.fem.dirichletbc(bc1_fun, bdofs, V.sub(i)))
+                bdofs = locate_boundary_dofs(Vi, bc1_fun.function_space)
+                bc1.append(dolfinx.fem.dirichletbc(bc1_fun, bdofs, Vi))
             return bc1
 
     return (lambda _: [],

diff --git a/tutorials/02_navier_stokes/tutorial_navier_stokes.ipynb b/tutorials/02_navier_stokes/tutorial_navier_stokes.ipynb
@@ -257,16 +257,16 @@
     "    J = ufl.derivative(F, up, dup)\n",
     "\n",
     "    # Boundary conditions\n",
-    "    u_in = dolfinx.fem.Function(W.sub(0).collapse()[0])\n",
+    "    W0 = W.sub(0)\n",
+    "    V, _ = W0.collapse()\n",
+    "    u_in = dolfinx.fem.Function(V)\n",
     "    u_in.interpolate(u_in_eval)\n",
-    "    u_wall = dolfinx.fem.Function(W.sub(0).collapse()[0])\n",
+    "    u_wall = dolfinx.fem.Function(V)\n",
     "    u_wall.interpolate(u_wall_eval)\n",
-    "    bdofs_V_1 = dolfinx.fem.locate_dofs_topological(\n",
-    "        (W.sub(0), W.sub(0).collapse()[0]), mesh.topology.dim - 1, boundaries_1)\n",
-    "    bdofs_V_2 = dolfinx.fem.locate_dofs_topological(\n",
-    "        (W.sub(0), W.sub(0).collapse()[0]), mesh.topology.dim - 1, boundaries_2)\n",
-    "    inlet_bc = dolfinx.fem.dirichletbc(u_in, bdofs_V_1, W.sub(0))\n",
-    "    wall_bc = dolfinx.fem.dirichletbc(u_wall, bdofs_V_2, W.sub(0))\n",
+    "    bdofs_V_1 = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, boundaries_1)\n",
+    "    bdofs_V_2 = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, boundaries_2)\n",
+    "    inlet_bc = dolfinx.fem.dirichletbc(u_in, bdofs_V_1, W0)\n",
+    "    wall_bc = dolfinx.fem.dirichletbc(u_wall, bdofs_V_2, W0)\n",
     "    bc = [inlet_bc, wall_bc]\n",
     "\n",
     "    # Class for interfacing with SNES\n",

diff --git a/tutorials/04_infsup_stokes/tutorial_infsup_stokes.ipynb b/tutorials/04_infsup_stokes/tutorial_infsup_stokes.ipynb
@@ -166,8 +166,9 @@
     "\n",
     "    # Define restriction for DOFs associated to homogenous Dirichlet boundary conditions\n",
     "    dofs_W = np.arange(0, W.dofmap.index_map.size_local + W.dofmap.index_map.num_ghosts)\n",
-    "    bdofs_V = dolfinx.fem.locate_dofs_topological(\n",
-    "        (W.sub(0), W.sub(0).collapse()[0]), mesh.topology.dim - 1, boundary_facets)[0]\n",
+    "    W0 = W.sub(0)\n",
+    "    V, _ = W0.collapse()\n",
+    "    bdofs_V = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, boundary_facets)[0]\n",
     "    restriction = multiphenicsx.fem.DofMapRestriction(W.dofmap, np.setdiff1d(dofs_W, bdofs_V))\n",
     "\n",
     "    # Assemble lhs and rhs matrices\n",
@@ -207,7 +208,8 @@
     "    with r_fun.vector.localForm() as r_fun_local, \\\n",
     "            multiphenicsx.fem.petsc.VecSubVectorWrapper(vr, W.dofmap, restriction) as vr_wrapper:\n",
     "        r_fun_local[:] = vr_wrapper\n",
-    "    (u_fun_1, u_fun_2) = (r_fun.sub(0).sub(0).collapse(), r_fun.sub(0).sub(1).collapse())\n",
+    "    u_fun = r_fun.sub(0).collapse()\n",
+    "    (u_fun_1, u_fun_2) = (u_fun.sub(0).collapse(), u_fun.sub(1).collapse())\n",
     "    p_fun = r_fun.sub(1).collapse()\n",
     "    normalize(u_fun_1, u_fun_2, p_fun)\n",
     "\n",

diff --git a/tutorials/06_optimal_control/tutorial_7a_stokes_dirichlet_control.ipynb b/tutorials/06_optimal_control/tutorial_7a_stokes_dirichlet_control.ipynb
@@ -275,6 +275,18 @@
     "Q_pressure = Y_pressure.clone()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Y_velocity_0 = Y_velocity.sub(0)\n",
+    "Y_velocity_1 = Y_velocity.sub(1)\n",
+    "Q_velocity_0 = Q_velocity.sub(0)\n",
+    "Q_velocity_1 = Q_velocity.sub(1)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -351,7 +363,7 @@
     "alpha = 1.e-2\n",
     "ff = dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2)))\n",
     "bc0 = dolfinx.fem.Function(Y_velocity)\n",
-    "bc0_component = dolfinx.fem.Function(Y_velocity.sub(0).collapse()[0])\n",
+    "bc0_component = dolfinx.fem.Function(Y_velocity_0.collapse()[0])\n",
     "bc1 = dolfinx.fem.Function(Y_velocity)\n",
     "bc1.interpolate(non_zero_eval)"
    ]
@@ -415,15 +427,15 @@
     "    dolfinx.fem.dirichletbc(\n",
     "        bc1, bdofs(Y_velocity, bc1.function_space, 1), Y_velocity),\n",
     "    dolfinx.fem.dirichletbc(\n",
-    "        bc0_component, bdofs(Y_velocity.sub(1), bc0_component.function_space, 2), Y_velocity.sub(1)),\n",
+    "        bc0_component, bdofs(Y_velocity_1, bc0_component.function_space, 2), Y_velocity_1),\n",
     "    dolfinx.fem.dirichletbc(\n",
-    "        bc0_component, bdofs(Y_velocity.sub(0), bc0_component.function_space, 4), Y_velocity.sub(0)),\n",
+    "        bc0_component, bdofs(Y_velocity_0, bc0_component.function_space, 4), Y_velocity_0),\n",
     "    dolfinx.fem.dirichletbc(\n",
     "        bc0, bdofs(Y_velocity, bc0.function_space, 5), Y_velocity),\n",
     "    dolfinx.fem.dirichletbc(\n",
     "        bc0, bdofs(Q_velocity, bc0.function_space, 1), Q_velocity),\n",
     "    dolfinx.fem.dirichletbc(\n",
-    "        bc0_component, bdofs(Q_velocity.sub(1), bc0_component.function_space, 2), Q_velocity.sub(1)),\n",
+    "        bc0_component, bdofs(Q_velocity_1, bc0_component.function_space, 2), Q_velocity_1),\n",
     "    dolfinx.fem.dirichletbc(\n",
     "        bc0, bdofs(Q_velocity, bc0.function_space, 4), Q_velocity),\n",
     "    dolfinx.fem.dirichletbc(\n",
@@ -488,7 +500,7 @@
     "    dolfinx.fem.dirichletbc(\n",
     "        bc1, bdofs(Y_velocity, bc1.function_space, 1), Y_velocity),\n",
     "    dolfinx.fem.dirichletbc(\n",
-    "        bc0_component, bdofs(Y_velocity.sub(1), bc0_component.function_space, 2), Y_velocity.sub(1)),\n",
+    "        bc0_component, bdofs(Y_velocity_1, bc0_component.function_space, 2), Y_velocity_1),\n",
     "    dolfinx.fem.dirichletbc(\n",
     "        bc0, bdofs(Y_velocity, bc0.function_space, 4), Y_velocity),\n",
     "    dolfinx.fem.dirichletbc(\n",