docstrings

environment.yml

@@ -8,5 +8,6 @@ dependencies:
   - numpy
   - pip
   - pytest
+  - pytest-mpi
   - pip:
     - -e .

examples/linear_elasticity/test_elasticity.py

@@ -43,18 +43,24 @@ def right_boundary(x):
 
     solver = NewtonSolver(MPI.COMM_WORLD, problem)
     n, converged = solver.solve(u)
+
+    # Compare the result with the analytical solution
     assert abs(problem.stress_1.x.array[0] - youngs_modulus * 0.01) < 1e-10 / (
         youngs_modulus * 0.01
     )
 
     problem.update()
+    # Check that the stress is updated correctly
     assert abs(problem.stress_0.x.array[0] - youngs_modulus * 0.01) < 1e-10 / (
         youngs_modulus * 0.01
     )
+    # Check that the displacement is updated correctly
     assert np.max(problem._u0.x.array) == displacement.value
 
     displacement.value = 0.02
     n, converged = solver.solve(u)
+
+    # Compare the result of the updated problem with new BC with the analytical solution
     assert abs(problem.stress_1.x.array[0] - youngs_modulus * 0.02) < 1e-10 / (
         youngs_modulus * 0.02
     )

src/fenics_constitutive/interfaces.py

@@ -83,7 +83,8 @@ def evaluate(
         tangent: np.ndarray,
         history: np.ndarray | dict[str, np.ndarray] | None,
     ) -> None:
-        """Evaluate the constitutive model and overwrite the stress, tangent and history.
+        """
+        Evaluate the constitutive model and overwrite the stress, tangent and history.
 
         Args:
             del_t : The time increment.

src/fenics_constitutive/maps.py

@@ -10,12 +10,30 @@
 
 @dataclass
 class SubSpaceMap:
+    """
+    Map between a subspace and a parent space.
+
+    Args:
+        parent: The dof indices of the parent space.
+        child: The dof indices of the child space.
+        sub_mesh: The mesh of the subspace.
+        parent_mesh: The mesh of the parent space.
+
+    """
+
     parent: np.ndarray
     child: np.ndarray
     sub_mesh: df.mesh.Mesh
     parent_mesh: df.mesh.Mesh
 
     def map_to_parent(self, sub: df.fem.Function, parent: df.fem.Function) -> None:
+        """
+        Map the values from the subspace to the parent space.
+
+        Args:
+            sub: The function in the subspace.
+            parent: The function in the parent space.
+        """
         assert sub.function_space.mesh == self.sub_mesh, "Subspace mesh does not match"
         assert (
             parent.function_space.mesh == self.parent_mesh
@@ -30,6 +48,13 @@ def map_to_parent(self, sub: df.fem.Function, parent: df.fem.Function) -> None:
         parent.x.scatter_forward()
 
     def map_to_child(self, parent: df.fem.Function, sub: df.fem.Function) -> None:
+        """
+        Map the values from the parent space to the subspace.
+
+        Args:
+            parent: The function in the parent space.
+            sub: The function in the subspace.
+        """
         assert sub.function_space.mesh == self.sub_mesh, "Subspace mesh does not match"
         assert (
             parent.function_space.mesh == self.parent_mesh
@@ -50,6 +75,18 @@ def build_subspace_map(
     tuple[SubSpaceMap, df.mesh.Mesh]
     | tuple[SubSpaceMap, df.mesh.Mesh, df.fem.FunctionSpace]
 ):
+    """
+    Build a map between a subspace and a parent space. This currently needs
+    to build a functionspace which can optionally be returned.
+
+    Args:
+        cells: The cells of the subspace.
+        V: The parent function space.
+        return_subspace: Return the subspace function space.
+
+    Returns:
+        The subspace map, the submesh and optionally the subspace function space.
+    """
     mesh = V.mesh
     submesh, cell_map, _, _ = df.mesh.create_submesh(mesh, mesh.topology.dim, cells)
     fe = V.ufl_element()

src/fenics_constitutive/solver.py

@@ -63,6 +63,28 @@ def build_history(
 
 
 class IncrSmallStrainProblem(df.fem.petsc.NonlinearProblem):
+    """
+    A nonlinear problem for incremental small strain models. To be used with
+    the dolfinx NewtonSolver.
+
+    Args:
+        laws: A list of tuples where the first element is the constitutive law and the second
+            element is the cells for the submesh. If only one law is provided, it is assumed
+            that the domain is homogenous.
+        u: The displacement field. This is the unknown in the nonlinear problem.
+        bcs: The Dirichlet boundary conditions.
+        q_degree: The quadrature degree (Polynomial degree which the quadrature rule needs to integrate exactly).
+        form_compiler_options: The options for the form compiler.
+        jit_options: The options for the JIT compiler.
+
+    Note:
+        If `super().__init__(R, u, bcs, dR)` is called within the __init__ method,
+        the user cannot add Neumann BCs. Therefore, the compilation (i.e. call to
+        `super().__init__()`) is done when `df.nls.petsc.NewtonSolver` is initialized.
+        The solver will call `self._A = fem.petsc.create_matrix(problem.a)` and hence
+        we override the property ``a`` of NonlinearProblem to ensure that the form is compiled.
+    """
+
     def __init__(
         self,
         laws: list[tuple[IncrSmallStrainModel, np.ndarray]] | IncrSmallStrainModel,
@@ -204,15 +226,6 @@ def __init__(
             ufl.nabla_grad(self._u - self._u0), self.q_points
         )
 
-        # ### Note on JIT compilation of UFL forms
-        # if super().__init__(R, u, bcs, dR) is called within ElasticityProblem.__init__
-        # the user cannot add Neumann BCs.
-        # Therefore, the compilation (i.e. call to super().__init__()) is done when
-        # df.nls.petsc.NewtonSolver is initialized.
-        # df.nls.petsc.NewtonSolver will call
-        # self._A = fem.petsc.create_matrix(problem.a) and hence it would suffice
-        # to override the property ``a`` of NonlinearProblem.
-
     @property
     def a(self) -> df.fem.FormMetaClass:
         """Compiled bilinear form (the Jacobian form)"""
@@ -232,13 +245,14 @@ def a(self) -> df.fem.FormMetaClass:
 
         return self._a
 
-    def form(self, x: PETSc.Vec):
+    def form(self, x: PETSc.Vec) -> None:
         """This function is called before the residual or Jacobian is
-        computed. This is usually used to update ghost values.
-        Parameters
-        ----------
-        x
-            The vector containing the latest solution
+        computed. This is usually used to update ghost values, but here
+        we use it to update the stress, tangent and history.
+
+        Args:
+            x: The vector containing the latest solution
+
         """
         super().form(x)
 
@@ -286,17 +300,16 @@ def form(self, x: PETSc.Vec):
                     self._del_grad_u[0].x.array,
                     self.stress_1.x.array,
                     self.tangent.x.array,
-                    self._history_1[0].x.array
-                    if law.history_dim is not None
-                    else None,  # history,
+                    self._history_1[0].x.array if law.history_dim is not None else None,
                 )
 
         self.stress_1.x.scatter_forward()
         self.tangent.x.scatter_forward()
 
-    def update(self):
-        # update the current displacement, stress and history
-        # works for both homogeneous and inhomogeneous domains
+    def update(self) -> None:
+        """
+        Update the current displacement, stress and history.
+        """
         self._u0.x.array[:] = self._u.x.array
         self._u0.x.scatter_forward()