Merge pull request #41 from jorgensd/dokken/main-typing

Typing for main branch

Author: jorgensd

demos/poisson_mother.py

@@ -140,7 +140,7 @@ def refinement_region(x, tol=1e-14):
 
 u = ufl.TrialFunction(V)
 v = ufl.TestFunction(V)
-kappa = dolfinx.fem.Constant(refined_mesh, dolfinx.default_scalar_type(1.0))  # Thermal diffusivity
+kappa = dolfinx.fem.Constant(refined_mesh, np.dtype(dolfinx.default_scalar_type).type(1.0))  # Thermal diffusivity
 F = ufl.inner(kappa * ufl.grad(u), ufl.grad(v)) * ufl.dx - f * v * ufl.dx
 a, L = ufl.system(F)
 
@@ -149,7 +149,7 @@ def refinement_region(x, tol=1e-14):
 refined_mesh.topology.create_connectivity(tdim - 1, tdim)
 exterior_facets = dolfinx.mesh.exterior_facet_indices(refined_mesh.topology)
 exterior_dofs = dolfinx.fem.locate_dofs_topological(V, tdim - 1, exterior_facets)
-zero = dolfinx.fem.Constant(refined_mesh, dolfinx.default_scalar_type(0.0))
+zero = dolfinx.fem.Constant(refined_mesh, np.dtype(dolfinx.default_scalar_type).type(0.0))
 bc = dolfinx.fem.dirichletbc(zero, exterior_dofs, V)
 
 # Next, we define a `dolfinx_adjoint.LinearProblem` instance, which overloads
@@ -194,7 +194,9 @@ def refinement_region(x, tol=1e-14):
 
 # The functional is written out in `ufl` and assembled with `dolfinx_adjoint.assemble_scalar`
 
-alpha = dolfinx.fem.Constant(refined_mesh, dolfinx.default_scalar_type(1.0e-6))  # Tikhonov regularization parameter
+alpha = dolfinx.fem.Constant(
+    refined_mesh, np.dtype(dolfinx.default_scalar_type).type(1.0e-6)
+)  # Tikhonov regularization parameter
 alpha.name = "alpha"  # type: ignore
 J_symbolic = 0.5 * ufl.inner(uh - d, uh - d) * ufl.dx + 0.5 * alpha * ufl.inner(f, f) * ufl.dx
 J = dolfinx_adjoint.assemble_scalar(J_symbolic)
@@ -354,14 +356,14 @@ def solve_optimal_problem(N: int, use_newton: bool = False) -> dict[str, float |
 
     u = ufl.TrialFunction(V)
     v = ufl.TestFunction(V)
-    kappa = dolfinx.fem.Constant(refined_mesh, dolfinx.default_scalar_type(1.0))  # Thermal diffusivity
+    kappa = dolfinx.fem.Constant(refined_mesh, np.dtype(dolfinx.default_scalar_type).type(1.0))  # Thermal diffusivity
     F = ufl.inner(kappa * ufl.grad(u), ufl.grad(v)) * ufl.dx - f * v * ufl.dx
     a, L = ufl.system(F)
 
     refined_mesh.topology.create_connectivity(tdim - 1, tdim)
     exterior_facets = dolfinx.mesh.exterior_facet_indices(refined_mesh.topology)
     exterior_dofs = dolfinx.fem.locate_dofs_topological(V, tdim - 1, exterior_facets)
-    zero = dolfinx.fem.Constant(refined_mesh, dolfinx.default_scalar_type(0.0))
+    zero = dolfinx.fem.Constant(refined_mesh, np.dtype(dolfinx.default_scalar_type).type(0.0))
     bc = dolfinx.fem.dirichletbc(zero, exterior_dofs, V)
 
     problem = dolfinx_adjoint.LinearProblem(
@@ -372,7 +374,9 @@ def solve_optimal_problem(N: int, use_newton: bool = False) -> dict[str, float |
     x, y = ufl.SpatialCoordinate(refined_mesh)
     d = 1 / (2 * ufl.pi**2) * ufl.sin(ufl.pi * x) * ufl.sin(ufl.pi * y)
 
-    alpha = dolfinx.fem.Constant(refined_mesh, dolfinx.default_scalar_type(1.0e-6))  # Tikhonov regularization parameter
+    alpha = dolfinx.fem.Constant(
+        refined_mesh, np.dtype(dolfinx.default_scalar_type).type(1.0e-6)
+    )  # Tikhonov regularization parameter
     J_symbolic = 0.5 * ufl.inner(uh - d, uh - d) * ufl.dx + 0.5 * alpha * ufl.inner(f, f) * ufl.dx
     J = dolfinx_adjoint.assemble_scalar(J_symbolic)
 

src/dolfinx_adjoint/blocks/solvers.py

@@ -363,11 +363,11 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         if bcs is not None:
             # This system should never be "blocked"
             dolfinx.fem.petsc.apply_lifting(b_tlm.petsc_vec, [dFdu], bcs=[bcs], alpha=0)
-            dolfinx.la.petsc._ghost_update(b_tlm.petsc_vec, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore [attr-defined]
+            dolfinx.la.petsc._ghost_update(b_tlm.petsc_vec, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[arg-type]
             for bc in bcs:
                 bc.set(b_tlm.array, alpha=0)
         else:
-            dolfinx.la.petsc._ghost_update(b_tlm, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore [attr-defined]
+            dolfinx.la.petsc._ghost_update(b_tlm, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[arg-type]
         solve_linear_problem(A_tlm, dudm.x, b_tlm, petsc_options=self._tlm_petsc_options)
         return dudm
 
@@ -929,11 +929,11 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         if bcs is not None:
             # This system should never be "blocked"
             dolfinx.fem.petsc.apply_lifting(b_tlm.petsc_vec, [dFdu], bcs=[bcs], alpha=0)
-            dolfinx.la.petsc._ghost_update(b_tlm.petsc_vec, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore [attr-defined]
+            dolfinx.la.petsc._ghost_update(b_tlm.petsc_vec, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore [arg-type]
             for bc in bcs:
                 bc.set(b_tlm.array, alpha=0)
         else:
-            dolfinx.la.petsc._ghost_update(b_tlm, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore [attr-defined]
+            dolfinx.la.petsc._ghost_update(b_tlm, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore [arg-type]
         solve_linear_problem(A_tlm, dudm.x, b_tlm, petsc_options=self._tlm_petsc_options)
         return dudm
 

src/dolfinx_adjoint/petsc_utils.py

@@ -37,7 +37,7 @@ def solve_linear_problem(
     opts = PETSc.Options()  # type: ignore [attr-defined]
     opts.prefixPush(problem_prefix)
     for k, v in petsc_options.items():
-        opts[k] = v
+        opts.setValue(k, v)
     opts.prefixPop()
     ksp.setFromOptions()
 
@@ -49,7 +49,7 @@ def solve_linear_problem(
 
     # Free option space post setting
     for k in petsc_options.keys():
-        del opts[k]
+        opts.delValue(k)
     ksp.solve(b.petsc_vec, x.petsc_vec)
     ksp.destroy()
     x.scatter_forward()
@@ -85,5 +85,5 @@ def solve(
         # Solve linear system and update ghost values in the solution
         self._solver.solve(self._b, self._x)
         dolfinx.la.petsc._ghost_update(self._x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore
-        dolfinx.fem.petsc.assign(self._x, self._u)
+        dolfinx.fem.petsc.assign(self._x, self._u)  # type: ignore
         return self._u

tests/test_linear_solver.py

@@ -48,7 +48,7 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int,
     mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
     boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
     boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets)
-    bc_val = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0))
+    bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0))
     bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V)
 
     options = {

tests/test_poisson_mother.py

@@ -73,7 +73,9 @@ def reference_solution(
     mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
     exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
     exterior_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, exterior_facets)
-    bc = dolfinx.fem.dirichletbc(dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0)), exterior_dofs, V)
+    bc = dolfinx.fem.dirichletbc(
+        dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(0.0)), exterior_dofs, V
+    )
     petsc_options = {
         "ksp_type": "preonly",
         "pc_type": "lu",
@@ -182,7 +184,9 @@ def step_func(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
     x, y = ufl.SpatialCoordinate(mesh)
     d = 1 / (2 * ufl.pi**2) * ufl.sin(ufl.pi * x) * ufl.sin(ufl.pi * y)
 
-    alpha = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0e-6))  # Tikhonov regularization parameter
+    alpha = dolfinx.fem.Constant(
+        mesh, np.dtype(dolfinx.default_scalar_type).type(1.0e-6)
+    )  # Tikhonov regularization parameter
     alpha.name = "alpha"  # type: ignore
 
     # Get reference values
@@ -206,7 +210,7 @@ def step_func(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
     mesh.topology.create_connectivity(tdim - 1, tdim)
     exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
     exterior_dofs = dolfinx.fem.locate_dofs_topological(V, tdim - 1, exterior_facets)
-    zero = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0))
+    zero = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(0.0))
     bc = dolfinx.fem.dirichletbc(zero, exterior_dofs, V)
 
     petsc_options = {