From b1950723f9a13411e7d91017a70225b379d085f5 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 14 Aug 2024 16:39:08 +0200
Subject: [PATCH 001/134] started a generic script

---
 discontinuity_generic.py        | 316 ++++++++++++++++++++++++++++++++
 festim/helpers_discontinuity.py | 179 ++++++++++++++++++
 2 files changed, 495 insertions(+)
 create mode 100644 discontinuity_generic.py
 create mode 100644 festim/helpers_discontinuity.py

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
new file mode 100644
index 000000000..78e96d6b6
--- /dev/null
+++ b/discontinuity_generic.py
@@ -0,0 +1,316 @@
+from mpi4py import MPI
+import dolfinx
+import dolfinx.fem.petsc
+import ufl
+import numpy as np
+from petsc4py import PETSc
+import basix
+from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
+import festim as F
+
+
+# ---------------- Generate a mesh ----------------
+def generate_mesh():
+    def bottom_boundary(x):
+        return np.isclose(x[1], 0.0)
+
+    def top_boundary(x):
+        return np.isclose(x[1], 1.0)
+
+    def half(x):
+        return x[1] <= 0.5 + 1e-14
+
+    mesh = dolfinx.mesh.create_unit_square(
+        MPI.COMM_WORLD, 20, 20, dolfinx.mesh.CellType.triangle
+    )
+
+    # Split domain in half and set an interface tag of 5
+    gdim = mesh.geometry.dim
+    tdim = mesh.topology.dim
+    fdim = tdim - 1
+    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    num_facets_local = (
+        mesh.topology.index_map(fdim).size_local
+        + mesh.topology.index_map(fdim).num_ghosts
+    )
+    facets = np.arange(num_facets_local, dtype=np.int32)
+    values = np.full_like(facets, 0, dtype=np.int32)
+    values[top_facets] = 1
+    values[bottom_facets] = 2
+
+    bottom_cells = dolfinx.mesh.locate_entities(mesh, tdim, half)
+    num_cells_local = (
+        mesh.topology.index_map(tdim).size_local
+        + mesh.topology.index_map(tdim).num_ghosts
+    )
+    cells = np.full(num_cells_local, 4, dtype=np.int32)
+    cells[bottom_cells] = 3
+    ct = dolfinx.mesh.meshtags(
+        mesh, tdim, np.arange(num_cells_local, dtype=np.int32), cells
+    )
+    all_b_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(3), tdim, fdim
+    )
+    all_t_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(4), tdim, fdim
+    )
+    interface = np.intersect1d(all_b_facets, all_t_facets)
+    values[interface] = 5
+
+    mt = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, values)
+    return mesh, mt, ct
+
+
+mesh, mt, ct = generate_mesh()
+
+top_domain = F.VolumeSubdomain(4, material=None)
+bottom_domain = F.VolumeSubdomain(3, material=None)
+list_of_subdomains = [bottom_domain, top_domain]
+list_of_interfaces = {5: [bottom_domain, top_domain]}
+
+gdim = mesh.geometry.dim
+tdim = mesh.topology.dim
+fdim = tdim - 1
+
+num_facets_local = (
+    mesh.topology.index_map(fdim).size_local + mesh.topology.index_map(fdim).num_ghosts
+)
+
+for subdomain in list_of_subdomains:
+    subdomain.submesh, subdomain.submesh_to_mesh, subdomain.v_map = (
+        dolfinx.mesh.create_submesh(mesh, tdim, ct.find(subdomain.id))[0:3]
+    )
+
+    subdomain.parent_to_submesh = np.full(num_facets_local, -1, dtype=np.int32)
+    subdomain.parent_to_submesh[subdomain.submesh_to_mesh] = np.arange(
+        len(subdomain.submesh_to_mesh), dtype=np.int32
+    )
+
+    # We need to modify the cell maps, as for `dS` integrals of interfaces between submeshes, there is no entity to map to.
+    # We use the entity on the same side to fix this (as all restrictions are one-sided)
+
+    # Transfer meshtags to submesh
+    subdomain.ft, subdomain.facet_to_parent = transfer_meshtags_to_submesh(
+        mesh, mt, subdomain.submesh, subdomain.v_map, subdomain.submesh_to_mesh
+    )
+
+# this seems to be not needed
+# t_parent_to_facet = np.full(num_facets_local, -1)
+# t_parent_to_facet[t_facet_to_parent] = np.arange(
+#     len(t_facet_to_parent), dtype=np.int32
+# )
+
+# Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
+# TODO ask Jorgen what this is for
+mesh.topology.create_connectivity(fdim, tdim)
+f_to_c = mesh.topology.connectivity(fdim, tdim)
+for interface in list_of_interfaces:
+    for facet in mt.find(interface):
+        cells = f_to_c.links(facet)
+        assert len(cells) == 2
+        for domain in list_of_interfaces[interface]:
+            map = domain.parent_to_submesh[cells]
+            domain.parent_to_submesh[cells] = max(map)
+
+entity_maps = {
+    subdomain.submesh: subdomain.parent_to_submesh for subdomain in list_of_subdomains
+}
+exit()
+
+
+def D(T):
+    k_B = 8.6173303e-5
+    return 2 * ufl.exp(-0.1 / k_B / T)
+
+
+def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
+    # Compute map from parent entity to submesh cell
+    codim = mesh.topology.dim - submesh.topology.dim
+    ptdim = mesh.topology.dim - codim
+    num_entities = (
+        mesh.topology.index_map(ptdim).size_local
+        + mesh.topology.index_map(ptdim).num_ghosts
+    )
+    mesh_to_submesh = np.full(num_entities, -1)
+    mesh_to_submesh[submesh_to_mesh] = np.arange(len(submesh_to_mesh), dtype=np.int32)
+
+    degree = 1
+    element_CG = basix.ufl.element(
+        basix.ElementFamily.P,
+        submesh.basix_cell(),
+        degree,
+        basix.LagrangeVariant.equispaced,
+    )
+    element = basix.ufl.mixed_element([element_CG, element_CG])
+    V = dolfinx.fem.functionspace(submesh, element)
+    u = dolfinx.fem.Function(V)
+    us = list(ufl.split(u))
+    vs = list(ufl.TestFunctions(V))
+    ct_r = dolfinx.mesh.meshtags(
+        mesh,
+        mesh.topology.dim,
+        submesh_to_mesh,
+        np.full_like(submesh_to_mesh, 1, dtype=np.int32),
+    )
+    val = dolfinx.fem.Constant(submesh, value)
+    dx_r = ufl.Measure("dx", domain=mesh, subdomain_data=ct_r, subdomain_id=1)
+    F = ufl.inner(ufl.grad(us[0]), ufl.grad(vs[0])) * dx_r - val * vs[0] * dx_r
+    k = 2
+    p = 0.1
+    n = 0.5
+    F += k * us[0] * (n - us[1]) * vs[1] * dx_r - p * us[1] * vs[1] * dx_r
+    return u, vs, F, mesh_to_submesh
+
+
+# for each subdomain, define the interior equation
+u_0, v_0s, F_00, m_to_b = define_interior_eq(mesh, 2, submesh_b, submesh_b_to_mesh, 0.0)
+u_1, v_1s, F_11, m_to_t = define_interior_eq(mesh, 1, submesh_t, submesh_t_to_mesh, 0.0)
+u_0.name = "u_b"
+u_1.name = "u_t"
+
+
+# Add coupling term to the interface
+# Get interface markers on submesh b
+dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=mt, subdomain_id=5)
+b_res = "+"
+t_res = "-"
+
+v_b = v_0s[0](b_res)
+v_t = v_1s[0](t_res)
+
+u_bs = list(ufl.split(u_0))
+u_ts = list(ufl.split(u_1))
+u_b = u_bs[0](b_res)
+u_t = u_ts[0](t_res)
+
+
+def mixed_term(u, v, n):
+    return ufl.dot(ufl.grad(u), n) * v
+
+
+n = ufl.FacetNormal(mesh)
+n_b = n(b_res)
+n_t = n(t_res)
+cr = ufl.Circumradius(mesh)
+h_b = 2 * cr(b_res)
+h_t = 2 * cr(t_res)
+gamma = 10.0
+
+W_0 = dolfinx.fem.functionspace(submesh_b, ("DG", 0))
+K_0 = dolfinx.fem.Function(W_0)
+K_0.x.array[:] = 2
+W_1 = dolfinx.fem.functionspace(submesh_t, ("DG", 0))
+K_1 = dolfinx.fem.Function(W_1)
+K_1.x.array[:] = 4
+
+K_b = K_0(b_res)
+K_t = K_1(t_res)
+
+
+F_0 = (
+    -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
+    - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
+)
+
+F_1 = (
+    +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
+    - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
+)
+F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
+F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
+
+F_0 += F_00
+F_1 += F_11
+
+jac00 = ufl.derivative(F_0, u_0)
+
+jac01 = ufl.derivative(F_0, u_1)
+
+jac10 = ufl.derivative(F_1, u_0)
+jac11 = ufl.derivative(F_1, u_1)
+
+J00 = dolfinx.fem.form(jac00, entity_maps=entity_maps)
+J01 = dolfinx.fem.form(jac01, entity_maps=entity_maps)
+J10 = dolfinx.fem.form(jac10, entity_maps=entity_maps)
+J11 = dolfinx.fem.form(jac11, entity_maps=entity_maps)
+J = [[J00, J01], [J10, J11]]
+F = [
+    dolfinx.fem.form(F_0, entity_maps=entity_maps),
+    dolfinx.fem.form(F_1, entity_maps=entity_maps),
+]
+
+
+# boundary conditions
+b_bc = dolfinx.fem.Function(u_0.function_space)
+b_bc.x.array[:] = 0.2
+submesh_b.topology.create_connectivity(
+    submesh_b.topology.dim - 1, submesh_b.topology.dim
+)
+bc_b = dolfinx.fem.dirichletbc(
+    b_bc,
+    dolfinx.fem.locate_dofs_topological(u_0.function_space.sub(0), fdim, ft_b.find(2)),
+)
+
+
+t_bc = dolfinx.fem.Function(u_1.function_space)
+t_bc.x.array[:] = 0.05
+submesh_t.topology.create_connectivity(
+    submesh_t.topology.dim - 1, submesh_t.topology.dim
+)
+bc_t = dolfinx.fem.dirichletbc(
+    t_bc,
+    dolfinx.fem.locate_dofs_topological(u_1.function_space.sub(0), fdim, ft_t.find(1)),
+)
+bcs = [bc_b, bc_t]
+
+
+solver = NewtonSolver(
+    F,
+    J,
+    [u_0, u_1],
+    bcs=bcs,
+    max_iterations=10,
+    petsc_options={
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "pc_factor_mat_solver_type": "mumps",
+    },
+)
+solver.solve(1e-5)
+
+# bp = dolfinx.io.VTXWriter(mesh.comm, "u_b.bp", [u_0.sub(0).collapse()], engine="BP4")
+bp = dolfinx.io.VTXWriter(mesh.comm, "u_b_0.bp", [u_0.sub(0).collapse()], engine="BP4")
+bp.write(0)
+bp.close()
+bp = dolfinx.io.VTXWriter(mesh.comm, "u_t_0.bp", [u_1.sub(0).collapse()], engine="BP4")
+bp.write(0)
+bp.close()
+bp = dolfinx.io.VTXWriter(mesh.comm, "u_b_1.bp", [u_0.sub(1).collapse()], engine="BP4")
+bp.write(0)
+bp.close()
+bp = dolfinx.io.VTXWriter(mesh.comm, "u_t_1.bp", [u_1.sub(1).collapse()], engine="BP4")
+bp.write(0)
+bp.close()
+
+
+# derived quantities
+V = dolfinx.fem.functionspace(mesh, ("CG", 1))
+T = dolfinx.fem.Function(V)
+T.interpolate(lambda x: 200 + x[1])
+
+
+T_b = dolfinx.fem.Function(u_0.sub(0).collapse().function_space)
+T_b.interpolate(T)
+
+ds_b = ufl.Measure("ds", domain=submesh_b)
+dx_b = ufl.Measure("dx", domain=submesh_b)
+dx = ufl.Measure("dx", domain=mesh)
+
+n_b = ufl.FacetNormal(submesh_b)
+
+form = dolfinx.fem.form(u_0.sub(0) * dx_b)
+print(dolfinx.fem.assemble_scalar(form))
+
+form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(u_0.sub(0)), n_b) * ds_b)
+print(dolfinx.fem.assemble_scalar(form))
diff --git a/festim/helpers_discontinuity.py b/festim/helpers_discontinuity.py
new file mode 100644
index 000000000..2ae19ccc0
--- /dev/null
+++ b/festim/helpers_discontinuity.py
@@ -0,0 +1,179 @@
+from mpi4py import MPI
+import dolfinx
+import dolfinx.fem.petsc
+import ufl
+import numpy as np
+from petsc4py import PETSc
+import basix
+
+
+class NewtonSolver:
+    max_iterations: int
+    bcs: list[dolfinx.fem.DirichletBC]
+    A: PETSc.Mat
+    b: PETSc.Vec
+    J: dolfinx.fem.Form
+    b: dolfinx.fem.Form
+    dx: PETSc.Vec
+
+    def __init__(
+        self,
+        F: list[dolfinx.fem.form],
+        J: list[list[dolfinx.fem.form]],
+        w: list[dolfinx.fem.Function],
+        bcs: list[dolfinx.fem.DirichletBC] | None = None,
+        max_iterations: int = 5,
+        petsc_options: dict[str, str | float | int | None] = None,
+        problem_prefix="newton",
+    ):
+        self.max_iterations = max_iterations
+        self.bcs = [] if bcs is None else bcs
+        self.b = dolfinx.fem.petsc.create_vector_block(F)
+        self.F = F
+        self.J = J
+        self.A = dolfinx.fem.petsc.create_matrix_block(J)
+        self.dx = self.A.createVecLeft()
+        self.w = w
+        self.x = dolfinx.fem.petsc.create_vector_block(F)
+
+        # Set PETSc options
+        opts = PETSc.Options()
+        if petsc_options is not None:
+            for k, v in petsc_options.items():
+                opts[k] = v
+
+        # Define KSP solver
+        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
+        self._solver.setOperators(self.A)
+        self._solver.setFromOptions()
+
+        # Set matrix and vector PETSc options
+        self.A.setFromOptions()
+        self.b.setFromOptions()
+
+    def solve(self, tol=1e-6, beta=1.0):
+        i = 0
+
+        while i < self.max_iterations:
+            dolfinx.cpp.la.petsc.scatter_local_vectors(
+                self.x,
+                [si.x.petsc_vec.array_r for si in self.w],
+                [
+                    (
+                        si.function_space.dofmap.index_map,
+                        si.function_space.dofmap.index_map_bs,
+                    )
+                    for si in self.w
+                ],
+            )
+            self.x.ghostUpdate(
+                addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
+            )
+
+            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
+            with self.b.localForm() as b_local:
+                b_local.set(0.0)
+            dolfinx.fem.petsc.assemble_vector_block(
+                self.b, self.F, self.J, bcs=self.bcs, x0=self.x, scale=-1.0
+            )
+            self.b.ghostUpdate(
+                PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
+            )
+
+            # Assemble Jacobian
+            self.A.zeroEntries()
+            dolfinx.fem.petsc.assemble_matrix_block(self.A, self.J, bcs=self.bcs)
+            self.A.assemble()
+
+            self._solver.solve(self.b, self.dx)
+            # self._solver.view()
+            assert (
+                self._solver.getConvergedReason() > 0
+            ), "Linear solver did not converge"
+            offset_start = 0
+            for s in self.w:
+                num_sub_dofs = (
+                    s.function_space.dofmap.index_map.size_local
+                    * s.function_space.dofmap.index_map_bs
+                )
+                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
+                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
+                )
+                s.x.petsc_vec.ghostUpdate(
+                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
+                )
+                offset_start += num_sub_dofs
+            # Compute norm of update
+
+            correction_norm = self.dx.norm(0)
+            print(f"Iteration {i}: Correction norm {correction_norm}")
+            if correction_norm < tol:
+                break
+            i += 1
+
+    def __del__(self):
+        self.A.destroy()
+        self.b.destroy()
+        self.dx.destroy()
+        self._solver.destroy()
+        self.x.destroy()
+
+
+def transfer_meshtags_to_submesh(
+    mesh, entity_tag, submesh, sub_vertex_to_parent, sub_cell_to_parent
+):
+    """
+    Transfer a meshtag from a parent mesh to a sub-mesh.
+    """
+
+    tdim = mesh.topology.dim
+    cell_imap = mesh.topology.index_map(tdim)
+    num_cells = cell_imap.size_local + cell_imap.num_ghosts
+    mesh_to_submesh = np.full(num_cells, -1)
+    mesh_to_submesh[sub_cell_to_parent] = np.arange(
+        len(sub_cell_to_parent), dtype=np.int32
+    )
+    sub_vertex_to_parent = np.asarray(sub_vertex_to_parent)
+
+    submesh.topology.create_connectivity(entity_tag.dim, 0)
+
+    num_child_entities = (
+        submesh.topology.index_map(entity_tag.dim).size_local
+        + submesh.topology.index_map(entity_tag.dim).num_ghosts
+    )
+    submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)
+
+    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
+    c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)
+
+    child_markers = np.full(num_child_entities, 0, dtype=np.int32)
+
+    mesh.topology.create_connectivity(entity_tag.dim, 0)
+    mesh.topology.create_connectivity(entity_tag.dim, mesh.topology.dim)
+    p_f_to_v = mesh.topology.connectivity(entity_tag.dim, 0)
+    p_f_to_c = mesh.topology.connectivity(entity_tag.dim, mesh.topology.dim)
+    sub_to_parent_entity_map = np.full(num_child_entities, -1, dtype=np.int32)
+    for facet, value in zip(entity_tag.indices, entity_tag.values):
+        facet_found = False
+        for cell in p_f_to_c.links(facet):
+            if facet_found:
+                break
+            if (child_cell := mesh_to_submesh[cell]) != -1:
+                for child_facet in c_c_to_e.links(child_cell):
+                    child_vertices = c_e_to_v.links(child_facet)
+                    child_vertices_as_parent = sub_vertex_to_parent[child_vertices]
+                    is_facet = np.isin(
+                        child_vertices_as_parent, p_f_to_v.links(facet)
+                    ).all()
+                    if is_facet:
+                        child_markers[child_facet] = value
+                        facet_found = True
+                        sub_to_parent_entity_map[child_facet] = facet
+    tags = dolfinx.mesh.meshtags(
+        submesh,
+        entity_tag.dim,
+        np.arange(num_child_entities, dtype=np.int32),
+        child_markers,
+    )
+    tags.name = entity_tag.name
+    return tags, sub_to_parent_entity_map

From 4034cc6e70b6be5a6379277734da44a18910483a Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 08:17:04 +0000
Subject: [PATCH 002/134] progress (runs on dolfinx nightly)

---
 .devcontainer/devcontainer.json |   3 +
 discontinuity_generic.py        | 193 ++++++++++++++++++--------------
 2 files changed, 113 insertions(+), 83 deletions(-)
 create mode 100644 .devcontainer/devcontainer.json

diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
new file mode 100644
index 000000000..dea4d4dd9
--- /dev/null
+++ b/.devcontainer/devcontainer.json
@@ -0,0 +1,3 @@
+{
+    "image": "dolfinx/dolfinx:nightly"
+}
\ No newline at end of file
diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 78e96d6b6..7b712a099 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -113,10 +113,11 @@ def half(x):
             map = domain.parent_to_submesh[cells]
             domain.parent_to_submesh[cells] = max(map)
 
+# ._cpp_object needed on dolfinx 0.8.0
 entity_maps = {
-    subdomain.submesh: subdomain.parent_to_submesh for subdomain in list_of_subdomains
+    subdomain.submesh: subdomain.parent_to_submesh
+    for subdomain in list_of_subdomains
 }
-exit()
 
 
 def D(T):
@@ -164,111 +165,137 @@ def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
 
 
 # for each subdomain, define the interior equation
-u_0, v_0s, F_00, m_to_b = define_interior_eq(mesh, 2, submesh_b, submesh_b_to_mesh, 0.0)
-u_1, v_1s, F_11, m_to_t = define_interior_eq(mesh, 1, submesh_t, submesh_t_to_mesh, 0.0)
-u_0.name = "u_b"
-u_1.name = "u_t"
-
+for subdomain in list_of_subdomains:
+    degree = 1
+    subdomain.u, subdomain.vs, subdomain.F, subdomain.m_to_s = define_interior_eq(
+        mesh, degree, subdomain.submesh, subdomain.submesh_to_mesh, 0.0
+    )
+    subdomain.u.name = f"u_{subdomain.id}"
 
 # Add coupling term to the interface
 # Get interface markers on submesh b
-dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=mt, subdomain_id=5)
-b_res = "+"
-t_res = "-"
-
-v_b = v_0s[0](b_res)
-v_t = v_1s[0](t_res)
-
-u_bs = list(ufl.split(u_0))
-u_ts = list(ufl.split(u_1))
-u_b = u_bs[0](b_res)
-u_t = u_ts[0](t_res)
-
-
-def mixed_term(u, v, n):
-    return ufl.dot(ufl.grad(u), n) * v
-
-
-n = ufl.FacetNormal(mesh)
-n_b = n(b_res)
-n_t = n(t_res)
-cr = ufl.Circumradius(mesh)
-h_b = 2 * cr(b_res)
-h_t = 2 * cr(t_res)
-gamma = 10.0
-
-W_0 = dolfinx.fem.functionspace(submesh_b, ("DG", 0))
-K_0 = dolfinx.fem.Function(W_0)
-K_0.x.array[:] = 2
-W_1 = dolfinx.fem.functionspace(submesh_t, ("DG", 0))
-K_1 = dolfinx.fem.Function(W_1)
-K_1.x.array[:] = 4
-
-K_b = K_0(b_res)
-K_t = K_1(t_res)
-
-
-F_0 = (
-    -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
-    - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
-)
-
-F_1 = (
-    +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
-    - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
-)
-F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
-F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
-
-F_0 += F_00
-F_1 += F_11
-
-jac00 = ufl.derivative(F_0, u_0)
-
-jac01 = ufl.derivative(F_0, u_1)
+for interface in list_of_interfaces:
+    subdomain_1 = list_of_interfaces[interface][0]
+    subdomain_2 = list_of_interfaces[interface][1]
 
-jac10 = ufl.derivative(F_1, u_0)
-jac11 = ufl.derivative(F_1, u_1)
+    dInterface = ufl.Measure(
+        "dS", domain=mesh, subdomain_data=mt, subdomain_id=interface
+    )
+    b_res = "+"
+    t_res = "-"
+
+    v_b = subdomain_1.vs[0](b_res)
+    v_t = subdomain_2.vs[0](t_res)
+
+    u_bs = list(ufl.split(subdomain_1.u))
+    u_ts = list(ufl.split(subdomain_2.u))
+    u_b = u_bs[0](b_res)
+    u_t = u_ts[0](t_res)
+
+    def mixed_term(u, v, n):
+        return ufl.dot(ufl.grad(u), n) * v
+
+    n = ufl.FacetNormal(mesh)
+    n_b = n(b_res)
+    n_t = n(t_res)
+    cr = ufl.Circumradius(mesh)
+    h_b = 2 * cr(b_res)
+    h_t = 2 * cr(t_res)
+    gamma = 10.0
+
+    # fabricate K
+    W_0 = dolfinx.fem.functionspace(subdomain_1.submesh, ("DG", 0))
+    K_0 = dolfinx.fem.Function(W_0)
+    K_0.x.array[:] = 2
+    W_1 = dolfinx.fem.functionspace(subdomain_2.submesh, ("DG", 0))
+    K_1 = dolfinx.fem.Function(W_1)
+    K_1.x.array[:] = 4
+
+    K_b = K_0(b_res)
+    K_t = K_1(t_res)
+
+    F_0 = (
+        -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
+        - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
+    )
 
-J00 = dolfinx.fem.form(jac00, entity_maps=entity_maps)
-J01 = dolfinx.fem.form(jac01, entity_maps=entity_maps)
-J10 = dolfinx.fem.form(jac10, entity_maps=entity_maps)
-J11 = dolfinx.fem.form(jac11, entity_maps=entity_maps)
-J = [[J00, J01], [J10, J11]]
-F = [
-    dolfinx.fem.form(F_0, entity_maps=entity_maps),
-    dolfinx.fem.form(F_1, entity_maps=entity_maps),
-]
+    F_1 = (
+        +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
+        - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
+    )
+    F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
+    F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
+
+    # F_0 += F_00
+    # F_1 += F_11
+    subdomain_1.F += F_0
+    subdomain_2.F += F_1
+
+J = []
+forms = []
+for subdomain1 in list_of_subdomains:
+    jac = []
+    form = subdomain1.F
+    for subdomain2 in list_of_subdomains:
+        jac.append(
+            dolfinx.fem.form(
+                ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
+            )
+        )
+    J.append(jac)
+    forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
+
+
+# jac00 = ufl.derivative(F_0, u_0)
+
+# jac01 = ufl.derivative(F_0, u_1)
+
+# jac10 = ufl.derivative(F_1, u_0)
+# jac11 = ufl.derivative(F_1, u_1)
+
+# J00 = dolfinx.fem.form(jac00, entity_maps=entity_maps)
+# J01 = dolfinx.fem.form(jac01, entity_maps=entity_maps)
+# J10 = dolfinx.fem.form(jac10, entity_maps=entity_maps)
+# J11 = dolfinx.fem.form(jac11, entity_maps=entity_maps)
+# J = [[J00, J01], [J10, J11]]
+# F = [
+#     dolfinx.fem.form(F_0, entity_maps=entity_maps),
+#     dolfinx.fem.form(F_1, entity_maps=entity_maps),
+# ]
 
 
 # boundary conditions
-b_bc = dolfinx.fem.Function(u_0.function_space)
+b_bc = dolfinx.fem.Function(bottom_domain.u.function_space)
 b_bc.x.array[:] = 0.2
-submesh_b.topology.create_connectivity(
-    submesh_b.topology.dim - 1, submesh_b.topology.dim
+bottom_domain.submesh.topology.create_connectivity(
+    bottom_domain.submesh.topology.dim - 1, bottom_domain.submesh.topology.dim
 )
 bc_b = dolfinx.fem.dirichletbc(
     b_bc,
-    dolfinx.fem.locate_dofs_topological(u_0.function_space.sub(0), fdim, ft_b.find(2)),
+    dolfinx.fem.locate_dofs_topological(
+        bottom_domain.u.function_space.sub(0), fdim, bottom_domain.ft.find(2)
+    ),
 )
 
 
-t_bc = dolfinx.fem.Function(u_1.function_space)
+t_bc = dolfinx.fem.Function(top_domain.u.function_space)
 t_bc.x.array[:] = 0.05
-submesh_t.topology.create_connectivity(
-    submesh_t.topology.dim - 1, submesh_t.topology.dim
+top_domain.submesh.topology.create_connectivity(
+    top_domain.submesh.topology.dim - 1, top_domain.submesh.topology.dim
 )
 bc_t = dolfinx.fem.dirichletbc(
     t_bc,
-    dolfinx.fem.locate_dofs_topological(u_1.function_space.sub(0), fdim, ft_t.find(1)),
+    dolfinx.fem.locate_dofs_topological(
+        top_domain.u.function_space.sub(0), fdim, top_domain.ft.find(1)
+    ),
 )
 bcs = [bc_b, bc_t]
 
 
 solver = NewtonSolver(
-    F,
+    forms,
     J,
-    [u_0, u_1],
+    [subdomain.u for subdomain in list_of_subdomains],
     bcs=bcs,
     max_iterations=10,
     petsc_options={
@@ -278,7 +305,7 @@ def mixed_term(u, v, n):
     },
 )
 solver.solve(1e-5)
-
+exit()
 # bp = dolfinx.io.VTXWriter(mesh.comm, "u_b.bp", [u_0.sub(0).collapse()], engine="BP4")
 bp = dolfinx.io.VTXWriter(mesh.comm, "u_b_0.bp", [u_0.sub(0).collapse()], engine="BP4")
 bp.write(0)

From 73fe94b268e81715d3d81927f7bf20edfa67bfb4 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 08:34:46 +0000
Subject: [PATCH 003/134] removed unneeded stuff

---
 discontinuity_generic.py | 64 +++++++++++-----------------------------
 1 file changed, 18 insertions(+), 46 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 7b712a099..eac3f31d0 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -95,11 +95,6 @@ def half(x):
         mesh, mt, subdomain.submesh, subdomain.v_map, subdomain.submesh_to_mesh
     )
 
-# this seems to be not needed
-# t_parent_to_facet = np.full(num_facets_local, -1)
-# t_parent_to_facet[t_facet_to_parent] = np.arange(
-#     len(t_facet_to_parent), dtype=np.int32
-# )
 
 # Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
 # TODO ask Jorgen what this is for
@@ -226,8 +221,6 @@ def mixed_term(u, v, n):
     F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
     F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
 
-    # F_0 += F_00
-    # F_1 += F_11
     subdomain_1.F += F_0
     subdomain_2.F += F_1
 
@@ -245,25 +238,6 @@ def mixed_term(u, v, n):
     J.append(jac)
     forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
-
-# jac00 = ufl.derivative(F_0, u_0)
-
-# jac01 = ufl.derivative(F_0, u_1)
-
-# jac10 = ufl.derivative(F_1, u_0)
-# jac11 = ufl.derivative(F_1, u_1)
-
-# J00 = dolfinx.fem.form(jac00, entity_maps=entity_maps)
-# J01 = dolfinx.fem.form(jac01, entity_maps=entity_maps)
-# J10 = dolfinx.fem.form(jac10, entity_maps=entity_maps)
-# J11 = dolfinx.fem.form(jac11, entity_maps=entity_maps)
-# J = [[J00, J01], [J10, J11]]
-# F = [
-#     dolfinx.fem.form(F_0, entity_maps=entity_maps),
-#     dolfinx.fem.form(F_1, entity_maps=entity_maps),
-# ]
-
-
 # boundary conditions
 b_bc = dolfinx.fem.Function(bottom_domain.u.function_space)
 b_bc.x.array[:] = 0.2
@@ -305,20 +279,18 @@ def mixed_term(u, v, n):
     },
 )
 solver.solve(1e-5)
-exit()
-# bp = dolfinx.io.VTXWriter(mesh.comm, "u_b.bp", [u_0.sub(0).collapse()], engine="BP4")
-bp = dolfinx.io.VTXWriter(mesh.comm, "u_b_0.bp", [u_0.sub(0).collapse()], engine="BP4")
-bp.write(0)
-bp.close()
-bp = dolfinx.io.VTXWriter(mesh.comm, "u_t_0.bp", [u_1.sub(0).collapse()], engine="BP4")
-bp.write(0)
-bp.close()
-bp = dolfinx.io.VTXWriter(mesh.comm, "u_b_1.bp", [u_0.sub(1).collapse()], engine="BP4")
-bp.write(0)
-bp.close()
-bp = dolfinx.io.VTXWriter(mesh.comm, "u_t_1.bp", [u_1.sub(1).collapse()], engine="BP4")
-bp.write(0)
-bp.close()
+
+for subdomain in list_of_subdomains:
+    u_sub_0 = subdomain.u.sub(0).collapse()
+    u_sub_0.name = "u_sub_0"
+
+    u_sub_1 = subdomain.u.sub(1).collapse()
+    u_sub_1.name = "u_sub_1"
+    bp = dolfinx.io.VTXWriter(
+        mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
+    )
+    bp.write(0)
+    bp.close()
 
 
 # derived quantities
@@ -327,17 +299,17 @@ def mixed_term(u, v, n):
 T.interpolate(lambda x: 200 + x[1])
 
 
-T_b = dolfinx.fem.Function(u_0.sub(0).collapse().function_space)
+T_b = dolfinx.fem.Function(top_domain.u.sub(0).collapse().function_space)
 T_b.interpolate(T)
 
-ds_b = ufl.Measure("ds", domain=submesh_b)
-dx_b = ufl.Measure("dx", domain=submesh_b)
+ds_b = ufl.Measure("ds", domain=top_domain.submesh)
+dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
 
-n_b = ufl.FacetNormal(submesh_b)
+n_b = ufl.FacetNormal(bottom_domain.submesh)
 
-form = dolfinx.fem.form(u_0.sub(0) * dx_b)
+form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b, entity_maps=entity_maps)
 print(dolfinx.fem.assemble_scalar(form))
 
-form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(u_0.sub(0)), n_b) * ds_b)
+form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
 print(dolfinx.fem.assemble_scalar(form))

From 3f2f8c6715a2782fafd3e905d8eaaabce4296bd3 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 09:29:54 +0000
Subject: [PATCH 004/134] bcs generic

---
 discontinuity_generic.py | 56 +++++++++++++++++++++-------------------
 1 file changed, 30 insertions(+), 26 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index eac3f31d0..81b2946c5 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -69,6 +69,18 @@ def half(x):
 list_of_subdomains = [bottom_domain, top_domain]
 list_of_interfaces = {5: [bottom_domain, top_domain]}
 
+top_surface = F.SurfaceSubdomain(id=1)
+bottom_surface = F.SurfaceSubdomain(id=2)
+
+H = F.Species("H", mobile=True)
+
+list_of_bcs = [
+    F.DirichletBC(top_surface, value=0.05, species=H),
+    F.DirichletBC(bottom_surface, value=0.2, species=H)
+    ]
+
+surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+
 gdim = mesh.geometry.dim
 tdim = mesh.topology.dim
 fdim = tdim - 1
@@ -239,32 +251,24 @@ def mixed_term(u, v, n):
     forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
 # boundary conditions
-b_bc = dolfinx.fem.Function(bottom_domain.u.function_space)
-b_bc.x.array[:] = 0.2
-bottom_domain.submesh.topology.create_connectivity(
-    bottom_domain.submesh.topology.dim - 1, bottom_domain.submesh.topology.dim
-)
-bc_b = dolfinx.fem.dirichletbc(
-    b_bc,
-    dolfinx.fem.locate_dofs_topological(
-        bottom_domain.u.function_space.sub(0), fdim, bottom_domain.ft.find(2)
-    ),
-)
-
-
-t_bc = dolfinx.fem.Function(top_domain.u.function_space)
-t_bc.x.array[:] = 0.05
-top_domain.submesh.topology.create_connectivity(
-    top_domain.submesh.topology.dim - 1, top_domain.submesh.topology.dim
-)
-bc_t = dolfinx.fem.dirichletbc(
-    t_bc,
-    dolfinx.fem.locate_dofs_topological(
-        top_domain.u.function_space.sub(0), fdim, top_domain.ft.find(1)
-    ),
-)
-bcs = [bc_b, bc_t]
-
+bcs = []
+for boundary_condition in list_of_bcs:
+    volume_subdomain = surface_to_volume[boundary_condition.subdomain]
+    bc = dolfinx.fem.Function(volume_subdomain.u.function_space)
+    bc.x.array[:] = boundary_condition.value
+    volume_subdomain.submesh.topology.create_connectivity(
+        volume_subdomain.submesh.topology.dim - 1,
+        volume_subdomain.submesh.topology.dim,
+    )
+    bc = dolfinx.fem.dirichletbc(
+        bc,
+        dolfinx.fem.locate_dofs_topological(
+            volume_subdomain.u.function_space.sub(0),
+            fdim,
+            volume_subdomain.ft.find(boundary_condition.subdomain.id),
+        ),
+    )
+    bcs.append(bc)
 
 solver = NewtonSolver(
     forms,

From a50810f497b1e41bb745421d83fa455932879c1a Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 09:30:58 +0000
Subject: [PATCH 005/134] bcs before jacobian

---
 discontinuity_generic.py | 41 +++++++++++++++++++++-------------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 81b2946c5..dfa7fb5f1 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -179,6 +179,27 @@ def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
     )
     subdomain.u.name = f"u_{subdomain.id}"
 
+
+# boundary conditions
+bcs = []
+for boundary_condition in list_of_bcs:
+    volume_subdomain = surface_to_volume[boundary_condition.subdomain]
+    bc = dolfinx.fem.Function(volume_subdomain.u.function_space)
+    bc.x.array[:] = boundary_condition.value
+    volume_subdomain.submesh.topology.create_connectivity(
+        volume_subdomain.submesh.topology.dim - 1,
+        volume_subdomain.submesh.topology.dim,
+    )
+    bc = dolfinx.fem.dirichletbc(
+        bc,
+        dolfinx.fem.locate_dofs_topological(
+            volume_subdomain.u.function_space.sub(0),
+            fdim,
+            volume_subdomain.ft.find(boundary_condition.subdomain.id),
+        ),
+    )
+    bcs.append(bc)
+
 # Add coupling term to the interface
 # Get interface markers on submesh b
 for interface in list_of_interfaces:
@@ -250,25 +271,7 @@ def mixed_term(u, v, n):
     J.append(jac)
     forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
-# boundary conditions
-bcs = []
-for boundary_condition in list_of_bcs:
-    volume_subdomain = surface_to_volume[boundary_condition.subdomain]
-    bc = dolfinx.fem.Function(volume_subdomain.u.function_space)
-    bc.x.array[:] = boundary_condition.value
-    volume_subdomain.submesh.topology.create_connectivity(
-        volume_subdomain.submesh.topology.dim - 1,
-        volume_subdomain.submesh.topology.dim,
-    )
-    bc = dolfinx.fem.dirichletbc(
-        bc,
-        dolfinx.fem.locate_dofs_topological(
-            volume_subdomain.u.function_space.sub(0),
-            fdim,
-            volume_subdomain.ft.find(boundary_condition.subdomain.id),
-        ),
-    )
-    bcs.append(bc)
+
 
 solver = NewtonSolver(
     forms,

From c59a875131bb126ee0fd2b640862cefbd6fd066e Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 13:35:28 +0000
Subject: [PATCH 006/134] issue with results

---
 discontinuity_generic.py | 117 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 109 insertions(+), 8 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index dfa7fb5f1..f79796cdc 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -73,6 +73,21 @@ def half(x):
 bottom_surface = F.SurfaceSubdomain(id=2)
 
 H = F.Species("H", mobile=True)
+trapped_H = F.Species("H_trapped", mobile=False)
+empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
+
+for species in [H, trapped_H]:
+    species.subdomains = [top_domain, bottom_domain]
+    species.subdomain_to_solution = {}
+    species.subdomain_to_prev_solution = {}
+    species.subdomain_to_test_function = {}
+
+list_of_species = [H, trapped_H]
+
+list_of_reactions = [
+    F.Reaction(reactant=[H, empty_trap], product=[trapped_H], k_0=2, E_k=0, p_0=0.1, E_p=0, volume=top_domain),
+    F.Reaction(reactant=[H, empty_trap], product=[trapped_H], k_0=2, E_k=0, p_0=0.1, E_p=0, volume=bottom_domain),
+    ]
 
 list_of_bcs = [
     F.DirichletBC(top_surface, value=0.05, species=H),
@@ -132,7 +147,11 @@ def D(T):
     return 2 * ufl.exp(-0.1 / k_B / T)
 
 
-def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
+def define_interior_eq(mesh, degree, subdomain, value):
+    submesh = subdomain.submesh
+    submesh_to_mesh = subdomain.submesh_to_mesh
+
+    # TODO mesh_to_submesh isn't used anywhere
     # Compute map from parent entity to submesh cell
     codim = mesh.topology.dim - submesh.topology.dim
     ptdim = mesh.topology.dim - codim
@@ -143,7 +162,6 @@ def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
     mesh_to_submesh = np.full(num_entities, -1)
     mesh_to_submesh[submesh_to_mesh] = np.arange(len(submesh_to_mesh), dtype=np.int32)
 
-    degree = 1
     element_CG = basix.ufl.element(
         basix.ElementFamily.P,
         submesh.basix_cell(),
@@ -171,15 +189,96 @@ def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
     return u, vs, F, mesh_to_submesh
 
 
+def define_function_spaces(subdomain: F.VolumeSubdomain):
+    # get number of species defined in the subdomain
+    all_species = [species for species in list_of_species if subdomain in species.subdomains]
+    unique_species = list(set(all_species))
+    nb_species = len(unique_species)
+
+    degree = 1
+    element_CG = basix.ufl.element(
+        basix.ElementFamily.P,
+        subdomain.submesh.basix_cell(),
+        degree,
+        basix.LagrangeVariant.equispaced,
+    )
+    element = basix.ufl.mixed_element([element_CG] * nb_species)
+    V = dolfinx.fem.functionspace(subdomain.submesh, element)
+    u = dolfinx.fem.Function(V)
+    u_n = dolfinx.fem.Function(V)
+
+    us = list(ufl.split(u))
+    u_ns = list(ufl.split(u_n))
+    vs = list(ufl.TestFunctions(V))
+    for i, species in enumerate(unique_species):
+        species.subdomain_to_solution[subdomain] = us[i]
+        species.subdomain_to_prev_solution[subdomain] = u_ns[i]
+        species.subdomain_to_test_function[subdomain] = vs[i]
+    subdomain.u = u
+
+def define_formulation(subdomain: F.VolumeSubdomain):
+    form = 0
+    T = dolfinx.fem.Constant(subdomain.submesh, 300.0)  # FIXME temperature is ignored for now
+    # add diffusion and time derivative for each species
+    for spe in list_of_species:
+        u = spe.subdomain_to_solution[subdomain]
+        u_n = spe.subdomain_to_prev_solution[subdomain]
+        v = spe.subdomain_to_test_function[subdomain]
+        dx = subdomain.dx
+
+        D = dolfinx.fem.Constant(subdomain.submesh, 1.0)  # TODO change this
+
+        if spe.mobile:
+            form += ufl.dot(D * ufl.grad(u), ufl.grad(v)) * dx
+
+    for reaction in list_of_reactions:
+        for species in reaction.reactant + reaction.product:
+            if isinstance(species, F.Species):
+                # TODO remove
+                # temporarily overide the solution and test function to the one of the subdomain
+                species.solution = species.subdomain_to_solution[subdomain]
+                species.test_function = species.subdomain_to_test_function[subdomain]
+
+        for reactant in reaction.reactant:
+            if isinstance(reactant, F.Species):
+                form += (
+                    reaction.reaction_term(T)  # FIXME temperature is ignored for now
+                    * reactant.test_function
+                    * dx
+                )
+
+        # product
+        if isinstance(reaction.product, list):
+            products = reaction.product
+        else:
+            products = [reaction.product]
+        for product in products:
+            form += (
+                -reaction.reaction_term(T)  # FIXME temperature is ignored for now
+                * product.subdomain_to_test_function[subdomain]
+                * dx
+            )
+
+    subdomain.F = form
+
+
 # for each subdomain, define the interior equation
 for subdomain in list_of_subdomains:
-    degree = 1
-    subdomain.u, subdomain.vs, subdomain.F, subdomain.m_to_s = define_interior_eq(
-        mesh, degree, subdomain.submesh, subdomain.submesh_to_mesh, 0.0
+    # degree = 1
+    # subdomain.u, subdomain.vs, subdomain.F, subdomain.m_to_s = define_interior_eq(
+    #     mesh, degree, subdomain, 0.0
+    # )
+    define_function_spaces(subdomain)
+    ct_r = dolfinx.mesh.meshtags(
+        mesh,
+        mesh.topology.dim,
+        subdomain.submesh_to_mesh,
+        np.full_like(subdomain.submesh_to_mesh, 1, dtype=np.int32),
     )
+    subdomain.dx = ufl.Measure("dx", domain=mesh, subdomain_data=ct_r, subdomain_id=1)
+    define_formulation(subdomain)
     subdomain.u.name = f"u_{subdomain.id}"
 
-
 # boundary conditions
 bcs = []
 for boundary_condition in list_of_bcs:
@@ -212,8 +311,10 @@ def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
     b_res = "+"
     t_res = "-"
 
-    v_b = subdomain_1.vs[0](b_res)
-    v_t = subdomain_2.vs[0](t_res)
+    v_b = H.subdomain_to_test_function[subdomain_1](b_res)
+    v_t = H.subdomain_to_test_function[subdomain_2](t_res)
+    # v_b = subdomain_1.vs[0](b_res)
+    # v_t = subdomain_2.vs[0](t_res)
 
     u_bs = list(ufl.split(subdomain_1.u))
     u_ts = list(ufl.split(subdomain_2.u))

From c0f9d09cd1ad7d498923fbcf8cf637e571f47a10 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 14:03:27 +0000
Subject: [PATCH 007/134] fixed double reaction + inner instead of dot

---
 discontinuity_generic.py | 52 ++++------------------------------------
 1 file changed, 4 insertions(+), 48 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index f79796cdc..ef810ad08 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -147,48 +147,6 @@ def D(T):
     return 2 * ufl.exp(-0.1 / k_B / T)
 
 
-def define_interior_eq(mesh, degree, subdomain, value):
-    submesh = subdomain.submesh
-    submesh_to_mesh = subdomain.submesh_to_mesh
-
-    # TODO mesh_to_submesh isn't used anywhere
-    # Compute map from parent entity to submesh cell
-    codim = mesh.topology.dim - submesh.topology.dim
-    ptdim = mesh.topology.dim - codim
-    num_entities = (
-        mesh.topology.index_map(ptdim).size_local
-        + mesh.topology.index_map(ptdim).num_ghosts
-    )
-    mesh_to_submesh = np.full(num_entities, -1)
-    mesh_to_submesh[submesh_to_mesh] = np.arange(len(submesh_to_mesh), dtype=np.int32)
-
-    element_CG = basix.ufl.element(
-        basix.ElementFamily.P,
-        submesh.basix_cell(),
-        degree,
-        basix.LagrangeVariant.equispaced,
-    )
-    element = basix.ufl.mixed_element([element_CG, element_CG])
-    V = dolfinx.fem.functionspace(submesh, element)
-    u = dolfinx.fem.Function(V)
-    us = list(ufl.split(u))
-    vs = list(ufl.TestFunctions(V))
-    ct_r = dolfinx.mesh.meshtags(
-        mesh,
-        mesh.topology.dim,
-        submesh_to_mesh,
-        np.full_like(submesh_to_mesh, 1, dtype=np.int32),
-    )
-    val = dolfinx.fem.Constant(submesh, value)
-    dx_r = ufl.Measure("dx", domain=mesh, subdomain_data=ct_r, subdomain_id=1)
-    F = ufl.inner(ufl.grad(us[0]), ufl.grad(vs[0])) * dx_r - val * vs[0] * dx_r
-    k = 2
-    p = 0.1
-    n = 0.5
-    F += k * us[0] * (n - us[1]) * vs[1] * dx_r - p * us[1] * vs[1] * dx_r
-    return u, vs, F, mesh_to_submesh
-
-
 def define_function_spaces(subdomain: F.VolumeSubdomain):
     # get number of species defined in the subdomain
     all_species = [species for species in list_of_species if subdomain in species.subdomains]
@@ -229,9 +187,12 @@ def define_formulation(subdomain: F.VolumeSubdomain):
         D = dolfinx.fem.Constant(subdomain.submesh, 1.0)  # TODO change this
 
         if spe.mobile:
-            form += ufl.dot(D * ufl.grad(u), ufl.grad(v)) * dx
+            # I noticed that if we use dot here it doesn't work....
+            form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
 
     for reaction in list_of_reactions:
+        if reaction.volume != subdomain:
+            continue
         for species in reaction.reactant + reaction.product:
             if isinstance(species, F.Species):
                 # TODO remove
@@ -262,12 +223,7 @@ def define_formulation(subdomain: F.VolumeSubdomain):
     subdomain.F = form
 
 
-# for each subdomain, define the interior equation
 for subdomain in list_of_subdomains:
-    # degree = 1
-    # subdomain.u, subdomain.vs, subdomain.F, subdomain.m_to_s = define_interior_eq(
-    #     mesh, degree, subdomain, 0.0
-    # )
     define_function_spaces(subdomain)
     ct_r = dolfinx.mesh.meshtags(
         mesh,

From 994931bd4babb190bcbf6caa0c505975eb0eea3e Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 15:02:46 +0000
Subject: [PATCH 008/134] fixed test function for reactant

---
 discontinuity_generic.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index ef810ad08..92c20129e 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -204,7 +204,7 @@ def define_formulation(subdomain: F.VolumeSubdomain):
             if isinstance(reactant, F.Species):
                 form += (
                     reaction.reaction_term(T)  # FIXME temperature is ignored for now
-                    * reactant.test_function
+                    * reactant.subdomain_to_test_function[subdomain]
                     * dx
                 )
 

From 66adf7c1a77d5e8db0a0a09e05ede5ad7f7fde01 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 15 Aug 2024 15:31:49 +0000
Subject: [PATCH 009/134] fixed a few random bugs

---
 discontinuity_generic.py | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 92c20129e..128353f8b 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -77,7 +77,7 @@ def half(x):
 empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
 
 for species in [H, trapped_H]:
-    species.subdomains = [top_domain, bottom_domain]
+    species.subdomains = [bottom_domain, top_domain]
     species.subdomain_to_solution = {}
     species.subdomain_to_prev_solution = {}
     species.subdomain_to_test_function = {}
@@ -150,7 +150,12 @@ def D(T):
 def define_function_spaces(subdomain: F.VolumeSubdomain):
     # get number of species defined in the subdomain
     all_species = [species for species in list_of_species if subdomain in species.subdomains]
-    unique_species = list(set(all_species))
+
+    # instead of using the set function we use a list to keep the order
+    unique_species = []
+    for species in all_species:
+        if species not in unique_species:
+            unique_species.append(species)
     nb_species = len(unique_species)
 
     degree = 1
@@ -269,13 +274,9 @@ def define_formulation(subdomain: F.VolumeSubdomain):
 
     v_b = H.subdomain_to_test_function[subdomain_1](b_res)
     v_t = H.subdomain_to_test_function[subdomain_2](t_res)
-    # v_b = subdomain_1.vs[0](b_res)
-    # v_t = subdomain_2.vs[0](t_res)
 
-    u_bs = list(ufl.split(subdomain_1.u))
-    u_ts = list(ufl.split(subdomain_2.u))
-    u_b = u_bs[0](b_res)
-    u_t = u_ts[0](t_res)
+    u_b = H.subdomain_to_solution[subdomain_1](b_res)
+    u_t = H.subdomain_to_solution[subdomain_2](t_res)
 
     def mixed_term(u, v, n):
         return ufl.dot(ufl.grad(u), n) * v
@@ -363,10 +364,10 @@ def mixed_term(u, v, n):
 T.interpolate(lambda x: 200 + x[1])
 
 
-T_b = dolfinx.fem.Function(top_domain.u.sub(0).collapse().function_space)
+T_b = dolfinx.fem.Function(bottom_domain.u.sub(0).collapse().function_space)
 T_b.interpolate(T)
 
-ds_b = ufl.Measure("ds", domain=top_domain.submesh)
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
 
@@ -375,5 +376,12 @@ def mixed_term(u, v, n):
 form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b, entity_maps=entity_maps)
 print(dolfinx.fem.assemble_scalar(form))
 
-form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
+form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b, entity_maps=entity_maps)
+print(dolfinx.fem.assemble_scalar(form))
+
+id_interface = 5
+form = dolfinx.fem.form(ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
 print(dolfinx.fem.assemble_scalar(form))
+
+# form = dolfinx.fem.form(T * dx_b, entity_maps=entity_maps)
+# print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file

From 0e5f5baa4b1075749cf54b22361e69298df852b1 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 07:15:46 +0000
Subject: [PATCH 010/134] correct way of making forms with parent functions

---
 discontinuity_generic.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 128353f8b..0d9a96851 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -359,29 +359,32 @@ def mixed_term(u, v, n):
 
 
 # derived quantities
+entity_maps[mesh] = bottom_domain.submesh_to_mesh
+
 V = dolfinx.fem.functionspace(mesh, ("CG", 1))
 T = dolfinx.fem.Function(V)
 T.interpolate(lambda x: 200 + x[1])
 
 
-T_b = dolfinx.fem.Function(bottom_domain.u.sub(0).collapse().function_space)
-T_b.interpolate(T)
-
-ds_b = ufl.Measure("ds", domain=bottom_domain.submesh)
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
+ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
 
 n_b = ufl.FacetNormal(bottom_domain.submesh)
+n_t = ufl.FacetNormal(top_domain.submesh)
 
-form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b, entity_maps=entity_maps)
+form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
 print(dolfinx.fem.assemble_scalar(form))
 
-form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b, entity_maps=entity_maps)
+form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b)
 print(dolfinx.fem.assemble_scalar(form))
 
-id_interface = 5
-form = dolfinx.fem.form(ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
+form = dolfinx.fem.form(T * dx_b, entity_maps={mesh: bottom_domain.submesh_to_mesh})
 print(dolfinx.fem.assemble_scalar(form))
 
-# form = dolfinx.fem.form(T * dx_b, entity_maps=entity_maps)
-# print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file
+id_interface = 5
+form = dolfinx.fem.form(T*ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b(id_interface), entity_maps={mesh: bottom_domain.submesh_to_mesh})
+print(dolfinx.fem.assemble_scalar(form))
+form = dolfinx.fem.form(T*ufl.dot(ufl.grad(top_domain.u.sub(0)), n_t) * ds_t(id_interface), entity_maps={mesh: top_domain.submesh_to_mesh})
+print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file

From ac07643469bd3d46738cde905a3da0c965be5a8d Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 07:15:59 +0000
Subject: [PATCH 011/134] formatting

---
 discontinuity_generic.py | 51 ++++++++++++++++++++++++++++++----------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 0d9a96851..02a0a9313 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -85,14 +85,30 @@ def half(x):
 list_of_species = [H, trapped_H]
 
 list_of_reactions = [
-    F.Reaction(reactant=[H, empty_trap], product=[trapped_H], k_0=2, E_k=0, p_0=0.1, E_p=0, volume=top_domain),
-    F.Reaction(reactant=[H, empty_trap], product=[trapped_H], k_0=2, E_k=0, p_0=0.1, E_p=0, volume=bottom_domain),
-    ]
+    F.Reaction(
+        reactant=[H, empty_trap],
+        product=[trapped_H],
+        k_0=2,
+        E_k=0,
+        p_0=0.1,
+        E_p=0,
+        volume=top_domain,
+    ),
+    F.Reaction(
+        reactant=[H, empty_trap],
+        product=[trapped_H],
+        k_0=2,
+        E_k=0,
+        p_0=0.1,
+        E_p=0,
+        volume=bottom_domain,
+    ),
+]
 
 list_of_bcs = [
     F.DirichletBC(top_surface, value=0.05, species=H),
-    F.DirichletBC(bottom_surface, value=0.2, species=H)
-    ]
+    F.DirichletBC(bottom_surface, value=0.2, species=H),
+]
 
 surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
@@ -137,8 +153,7 @@ def half(x):
 
 # ._cpp_object needed on dolfinx 0.8.0
 entity_maps = {
-    subdomain.submesh: subdomain.parent_to_submesh
-    for subdomain in list_of_subdomains
+    subdomain.submesh: subdomain.parent_to_submesh for subdomain in list_of_subdomains
 }
 
 
@@ -149,7 +164,9 @@ def D(T):
 
 def define_function_spaces(subdomain: F.VolumeSubdomain):
     # get number of species defined in the subdomain
-    all_species = [species for species in list_of_species if subdomain in species.subdomains]
+    all_species = [
+        species for species in list_of_species if subdomain in species.subdomains
+    ]
 
     # instead of using the set function we use a list to keep the order
     unique_species = []
@@ -179,9 +196,12 @@ def define_function_spaces(subdomain: F.VolumeSubdomain):
         species.subdomain_to_test_function[subdomain] = vs[i]
     subdomain.u = u
 
+
 def define_formulation(subdomain: F.VolumeSubdomain):
     form = 0
-    T = dolfinx.fem.Constant(subdomain.submesh, 300.0)  # FIXME temperature is ignored for now
+    T = dolfinx.fem.Constant(
+        subdomain.submesh, 300.0
+    )  # FIXME temperature is ignored for now
     # add diffusion and time derivative for each species
     for spe in list_of_species:
         u = spe.subdomain_to_solution[subdomain]
@@ -330,7 +350,6 @@ def mixed_term(u, v, n):
     forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
 
-
 solver = NewtonSolver(
     forms,
     J,
@@ -384,7 +403,13 @@ def mixed_term(u, v, n):
 print(dolfinx.fem.assemble_scalar(form))
 
 id_interface = 5
-form = dolfinx.fem.form(T*ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b(id_interface), entity_maps={mesh: bottom_domain.submesh_to_mesh})
+form = dolfinx.fem.form(
+    T * ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b(id_interface),
+    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+)
+print(dolfinx.fem.assemble_scalar(form))
+form = dolfinx.fem.form(
+    T * ufl.dot(ufl.grad(top_domain.u.sub(0)), n_t) * ds_t(id_interface),
+    entity_maps={mesh: top_domain.submesh_to_mesh},
+)
 print(dolfinx.fem.assemble_scalar(form))
-form = dolfinx.fem.form(T*ufl.dot(ufl.grad(top_domain.u.sub(0)), n_t) * ds_t(id_interface), entity_maps={mesh: top_domain.submesh_to_mesh})
-print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file

From 659cc7bb0ad071c09984a83678d289d2fd42e62b Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 08:49:56 +0000
Subject: [PATCH 012/134] added integration with parent mesh coefficients

---
 discontinuity_generic.py | 39 +++++++++++++++++++--------------------
 1 file changed, 19 insertions(+), 20 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 02a0a9313..2f0c4964c 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -112,6 +112,16 @@ def half(x):
 
 surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
+V = dolfinx.fem.functionspace(mesh, ("CG", 1))
+T = dolfinx.fem.Function(V)
+T.interpolate(lambda x: 300 + 10 * x[1])
+
+
+def D_fun(T):
+    k_B = 8.6173303e-5
+    return 2 * ufl.exp(-0.1 / k_B / T)
+
+
 gdim = mesh.geometry.dim
 tdim = mesh.topology.dim
 fdim = tdim - 1
@@ -157,11 +167,6 @@ def half(x):
 }
 
 
-def D(T):
-    k_B = 8.6173303e-5
-    return 2 * ufl.exp(-0.1 / k_B / T)
-
-
 def define_function_spaces(subdomain: F.VolumeSubdomain):
     # get number of species defined in the subdomain
     all_species = [
@@ -199,9 +204,6 @@ def define_function_spaces(subdomain: F.VolumeSubdomain):
 
 def define_formulation(subdomain: F.VolumeSubdomain):
     form = 0
-    T = dolfinx.fem.Constant(
-        subdomain.submesh, 300.0
-    )  # FIXME temperature is ignored for now
     # add diffusion and time derivative for each species
     for spe in list_of_species:
         u = spe.subdomain_to_solution[subdomain]
@@ -209,10 +211,9 @@ def define_formulation(subdomain: F.VolumeSubdomain):
         v = spe.subdomain_to_test_function[subdomain]
         dx = subdomain.dx
 
-        D = dolfinx.fem.Constant(subdomain.submesh, 1.0)  # TODO change this
+        D = D_fun(T)
 
         if spe.mobile:
-            # I noticed that if we use dot here it doesn't work....
             form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
 
     for reaction in list_of_reactions:
@@ -307,7 +308,7 @@ def mixed_term(u, v, n):
     cr = ufl.Circumradius(mesh)
     h_b = 2 * cr(b_res)
     h_t = 2 * cr(t_res)
-    gamma = 10.0
+    gamma = 400.0  # this needs to be "sufficiently large"
 
     # fabricate K
     W_0 = dolfinx.fem.functionspace(subdomain_1.submesh, ("DG", 0))
@@ -340,14 +341,17 @@ def mixed_term(u, v, n):
 for subdomain1 in list_of_subdomains:
     jac = []
     form = subdomain1.F
+    # copy entity_maps
+    entity_maps_ = entity_maps.copy()
+    entity_maps_[mesh] = subdomain1.submesh_to_mesh
     for subdomain2 in list_of_subdomains:
         jac.append(
             dolfinx.fem.form(
-                ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
+                ufl.derivative(form, subdomain2.u), entity_maps=entity_maps_
             )
         )
     J.append(jac)
-    forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
+    forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps_))
 
 
 solver = NewtonSolver(
@@ -380,11 +384,6 @@ def mixed_term(u, v, n):
 # derived quantities
 entity_maps[mesh] = bottom_domain.submesh_to_mesh
 
-V = dolfinx.fem.functionspace(mesh, ("CG", 1))
-T = dolfinx.fem.Function(V)
-T.interpolate(lambda x: 200 + x[1])
-
-
 ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
 ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
@@ -404,12 +403,12 @@ def mixed_term(u, v, n):
 
 id_interface = 5
 form = dolfinx.fem.form(
-    T * ufl.dot(ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b(id_interface),
+    ufl.dot(D_fun(T) * ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b(id_interface),
     entity_maps={mesh: bottom_domain.submesh_to_mesh},
 )
 print(dolfinx.fem.assemble_scalar(form))
 form = dolfinx.fem.form(
-    T * ufl.dot(ufl.grad(top_domain.u.sub(0)), n_t) * ds_t(id_interface),
+    ufl.dot(D_fun(T) * ufl.grad(top_domain.u.sub(0)), n_t) * ds_t(id_interface),
     entity_maps={mesh: top_domain.submesh_to_mesh},
 )
 print(dolfinx.fem.assemble_scalar(form))

From fb431516fccd8e26cb122d02a8118274e9679952 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 09:04:39 +0000
Subject: [PATCH 013/134] now solubility is generic

---
 discontinuity_generic.py | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 2f0c4964c..1f8a930a7 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -122,6 +122,16 @@ def D_fun(T):
     return 2 * ufl.exp(-0.1 / k_B / T)
 
 
+def K_1_fun(T):
+    k_B = 8.6173303e-5
+    return 4 * ufl.exp(-0.1 / k_B / T)
+
+
+def K_2_fun(T):
+    k_B = 8.6173303e-5
+    return 4 * ufl.exp(-0.1 / k_B / T)
+
+
 gdim = mesh.geometry.dim
 tdim = mesh.topology.dim
 fdim = tdim - 1
@@ -310,16 +320,8 @@ def mixed_term(u, v, n):
     h_t = 2 * cr(t_res)
     gamma = 400.0  # this needs to be "sufficiently large"
 
-    # fabricate K
-    W_0 = dolfinx.fem.functionspace(subdomain_1.submesh, ("DG", 0))
-    K_0 = dolfinx.fem.Function(W_0)
-    K_0.x.array[:] = 2
-    W_1 = dolfinx.fem.functionspace(subdomain_2.submesh, ("DG", 0))
-    K_1 = dolfinx.fem.Function(W_1)
-    K_1.x.array[:] = 4
-
-    K_b = K_0(b_res)
-    K_t = K_1(t_res)
+    K_b = K_1_fun(T(b_res))
+    K_t = K_2_fun(T(t_res))
 
     F_0 = (
         -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface

From 91a6a737ffcfaff2f73d1b2ada07e83d46871bee Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 09:05:29 +0000
Subject: [PATCH 014/134] back to solubility of 2

---
 discontinuity_generic.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 1f8a930a7..d1f10740c 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -124,7 +124,7 @@ def D_fun(T):
 
 def K_1_fun(T):
     k_B = 8.6173303e-5
-    return 4 * ufl.exp(-0.1 / k_B / T)
+    return 2 * ufl.exp(-0.1 / k_B / T)
 
 
 def K_2_fun(T):

From 53e74306548f75cd25b565871945b75808f00ef7 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 09:13:02 +0000
Subject: [PATCH 015/134] T varies in x

---
 discontinuity_generic.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index d1f10740c..ceffb02ad 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -114,7 +114,7 @@ def half(x):
 
 V = dolfinx.fem.functionspace(mesh, ("CG", 1))
 T = dolfinx.fem.Function(V)
-T.interpolate(lambda x: 300 + 10 * x[1])
+T.interpolate(lambda x: 300 + 10 * x[1] + 100 * x[0])
 
 
 def D_fun(T):
@@ -129,7 +129,7 @@ def K_1_fun(T):
 
 def K_2_fun(T):
     k_B = 8.6173303e-5
-    return 4 * ufl.exp(-0.1 / k_B / T)
+    return 4 * ufl.exp(-0.12 / k_B / T)
 
 
 gdim = mesh.geometry.dim
@@ -239,7 +239,7 @@ def define_formulation(subdomain: F.VolumeSubdomain):
         for reactant in reaction.reactant:
             if isinstance(reactant, F.Species):
                 form += (
-                    reaction.reaction_term(T)  # FIXME temperature is ignored for now
+                    reaction.reaction_term(T)
                     * reactant.subdomain_to_test_function[subdomain]
                     * dx
                 )
@@ -251,7 +251,7 @@ def define_formulation(subdomain: F.VolumeSubdomain):
             products = [reaction.product]
         for product in products:
             form += (
-                -reaction.reaction_term(T)  # FIXME temperature is ignored for now
+                -reaction.reaction_term(T)
                 * product.subdomain_to_test_function[subdomain]
                 * dx
             )

From f1f6d65c31e5ad5bb012741ea0da538aef1bbffe Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 09:21:57 +0000
Subject: [PATCH 016/134] K_S no longer hard coded

---
 discontinuity_generic.py | 46 ++++++++++++++++++++++++++--------------
 1 file changed, 30 insertions(+), 16 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index ceffb02ad..b3790b658 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -64,8 +64,16 @@ def half(x):
 
 mesh, mt, ct = generate_mesh()
 
-top_domain = F.VolumeSubdomain(4, material=None)
-bottom_domain = F.VolumeSubdomain(3, material=None)
+material_bottom = F.Material(D_0=2.0, E_D=0.1)
+material_top = F.Material(D_0=2.0, E_D=0.1)
+
+material_bottom.K_S_0 = 2.0
+material_bottom.E_K_S = 0.1
+material_top.K_S_0 = 4.0
+material_top.E_K_S = 0.12
+
+top_domain = F.VolumeSubdomain(4, material=material_top)
+bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
 list_of_subdomains = [bottom_domain, top_domain]
 list_of_interfaces = {5: [bottom_domain, top_domain]}
 
@@ -117,19 +125,14 @@ def half(x):
 T.interpolate(lambda x: 300 + 10 * x[1] + 100 * x[0])
 
 
-def D_fun(T):
+def D_fun(T, D_0, E_D):
     k_B = 8.6173303e-5
-    return 2 * ufl.exp(-0.1 / k_B / T)
+    return D_0 * ufl.exp(-E_D / k_B / T)
 
 
-def K_1_fun(T):
+def K_S_fun(T, K_S_0, E_K_S):
     k_B = 8.6173303e-5
-    return 2 * ufl.exp(-0.1 / k_B / T)
-
-
-def K_2_fun(T):
-    k_B = 8.6173303e-5
-    return 4 * ufl.exp(-0.12 / k_B / T)
+    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
 
 
 gdim = mesh.geometry.dim
@@ -221,7 +224,7 @@ def define_formulation(subdomain: F.VolumeSubdomain):
         v = spe.subdomain_to_test_function[subdomain]
         dx = subdomain.dx
 
-        D = D_fun(T)
+        D = D_fun(T, subdomain.material.D_0, subdomain.material.E_D)
 
         if spe.mobile:
             form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
@@ -320,8 +323,9 @@ def mixed_term(u, v, n):
     h_t = 2 * cr(t_res)
     gamma = 400.0  # this needs to be "sufficiently large"
 
-    K_b = K_1_fun(T(b_res))
-    K_t = K_2_fun(T(t_res))
+    # TODO how do we know if b_res corresponds to the bottom or top domain?
+    K_b = K_S_fun(T(b_res), subdomain_1.material.K_S_0, subdomain_1.material.E_K_S)
+    K_t = K_S_fun(T(t_res), subdomain_2.material.K_S_0, subdomain_2.material.E_K_S)
 
     F_0 = (
         -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
@@ -405,12 +409,22 @@ def mixed_term(u, v, n):
 
 id_interface = 5
 form = dolfinx.fem.form(
-    ufl.dot(D_fun(T) * ufl.grad(bottom_domain.u.sub(0)), n_b) * ds_b(id_interface),
+    ufl.dot(
+        D_fun(T, bottom_domain.material.D_0, bottom_domain.material.E_D)
+        * ufl.grad(bottom_domain.u.sub(0)),
+        n_b,
+    )
+    * ds_b(id_interface),
     entity_maps={mesh: bottom_domain.submesh_to_mesh},
 )
 print(dolfinx.fem.assemble_scalar(form))
 form = dolfinx.fem.form(
-    ufl.dot(D_fun(T) * ufl.grad(top_domain.u.sub(0)), n_t) * ds_t(id_interface),
+    ufl.dot(
+        D_fun(T, top_domain.material.D_0, top_domain.material.E_D)
+        * ufl.grad(top_domain.u.sub(0)),
+        n_t,
+    )
+    * ds_t(id_interface),
     entity_maps={mesh: top_domain.submesh_to_mesh},
 )
 print(dolfinx.fem.assemble_scalar(form))

From 9f05c8278143233a3f12eddc9f99d012d0d91c56 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 09:43:40 +0000
Subject: [PATCH 017/134] + and - restrictions correspond to the correct
 solubilities

---
 discontinuity_generic.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index b3790b658..ce731bcd0 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -68,9 +68,9 @@ def half(x):
 material_top = F.Material(D_0=2.0, E_D=0.1)
 
 material_bottom.K_S_0 = 2.0
-material_bottom.E_K_S = 0.1
+material_bottom.E_K_S = 0 * 0.1
 material_top.K_S_0 = 4.0
-material_top.E_K_S = 0.12
+material_top.E_K_S = 0 * 0.12
 
 top_domain = F.VolumeSubdomain(4, material=material_top)
 bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
@@ -297,8 +297,6 @@ def define_formulation(subdomain: F.VolumeSubdomain):
 # Add coupling term to the interface
 # Get interface markers on submesh b
 for interface in list_of_interfaces:
-    subdomain_1 = list_of_interfaces[interface][0]
-    subdomain_2 = list_of_interfaces[interface][1]
 
     dInterface = ufl.Measure(
         "dS", domain=mesh, subdomain_data=mt, subdomain_id=interface
@@ -306,6 +304,22 @@ def define_formulation(subdomain: F.VolumeSubdomain):
     b_res = "+"
     t_res = "-"
 
+    # look at the first facet on interface
+    # and get the two cells that are connected to it
+    # and get the material properties of these cells
+    first_facet_interface = mt.find(interface)[0]
+    c_plus, c_minus = (
+        f_to_c.links(first_facet_interface)[0],
+        f_to_c.links(first_facet_interface)[1],
+    )
+    id_minus, id_plus = ct.values[c_minus], ct.values[c_plus]
+
+    for subdomain in list_of_interfaces[interface]:
+        if subdomain.id == id_plus:
+            subdomain_1 = subdomain
+        if subdomain.id == id_minus:
+            subdomain_2 = subdomain
+
     v_b = H.subdomain_to_test_function[subdomain_1](b_res)
     v_t = H.subdomain_to_test_function[subdomain_2](t_res)
 
@@ -323,7 +337,6 @@ def mixed_term(u, v, n):
     h_t = 2 * cr(t_res)
     gamma = 400.0  # this needs to be "sufficiently large"
 
-    # TODO how do we know if b_res corresponds to the bottom or top domain?
     K_b = K_S_fun(T(b_res), subdomain_1.material.K_S_0, subdomain_1.material.E_K_S)
     K_t = K_S_fun(T(t_res), subdomain_2.material.K_S_0, subdomain_2.material.E_K_S)
 

From de79789c3862881ffae87ca1f62a51a6e5a80070 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 13:11:16 +0000
Subject: [PATCH 018/134] simplified code

---
 discontinuity_generic.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index ce731bcd0..ca43331b2 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -224,8 +224,7 @@ def define_formulation(subdomain: F.VolumeSubdomain):
         v = spe.subdomain_to_test_function[subdomain]
         dx = subdomain.dx
 
-        D = D_fun(T, subdomain.material.D_0, subdomain.material.E_D)
-
+        D = subdomain.material.get_diffusion_coefficient(mesh, T, spe)
         if spe.mobile:
             form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
 
@@ -360,17 +359,14 @@ def mixed_term(u, v, n):
 for subdomain1 in list_of_subdomains:
     jac = []
     form = subdomain1.F
-    # copy entity_maps
-    entity_maps_ = entity_maps.copy()
-    entity_maps_[mesh] = subdomain1.submesh_to_mesh
     for subdomain2 in list_of_subdomains:
         jac.append(
             dolfinx.fem.form(
-                ufl.derivative(form, subdomain2.u), entity_maps=entity_maps_
+                ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
             )
         )
     J.append(jac)
-    forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps_))
+    forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
 
 solver = NewtonSolver(

From dab4590507b6c02ed1bfa05c97bf9fc4d219cf00 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 13:45:00 +0000
Subject: [PATCH 019/134] HTransportProblemDiscontinuous class

---
 problem.py             | 327 +++++++++++++++++++++++++++++++++++++++++
 test_implementation.py | 201 +++++++++++++++++++++++++
 2 files changed, 528 insertions(+)
 create mode 100644 problem.py
 create mode 100644 test_implementation.py

diff --git a/problem.py b/problem.py
new file mode 100644
index 000000000..4a7c63a56
--- /dev/null
+++ b/problem.py
@@ -0,0 +1,327 @@
+import festim as F
+import dolfinx
+from dolfinx import fem
+import numpy as np
+import basix
+import ufl
+from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
+
+
+def K_S_fun(T, K_S_0, E_K_S):
+    k_B = 8.6173303e-5
+    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
+
+
+class HTransportProblemDiscontinuous(F.HydrogenTransportProblem):
+
+    def initialise(self):
+        self.create_submeshes()
+        self.create_species_from_traps()
+        self.define_temperature()
+
+        self.entity_maps = {
+            subdomain.submesh: subdomain.parent_to_submesh
+            for subdomain in self.volume_subdomains
+        }
+
+        for subdomain in self.volume_subdomains:
+            self.define_function_spaces(subdomain)
+            ct_r = dolfinx.mesh.meshtags(
+                self.mesh.mesh,
+                self.mesh.mesh.topology.dim,
+                subdomain.submesh_to_mesh,
+                np.full_like(subdomain.submesh_to_mesh, 1, dtype=np.int32),
+            )
+            subdomain.dx = ufl.Measure(
+                "dx", domain=self.mesh.mesh, subdomain_data=ct_r, subdomain_id=1
+            )
+            self.create_subdomain_formulation(subdomain)
+            subdomain.u.name = f"u_{subdomain.id}"
+
+        self.define_meshtags_and_measures()
+        # self.assign_functions_to_species()
+
+        self.t = fem.Constant(self.mesh.mesh, 0.0)
+        if self.settings.transient:
+            # TODO should raise error if no stepsize is provided
+            # TODO Should this be an attribute of festim.Stepsize?
+            self.dt = F.as_fenics_constant(
+                self.settings.stepsize.initial_value, self.mesh.mesh
+            )
+
+        self.define_boundary_conditions()
+        self.create_source_values_fenics()
+        self.create_flux_values_fenics()
+        self.create_initial_conditions()
+        self.create_formulation()
+        self.create_solver()
+        self.initialise_exports()
+
+    def create_dirichletbc_form(self, bc):
+        fdim = self.mesh.mesh.topology.dim - 1
+        volume_subdomain = self.surface_to_volume[bc.subdomain]
+        bc_function = dolfinx.fem.Function(volume_subdomain.u.function_space)
+        bc_function.x.array[:] = bc.value
+        volume_subdomain.submesh.topology.create_connectivity(
+            volume_subdomain.submesh.topology.dim - 1,
+            volume_subdomain.submesh.topology.dim,
+        )
+        form = dolfinx.fem.dirichletbc(
+            bc_function,
+            dolfinx.fem.locate_dofs_topological(
+                volume_subdomain.u.function_space.sub(0),
+                fdim,
+                volume_subdomain.ft.find(bc.subdomain.id),
+            ),
+        )
+        return form
+
+    def create_submeshes(self):
+        mesh = self.mesh.mesh
+        ct = self.volume_meshtags
+        mt = self.facet_meshtags
+
+        gdim = mesh.geometry.dim
+        tdim = mesh.topology.dim
+        fdim = tdim - 1
+
+        num_facets_local = (
+            mesh.topology.index_map(fdim).size_local
+            + mesh.topology.index_map(fdim).num_ghosts
+        )
+
+        for subdomain in self.volume_subdomains:
+            subdomain.submesh, subdomain.submesh_to_mesh, subdomain.v_map = (
+                dolfinx.mesh.create_submesh(mesh, tdim, ct.find(subdomain.id))[0:3]
+            )
+
+            subdomain.parent_to_submesh = np.full(num_facets_local, -1, dtype=np.int32)
+            subdomain.parent_to_submesh[subdomain.submesh_to_mesh] = np.arange(
+                len(subdomain.submesh_to_mesh), dtype=np.int32
+            )
+
+            # We need to modify the cell maps, as for `dS` integrals of interfaces between submeshes, there is no entity to map to.
+            # We use the entity on the same side to fix this (as all restrictions are one-sided)
+
+            # Transfer meshtags to submesh
+            subdomain.ft, subdomain.facet_to_parent = transfer_meshtags_to_submesh(
+                mesh, mt, subdomain.submesh, subdomain.v_map, subdomain.submesh_to_mesh
+            )
+
+        # Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
+        # TODO ask Jorgen what this is for
+        mesh.topology.create_connectivity(fdim, tdim)
+        f_to_c = mesh.topology.connectivity(fdim, tdim)
+        for interface in self.interfaces:
+            for facet in mt.find(interface):
+                cells = f_to_c.links(facet)
+                assert len(cells) == 2
+                for domain in self.interfaces[interface]:
+                    map = domain.parent_to_submesh[cells]
+                    domain.parent_to_submesh[cells] = max(map)
+
+        self.f_to_c = f_to_c
+
+    def define_function_spaces(self, subdomain: F.VolumeSubdomain):
+        # get number of species defined in the subdomain
+        all_species = [
+            species for species in self.species if subdomain in species.subdomains
+        ]
+
+        # instead of using the set function we use a list to keep the order
+        unique_species = []
+        for species in all_species:
+            if species not in unique_species:
+                unique_species.append(species)
+        nb_species = len(unique_species)
+
+        degree = 1
+        element_CG = basix.ufl.element(
+            basix.ElementFamily.P,
+            subdomain.submesh.basix_cell(),
+            degree,
+            basix.LagrangeVariant.equispaced,
+        )
+        element = basix.ufl.mixed_element([element_CG] * nb_species)
+        V = dolfinx.fem.functionspace(subdomain.submesh, element)
+        u = dolfinx.fem.Function(V)
+        u_n = dolfinx.fem.Function(V)
+
+        us = list(ufl.split(u))
+        u_ns = list(ufl.split(u_n))
+        vs = list(ufl.TestFunctions(V))
+        for i, species in enumerate(unique_species):
+            species.subdomain_to_solution[subdomain] = us[i]
+            species.subdomain_to_prev_solution[subdomain] = u_ns[i]
+            species.subdomain_to_test_function[subdomain] = vs[i]
+        subdomain.u = u
+
+    def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
+        form = 0
+        # add diffusion and time derivative for each species
+        for spe in self.species:
+            u = spe.subdomain_to_solution[subdomain]
+            u_n = spe.subdomain_to_prev_solution[subdomain]
+            v = spe.subdomain_to_test_function[subdomain]
+            dx = subdomain.dx
+
+            D = subdomain.material.get_diffusion_coefficient(
+                self.mesh.mesh, self.temperature_fenics, spe
+            )
+            if self.settings.transient:
+                raise NotImplementedError("Transient not implemented")
+
+            if spe.mobile:
+                form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
+
+        for reaction in self.reactions:
+            if reaction.volume != subdomain:
+                continue
+            for species in reaction.reactant + reaction.product:
+                if isinstance(species, F.Species):
+                    # TODO remove
+                    # temporarily overide the solution and test function to the one of the subdomain
+                    species.solution = species.subdomain_to_solution[subdomain]
+                    species.test_function = species.subdomain_to_test_function[
+                        subdomain
+                    ]
+
+            for reactant in reaction.reactant:
+                if isinstance(reactant, F.Species):
+                    form += (
+                        reaction.reaction_term(self.temperature_fenics)
+                        * reactant.subdomain_to_test_function[subdomain]
+                        * dx
+                    )
+
+            # product
+            if isinstance(reaction.product, list):
+                products = reaction.product
+            else:
+                products = [reaction.product]
+            for product in products:
+                form += (
+                    -reaction.reaction_term(self.temperature_fenics)
+                    * product.subdomain_to_test_function[subdomain]
+                    * dx
+                )
+
+        subdomain.F = form
+
+    def create_formulation(self):
+        # Add coupling term to the interface
+        # Get interface markers on submesh b
+        mesh = self.mesh.mesh
+        ct = self.volume_meshtags
+        mt = self.facet_meshtags
+        f_to_c = mesh.topology.connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+
+        for interface in self.interfaces:
+
+            dInterface = ufl.Measure(
+                "dS", domain=mesh, subdomain_data=mt, subdomain_id=interface
+            )
+            b_res = "+"
+            t_res = "-"
+
+            # look at the first facet on interface
+            # and get the two cells that are connected to it
+            # and get the material properties of these cells
+            first_facet_interface = mt.find(interface)[0]
+            c_plus, c_minus = (
+                f_to_c.links(first_facet_interface)[0],
+                f_to_c.links(first_facet_interface)[1],
+            )
+            id_minus, id_plus = ct.values[c_minus], ct.values[c_plus]
+
+            for subdomain in self.interfaces[interface]:
+                if subdomain.id == id_plus:
+                    subdomain_1 = subdomain
+                if subdomain.id == id_minus:
+                    subdomain_2 = subdomain
+
+            all_mobile_species = [spe for spe in self.species if spe.mobile]
+            if len(all_mobile_species) > 1:
+                raise NotImplementedError("Multiple mobile species not implemented")
+            H = all_mobile_species[0]
+
+            v_b = H.subdomain_to_test_function[subdomain_1](b_res)
+            v_t = H.subdomain_to_test_function[subdomain_2](t_res)
+
+            u_b = H.subdomain_to_solution[subdomain_1](b_res)
+            u_t = H.subdomain_to_solution[subdomain_2](t_res)
+
+            def mixed_term(u, v, n):
+                return ufl.dot(ufl.grad(u), n) * v
+
+            n = ufl.FacetNormal(mesh)
+            n_b = n(b_res)
+            n_t = n(t_res)
+            cr = ufl.Circumradius(mesh)
+            h_b = 2 * cr(b_res)
+            h_t = 2 * cr(t_res)
+            gamma = 400.0  # this needs to be "sufficiently large"
+
+            K_b = K_S_fun(
+                self.temperature_fenics(b_res),
+                subdomain_1.material.K_S_0,
+                subdomain_1.material.E_K_S,
+            )
+            K_t = K_S_fun(
+                self.temperature_fenics(t_res),
+                subdomain_2.material.K_S_0,
+                subdomain_2.material.E_K_S,
+            )
+
+            F_0 = (
+                -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
+                - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
+            )
+
+            F_1 = (
+                +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
+                - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
+            )
+            F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
+            F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
+
+            subdomain_1.F += F_0
+            subdomain_2.F += F_1
+
+        J = []
+        forms = []
+        for subdomain1 in self.volume_subdomains:
+            jac = []
+            form = subdomain1.F
+            for subdomain2 in self.volume_subdomains:
+                jac.append(
+                    dolfinx.fem.form(
+                        ufl.derivative(form, subdomain2.u), entity_maps=self.entity_maps
+                    )
+                )
+            J.append(jac)
+            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=self.entity_maps))
+        self.forms = forms
+        self.J = J
+
+    def create_solver(self):
+        self.solver = NewtonSolver(
+            self.forms,
+            self.J,
+            [subdomain.u for subdomain in self.volume_subdomains],
+            bcs=self.bc_forms,
+            max_iterations=10,
+            petsc_options={
+                "ksp_type": "preonly",
+                "pc_type": "lu",
+                "pc_factor_mat_solver_type": "mumps",
+            },
+        )
+
+    def run(self):
+        if self.settings.transient:
+            raise NotImplementedError("Transient not implemented")
+        else:
+            # Solve steady-state
+            self.solver.solve(1e-5)
+            self.post_processing()
diff --git a/test_implementation.py b/test_implementation.py
new file mode 100644
index 000000000..933ec6e58
--- /dev/null
+++ b/test_implementation.py
@@ -0,0 +1,201 @@
+import festim as F
+from problem import HTransportProblemDiscontinuous
+
+from mpi4py import MPI
+import dolfinx
+import dolfinx.fem.petsc
+import numpy as np
+import festim as F
+import ufl
+from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
+import basix
+
+
+def K_S_fun(T, K_S_0, E_K_S):
+    k_B = 8.6173303e-5
+    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
+
+
+# ---------------- Generate a mesh ----------------
+def generate_mesh():
+    def bottom_boundary(x):
+        return np.isclose(x[1], 0.0)
+
+    def top_boundary(x):
+        return np.isclose(x[1], 1.0)
+
+    def half(x):
+        return x[1] <= 0.5 + 1e-14
+
+    mesh = dolfinx.mesh.create_unit_square(
+        MPI.COMM_WORLD, 20, 20, dolfinx.mesh.CellType.triangle
+    )
+
+    # Split domain in half and set an interface tag of 5
+    gdim = mesh.geometry.dim
+    tdim = mesh.topology.dim
+    fdim = tdim - 1
+    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    num_facets_local = (
+        mesh.topology.index_map(fdim).size_local
+        + mesh.topology.index_map(fdim).num_ghosts
+    )
+    facets = np.arange(num_facets_local, dtype=np.int32)
+    values = np.full_like(facets, 0, dtype=np.int32)
+    values[top_facets] = 1
+    values[bottom_facets] = 2
+
+    bottom_cells = dolfinx.mesh.locate_entities(mesh, tdim, half)
+    num_cells_local = (
+        mesh.topology.index_map(tdim).size_local
+        + mesh.topology.index_map(tdim).num_ghosts
+    )
+    cells = np.full(num_cells_local, 4, dtype=np.int32)
+    cells[bottom_cells] = 3
+    ct = dolfinx.mesh.meshtags(
+        mesh, tdim, np.arange(num_cells_local, dtype=np.int32), cells
+    )
+    all_b_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(3), tdim, fdim
+    )
+    all_t_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(4), tdim, fdim
+    )
+    interface = np.intersect1d(all_b_facets, all_t_facets)
+    values[interface] = 5
+
+    mt = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, values)
+    return mesh, mt, ct
+
+
+mesh, mt, ct = generate_mesh()
+
+my_model = HTransportProblemDiscontinuous()
+my_model.mesh = F.Mesh(mesh)
+my_model.volume_meshtags = ct
+my_model.facet_meshtags = mt
+
+material_bottom = F.Material(D_0=2.0, E_D=0.1)
+material_top = F.Material(D_0=2.0, E_D=0.1)
+
+material_bottom.K_S_0 = 2.0
+material_bottom.E_K_S = 0 * 0.1
+material_top.K_S_0 = 4.0
+material_top.E_K_S = 0 * 0.12
+
+top_domain = F.VolumeSubdomain(4, material=material_top)
+bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
+
+# we should be able to automate this
+my_model.interfaces = {5: [bottom_domain, top_domain]}
+
+top_surface = F.SurfaceSubdomain(id=1)
+bottom_surface = F.SurfaceSubdomain(id=2)
+
+my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+
+H = F.Species("H", mobile=True)
+trapped_H = F.Species("H_trapped", mobile=False)
+empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
+
+my_model.species = [H, trapped_H]
+
+for species in [H, trapped_H]:
+    species.subdomains = [bottom_domain, top_domain]
+    species.subdomain_to_solution = {}
+    species.subdomain_to_prev_solution = {}
+    species.subdomain_to_test_function = {}
+
+my_model.reactions = [
+    F.Reaction(
+        reactant=[H, empty_trap],
+        product=[trapped_H],
+        k_0=2,
+        E_k=0,
+        p_0=0.1,
+        E_p=0,
+        volume=top_domain,
+    ),
+    F.Reaction(
+        reactant=[H, empty_trap],
+        product=[trapped_H],
+        k_0=2,
+        E_k=0,
+        p_0=0.1,
+        E_p=0,
+        volume=bottom_domain,
+    ),
+]
+
+my_model.boundary_conditions = [
+    F.DirichletBC(top_surface, value=0.05, species=H),
+    F.DirichletBC(bottom_surface, value=0.2, species=H),
+]
+
+my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+
+my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
+
+my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+
+
+my_model.initialise()
+my_model.run()
+
+list_of_subdomains = my_model.volume_subdomains
+
+for subdomain in list_of_subdomains:
+    u_sub_0 = subdomain.u.sub(0).collapse()
+    u_sub_0.name = "u_sub_0"
+
+    u_sub_1 = subdomain.u.sub(1).collapse()
+    u_sub_1.name = "u_sub_1"
+    bp = dolfinx.io.VTXWriter(
+        mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
+    )
+    bp.write(0)
+    bp.close()
+
+
+# # derived quantities
+# entity_maps[mesh] = bottom_domain.submesh_to_mesh
+
+# ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
+# ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
+# dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
+# dx = ufl.Measure("dx", domain=mesh)
+
+# n_b = ufl.FacetNormal(bottom_domain.submesh)
+# n_t = ufl.FacetNormal(top_domain.submesh)
+
+# form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
+# print(dolfinx.fem.assemble_scalar(form))
+
+# form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b)
+# print(dolfinx.fem.assemble_scalar(form))
+
+# form = dolfinx.fem.form(T * dx_b, entity_maps={mesh: bottom_domain.submesh_to_mesh})
+# print(dolfinx.fem.assemble_scalar(form))
+
+# id_interface = 5
+# form = dolfinx.fem.form(
+#     ufl.dot(
+#         D_fun(T, bottom_domain.material.D_0, bottom_domain.material.E_D)
+#         * ufl.grad(bottom_domain.u.sub(0)),
+#         n_b,
+#     )
+#     * ds_b(id_interface),
+#     entity_maps={mesh: bottom_domain.submesh_to_mesh},
+# )
+# print(dolfinx.fem.assemble_scalar(form))
+# form = dolfinx.fem.form(
+#     ufl.dot(
+#         D_fun(T, top_domain.material.D_0, top_domain.material.E_D)
+#         * ufl.grad(top_domain.u.sub(0)),
+#         n_t,
+#     )
+#     * ds_t(id_interface),
+#     entity_maps={mesh: top_domain.submesh_to_mesh},
+# )
+# print(dolfinx.fem.assemble_scalar(form))

From 6a700b0eac1e3b973e3413b81466c012b0683be0 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 16 Aug 2024 13:49:20 +0000
Subject: [PATCH 020/134] added back the derived quantities + removed unused
 stuff

---
 test_implementation.py | 93 ++++++++++++++++++++----------------------
 1 file changed, 45 insertions(+), 48 deletions(-)

diff --git a/test_implementation.py b/test_implementation.py
index 933ec6e58..34ce1c389 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -7,13 +7,6 @@
 import numpy as np
 import festim as F
 import ufl
-from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
-import basix
-
-
-def K_S_fun(T, K_S_0, E_K_S):
-    k_B = 8.6173303e-5
-    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
 
 
 # ---------------- Generate a mesh ----------------
@@ -158,44 +151,48 @@ def half(x):
     bp.close()
 
 
-# # derived quantities
-# entity_maps[mesh] = bottom_domain.submesh_to_mesh
-
-# ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
-# ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
-# dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
-# dx = ufl.Measure("dx", domain=mesh)
-
-# n_b = ufl.FacetNormal(bottom_domain.submesh)
-# n_t = ufl.FacetNormal(top_domain.submesh)
-
-# form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
-# print(dolfinx.fem.assemble_scalar(form))
-
-# form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b)
-# print(dolfinx.fem.assemble_scalar(form))
-
-# form = dolfinx.fem.form(T * dx_b, entity_maps={mesh: bottom_domain.submesh_to_mesh})
-# print(dolfinx.fem.assemble_scalar(form))
-
-# id_interface = 5
-# form = dolfinx.fem.form(
-#     ufl.dot(
-#         D_fun(T, bottom_domain.material.D_0, bottom_domain.material.E_D)
-#         * ufl.grad(bottom_domain.u.sub(0)),
-#         n_b,
-#     )
-#     * ds_b(id_interface),
-#     entity_maps={mesh: bottom_domain.submesh_to_mesh},
-# )
-# print(dolfinx.fem.assemble_scalar(form))
-# form = dolfinx.fem.form(
-#     ufl.dot(
-#         D_fun(T, top_domain.material.D_0, top_domain.material.E_D)
-#         * ufl.grad(top_domain.u.sub(0)),
-#         n_t,
-#     )
-#     * ds_t(id_interface),
-#     entity_maps={mesh: top_domain.submesh_to_mesh},
-# )
-# print(dolfinx.fem.assemble_scalar(form))
+# derived quantities
+my_model.entity_maps[mesh] = bottom_domain.submesh_to_mesh
+
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
+ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
+dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
+dx = ufl.Measure("dx", domain=mesh)
+
+n_b = ufl.FacetNormal(bottom_domain.submesh)
+n_t = ufl.FacetNormal(top_domain.submesh)
+
+form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
+print(dolfinx.fem.assemble_scalar(form))
+
+form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b)
+print(dolfinx.fem.assemble_scalar(form))
+
+form = dolfinx.fem.form(
+    my_model.temperature_fenics * dx_b,
+    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+)
+print(dolfinx.fem.assemble_scalar(form))
+
+D = subdomain.material.get_diffusion_coefficient(
+    my_model.mesh.mesh, my_model.temperature_fenics, H
+)
+id_interface = 5
+form = dolfinx.fem.form(
+    ufl.dot(
+        D * ufl.grad(bottom_domain.u.sub(0)),
+        n_b,
+    )
+    * ds_b(id_interface),
+    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+)
+print(dolfinx.fem.assemble_scalar(form))
+form = dolfinx.fem.form(
+    ufl.dot(
+        D * ufl.grad(top_domain.u.sub(0)),
+        n_t,
+    )
+    * ds_t(id_interface),
+    entity_maps={mesh: top_domain.submesh_to_mesh},
+)
+print(dolfinx.fem.assemble_scalar(form))

From 34f82e3a89298f033fa7235833e9bdd0aa251329 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 07:47:25 +0000
Subject: [PATCH 021/134] correspondance dicts are initialised in species

---
 festim/species.py      | 4 ++++
 test_implementation.py | 4 +---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/festim/species.py b/festim/species.py
index 04ac3f25f..19167642e 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -47,6 +47,10 @@ def __init__(self, name: str = None, mobile=True) -> None:
         self.post_processing_solution = None
         self.collapsed_function_space = None
 
+        self.subdomain_to_solution = {}
+        self.subdomain_to_prev_solution = {}
+        self.subdomain_to_test_function = {}
+
     def __repr__(self) -> str:
         return f"Species({self.name})"
 
diff --git a/test_implementation.py b/test_implementation.py
index 34ce1c389..d1b698d0a 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -96,9 +96,7 @@ def half(x):
 
 for species in [H, trapped_H]:
     species.subdomains = [bottom_domain, top_domain]
-    species.subdomain_to_solution = {}
-    species.subdomain_to_prev_solution = {}
-    species.subdomain_to_test_function = {}
+
 
 my_model.reactions = [
     F.Reaction(

From e5c6793a9589dfbf7f9c0d0852f481c171989543 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 07:56:50 +0000
Subject: [PATCH 022/134] add subdomains arg and attribute to Species +
 documentation

---
 festim/__init__.py |  4 ++--
 festim/species.py  | 14 +++++++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/festim/__init__.py b/festim/__init__.py
index a294d9752..9c8120942 100644
--- a/festim/__init__.py
+++ b/festim/__init__.py
@@ -43,13 +43,13 @@
 
 from .source import ParticleSource, HeatSource, SourceBase
 
-from .species import Species, Trap, ImplicitSpecies, find_species_from_name
-
 from .subdomain.surface_subdomain import SurfaceSubdomain, find_surface_from_id
 from .subdomain.surface_subdomain_1d import SurfaceSubdomain1D
 from .subdomain.volume_subdomain import VolumeSubdomain, find_volume_from_id
 from .subdomain.volume_subdomain_1d import VolumeSubdomain1D
 
+from .species import Species, Trap, ImplicitSpecies, find_species_from_name
+
 from .stepsize import Stepsize
 
 from .exports.surface_quantity import SurfaceQuantity
diff --git a/festim/species.py b/festim/species.py
index 19167642e..bea839a51 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -9,6 +9,9 @@ class Species:
     Args:
         name (str, optional): a name given to the species. Defaults to None.
         mobile (bool, optional): whether the species is mobile or not.
+            Defaults to True.
+        subdomain (F.VolumeSubdomain, optional): the volume subdomain where the
+            species is. Defaults to None.
 
     Attributes:
         name (str): a name given to the species.
@@ -26,6 +29,12 @@ class Species:
         post_processing_solution (dolfinx.fem.Function): the solution for post
             processing
         concentration (dolfinx.fem.Function): the concentration of the species
+        subdomains (F.VolumeSubdomain): the volume subdomains where the species is
+        subdomain_to_solution (dict): a dictionary mapping subdomains to solutions
+        subdomain_to_prev_solution (dict): a dictionary mapping subdomains to
+            previous solutions
+        subdomain_to_test_function (dict): a dictionary mapping subdomains to
+            test functions
 
     Usage:
         >>> from festim import Species, HTransportProblem
@@ -37,7 +46,9 @@ class Species:
 
     """
 
-    def __init__(self, name: str = None, mobile=True) -> None:
+    def __init__(
+        self, name: str = None, mobile=True, subdomains: F.VolumeSubdomain = None
+    ) -> None:
         self.name = name
         self.mobile = mobile
         self.solution = None
@@ -47,6 +58,7 @@ def __init__(self, name: str = None, mobile=True) -> None:
         self.post_processing_solution = None
         self.collapsed_function_space = None
 
+        self.subdomains = subdomains
         self.subdomain_to_solution = {}
         self.subdomain_to_prev_solution = {}
         self.subdomain_to_test_function = {}

From ce31d38ba3eb046858aa760296d740632099e9bc Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 07:57:28 +0000
Subject: [PATCH 023/134] gather connectivity dicts

---
 test_implementation.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/test_implementation.py b/test_implementation.py
index d1b698d0a..a49ce6ab7 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -80,14 +80,14 @@ def half(x):
 top_domain = F.VolumeSubdomain(4, material=material_top)
 bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
 
-# we should be able to automate this
-my_model.interfaces = {5: [bottom_domain, top_domain]}
-
 top_surface = F.SurfaceSubdomain(id=1)
 bottom_surface = F.SurfaceSubdomain(id=2)
-
 my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
+# we should be able to automate this
+my_model.interfaces = {5: [bottom_domain, top_domain]}
+my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
 empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
@@ -124,7 +124,6 @@ def half(x):
     F.DirichletBC(bottom_surface, value=0.2, species=H),
 ]
 
-my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
 

From 4f075beff6358f6de3ac2b63ae7b87ab9dc33f75 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 07:57:46 +0000
Subject: [PATCH 024/134] NotImplementedError for exports for now

---
 problem.py | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/problem.py b/problem.py
index 4a7c63a56..20ed63d0f 100644
--- a/problem.py
+++ b/problem.py
@@ -14,6 +14,38 @@ def K_S_fun(T, K_S_0, E_K_S):
 
 class HTransportProblemDiscontinuous(F.HydrogenTransportProblem):
 
+    def __init__(
+        self,
+        mesh=None,
+        subdomains=None,
+        species=None,
+        reactions=None,
+        temperature=None,
+        sources=None,
+        initial_conditions=None,
+        boundary_conditions=None,
+        settings=None,
+        exports=None,
+        traps=None,
+        interfaces=None,
+        surface_to_volume=None,
+    ):
+        super().__init__(
+            mesh,
+            subdomains,
+            species,
+            reactions,
+            temperature,
+            sources,
+            initial_conditions,
+            boundary_conditions,
+            settings,
+            exports,
+            traps,
+        )
+        self.interfaces = interfaces or {}
+        self.surface_to_volume = surface_to_volume or {}
+
     def initialise(self):
         self.create_submeshes()
         self.create_species_from_traps()
@@ -318,6 +350,12 @@ def create_solver(self):
             },
         )
 
+    def initialise_exports(self):
+        if self.exports:
+            raise NotImplementedError(
+                "Exports not implemented for HTransportProblemDiscontinuous"
+            )
+
     def run(self):
         if self.settings.transient:
             raise NotImplementedError("Transient not implemented")

From ea3da5da138966451dc0404ec54dc39d497830b0 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 08:14:01 +0000
Subject: [PATCH 025/134] Refactor code to import HydrogenTransportProblem and
 HTransportProblemDiscontinuous from the correct module

---
 festim/__init__.py                   |   8 +-
 festim/hydrogen_transport_problem.py | 359 +++++++++++++++++++++++++++
 2 files changed, 365 insertions(+), 2 deletions(-)

diff --git a/festim/__init__.py b/festim/__init__.py
index 9c8120942..c20a44352 100644
--- a/festim/__init__.py
+++ b/festim/__init__.py
@@ -34,8 +34,6 @@
 from .mesh.mesh_1d import Mesh1D
 from .mesh.mesh_from_xdmf import MeshFromXDMF
 
-from .hydrogen_transport_problem import HydrogenTransportProblem
-from .heat_transfer_problem import HeatTransferProblem
 
 from .initial_condition import InitialCondition, InitialTemperature
 
@@ -60,3 +58,9 @@
 from .exports.xdmf import XDMFExport
 
 from .reaction import Reaction
+
+from .hydrogen_transport_problem import (
+    HydrogenTransportProblem,
+    HTransportProblemDiscontinuous,
+)
+from .heat_transfer_problem import HeatTransferProblem
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 7207cac54..5710cb7aa 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -8,6 +8,7 @@
 import numpy as np
 import tqdm.auto as tqdm
 import festim as F
+from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
 
 
 class HydrogenTransportProblem:
@@ -806,3 +807,361 @@ def post_processing(self):
                     export.write(t=float(self.t))
             if isinstance(export, (F.VTXExport, F.XDMFExport)):
                 export.write(float(self.t))
+
+
+def K_S_fun(T, K_S_0, E_K_S):
+    k_B = 8.6173303e-5
+    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
+
+
+class HTransportProblemDiscontinuous(HydrogenTransportProblem):
+
+    def __init__(
+        self,
+        mesh=None,
+        subdomains=None,
+        species=None,
+        reactions=None,
+        temperature=None,
+        sources=None,
+        initial_conditions=None,
+        boundary_conditions=None,
+        settings=None,
+        exports=None,
+        traps=None,
+        interfaces=None,
+        surface_to_volume=None,
+    ):
+        super().__init__(
+            mesh,
+            subdomains,
+            species,
+            reactions,
+            temperature,
+            sources,
+            initial_conditions,
+            boundary_conditions,
+            settings,
+            exports,
+            traps,
+        )
+        self.interfaces = interfaces or {}
+        self.surface_to_volume = surface_to_volume or {}
+
+    def initialise(self):
+        self.create_submeshes()
+        self.create_species_from_traps()
+        self.define_temperature()
+
+        self.entity_maps = {
+            subdomain.submesh: subdomain.parent_to_submesh
+            for subdomain in self.volume_subdomains
+        }
+
+        for subdomain in self.volume_subdomains:
+            self.define_function_spaces(subdomain)
+            ct_r = dolfinx.mesh.meshtags(
+                self.mesh.mesh,
+                self.mesh.mesh.topology.dim,
+                subdomain.submesh_to_mesh,
+                np.full_like(subdomain.submesh_to_mesh, 1, dtype=np.int32),
+            )
+            subdomain.dx = ufl.Measure(
+                "dx", domain=self.mesh.mesh, subdomain_data=ct_r, subdomain_id=1
+            )
+            self.create_subdomain_formulation(subdomain)
+            subdomain.u.name = f"u_{subdomain.id}"
+
+        self.define_meshtags_and_measures()
+        # self.assign_functions_to_species()
+
+        self.t = fem.Constant(self.mesh.mesh, 0.0)
+        if self.settings.transient:
+            # TODO should raise error if no stepsize is provided
+            # TODO Should this be an attribute of festim.Stepsize?
+            self.dt = F.as_fenics_constant(
+                self.settings.stepsize.initial_value, self.mesh.mesh
+            )
+
+        self.define_boundary_conditions()
+        self.create_source_values_fenics()
+        self.create_flux_values_fenics()
+        self.create_initial_conditions()
+        self.create_formulation()
+        self.create_solver()
+        self.initialise_exports()
+
+    def create_dirichletbc_form(self, bc):
+        fdim = self.mesh.mesh.topology.dim - 1
+        volume_subdomain = self.surface_to_volume[bc.subdomain]
+        bc_function = dolfinx.fem.Function(volume_subdomain.u.function_space)
+        bc_function.x.array[:] = bc.value
+        volume_subdomain.submesh.topology.create_connectivity(
+            volume_subdomain.submesh.topology.dim - 1,
+            volume_subdomain.submesh.topology.dim,
+        )
+        form = dolfinx.fem.dirichletbc(
+            bc_function,
+            dolfinx.fem.locate_dofs_topological(
+                volume_subdomain.u.function_space.sub(0),
+                fdim,
+                volume_subdomain.ft.find(bc.subdomain.id),
+            ),
+        )
+        return form
+
+    def create_submeshes(self):
+        mesh = self.mesh.mesh
+        ct = self.volume_meshtags
+        mt = self.facet_meshtags
+
+        gdim = mesh.geometry.dim
+        tdim = mesh.topology.dim
+        fdim = tdim - 1
+
+        num_facets_local = (
+            mesh.topology.index_map(fdim).size_local
+            + mesh.topology.index_map(fdim).num_ghosts
+        )
+
+        for subdomain in self.volume_subdomains:
+            subdomain.submesh, subdomain.submesh_to_mesh, subdomain.v_map = (
+                dolfinx.mesh.create_submesh(mesh, tdim, ct.find(subdomain.id))[0:3]
+            )
+
+            subdomain.parent_to_submesh = np.full(num_facets_local, -1, dtype=np.int32)
+            subdomain.parent_to_submesh[subdomain.submesh_to_mesh] = np.arange(
+                len(subdomain.submesh_to_mesh), dtype=np.int32
+            )
+
+            # We need to modify the cell maps, as for `dS` integrals of interfaces between submeshes, there is no entity to map to.
+            # We use the entity on the same side to fix this (as all restrictions are one-sided)
+
+            # Transfer meshtags to submesh
+            subdomain.ft, subdomain.facet_to_parent = transfer_meshtags_to_submesh(
+                mesh, mt, subdomain.submesh, subdomain.v_map, subdomain.submesh_to_mesh
+            )
+
+        # Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
+        # TODO ask Jorgen what this is for
+        mesh.topology.create_connectivity(fdim, tdim)
+        f_to_c = mesh.topology.connectivity(fdim, tdim)
+        for interface in self.interfaces:
+            for facet in mt.find(interface):
+                cells = f_to_c.links(facet)
+                assert len(cells) == 2
+                for domain in self.interfaces[interface]:
+                    map = domain.parent_to_submesh[cells]
+                    domain.parent_to_submesh[cells] = max(map)
+
+        self.f_to_c = f_to_c
+
+    def define_function_spaces(self, subdomain: F.VolumeSubdomain):
+        # get number of species defined in the subdomain
+        all_species = [
+            species for species in self.species if subdomain in species.subdomains
+        ]
+
+        # instead of using the set function we use a list to keep the order
+        unique_species = []
+        for species in all_species:
+            if species not in unique_species:
+                unique_species.append(species)
+        nb_species = len(unique_species)
+
+        degree = 1
+        element_CG = basix.ufl.element(
+            basix.ElementFamily.P,
+            subdomain.submesh.basix_cell(),
+            degree,
+            basix.LagrangeVariant.equispaced,
+        )
+        element = basix.ufl.mixed_element([element_CG] * nb_species)
+        V = dolfinx.fem.functionspace(subdomain.submesh, element)
+        u = dolfinx.fem.Function(V)
+        u_n = dolfinx.fem.Function(V)
+
+        us = list(ufl.split(u))
+        u_ns = list(ufl.split(u_n))
+        vs = list(ufl.TestFunctions(V))
+        for i, species in enumerate(unique_species):
+            species.subdomain_to_solution[subdomain] = us[i]
+            species.subdomain_to_prev_solution[subdomain] = u_ns[i]
+            species.subdomain_to_test_function[subdomain] = vs[i]
+        subdomain.u = u
+
+    def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
+        form = 0
+        # add diffusion and time derivative for each species
+        for spe in self.species:
+            u = spe.subdomain_to_solution[subdomain]
+            u_n = spe.subdomain_to_prev_solution[subdomain]
+            v = spe.subdomain_to_test_function[subdomain]
+            dx = subdomain.dx
+
+            D = subdomain.material.get_diffusion_coefficient(
+                self.mesh.mesh, self.temperature_fenics, spe
+            )
+            if self.settings.transient:
+                raise NotImplementedError("Transient not implemented")
+
+            if spe.mobile:
+                form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
+
+        for reaction in self.reactions:
+            if reaction.volume != subdomain:
+                continue
+            for species in reaction.reactant + reaction.product:
+                if isinstance(species, F.Species):
+                    # TODO remove
+                    # temporarily overide the solution and test function to the one of the subdomain
+                    species.solution = species.subdomain_to_solution[subdomain]
+                    species.test_function = species.subdomain_to_test_function[
+                        subdomain
+                    ]
+
+            for reactant in reaction.reactant:
+                if isinstance(reactant, F.Species):
+                    form += (
+                        reaction.reaction_term(self.temperature_fenics)
+                        * reactant.subdomain_to_test_function[subdomain]
+                        * dx
+                    )
+
+            # product
+            if isinstance(reaction.product, list):
+                products = reaction.product
+            else:
+                products = [reaction.product]
+            for product in products:
+                form += (
+                    -reaction.reaction_term(self.temperature_fenics)
+                    * product.subdomain_to_test_function[subdomain]
+                    * dx
+                )
+
+        subdomain.F = form
+
+    def create_formulation(self):
+        # Add coupling term to the interface
+        # Get interface markers on submesh b
+        mesh = self.mesh.mesh
+        ct = self.volume_meshtags
+        mt = self.facet_meshtags
+        f_to_c = mesh.topology.connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+
+        for interface in self.interfaces:
+
+            dInterface = ufl.Measure(
+                "dS", domain=mesh, subdomain_data=mt, subdomain_id=interface
+            )
+            b_res = "+"
+            t_res = "-"
+
+            # look at the first facet on interface
+            # and get the two cells that are connected to it
+            # and get the material properties of these cells
+            first_facet_interface = mt.find(interface)[0]
+            c_plus, c_minus = (
+                f_to_c.links(first_facet_interface)[0],
+                f_to_c.links(first_facet_interface)[1],
+            )
+            id_minus, id_plus = ct.values[c_minus], ct.values[c_plus]
+
+            for subdomain in self.interfaces[interface]:
+                if subdomain.id == id_plus:
+                    subdomain_1 = subdomain
+                if subdomain.id == id_minus:
+                    subdomain_2 = subdomain
+
+            all_mobile_species = [spe for spe in self.species if spe.mobile]
+            if len(all_mobile_species) > 1:
+                raise NotImplementedError("Multiple mobile species not implemented")
+            H = all_mobile_species[0]
+
+            v_b = H.subdomain_to_test_function[subdomain_1](b_res)
+            v_t = H.subdomain_to_test_function[subdomain_2](t_res)
+
+            u_b = H.subdomain_to_solution[subdomain_1](b_res)
+            u_t = H.subdomain_to_solution[subdomain_2](t_res)
+
+            def mixed_term(u, v, n):
+                return ufl.dot(ufl.grad(u), n) * v
+
+            n = ufl.FacetNormal(mesh)
+            n_b = n(b_res)
+            n_t = n(t_res)
+            cr = ufl.Circumradius(mesh)
+            h_b = 2 * cr(b_res)
+            h_t = 2 * cr(t_res)
+            gamma = 400.0  # this needs to be "sufficiently large"
+
+            K_b = K_S_fun(
+                self.temperature_fenics(b_res),
+                subdomain_1.material.K_S_0,
+                subdomain_1.material.E_K_S,
+            )
+            K_t = K_S_fun(
+                self.temperature_fenics(t_res),
+                subdomain_2.material.K_S_0,
+                subdomain_2.material.E_K_S,
+            )
+
+            F_0 = (
+                -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
+                - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
+            )
+
+            F_1 = (
+                +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
+                - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
+            )
+            F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
+            F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
+
+            subdomain_1.F += F_0
+            subdomain_2.F += F_1
+
+        J = []
+        forms = []
+        for subdomain1 in self.volume_subdomains:
+            jac = []
+            form = subdomain1.F
+            for subdomain2 in self.volume_subdomains:
+                jac.append(
+                    dolfinx.fem.form(
+                        ufl.derivative(form, subdomain2.u), entity_maps=self.entity_maps
+                    )
+                )
+            J.append(jac)
+            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=self.entity_maps))
+        self.forms = forms
+        self.J = J
+
+    def create_solver(self):
+        self.solver = NewtonSolver(
+            self.forms,
+            self.J,
+            [subdomain.u for subdomain in self.volume_subdomains],
+            bcs=self.bc_forms,
+            max_iterations=10,
+            petsc_options={
+                "ksp_type": "preonly",
+                "pc_type": "lu",
+                "pc_factor_mat_solver_type": "mumps",
+            },
+        )
+
+    def initialise_exports(self):
+        if self.exports:
+            raise NotImplementedError(
+                "Exports not implemented for HTransportProblemDiscontinuous"
+            )
+
+    def run(self):
+        if self.settings.transient:
+            raise NotImplementedError("Transient not implemented")
+        else:
+            # Solve steady-state
+            self.solver.solve(1e-5)
+            self.post_processing()

From 76d1ac26ec73814e76c0cafa39179db63030e2b8 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 08:15:09 +0000
Subject: [PATCH 026/134] remove problem file + comment

---
 problem.py             | 365 -----------------------------------------
 test_implementation.py |   5 +-
 2 files changed, 3 insertions(+), 367 deletions(-)
 delete mode 100644 problem.py

diff --git a/problem.py b/problem.py
deleted file mode 100644
index 20ed63d0f..000000000
--- a/problem.py
+++ /dev/null
@@ -1,365 +0,0 @@
-import festim as F
-import dolfinx
-from dolfinx import fem
-import numpy as np
-import basix
-import ufl
-from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
-
-
-def K_S_fun(T, K_S_0, E_K_S):
-    k_B = 8.6173303e-5
-    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
-
-
-class HTransportProblemDiscontinuous(F.HydrogenTransportProblem):
-
-    def __init__(
-        self,
-        mesh=None,
-        subdomains=None,
-        species=None,
-        reactions=None,
-        temperature=None,
-        sources=None,
-        initial_conditions=None,
-        boundary_conditions=None,
-        settings=None,
-        exports=None,
-        traps=None,
-        interfaces=None,
-        surface_to_volume=None,
-    ):
-        super().__init__(
-            mesh,
-            subdomains,
-            species,
-            reactions,
-            temperature,
-            sources,
-            initial_conditions,
-            boundary_conditions,
-            settings,
-            exports,
-            traps,
-        )
-        self.interfaces = interfaces or {}
-        self.surface_to_volume = surface_to_volume or {}
-
-    def initialise(self):
-        self.create_submeshes()
-        self.create_species_from_traps()
-        self.define_temperature()
-
-        self.entity_maps = {
-            subdomain.submesh: subdomain.parent_to_submesh
-            for subdomain in self.volume_subdomains
-        }
-
-        for subdomain in self.volume_subdomains:
-            self.define_function_spaces(subdomain)
-            ct_r = dolfinx.mesh.meshtags(
-                self.mesh.mesh,
-                self.mesh.mesh.topology.dim,
-                subdomain.submesh_to_mesh,
-                np.full_like(subdomain.submesh_to_mesh, 1, dtype=np.int32),
-            )
-            subdomain.dx = ufl.Measure(
-                "dx", domain=self.mesh.mesh, subdomain_data=ct_r, subdomain_id=1
-            )
-            self.create_subdomain_formulation(subdomain)
-            subdomain.u.name = f"u_{subdomain.id}"
-
-        self.define_meshtags_and_measures()
-        # self.assign_functions_to_species()
-
-        self.t = fem.Constant(self.mesh.mesh, 0.0)
-        if self.settings.transient:
-            # TODO should raise error if no stepsize is provided
-            # TODO Should this be an attribute of festim.Stepsize?
-            self.dt = F.as_fenics_constant(
-                self.settings.stepsize.initial_value, self.mesh.mesh
-            )
-
-        self.define_boundary_conditions()
-        self.create_source_values_fenics()
-        self.create_flux_values_fenics()
-        self.create_initial_conditions()
-        self.create_formulation()
-        self.create_solver()
-        self.initialise_exports()
-
-    def create_dirichletbc_form(self, bc):
-        fdim = self.mesh.mesh.topology.dim - 1
-        volume_subdomain = self.surface_to_volume[bc.subdomain]
-        bc_function = dolfinx.fem.Function(volume_subdomain.u.function_space)
-        bc_function.x.array[:] = bc.value
-        volume_subdomain.submesh.topology.create_connectivity(
-            volume_subdomain.submesh.topology.dim - 1,
-            volume_subdomain.submesh.topology.dim,
-        )
-        form = dolfinx.fem.dirichletbc(
-            bc_function,
-            dolfinx.fem.locate_dofs_topological(
-                volume_subdomain.u.function_space.sub(0),
-                fdim,
-                volume_subdomain.ft.find(bc.subdomain.id),
-            ),
-        )
-        return form
-
-    def create_submeshes(self):
-        mesh = self.mesh.mesh
-        ct = self.volume_meshtags
-        mt = self.facet_meshtags
-
-        gdim = mesh.geometry.dim
-        tdim = mesh.topology.dim
-        fdim = tdim - 1
-
-        num_facets_local = (
-            mesh.topology.index_map(fdim).size_local
-            + mesh.topology.index_map(fdim).num_ghosts
-        )
-
-        for subdomain in self.volume_subdomains:
-            subdomain.submesh, subdomain.submesh_to_mesh, subdomain.v_map = (
-                dolfinx.mesh.create_submesh(mesh, tdim, ct.find(subdomain.id))[0:3]
-            )
-
-            subdomain.parent_to_submesh = np.full(num_facets_local, -1, dtype=np.int32)
-            subdomain.parent_to_submesh[subdomain.submesh_to_mesh] = np.arange(
-                len(subdomain.submesh_to_mesh), dtype=np.int32
-            )
-
-            # We need to modify the cell maps, as for `dS` integrals of interfaces between submeshes, there is no entity to map to.
-            # We use the entity on the same side to fix this (as all restrictions are one-sided)
-
-            # Transfer meshtags to submesh
-            subdomain.ft, subdomain.facet_to_parent = transfer_meshtags_to_submesh(
-                mesh, mt, subdomain.submesh, subdomain.v_map, subdomain.submesh_to_mesh
-            )
-
-        # Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
-        # TODO ask Jorgen what this is for
-        mesh.topology.create_connectivity(fdim, tdim)
-        f_to_c = mesh.topology.connectivity(fdim, tdim)
-        for interface in self.interfaces:
-            for facet in mt.find(interface):
-                cells = f_to_c.links(facet)
-                assert len(cells) == 2
-                for domain in self.interfaces[interface]:
-                    map = domain.parent_to_submesh[cells]
-                    domain.parent_to_submesh[cells] = max(map)
-
-        self.f_to_c = f_to_c
-
-    def define_function_spaces(self, subdomain: F.VolumeSubdomain):
-        # get number of species defined in the subdomain
-        all_species = [
-            species for species in self.species if subdomain in species.subdomains
-        ]
-
-        # instead of using the set function we use a list to keep the order
-        unique_species = []
-        for species in all_species:
-            if species not in unique_species:
-                unique_species.append(species)
-        nb_species = len(unique_species)
-
-        degree = 1
-        element_CG = basix.ufl.element(
-            basix.ElementFamily.P,
-            subdomain.submesh.basix_cell(),
-            degree,
-            basix.LagrangeVariant.equispaced,
-        )
-        element = basix.ufl.mixed_element([element_CG] * nb_species)
-        V = dolfinx.fem.functionspace(subdomain.submesh, element)
-        u = dolfinx.fem.Function(V)
-        u_n = dolfinx.fem.Function(V)
-
-        us = list(ufl.split(u))
-        u_ns = list(ufl.split(u_n))
-        vs = list(ufl.TestFunctions(V))
-        for i, species in enumerate(unique_species):
-            species.subdomain_to_solution[subdomain] = us[i]
-            species.subdomain_to_prev_solution[subdomain] = u_ns[i]
-            species.subdomain_to_test_function[subdomain] = vs[i]
-        subdomain.u = u
-
-    def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
-        form = 0
-        # add diffusion and time derivative for each species
-        for spe in self.species:
-            u = spe.subdomain_to_solution[subdomain]
-            u_n = spe.subdomain_to_prev_solution[subdomain]
-            v = spe.subdomain_to_test_function[subdomain]
-            dx = subdomain.dx
-
-            D = subdomain.material.get_diffusion_coefficient(
-                self.mesh.mesh, self.temperature_fenics, spe
-            )
-            if self.settings.transient:
-                raise NotImplementedError("Transient not implemented")
-
-            if spe.mobile:
-                form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
-
-        for reaction in self.reactions:
-            if reaction.volume != subdomain:
-                continue
-            for species in reaction.reactant + reaction.product:
-                if isinstance(species, F.Species):
-                    # TODO remove
-                    # temporarily overide the solution and test function to the one of the subdomain
-                    species.solution = species.subdomain_to_solution[subdomain]
-                    species.test_function = species.subdomain_to_test_function[
-                        subdomain
-                    ]
-
-            for reactant in reaction.reactant:
-                if isinstance(reactant, F.Species):
-                    form += (
-                        reaction.reaction_term(self.temperature_fenics)
-                        * reactant.subdomain_to_test_function[subdomain]
-                        * dx
-                    )
-
-            # product
-            if isinstance(reaction.product, list):
-                products = reaction.product
-            else:
-                products = [reaction.product]
-            for product in products:
-                form += (
-                    -reaction.reaction_term(self.temperature_fenics)
-                    * product.subdomain_to_test_function[subdomain]
-                    * dx
-                )
-
-        subdomain.F = form
-
-    def create_formulation(self):
-        # Add coupling term to the interface
-        # Get interface markers on submesh b
-        mesh = self.mesh.mesh
-        ct = self.volume_meshtags
-        mt = self.facet_meshtags
-        f_to_c = mesh.topology.connectivity(mesh.topology.dim - 1, mesh.topology.dim)
-
-        for interface in self.interfaces:
-
-            dInterface = ufl.Measure(
-                "dS", domain=mesh, subdomain_data=mt, subdomain_id=interface
-            )
-            b_res = "+"
-            t_res = "-"
-
-            # look at the first facet on interface
-            # and get the two cells that are connected to it
-            # and get the material properties of these cells
-            first_facet_interface = mt.find(interface)[0]
-            c_plus, c_minus = (
-                f_to_c.links(first_facet_interface)[0],
-                f_to_c.links(first_facet_interface)[1],
-            )
-            id_minus, id_plus = ct.values[c_minus], ct.values[c_plus]
-
-            for subdomain in self.interfaces[interface]:
-                if subdomain.id == id_plus:
-                    subdomain_1 = subdomain
-                if subdomain.id == id_minus:
-                    subdomain_2 = subdomain
-
-            all_mobile_species = [spe for spe in self.species if spe.mobile]
-            if len(all_mobile_species) > 1:
-                raise NotImplementedError("Multiple mobile species not implemented")
-            H = all_mobile_species[0]
-
-            v_b = H.subdomain_to_test_function[subdomain_1](b_res)
-            v_t = H.subdomain_to_test_function[subdomain_2](t_res)
-
-            u_b = H.subdomain_to_solution[subdomain_1](b_res)
-            u_t = H.subdomain_to_solution[subdomain_2](t_res)
-
-            def mixed_term(u, v, n):
-                return ufl.dot(ufl.grad(u), n) * v
-
-            n = ufl.FacetNormal(mesh)
-            n_b = n(b_res)
-            n_t = n(t_res)
-            cr = ufl.Circumradius(mesh)
-            h_b = 2 * cr(b_res)
-            h_t = 2 * cr(t_res)
-            gamma = 400.0  # this needs to be "sufficiently large"
-
-            K_b = K_S_fun(
-                self.temperature_fenics(b_res),
-                subdomain_1.material.K_S_0,
-                subdomain_1.material.E_K_S,
-            )
-            K_t = K_S_fun(
-                self.temperature_fenics(t_res),
-                subdomain_2.material.K_S_0,
-                subdomain_2.material.E_K_S,
-            )
-
-            F_0 = (
-                -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
-                - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
-            )
-
-            F_1 = (
-                +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
-                - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
-            )
-            F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
-            F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
-
-            subdomain_1.F += F_0
-            subdomain_2.F += F_1
-
-        J = []
-        forms = []
-        for subdomain1 in self.volume_subdomains:
-            jac = []
-            form = subdomain1.F
-            for subdomain2 in self.volume_subdomains:
-                jac.append(
-                    dolfinx.fem.form(
-                        ufl.derivative(form, subdomain2.u), entity_maps=self.entity_maps
-                    )
-                )
-            J.append(jac)
-            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=self.entity_maps))
-        self.forms = forms
-        self.J = J
-
-    def create_solver(self):
-        self.solver = NewtonSolver(
-            self.forms,
-            self.J,
-            [subdomain.u for subdomain in self.volume_subdomains],
-            bcs=self.bc_forms,
-            max_iterations=10,
-            petsc_options={
-                "ksp_type": "preonly",
-                "pc_type": "lu",
-                "pc_factor_mat_solver_type": "mumps",
-            },
-        )
-
-    def initialise_exports(self):
-        if self.exports:
-            raise NotImplementedError(
-                "Exports not implemented for HTransportProblemDiscontinuous"
-            )
-
-    def run(self):
-        if self.settings.transient:
-            raise NotImplementedError("Transient not implemented")
-        else:
-            # Solve steady-state
-            self.solver.solve(1e-5)
-            self.post_processing()
diff --git a/test_implementation.py b/test_implementation.py
index a49ce6ab7..b494ba120 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -1,5 +1,4 @@
 import festim as F
-from problem import HTransportProblemDiscontinuous
 
 from mpi4py import MPI
 import dolfinx
@@ -64,7 +63,7 @@ def half(x):
 
 mesh, mt, ct = generate_mesh()
 
-my_model = HTransportProblemDiscontinuous()
+my_model = F.HTransportProblemDiscontinuous()
 my_model.mesh = F.Mesh(mesh)
 my_model.volume_meshtags = ct
 my_model.facet_meshtags = mt
@@ -133,6 +132,8 @@ def half(x):
 my_model.initialise()
 my_model.run()
 
+# -------------------- post processing --------------------
+
 list_of_subdomains = my_model.volume_subdomains
 
 for subdomain in list_of_subdomains:

From 1665bcfc567bebf1c4243a39542276b51144629b Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 08:56:20 +0000
Subject: [PATCH 027/134] need to define the correct ds

---
 festim/hydrogen_transport_problem.py |  42 ++++---
 test_implementation_2.py             | 169 +++++++++++++++++++++++++++
 2 files changed, 197 insertions(+), 14 deletions(-)
 create mode 100644 test_implementation_2.py

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 5710cb7aa..2dc6444ef 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -851,6 +851,10 @@ def __init__(
     def initialise(self):
         self.create_submeshes()
         self.create_species_from_traps()
+
+        self.t = fem.Constant(self.mesh.mesh, 0.0)
+        if self.settings.transient:
+            raise NotImplementedError("Transient simulations not supported yet")
         self.define_temperature()
 
         self.entity_maps = {
@@ -858,6 +862,10 @@ def initialise(self):
             for subdomain in self.volume_subdomains
         }
 
+        self.create_source_values_fenics()
+        self.create_flux_values_fenics()
+        self.create_initial_conditions()
+
         for subdomain in self.volume_subdomains:
             self.define_function_spaces(subdomain)
             ct_r = dolfinx.mesh.meshtags(
@@ -869,24 +877,12 @@ def initialise(self):
             subdomain.dx = ufl.Measure(
                 "dx", domain=self.mesh.mesh, subdomain_data=ct_r, subdomain_id=1
             )
+            subdomain.ds = None
             self.create_subdomain_formulation(subdomain)
             subdomain.u.name = f"u_{subdomain.id}"
 
-        self.define_meshtags_and_measures()
-        # self.assign_functions_to_species()
-
-        self.t = fem.Constant(self.mesh.mesh, 0.0)
-        if self.settings.transient:
-            # TODO should raise error if no stepsize is provided
-            # TODO Should this be an attribute of festim.Stepsize?
-            self.dt = F.as_fenics_constant(
-                self.settings.stepsize.initial_value, self.mesh.mesh
-            )
-
         self.define_boundary_conditions()
-        self.create_source_values_fenics()
-        self.create_flux_values_fenics()
-        self.create_initial_conditions()
+
         self.create_formulation()
         self.create_solver()
         self.initialise_exports()
@@ -998,6 +994,7 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
             u_n = spe.subdomain_to_prev_solution[subdomain]
             v = spe.subdomain_to_test_function[subdomain]
             dx = subdomain.dx
+            ds = subdomain.ds
 
             D = subdomain.material.get_diffusion_coefficient(
                 self.mesh.mesh, self.temperature_fenics, spe
@@ -1040,6 +1037,11 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
                     * dx
                 )
 
+        # add fluxes
+        for bc in self.boundary_conditions:
+            if isinstance(bc, F.ParticleFluxBC):
+                form -= bc.value_fenics * v * ds(bc.subdomain.id)
+
         subdomain.F = form
 
     def create_formulation(self):
@@ -1152,6 +1154,18 @@ def create_solver(self):
             },
         )
 
+    def create_flux_values_fenics(self):
+        """For each particle flux create the value_fenics"""
+        for bc in self.boundary_conditions:
+            # create value_fenics for all F.ParticleFluxBC objects
+            if isinstance(bc, F.ParticleFluxBC):
+                volume_subdomain = self.surface_to_volume[bc.subdomain]
+                bc.create_value_fenics(
+                    mesh=volume_subdomain.submesh,
+                    temperature=self.temperature_fenics,
+                    t=self.t,
+                )
+
     def initialise_exports(self):
         if self.exports:
             raise NotImplementedError(
diff --git a/test_implementation_2.py b/test_implementation_2.py
new file mode 100644
index 000000000..1cbada6cb
--- /dev/null
+++ b/test_implementation_2.py
@@ -0,0 +1,169 @@
+import festim as F
+
+from mpi4py import MPI
+import dolfinx
+import dolfinx.fem.petsc
+import numpy as np
+import festim as F
+import ufl
+
+
+# ---------------- Generate a mesh ----------------
+def generate_mesh():
+    def bottom_boundary(x):
+        return np.isclose(x[1], 0.0)
+
+    def top_boundary(x):
+        return np.isclose(x[1], 1.0)
+
+    def half(x):
+        return x[1] <= 0.5 + 1e-14
+
+    mesh = dolfinx.mesh.create_unit_square(
+        MPI.COMM_WORLD, 20, 20, dolfinx.mesh.CellType.triangle
+    )
+
+    # Split domain in half and set an interface tag of 5
+    gdim = mesh.geometry.dim
+    tdim = mesh.topology.dim
+    fdim = tdim - 1
+    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    num_facets_local = (
+        mesh.topology.index_map(fdim).size_local
+        + mesh.topology.index_map(fdim).num_ghosts
+    )
+    facets = np.arange(num_facets_local, dtype=np.int32)
+    values = np.full_like(facets, 0, dtype=np.int32)
+    values[top_facets] = 1
+    values[bottom_facets] = 2
+
+    bottom_cells = dolfinx.mesh.locate_entities(mesh, tdim, half)
+    num_cells_local = (
+        mesh.topology.index_map(tdim).size_local
+        + mesh.topology.index_map(tdim).num_ghosts
+    )
+    cells = np.full(num_cells_local, 4, dtype=np.int32)
+    cells[bottom_cells] = 3
+    ct = dolfinx.mesh.meshtags(
+        mesh, tdim, np.arange(num_cells_local, dtype=np.int32), cells
+    )
+    all_b_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(3), tdim, fdim
+    )
+    all_t_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(4), tdim, fdim
+    )
+    interface = np.intersect1d(all_b_facets, all_t_facets)
+    values[interface] = 5
+
+    mt = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, values)
+    return mesh, mt, ct
+
+
+mesh, mt, ct = generate_mesh()
+
+my_model = F.HTransportProblemDiscontinuous()
+my_model.mesh = F.Mesh(mesh)
+my_model.volume_meshtags = ct
+my_model.facet_meshtags = mt
+
+material_bottom = F.Material(D_0=2.0, E_D=0.1)
+material_top = F.Material(D_0=2.0, E_D=0.1)
+
+material_bottom.K_S_0 = 2.0
+material_bottom.E_K_S = 0 * 0.1
+material_top.K_S_0 = 4.0
+material_top.E_K_S = 0 * 0.12
+
+top_domain = F.VolumeSubdomain(4, material=material_top)
+bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
+
+top_surface = F.SurfaceSubdomain(id=1)
+bottom_surface = F.SurfaceSubdomain(id=2)
+my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+
+# we should be able to automate this
+my_model.interfaces = {5: [bottom_domain, top_domain]}
+my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+
+H = F.Species("H", mobile=True)
+
+my_model.species = [H]
+
+for species in my_model.species:
+    species.subdomains = [bottom_domain, top_domain]
+
+
+my_model.boundary_conditions = [
+    F.DirichletBC(top_surface, value=0.05, species=H),
+    # F.DirichletBC(bottom_surface, value=0.2, species=H),
+    F.ParticleFluxBC(bottom_surface, value=1.0, species=H),
+]
+
+
+my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
+
+my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+
+
+my_model.initialise()
+my_model.run()
+
+# -------------------- post processing --------------------
+
+list_of_subdomains = my_model.volume_subdomains
+
+for subdomain in list_of_subdomains:
+    u_sub_0 = subdomain.u.sub(0).collapse()
+    u_sub_0.name = "u_sub_0"
+
+    bp = dolfinx.io.VTXWriter(
+        mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0], engine="BP4"
+    )
+    bp.write(0)
+    bp.close()
+
+
+# derived quantities
+my_model.entity_maps[mesh] = bottom_domain.submesh_to_mesh
+
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
+ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
+dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
+dx = ufl.Measure("dx", domain=mesh)
+
+n_b = ufl.FacetNormal(bottom_domain.submesh)
+n_t = ufl.FacetNormal(top_domain.submesh)
+
+form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
+print(dolfinx.fem.assemble_scalar(form))
+
+form = dolfinx.fem.form(
+    my_model.temperature_fenics * dx_b,
+    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+)
+print(dolfinx.fem.assemble_scalar(form))
+
+D = subdomain.material.get_diffusion_coefficient(
+    my_model.mesh.mesh, my_model.temperature_fenics, H
+)
+id_interface = 5
+form = dolfinx.fem.form(
+    ufl.dot(
+        D * ufl.grad(bottom_domain.u.sub(0)),
+        n_b,
+    )
+    * ds_b(id_interface),
+    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+)
+print(dolfinx.fem.assemble_scalar(form))
+form = dolfinx.fem.form(
+    ufl.dot(
+        D * ufl.grad(top_domain.u.sub(0)),
+        n_t,
+    )
+    * ds_t(id_interface),
+    entity_maps={mesh: top_domain.submesh_to_mesh},
+)
+print(dolfinx.fem.assemble_scalar(form))

From dd0f65c118f04aef709cb222f63a83dec871c0ac Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 10:27:17 +0000
Subject: [PATCH 028/134] particle fluxes work correctly

---
 festim/hydrogen_transport_problem.py | 12 +++++++-----
 test_implementation_2.py             | 10 +++++-----
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 46d1b3318..248010e21 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -754,6 +754,7 @@ def __init__(
         self.surface_to_volume = surface_to_volume or {}
 
     def initialise(self):
+        self.define_meshtags_and_measures()
         self.create_submeshes()
         self.create_species_from_traps()
 
@@ -782,7 +783,6 @@ def initialise(self):
             subdomain.dx = ufl.Measure(
                 "dx", domain=self.mesh.mesh, subdomain_data=ct_r, subdomain_id=1
             )
-            subdomain.ds = None
             self.create_subdomain_formulation(subdomain)
             subdomain.u.name = f"u_{subdomain.id}"
 
@@ -899,7 +899,7 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
             u_n = spe.subdomain_to_prev_solution[subdomain]
             v = spe.subdomain_to_test_function[subdomain]
             dx = subdomain.dx
-            ds = subdomain.ds
+            ds = self.ds
 
             D = subdomain.material.get_diffusion_coefficient(
                 self.mesh.mesh, self.temperature_fenics, spe
@@ -945,7 +945,8 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         # add fluxes
         for bc in self.boundary_conditions:
             if isinstance(bc, F.ParticleFluxBC):
-                form -= bc.value_fenics * v * ds(bc.subdomain.id)
+                if subdomain == self.surface_to_volume[bc.subdomain]:
+                    form -= bc.value_fenics * v * ds(bc.subdomain.id)
 
         subdomain.F = form
 
@@ -997,11 +998,12 @@ def mixed_term(u, v, n):
 
             n = ufl.FacetNormal(mesh)
             n_b = n(b_res)
-            n_t = n(t_res)
+            n_t = n(t_res)  # this doesn't seem to be used
             cr = ufl.Circumradius(mesh)
             h_b = 2 * cr(b_res)
             h_t = 2 * cr(t_res)
-            gamma = 400.0  # this needs to be "sufficiently large"
+            # TODO make this a user parameter
+            gamma = 10.0  # this needs to be "sufficiently large"
 
             K_b = K_S_fun(
                 self.temperature_fenics(b_res),
diff --git a/test_implementation_2.py b/test_implementation_2.py
index 1cbada6cb..51f9d5e97 100644
--- a/test_implementation_2.py
+++ b/test_implementation_2.py
@@ -68,8 +68,8 @@ def half(x):
 my_model.volume_meshtags = ct
 my_model.facet_meshtags = mt
 
-material_bottom = F.Material(D_0=2.0, E_D=0.1)
-material_top = F.Material(D_0=2.0, E_D=0.1)
+material_bottom = F.Material(D_0=1, E_D=0 * 0.1)
+material_top = F.Material(D_0=1, E_D=0 * 0.1)
 
 material_bottom.K_S_0 = 2.0
 material_bottom.E_K_S = 0 * 0.1
@@ -97,12 +97,12 @@ def half(x):
 
 my_model.boundary_conditions = [
     F.DirichletBC(top_surface, value=0.05, species=H),
-    # F.DirichletBC(bottom_surface, value=0.2, species=H),
-    F.ParticleFluxBC(bottom_surface, value=1.0, species=H),
+    # F.DirichletBC(bottom_surface, value=1, species=H),
+    F.ParticleFluxBC(bottom_surface, value=lambda x: 1.0 + x[0], species=H),
 ]
 
 
-my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
+my_model.temperature = 500.0  # lambda x: 300 + 10 * x[1] + 100 * x[0]
 
 my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
 

From 46b6c476358166690a4e05b4e630a28a1725a272 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 10:28:50 +0000
Subject: [PATCH 029/134] use festim k_B

---
 festim/hydrogen_transport_problem.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 248010e21..5a6c91254 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -715,8 +715,7 @@ def post_processing(self):
 
 
 def K_S_fun(T, K_S_0, E_K_S):
-    k_B = 8.6173303e-5
-    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
+    return K_S_0 * ufl.exp(-E_K_S / F.k_B / T)
 
 
 class HTransportProblemDiscontinuous(HydrogenTransportProblem):

From f7408fec240ea91bf58f45bc18b6dbc66a72257d Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 10:36:23 +0000
Subject: [PATCH 030/134] ICs not implemented

---
 festim/hydrogen_transport_problem.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 5a6c91254..6c13868a3 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -810,6 +810,10 @@ def create_dirichletbc_form(self, bc):
         )
         return form
 
+    def create_initial_conditions(self):
+        if self.initial_conditions:
+            raise NotImplementedError("initial conditions not yet implemented for discontinuous")
+
     def create_submeshes(self):
         mesh = self.mesh.mesh
         ct = self.volume_meshtags

From 0df93381a896b1b565503b5546837db02a8d9bdf Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 10:54:18 +0000
Subject: [PATCH 031/134] black formatted

---
 festim/hydrogen_transport_problem.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 6c13868a3..c8aaf6b92 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -812,7 +812,9 @@ def create_dirichletbc_form(self, bc):
 
     def create_initial_conditions(self):
         if self.initial_conditions:
-            raise NotImplementedError("initial conditions not yet implemented for discontinuous")
+            raise NotImplementedError(
+                "initial conditions not yet implemented for discontinuous"
+            )
 
     def create_submeshes(self):
         mesh = self.mesh.mesh

From 89fbef5108fbd406b7aadbe238f7d1c4f63f8cbb Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 10:55:58 +0000
Subject: [PATCH 032/134] added new dict

---
 festim/species.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/festim/species.py b/festim/species.py
index bea839a51..e63faed35 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -62,6 +62,7 @@ def __init__(
         self.subdomain_to_solution = {}
         self.subdomain_to_prev_solution = {}
         self.subdomain_to_test_function = {}
+        self.subdomain_to_post_processing_solution = {}
 
     def __repr__(self) -> str:
         return f"Species({self.name})"

From d382ad4219cd18bf3a54086d1912cba30906783d Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 12:38:01 +0000
Subject: [PATCH 033/134] added integration with Mesh1D

---
 festim/mesh/mesh.py      | 20 +++++++--
 festim/mesh/mesh_1d.py   |  4 +-
 festim/problem.py        |  2 +
 test_implementation_3.py | 93 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 114 insertions(+), 5 deletions(-)
 create mode 100644 test_implementation_3.py

diff --git a/festim/mesh/mesh.py b/festim/mesh/mesh.py
index 83537330a..e99d1094a 100644
--- a/festim/mesh/mesh.py
+++ b/festim/mesh/mesh.py
@@ -55,12 +55,14 @@ def fdim(self):
     def n(self):
         return ufl.FacetNormal(self.mesh)
 
-    def define_meshtags(self, surface_subdomains, volume_subdomains):
+    def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None):
         """Defines the facet and volume meshtags of the mesh
 
         Args:
             surface_subdomains (list of festim.SufaceSubdomains): the surface subdomains of the model
             volume_subdomains (list of festim.VolumeSubdomains): the volume subdomains of the model
+            interfaces (dict, optional): the interfaces between volume
+                subdomains {int: [VolumeSubdomain, VolumeSubdomain]}. Defaults to None.
 
         Returns:
             dolfinx.mesh.MeshTags: the facet meshtags
@@ -86,10 +88,22 @@ def define_meshtags(self, surface_subdomains, volume_subdomains):
             entities = vol.locate_subdomain_entities(self.mesh, self.vdim)
             tags_volumes[entities] = vol.id
 
-        # define mesh tags
-        facet_meshtags = meshtags(self.mesh, self.fdim, mesh_facet_indices, tags_facets)
         volume_meshtags = meshtags(
             self.mesh, self.vdim, mesh_cell_indices, tags_volumes
         )
 
+        # tag interfaces
+        interfaces = interfaces or {}  # if interfaces is None, set it to empty dict
+        for interface_id, (domain_1, domain_2) in interfaces.items():
+            all_1_facets = dolfinx.mesh.compute_incident_entities(
+                self.mesh.topology, volume_meshtags.find(domain_1.id), self.vdim, self.fdim
+            )
+            all_2_facets = dolfinx.mesh.compute_incident_entities(
+                self.mesh.topology, volume_meshtags.find(domain_2.id), self.vdim, self.fdim
+            )
+            interface_entities = np.intersect1d(all_1_facets, all_2_facets)
+            tags_facets[interface_entities] = interface_id
+
+        facet_meshtags = meshtags(self.mesh, self.fdim, mesh_facet_indices, tags_facets)
+
         return facet_meshtags, volume_meshtags
diff --git a/festim/mesh/mesh_1d.py b/festim/mesh/mesh_1d.py
index d9c5ae030..9e0c56802 100644
--- a/festim/mesh/mesh_1d.py
+++ b/festim/mesh/mesh_1d.py
@@ -72,7 +72,7 @@ def check_borders(self, volume_subdomains):
         ):
             raise ValueError("borders dont match domain borders")
 
-    def define_meshtags(self, surface_subdomains, volume_subdomains):
+    def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None):
         # check if all borders are defined
         self.check_borders(volume_subdomains)
-        return super().define_meshtags(surface_subdomains, volume_subdomains)
+        return super().define_meshtags(surface_subdomains, volume_subdomains, interfaces)
diff --git a/festim/problem.py b/festim/problem.py
index 79c73a110..344e6eaf3 100644
--- a/festim/problem.py
+++ b/festim/problem.py
@@ -63,6 +63,8 @@ def define_meshtags_and_measures(self):
             self.facet_meshtags, self.volume_meshtags = self.mesh.define_meshtags(
                 surface_subdomains=self.surface_subdomains,
                 volume_subdomains=self.volume_subdomains,
+                # if self has attribute interfaces pass it
+                interfaces=getattr(self, "interfaces", None),
             )
 
         # check volume ids are unique
diff --git a/test_implementation_3.py b/test_implementation_3.py
new file mode 100644
index 000000000..c261e654e
--- /dev/null
+++ b/test_implementation_3.py
@@ -0,0 +1,93 @@
+import festim as F
+
+import dolfinx
+import numpy as np
+import festim as F
+
+my_model = F.HTransportProblemDiscontinuous()
+
+interface_1 = 0.2
+interface_2 = 0.8
+
+vertices = np.concatenate(
+    [
+        np.linspace(0, interface_1, num=100),
+        np.linspace(interface_1, interface_2, num=100),
+        np.linspace(interface_2, 1, num=100),
+    ]
+)
+
+my_model.mesh = F.Mesh1D(vertices)
+
+material_left = F.Material(D_0=2.0, E_D=0.1)
+material_mid = F.Material(D_0=2.0, E_D=0.1)
+material_right = F.Material(D_0=2.0, E_D=0.1)
+
+material_left.K_S_0 = 2.0
+material_left.E_K_S = 0
+material_mid.K_S_0 = 4.0
+material_mid.E_K_S = 0
+material_right.K_S_0 = 6.0
+material_right.E_K_S = 0
+
+left_domain = F.VolumeSubdomain1D(3, borders=[0, interface_1], material=material_left)
+middle_domain = F.VolumeSubdomain1D(4, borders=[interface_1, interface_2], material=material_mid)
+right_domain = F.VolumeSubdomain1D(5, borders=[interface_2, 1], material=material_right)
+
+left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
+right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
+
+# the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
+my_model.interfaces = {6: [left_domain, middle_domain], 7: [middle_domain, right_domain]}
+my_model.subdomains = [left_domain, middle_domain, right_domain, left_surface, right_surface]
+my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
+
+H = F.Species("H", mobile=True)
+trapped_H = F.Species("H_trapped", mobile=False)
+empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
+
+my_model.species = [H, trapped_H]
+
+for species in [H, trapped_H]:
+    species.subdomains = [left_domain, middle_domain, right_domain]
+
+
+my_model.reactions = [
+    F.Reaction(
+        reactant=[H, empty_trap],
+        product=[trapped_H],
+        k_0=2,
+        E_k=0,
+        p_0=0.1,
+        E_p=0,
+        volume=domain,
+    )
+    for domain in [left_domain, middle_domain, right_domain]
+]
+
+my_model.boundary_conditions = [
+    F.DirichletBC(left_surface, value=0.05, species=H),
+    F.DirichletBC(right_surface, value=0.2, species=H),
+]
+
+
+my_model.temperature = lambda x: 300 + 100 * x[0]
+
+my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+
+
+my_model.initialise()
+my_model.run()
+
+
+for subdomain in my_model.volume_subdomains:
+    u_sub_0 = subdomain.u.sub(0).collapse()
+    u_sub_0.name = "u_sub_0"
+
+    u_sub_1 = subdomain.u.sub(1).collapse()
+    u_sub_1.name = "u_sub_1"
+    bp = dolfinx.io.VTXWriter(
+        my_model.mesh.mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
+    )
+    bp.write(0)
+    bp.close()
\ No newline at end of file

From 28c1d44e828d2250fe99f239a263505c46d38967 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 12:54:14 +0000
Subject: [PATCH 034/134] added usage in implementation script

---
 test_implementation_3.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/test_implementation_3.py b/test_implementation_3.py
index c261e654e..ab94da1b9 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -81,13 +81,16 @@
 
 
 for subdomain in my_model.volume_subdomains:
-    u_sub_0 = subdomain.u.sub(0).collapse()
-    u_sub_0.name = "u_sub_0"
+    us = []
+    for species in my_model.species:
+        if subdomain not in species.subdomains:
+            continue
+        u_sub_0 = species.subdomain_to_post_processing_solution[subdomain].collapse()
+        u_sub_0.name = species.name
+        us.append(u_sub_0)
 
-    u_sub_1 = subdomain.u.sub(1).collapse()
-    u_sub_1.name = "u_sub_1"
     bp = dolfinx.io.VTXWriter(
-        my_model.mesh.mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
+        my_model.mesh.mesh.comm, f"u_{subdomain.id}.bp", us, engine="BP4"
     )
     bp.write(0)
     bp.close()
\ No newline at end of file

From eb7885c30ad551cf738d6f992be62cbfb0b2ab47 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 14:42:10 +0000
Subject: [PATCH 035/134] small changes

---
 festim/hydrogen_transport_problem.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index c8aaf6b92..be56872ef 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -900,11 +900,12 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         form = 0
         # add diffusion and time derivative for each species
         for spe in self.species:
+            if subdomain not in spe.subdomains:
+                continue
             u = spe.subdomain_to_solution[subdomain]
             u_n = spe.subdomain_to_prev_solution[subdomain]
             v = spe.subdomain_to_test_function[subdomain]
             dx = subdomain.dx
-            ds = self.ds
 
             D = subdomain.material.get_diffusion_coefficient(
                 self.mesh.mesh, self.temperature_fenics, spe
@@ -951,7 +952,8 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         for bc in self.boundary_conditions:
             if isinstance(bc, F.ParticleFluxBC):
                 if subdomain == self.surface_to_volume[bc.subdomain]:
-                    form -= bc.value_fenics * v * ds(bc.subdomain.id)
+                    v = bc.species.subdomain_to_test_function[subdomain]
+                    form -= bc.value_fenics * v * self.ds(bc.subdomain.id)
 
         subdomain.F = form
 
@@ -1090,4 +1092,4 @@ def run(self):
         else:
             # Solve steady-state
             self.solver.solve(1e-5)
-            self.post_processing()
+            # self.post_processing()

From 2144141436a2308fcc26ddadea01d8cb70ae2ce8 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 15:08:25 +0000
Subject: [PATCH 036/134] a few changes + comment about random SEGV crash

---
 test_implementation_3.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/test_implementation_3.py b/test_implementation_3.py
index c261e654e..e44cd6768 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -1,14 +1,13 @@
-import festim as F
-
 import dolfinx
 import numpy as np
 import festim as F
 
 my_model = F.HTransportProblemDiscontinuous()
 
-interface_1 = 0.2
-interface_2 = 0.8
+interface_1 = 0.5
+interface_2 = 0.7
 
+# for some reason if the mesh isn't fine enough then I have a random SEGV error
 vertices = np.concatenate(
     [
         np.linspace(0, interface_1, num=100),
@@ -19,9 +18,9 @@
 
 my_model.mesh = F.Mesh1D(vertices)
 
-material_left = F.Material(D_0=2.0, E_D=0.1)
-material_mid = F.Material(D_0=2.0, E_D=0.1)
-material_right = F.Material(D_0=2.0, E_D=0.1)
+material_left = F.Material(D_0=2.0, E_D=0)
+material_mid = F.Material(D_0=2.0, E_D=0)
+material_right = F.Material(D_0=2.0, E_D=0)
 
 material_left.K_S_0 = 2.0
 material_left.E_K_S = 0
@@ -30,9 +29,9 @@
 material_right.K_S_0 = 6.0
 material_right.E_K_S = 0
 
-left_domain = F.VolumeSubdomain1D(3, borders=[0, interface_1], material=material_left)
+left_domain = F.VolumeSubdomain1D(3, borders=[vertices[0], interface_1], material=material_left)
 middle_domain = F.VolumeSubdomain1D(4, borders=[interface_1, interface_2], material=material_mid)
-right_domain = F.VolumeSubdomain1D(5, borders=[interface_2, 1], material=material_right)
+right_domain = F.VolumeSubdomain1D(5, borders=[interface_2, vertices[-1]], material=material_right)
 
 left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
@@ -48,8 +47,8 @@
 
 my_model.species = [H, trapped_H]
 
-for species in [H, trapped_H]:
-    species.subdomains = [left_domain, middle_domain, right_domain]
+for species in my_model.species:
+    species.subdomains = my_model.volume_subdomains
 
 
 my_model.reactions = [

From 283abd4d5b483a0a46528883c89d28e6c256338c Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 15:18:20 +0000
Subject: [PATCH 037/134] added a map to collapsed function space and parent
 map

---
 festim/species.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/festim/species.py b/festim/species.py
index e63faed35..89ef4b5e4 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -63,6 +63,7 @@ def __init__(
         self.subdomain_to_prev_solution = {}
         self.subdomain_to_test_function = {}
         self.subdomain_to_post_processing_solution = {}
+        self.subdomain_to_collapsed_function_space = {}
 
     def __repr__(self) -> str:
         return f"Species({self.name})"

From b90693129dd3b45917b2299ff689906fd4ea32fb Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 15:18:52 +0000
Subject: [PATCH 038/134] store collapsed functions and function space + update
 in post process

---
 festim/hydrogen_transport_problem.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index be56872ef..2fb3ebbe0 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -894,6 +894,8 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_solution[subdomain] = us[i]
             species.subdomain_to_prev_solution[subdomain] = u_ns[i]
             species.subdomain_to_test_function[subdomain] = vs[i]
+            species.subdomain_to_post_processing_solution[subdomain] = u.sub(i).collapse()
+            species.subdomain_to_collapsed_function_space[subdomain] = (V.sub(i).collapse())
         subdomain.u = u
 
     def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
@@ -1086,10 +1088,24 @@ def initialise_exports(self):
                 "Exports not implemented for HTransportProblemDiscontinuous"
             )
 
+    def post_processing(self):
+        for subdomain in self.volume_subdomains:
+            for species in self.species:
+                if subdomain not in species.subdomains:
+                    continue
+                collapsed_function = species.subdomain_to_post_processing_solution[
+                    subdomain
+                ]
+                u = subdomain.u
+                v0_to_V = species.subdomain_to_collapsed_function_space[
+                    subdomain
+                ][1]
+                collapsed_function.x.array[:] = u.x.array[v0_to_V]
+
     def run(self):
         if self.settings.transient:
             raise NotImplementedError("Transient not implemented")
         else:
             # Solve steady-state
             self.solver.solve(1e-5)
-            # self.post_processing()
+            self.post_processing()

From 256fe0628ecea58b882bded8dcac5a9fbb9e999f Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 15:30:04 +0000
Subject: [PATCH 039/134] don't need to collapse here

---
 test_implementation_3.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test_implementation_3.py b/test_implementation_3.py
index ba2dc0fe3..e4d02548a 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -10,9 +10,9 @@
 # for some reason if the mesh isn't fine enough then I have a random SEGV error
 vertices = np.concatenate(
     [
-        np.linspace(0, interface_1, num=100),
-        np.linspace(interface_1, interface_2, num=100),
-        np.linspace(interface_2, 1, num=100),
+        np.linspace(0, interface_1, num=1500),
+        np.linspace(interface_1, interface_2, num=1500),
+        np.linspace(interface_2, 1, num=1500),
     ]
 )
 
@@ -84,7 +84,7 @@
     for species in my_model.species:
         if subdomain not in species.subdomains:
             continue
-        u_sub_0 = species.subdomain_to_post_processing_solution[subdomain].collapse()
+        u_sub_0 = species.subdomain_to_post_processing_solution[subdomain]
         u_sub_0.name = species.name
         us.append(u_sub_0)
 

From d83f0b7b52a5e66e451236e8db7c82a9043812dd Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 15:42:02 +0000
Subject: [PATCH 040/134] integrated VTX Export

---
 festim/exports/vtx.py                |  3 ++-
 festim/hydrogen_transport_problem.py | 19 +++++++++++++++----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/festim/exports/vtx.py b/festim/exports/vtx.py
index afc1191e0..00d75eec0 100644
--- a/festim/exports/vtx.py
+++ b/festim/exports/vtx.py
@@ -63,8 +63,9 @@ class VTXExport(VTXExportBase):
         ...    my_export.write(t)
     """
 
-    def __init__(self, filename: str, field) -> None:
+    def __init__(self, filename: str, field, subdomain: F.VolumeSubdomain=None) -> None:
         self.field = field
+        self.subdomain = subdomain
         super().__init__(filename)
 
     @property
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 2fb3ebbe0..065504294 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -896,6 +896,7 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_test_function[subdomain] = vs[i]
             species.subdomain_to_post_processing_solution[subdomain] = u.sub(i).collapse()
             species.subdomain_to_collapsed_function_space[subdomain] = (V.sub(i).collapse())
+            species.subdomain_to_post_processing_solution[subdomain].name = f"{species.name}_{subdomain.id}"
         subdomain.u = u
 
     def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
@@ -1083,10 +1084,14 @@ def create_flux_values_fenics(self):
                 )
 
     def initialise_exports(self):
-        if self.exports:
-            raise NotImplementedError(
-                "Exports not implemented for HTransportProblemDiscontinuous"
-            )
+        for export in self.exports:
+            if isinstance(export, F.VTXExport):
+                species = export.field[0]
+                # override post_processing_solution attribute of species
+                species.post_processing_solution = species.subdomain_to_post_processing_solution[export.subdomain]
+                export.define_writer(MPI.COMM_WORLD)
+            else:
+                raise NotImplementedError("Export type not implemented")
 
     def post_processing(self):
         for subdomain in self.volume_subdomains:
@@ -1102,6 +1107,12 @@ def post_processing(self):
                 ][1]
                 collapsed_function.x.array[:] = u.x.array[v0_to_V]
 
+        for export in self.exports:
+            if isinstance(export, F.VTXExport):
+                export.write(float(self.t))
+            else:
+                raise NotImplementedError("Export type not implemented")
+
     def run(self):
         if self.settings.transient:
             raise NotImplementedError("Transient not implemented")

From 9b75dd6fa986fa26d97c7ec8bab07d5424f8697f Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 15:42:31 +0000
Subject: [PATCH 041/134] integrated export in example

---
 test_implementation_3.py | 20 +++-----------------
 1 file changed, 3 insertions(+), 17 deletions(-)

diff --git a/test_implementation_3.py b/test_implementation_3.py
index e4d02548a..3f21819b0 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -1,4 +1,3 @@
-import dolfinx
 import numpy as np
 import festim as F
 
@@ -74,22 +73,9 @@
 
 my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
 
+my_model.exports = [
+    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+]
 
 my_model.initialise()
 my_model.run()
-
-
-for subdomain in my_model.volume_subdomains:
-    us = []
-    for species in my_model.species:
-        if subdomain not in species.subdomains:
-            continue
-        u_sub_0 = species.subdomain_to_post_processing_solution[subdomain]
-        u_sub_0.name = species.name
-        us.append(u_sub_0)
-
-    bp = dolfinx.io.VTXWriter(
-        my_model.mesh.mesh.comm, f"u_{subdomain.id}.bp", us, engine="BP4"
-    )
-    bp.write(0)
-    bp.close()
\ No newline at end of file

From 3f2c79c5b5da6b403b65cd8362be6cfaeba63d81 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 16:11:37 +0000
Subject: [PATCH 042/134] integrated VTX export in examples

---
 test_implementation.py   | 29 ++++++++++-------------------
 test_implementation_2.py | 18 ++++--------------
 test_implementation_3.py |  7 ++++---
 3 files changed, 18 insertions(+), 36 deletions(-)

diff --git a/test_implementation.py b/test_implementation.py
index b494ba120..4e7df1395 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -72,9 +72,9 @@ def half(x):
 material_top = F.Material(D_0=2.0, E_D=0.1)
 
 material_bottom.K_S_0 = 2.0
-material_bottom.E_K_S = 0 * 0.1
+material_bottom.E_K_S = 0
 material_top.K_S_0 = 4.0
-material_top.E_K_S = 0 * 0.12
+material_top.E_K_S = 0
 
 top_domain = F.VolumeSubdomain(4, material=material_top)
 bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
@@ -93,7 +93,7 @@ def half(x):
 
 my_model.species = [H, trapped_H]
 
-for species in [H, trapped_H]:
+for species in my_model.species:
     species.subdomains = [bottom_domain, top_domain]
 
 
@@ -124,30 +124,21 @@ def half(x):
 ]
 
 
-my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
+my_model.temperature = lambda x: 400 + 10 * x[1] + 100 * x[0]
 
 my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
 
+my_model.exports = [
+    F.VTXExport(f"c_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+] + [
+    F.VTXExport(f"c_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+]
 
 my_model.initialise()
 my_model.run()
 
 # -------------------- post processing --------------------
 
-list_of_subdomains = my_model.volume_subdomains
-
-for subdomain in list_of_subdomains:
-    u_sub_0 = subdomain.u.sub(0).collapse()
-    u_sub_0.name = "u_sub_0"
-
-    u_sub_1 = subdomain.u.sub(1).collapse()
-    u_sub_1.name = "u_sub_1"
-    bp = dolfinx.io.VTXWriter(
-        mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
-    )
-    bp.write(0)
-    bp.close()
-
 
 # derived quantities
 my_model.entity_maps[mesh] = bottom_domain.submesh_to_mesh
@@ -172,7 +163,7 @@ def half(x):
 )
 print(dolfinx.fem.assemble_scalar(form))
 
-D = subdomain.material.get_diffusion_coefficient(
+D = bottom_domain.material.get_diffusion_coefficient(
     my_model.mesh.mesh, my_model.temperature_fenics, H
 )
 id_interface = 5
diff --git a/test_implementation_2.py b/test_implementation_2.py
index 51f9d5e97..28ab2d256 100644
--- a/test_implementation_2.py
+++ b/test_implementation_2.py
@@ -106,25 +106,15 @@ def half(x):
 
 my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
 
+my_model.exports = [
+    F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+]
 
 my_model.initialise()
 my_model.run()
 
 # -------------------- post processing --------------------
 
-list_of_subdomains = my_model.volume_subdomains
-
-for subdomain in list_of_subdomains:
-    u_sub_0 = subdomain.u.sub(0).collapse()
-    u_sub_0.name = "u_sub_0"
-
-    bp = dolfinx.io.VTXWriter(
-        mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0], engine="BP4"
-    )
-    bp.write(0)
-    bp.close()
-
-
 # derived quantities
 my_model.entity_maps[mesh] = bottom_domain.submesh_to_mesh
 
@@ -145,7 +135,7 @@ def half(x):
 )
 print(dolfinx.fem.assemble_scalar(form))
 
-D = subdomain.material.get_diffusion_coefficient(
+D = bottom_domain.material.get_diffusion_coefficient(
     my_model.mesh.mesh, my_model.temperature_fenics, H
 )
 id_interface = 5
diff --git a/test_implementation_3.py b/test_implementation_3.py
index 3f21819b0..568bd5cbc 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -7,11 +7,12 @@
 interface_2 = 0.7
 
 # for some reason if the mesh isn't fine enough then I have a random SEGV error
+N = 1500
 vertices = np.concatenate(
     [
-        np.linspace(0, interface_1, num=1500),
-        np.linspace(interface_1, interface_2, num=1500),
-        np.linspace(interface_2, 1, num=1500),
+        np.linspace(0, interface_1, num=N),
+        np.linspace(interface_1, interface_2, num=N),
+        np.linspace(interface_2, 1, num=N),
     ]
 )
 

From ffe8212173349136ea94133e8ae940e36a9049a9 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 16:16:19 +0000
Subject: [PATCH 043/134] added example for only one material (no interfaces)

---
 test_one_material.py | 62 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)
 create mode 100644 test_one_material.py

diff --git a/test_one_material.py b/test_one_material.py
new file mode 100644
index 000000000..da5391639
--- /dev/null
+++ b/test_one_material.py
@@ -0,0 +1,62 @@
+import numpy as np
+import festim as F
+
+my_model = F.HTransportProblemDiscontinuous()
+
+
+N = 1500
+vertices = np.linspace(0, 1, num=N)
+my_model.mesh = F.Mesh1D(vertices)
+
+material = F.Material(D_0=2.0, E_D=0)
+
+material.K_S_0 = 2.0
+material.E_K_S = 0
+
+subdomain = F.VolumeSubdomain1D(1, borders=[vertices[0], vertices[-1]], material=material)
+
+left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
+right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
+
+my_model.subdomains = [subdomain,left_surface, right_surface]
+my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
+
+H = F.Species("H", mobile=True)
+trapped_H = F.Species("H_trapped", mobile=False)
+empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
+
+my_model.species = [H, trapped_H]
+
+for species in my_model.species:
+    species.subdomains = my_model.volume_subdomains
+
+
+my_model.reactions = [
+    F.Reaction(
+        reactant=[H, empty_trap],
+        product=[trapped_H],
+        k_0=2,
+        E_k=0,
+        p_0=0.1,
+        E_p=0,
+        volume=domain,
+    )
+    for domain in my_model.volume_subdomains
+]
+
+my_model.boundary_conditions = [
+    F.DirichletBC(left_surface, value=0.05, species=H),
+    F.DirichletBC(right_surface, value=0.2, species=H),
+]
+
+
+my_model.temperature = lambda x: 300 + 100 * x[0]
+
+my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+
+my_model.exports = [
+    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+]
+
+my_model.initialise()
+my_model.run()

From 2231f89c896f9a8a6b6b45d6b6bd8fe1a4919f01 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 17:26:21 +0000
Subject: [PATCH 044/134] started implementing transient

---
 festim/hydrogen_transport_problem.py | 48 +++++++++++++---
 test_implementation_transient.py     | 85 ++++++++++++++++++++++++++++
 2 files changed, 126 insertions(+), 7 deletions(-)
 create mode 100644 test_implementation_transient.py

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 065504294..8c1d854f1 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -4,6 +4,7 @@
 import ufl
 from mpi4py import MPI
 import numpy as np
+import tqdm.autonotebook
 import festim as F
 from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
 
@@ -759,7 +760,11 @@ def initialise(self):
 
         self.t = fem.Constant(self.mesh.mesh, 0.0)
         if self.settings.transient:
-            raise NotImplementedError("Transient simulations not supported yet")
+            # TODO should raise error if no stepsize is provided
+            # TODO Should this be an attribute of festim.Stepsize?
+            self.dt = F.as_fenics_constant(
+                self.settings.stepsize.initial_value, self.mesh.mesh
+            )
         self.define_temperature()
 
         self.entity_maps = {
@@ -898,6 +903,7 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_collapsed_function_space[subdomain] = (V.sub(i).collapse())
             species.subdomain_to_post_processing_solution[subdomain].name = f"{species.name}_{subdomain.id}"
         subdomain.u = u
+        subdomain.u_n = u_n
 
     def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         form = 0
@@ -914,7 +920,7 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
                 self.mesh.mesh, self.temperature_fenics, spe
             )
             if self.settings.transient:
-                raise NotImplementedError("Transient not implemented")
+                form += ((u - u_n) / self.dt) * v * dx
 
             if spe.mobile:
                 form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
@@ -925,11 +931,8 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
             for species in reaction.reactant + reaction.product:
                 if isinstance(species, F.Species):
                     # TODO remove
-                    # temporarily overide the solution and test function to the one of the subdomain
+                    # temporarily overide the solution to the one of the subdomain
                     species.solution = species.subdomain_to_solution[subdomain]
-                    species.test_function = species.subdomain_to_test_function[
-                        subdomain
-                    ]
 
             for reactant in reaction.reactant:
                 if isinstance(reactant, F.Species):
@@ -1113,9 +1116,40 @@ def post_processing(self):
             else:
                 raise NotImplementedError("Export type not implemented")
 
+    def iterate(self):
+        """Iterates the model for a given time step"""
+        self.progress.update(
+            min(self.dt.value, abs(self.settings.final_time - self.t.value))
+        )
+        self.t.value += self.dt.value
+
+        self.update_time_dependent_values()
+
+        # solve main problem
+        self.solver.solve(1e-5)
+
+        # post processing
+        self.post_processing()
+
+        # update previous solution
+        for subdomain in self.volume_subdomains:
+            subdomain.u_n.x.array[:] = subdomain.u.x.array[:]
+
+        # adapt stepsize
+        if self.settings.stepsize.adaptive:
+            raise NotImplementedError("Adaptive stepsize not implemented")
+
     def run(self):
         if self.settings.transient:
-            raise NotImplementedError("Transient not implemented")
+            # Solve transient
+            self.progress = tqdm.autonotebook.tqdm(
+                desc=f"Solving {self.__class__.__name__}",
+                total=self.settings.final_time,
+                unit_scale=True,
+            )
+            while self.t.value < self.settings.final_time:
+                self.iterate()
+            self.progress.refresh()  # refresh progress bar to show 100%
         else:
             # Solve steady-state
             self.solver.solve(1e-5)
diff --git a/test_implementation_transient.py b/test_implementation_transient.py
new file mode 100644
index 000000000..2d359ab81
--- /dev/null
+++ b/test_implementation_transient.py
@@ -0,0 +1,85 @@
+import numpy as np
+import festim as F
+
+my_model = F.HTransportProblemDiscontinuous()
+
+interface_1 = 0.5
+interface_2 = 0.7
+
+# for some reason if the mesh isn't fine enough then I have a random SEGV error
+N = 1500
+vertices = np.concatenate(
+    [
+        np.linspace(0, interface_1, num=N),
+        np.linspace(interface_1, interface_2, num=N),
+        np.linspace(interface_2, 1, num=N),
+    ]
+)
+
+my_model.mesh = F.Mesh1D(vertices)
+
+material_left = F.Material(D_0=2.0, E_D=0)
+material_mid = F.Material(D_0=2.0, E_D=0)
+material_right = F.Material(D_0=2.0, E_D=0)
+
+material_left.K_S_0 = 2.0
+material_left.E_K_S = 0
+material_mid.K_S_0 = 4.0
+material_mid.E_K_S = 0
+material_right.K_S_0 = 6.0
+material_right.E_K_S = 0
+
+left_domain = F.VolumeSubdomain1D(3, borders=[vertices[0], interface_1], material=material_left)
+middle_domain = F.VolumeSubdomain1D(4, borders=[interface_1, interface_2], material=material_mid)
+right_domain = F.VolumeSubdomain1D(5, borders=[interface_2, vertices[-1]], material=material_right)
+
+left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
+right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
+
+# the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
+my_model.interfaces = {6: [left_domain, middle_domain], 7: [middle_domain, right_domain]}
+my_model.subdomains = [left_domain, middle_domain, right_domain, left_surface, right_surface]
+my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
+
+H = F.Species("H", mobile=True)
+trapped_H = F.Species("H_trapped", mobile=False)
+empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
+
+my_model.species = [H, trapped_H]
+
+for species in my_model.species:
+    species.subdomains = my_model.volume_subdomains
+
+
+my_model.reactions = [
+    F.Reaction(
+        reactant=[H, empty_trap],
+        product=[trapped_H],
+        k_0=2,
+        E_k=0,
+        p_0=0.1,
+        E_p=0,
+        volume=domain,
+    )
+    for domain in [left_domain, middle_domain, right_domain]
+]
+
+my_model.boundary_conditions = [
+    F.DirichletBC(left_surface, value=0.05, species=H),
+    F.DirichletBC(right_surface, value=0.2, species=H),
+]
+
+
+my_model.temperature = lambda x: 300 + 100 * x[0]
+
+my_model.settings = F.Settings(atol=None, rtol=None, transient=True, final_time=5)
+my_model.settings.stepsize = 1e-2
+
+my_model.exports = [
+    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+] + [
+    F.VTXExport(filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+]
+
+my_model.initialise()
+my_model.run()

From 2c47f522d9c33ded20f8f90be7e93b70373d5539 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 20 Aug 2024 17:27:43 +0000
Subject: [PATCH 045/134] updated times

---
 test_implementation_transient.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test_implementation_transient.py b/test_implementation_transient.py
index 2d359ab81..cdabe45d1 100644
--- a/test_implementation_transient.py
+++ b/test_implementation_transient.py
@@ -72,8 +72,8 @@
 
 my_model.temperature = lambda x: 300 + 100 * x[0]
 
-my_model.settings = F.Settings(atol=None, rtol=None, transient=True, final_time=5)
-my_model.settings.stepsize = 1e-2
+my_model.settings = F.Settings(atol=None, rtol=None, transient=True, final_time=100)
+my_model.settings.stepsize = 1
 
 my_model.exports = [
     F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains

From c91dd3d76797714a676d42a8198b97368354af78 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 22 Aug 2024 16:17:05 -0400
Subject: [PATCH 046/134] added K_S to material + black

---
 festim/material.py               | 12 +++++++++
 test_implementation_transient.py | 44 ++++++++++++++++++++------------
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/festim/material.py b/festim/material.py
index e93464d57..dded153a4 100644
--- a/festim/material.py
+++ b/festim/material.py
@@ -11,6 +11,10 @@ class Material:
             diffusion coefficient (m2/s)
         E_D (float or dict): the activation energy of the diffusion
             coeficient (eV)
+        K_S_0 (float or dict): the pre-exponential factor of the
+            solubility coefficient (H/m3/Pa0.5)
+        E_K_S (float or dict): the activation energy of the solubility
+            coeficient (eV)
         name (str): the name of the material
         thermal_conductivity (float, callable): the thermal conductivity of the material (W/m/K)
         density (float, callable): the density of the material (kg/m3)
@@ -21,6 +25,10 @@ class Material:
             diffusion coefficient (m2/s)
         E_D (float or dict): the activation energy of the diffusion
             coeficient (eV)
+        K_S_0 (float or dict): the pre-exponential factor of the
+            solubility coefficient (H/m3/Pa0.5)
+        E_K_S (float or dict): the activation energy of the solubility
+            coeficient (eV)
         name (str): the name of the material
         thermal_conductivity (float, callable): the thermal conductivity of the material (W/m/K)
         density (float, callable): the density of the material (kg/m3)
@@ -40,6 +48,8 @@ def __init__(
         self,
         D_0,
         E_D,
+        K_S_0=None,
+        E_K_S=None,
         thermal_conductivity=None,
         density=None,
         heat_capacity=None,
@@ -47,6 +57,8 @@ def __init__(
     ) -> None:
         self.D_0 = D_0
         self.E_D = E_D
+        self.K_S_0 = K_S_0
+        self.E_K_S = E_K_S
 
         self.thermal_conductivity = thermal_conductivity
         self.density = density
diff --git a/test_implementation_transient.py b/test_implementation_transient.py
index cdabe45d1..54589e940 100644
--- a/test_implementation_transient.py
+++ b/test_implementation_transient.py
@@ -18,27 +18,35 @@
 
 my_model.mesh = F.Mesh1D(vertices)
 
-material_left = F.Material(D_0=2.0, E_D=0)
-material_mid = F.Material(D_0=2.0, E_D=0)
-material_right = F.Material(D_0=2.0, E_D=0)
+material_left = F.Material(D_0=2.0, E_D=0, K_S_0=2.0, E_K_S=0)
+material_mid = F.Material(D_0=2.0, E_D=0, K_S_0=4.0, E_K_S=0)
+material_right = F.Material(D_0=2.0, E_D=0, K_S_0=6.0, E_K_S=0)
 
-material_left.K_S_0 = 2.0
-material_left.E_K_S = 0
-material_mid.K_S_0 = 4.0
-material_mid.E_K_S = 0
-material_right.K_S_0 = 6.0
-material_right.E_K_S = 0
-
-left_domain = F.VolumeSubdomain1D(3, borders=[vertices[0], interface_1], material=material_left)
-middle_domain = F.VolumeSubdomain1D(4, borders=[interface_1, interface_2], material=material_mid)
-right_domain = F.VolumeSubdomain1D(5, borders=[interface_2, vertices[-1]], material=material_right)
+left_domain = F.VolumeSubdomain1D(
+    3, borders=[vertices[0], interface_1], material=material_left
+)
+middle_domain = F.VolumeSubdomain1D(
+    4, borders=[interface_1, interface_2], material=material_mid
+)
+right_domain = F.VolumeSubdomain1D(
+    5, borders=[interface_2, vertices[-1]], material=material_right
+)
 
 left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
 # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
-my_model.interfaces = {6: [left_domain, middle_domain], 7: [middle_domain, right_domain]}
-my_model.subdomains = [left_domain, middle_domain, right_domain, left_surface, right_surface]
+my_model.interfaces = {
+    6: [left_domain, middle_domain],
+    7: [middle_domain, right_domain],
+}
+my_model.subdomains = [
+    left_domain,
+    middle_domain,
+    right_domain,
+    left_surface,
+    right_surface,
+]
 my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
 
 H = F.Species("H", mobile=True)
@@ -76,9 +84,11 @@
 my_model.settings.stepsize = 1
 
 my_model.exports = [
-    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    for subdomain in my_model.volume_subdomains
 ] + [
-    F.VTXExport(filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+    F.VTXExport(filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain)
+    for subdomain in my_model.volume_subdomains
 ]
 
 my_model.initialise()

From 2b58f0003402f12730dfbe7130ada9fe6ef980f3 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 22 Aug 2024 21:50:37 +0000
Subject: [PATCH 047/134] Implemented refactoring done by Jorgen

Co-authored-by: jorgensd <dokken@simula.no>
---
 discontinuity_generic.py             | 735 ++++++++++++++++-----------
 festim/__init__.py                   |   1 +
 festim/hydrogen_transport_problem.py | 107 ++--
 festim/mesh/mesh.py                  |   5 +-
 festim/subdomain/interface.py        |  65 +++
 festim/subdomain/volume_subdomain.py |  35 +-
 test_implementation.py               |   2 +-
 test_implementation_2.py             |   2 +-
 test_implementation_3.py             |   7 +-
 test_implementation_transient.py     |  10 +-
 10 files changed, 591 insertions(+), 378 deletions(-)
 create mode 100644 festim/subdomain/interface.py

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index ca43331b2..1277be7be 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -1,373 +1,533 @@
 from mpi4py import MPI
+from petsc4py import PETSc
 import dolfinx
-import dolfinx.fem.petsc
+from dolfinx.cpp.fem import compute_integration_domains
+
 import ufl
 import numpy as np
-from petsc4py import PETSc
 import basix
-from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
-import festim as F
+import dolfinx.fem.petsc
 
+class NewtonSolver:
+    max_iterations: int
+    bcs: list[dolfinx.fem.DirichletBC]
+    A: PETSc.Mat
+    b: PETSc.Vec
+    J: dolfinx.fem.Form
+    b: dolfinx.fem.Form
+    dx: PETSc.Vec
+
+    def __init__(
+        self,
+        F: list[dolfinx.fem.form],
+        J: list[list[dolfinx.fem.form]],
+        w: list[dolfinx.fem.Function],
+        bcs: list[dolfinx.fem.DirichletBC] | None = None,
+        max_iterations: int = 5,
+        petsc_options: dict[str, str | float | int | None] = None,
+        problem_prefix="newton",
+    ):
+        self.max_iterations = max_iterations
+        self.bcs = [] if bcs is None else bcs
+        self.b = dolfinx.fem.petsc.create_vector_block(F)
+        self.F = F
+        self.J = J
+        self.A = dolfinx.fem.petsc.create_matrix_block(J)
+        self.dx = self.A.createVecLeft()
+        self.w = w
+        self.x = dolfinx.fem.petsc.create_vector_block(F)
+
+        # Set PETSc options
+        opts = PETSc.Options()
+        if petsc_options is not None:
+            for k, v in petsc_options.items():
+                opts[k] = v
+
+        # Define KSP solver
+        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
+        self._solver.setOperators(self.A)
+        self._solver.setFromOptions()
+
+        # Set matrix and vector PETSc options
+        self.A.setFromOptions()
+        self.b.setFromOptions()
+
+    def solve(self, tol=1e-6, beta=1.0):
+        i = 0
+
+        while i < self.max_iterations:
+            dolfinx.cpp.la.petsc.scatter_local_vectors(
+                self.x,
+                [si.x.petsc_vec.array_r for si in self.w],
+                [
+                    (
+                        si.function_space.dofmap.index_map,
+                        si.function_space.dofmap.index_map_bs,
+                    )
+                    for si in self.w
+                ],
+            )
+            self.x.ghostUpdate(
+                addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
+            )
 
-# ---------------- Generate a mesh ----------------
-def generate_mesh():
-    def bottom_boundary(x):
-        return np.isclose(x[1], 0.0)
+            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
+            with self.b.localForm() as b_local:
+                b_local.set(0.0)
+            
+            dolfinx.fem.petsc.assemble_vector_block(
+                self.b, self.F, self.J, bcs=self.bcs, x0=self.x, scale=-1.0
+            )
+            self.b.ghostUpdate(
+                PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
+            )
 
-    def top_boundary(x):
-        return np.isclose(x[1], 1.0)
+            # Assemble Jacobian
+            self.A.zeroEntries()
+            dolfinx.fem.petsc.assemble_matrix_block(self.A, self.J, bcs=self.bcs)
+            self.A.assemble()
+
+            self._solver.solve(self.b, self.dx)
+            # self._solver.view()
+            assert (
+                self._solver.getConvergedReason() > 0
+            ), "Linear solver did not converge"
+            offset_start = 0
+            for s in self.w:
+                num_sub_dofs = (
+                    s.function_space.dofmap.index_map.size_local
+                    * s.function_space.dofmap.index_map_bs
+                )
+                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
+                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
+                )
+                s.x.petsc_vec.ghostUpdate(
+                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
+                )
+                offset_start += num_sub_dofs
+            # Compute norm of update
+
+            correction_norm = self.dx.norm(0)
+            print(f"Iteration {i}: Correction norm {correction_norm}")
+            if correction_norm < tol:
+                break
+            i += 1
+
+    def __del__(self):
+        self.A.destroy()
+        self.b.destroy()
+        self.dx.destroy()
+        self._solver.destroy()
+        self.x.destroy()
+
+
+def transfer_meshtags_to_submesh(
+    mesh, entity_tag, submesh, sub_vertex_to_parent, sub_cell_to_parent
+):
+    """
+    Transfer a meshtag from a parent mesh to a sub-mesh.
+    """
+
+    tdim = mesh.topology.dim
+    cell_imap = mesh.topology.index_map(tdim)
+    num_cells = cell_imap.size_local + cell_imap.num_ghosts
+    mesh_to_submesh = np.full(num_cells, -1)
+    mesh_to_submesh[sub_cell_to_parent] = np.arange(
+        len(sub_cell_to_parent), dtype=np.int32
+    )
+    sub_vertex_to_parent = np.asarray(sub_vertex_to_parent)
+
+    submesh.topology.create_connectivity(entity_tag.dim, 0)
 
-    def half(x):
-        return x[1] <= 0.5 + 1e-14
+    num_child_entities = (
+        submesh.topology.index_map(entity_tag.dim).size_local
+        + submesh.topology.index_map(entity_tag.dim).num_ghosts
+    )
+    submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)
+
+    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
+    c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)
+
+    child_markers = np.full(num_child_entities, 0, dtype=np.int32)
+
+    mesh.topology.create_connectivity(entity_tag.dim, 0)
+    mesh.topology.create_connectivity(entity_tag.dim, mesh.topology.dim)
+    p_f_to_v = mesh.topology.connectivity(entity_tag.dim, 0)
+    p_f_to_c = mesh.topology.connectivity(entity_tag.dim, mesh.topology.dim)
+    sub_to_parent_entity_map = np.full(num_child_entities, -1, dtype=np.int32)
+    for facet, value in zip(entity_tag.indices, entity_tag.values):
+        facet_found = False
+        for cell in p_f_to_c.links(facet):
+            if facet_found:
+                break
+            if (child_cell := mesh_to_submesh[cell]) != -1:
+                for child_facet in c_c_to_e.links(child_cell):
+                    child_vertices = c_e_to_v.links(child_facet)
+                    child_vertices_as_parent = sub_vertex_to_parent[child_vertices]
+                    is_facet = np.isin(
+                        child_vertices_as_parent, p_f_to_v.links(facet)
+                    ).all()
+                    if is_facet:
+                        child_markers[child_facet] = value
+                        facet_found = True
+                        sub_to_parent_entity_map[child_facet] = facet
+    tags = dolfinx.mesh.meshtags(
+        submesh,
+        entity_tag.dim,
+        np.arange(num_child_entities, dtype=np.int32),
+        child_markers,
+    )
+    tags.name = entity_tag.name
+    return tags, sub_to_parent_entity_map
+
+
+# ---------------- Generate a mesh ----------------
+def generate_mesh():
+    def left_boundary(x):
+        return np.isclose(x[0], 0.0)
+
+    def right_boundary(x):
+        return np.isclose(x[0], 1.0)
+
+    def left_domain(x):
+        return x[0] <= 0.5 + 1e-14
+    
+    def right_domain(x):
+        return x[0] >= 0.7 - 1e-14
+
+    interface_1 = 0.5
+    interface_2 = 0.7
+
+    # with N = 2000 I've never had the error but with N = 10 I had it
+    N = 3
+    vertices = np.concatenate(
+        [
+            np.linspace(0, interface_1, num=N),
+            np.linspace(interface_1, interface_2, num=N),
+            np.linspace(interface_2, 1, num=N),
+        ]
+    )
 
-    mesh = dolfinx.mesh.create_unit_square(
-        MPI.COMM_WORLD, 20, 20, dolfinx.mesh.CellType.triangle
+    vertices = np.sort(np.unique(vertices)).astype(float)
+    degree = 1
+    domain = ufl.Mesh(
+        basix.ufl.element(basix.ElementFamily.P, "interval", degree, shape=(1,))
     )
+    mesh_points = np.reshape(vertices, (len(vertices), 1))
+    indexes = np.arange(vertices.shape[0])
+    cells = np.stack((indexes[:-1], indexes[1:]), axis=-1)
+
+    mesh = dolfinx.mesh.create_mesh(MPI.COMM_WORLD, cells, mesh_points, domain)
 
     # Split domain in half and set an interface tag of 5
-    gdim = mesh.geometry.dim
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
-    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    left_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, left_boundary)
+    right_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, right_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
     )
     facets = np.arange(num_facets_local, dtype=np.int32)
     values = np.full_like(facets, 0, dtype=np.int32)
-    values[top_facets] = 1
-    values[bottom_facets] = 2
+    values[left_facets] = 1
+    values[right_facets] = 2
 
-    bottom_cells = dolfinx.mesh.locate_entities(mesh, tdim, half)
+    left_cells = dolfinx.mesh.locate_entities(mesh, tdim, left_domain)
+    right_cells = dolfinx.mesh.locate_entities(mesh, tdim, right_domain)
     num_cells_local = (
         mesh.topology.index_map(tdim).size_local
         + mesh.topology.index_map(tdim).num_ghosts
     )
     cells = np.full(num_cells_local, 4, dtype=np.int32)
-    cells[bottom_cells] = 3
+    cells[left_cells] = 3
+    cells[right_cells] = 5
     ct = dolfinx.mesh.meshtags(
         mesh, tdim, np.arange(num_cells_local, dtype=np.int32), cells
     )
-    all_b_facets = dolfinx.mesh.compute_incident_entities(
+    all_l_facets = dolfinx.mesh.compute_incident_entities(
         mesh.topology, ct.find(3), tdim, fdim
     )
-    all_t_facets = dolfinx.mesh.compute_incident_entities(
+    all_m_facets = dolfinx.mesh.compute_incident_entities(
         mesh.topology, ct.find(4), tdim, fdim
     )
-    interface = np.intersect1d(all_b_facets, all_t_facets)
-    values[interface] = 5
+    all_r_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(5), tdim, fdim
+    )
+    interface1 = np.intersect1d(all_l_facets, all_m_facets)
+    interface2 = np.intersect1d(all_m_facets, all_r_facets)
+    values[interface1] = 5
+    values[interface2] = 6
 
     mt = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, values)
     return mesh, mt, ct
 
 
-mesh, mt, ct = generate_mesh()
-
-material_bottom = F.Material(D_0=2.0, E_D=0.1)
-material_top = F.Material(D_0=2.0, E_D=0.1)
-
-material_bottom.K_S_0 = 2.0
-material_bottom.E_K_S = 0 * 0.1
-material_top.K_S_0 = 4.0
-material_top.E_K_S = 0 * 0.12
-
-top_domain = F.VolumeSubdomain(4, material=material_top)
-bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
-list_of_subdomains = [bottom_domain, top_domain]
-list_of_interfaces = {5: [bottom_domain, top_domain]}
-
-top_surface = F.SurfaceSubdomain(id=1)
-bottom_surface = F.SurfaceSubdomain(id=2)
-
-H = F.Species("H", mobile=True)
-trapped_H = F.Species("H_trapped", mobile=False)
-empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
-
-for species in [H, trapped_H]:
-    species.subdomains = [bottom_domain, top_domain]
-    species.subdomain_to_solution = {}
-    species.subdomain_to_prev_solution = {}
-    species.subdomain_to_test_function = {}
-
-list_of_species = [H, trapped_H]
-
-list_of_reactions = [
-    F.Reaction(
-        reactant=[H, empty_trap],
-        product=[trapped_H],
-        k_0=2,
-        E_k=0,
-        p_0=0.1,
-        E_p=0,
-        volume=top_domain,
-    ),
-    F.Reaction(
-        reactant=[H, empty_trap],
-        product=[trapped_H],
-        k_0=2,
-        E_k=0,
-        p_0=0.1,
-        E_p=0,
-        volume=bottom_domain,
-    ),
-]
-
-list_of_bcs = [
-    F.DirichletBC(top_surface, value=0.05, species=H),
-    F.DirichletBC(bottom_surface, value=0.2, species=H),
-]
-
-surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
-
-V = dolfinx.fem.functionspace(mesh, ("CG", 1))
-T = dolfinx.fem.Function(V)
-T.interpolate(lambda x: 300 + 10 * x[1] + 100 * x[0])
+class VolumeSubdomain:
+    id: int
+    submesh: dolfinx.mesh.Mesh
+    submesh_to_mesh: np.ndarray
+    parent_mesh: dolfinx.mesh.Mesh
+    parent_to_submesh: np.ndarray
+    v_map: np.ndarray
+    facet_to_parent: np.ndarray
+    ft: dolfinx.mesh.MeshTags
+    padded:bool
+    def __init__(self, id):
+        self.id = id
+
+    def create_subdomain(self, mesh, marker):
+        assert marker.dim == mesh.topology.dim
+        self.parent_mesh = mesh
+        self.submesh, self.submesh_to_mesh, self.v_map = (
+            dolfinx.mesh.create_submesh(mesh, marker.dim, marker.find(self.id))[0:3]
+        )
+        num_cells_local = (
+            mesh.topology.index_map(marker.dim).size_local
+            + mesh.topology.index_map(marker.dim).num_ghosts
+        )
+        self.parent_to_submesh = np.full(num_cells_local, -1, dtype=np.int32)
+        self.parent_to_submesh[self.submesh_to_mesh] = np.arange(
+            len(self.submesh_to_mesh), dtype=np.int32
+        )
+        self.padded=False
 
+    def transfer_meshtag(self, tag):
+        # Transfer meshtags to submesh
+        assert self.submesh is not None, "Need to call create_subdomain first"
+        self.ft, self.facet_to_parent = transfer_meshtags_to_submesh(
+            mesh, tag, self.submesh, self.v_map, self.submesh_to_mesh
+        )
 
-def D_fun(T, D_0, E_D):
-    k_B = 8.6173303e-5
-    return D_0 * ufl.exp(-E_D / k_B / T)
+class Interface():
+    id: int
+    subdomains: tuple[VolumeSubdomain, VolumeSubdomain]
+    parent_mesh: dolfinx.mesh.Mesh
+    restriction: [str, str] = ("+", "-")
+    padded: bool
+    def __init__(self, parent_mesh, mt,  id, subdomains):
+        self.id = id
+        self.subdomains = tuple(subdomains)
+        self.mt = mt
+        self.parent_mesh = parent_mesh
+    def pad_parent_maps(self):
+        """Workaround to make sparsity-pattern work without skips
+        """ 
+
+        integration_data = compute_integration_domains(
+                dolfinx.fem.IntegralType.interior_facet, self.parent_mesh.topology, self.mt.find(self.id), self.mt.dim).reshape(-1, 4)
+        for i in range(2):
+            # We pad the parent to submesh map to make sure that sparsity pattern is correct
+            mapped_cell_0 = self.subdomains[i].parent_to_submesh[integration_data[:, 0]]
+            mapped_cell_1 = self.subdomains[i].parent_to_submesh[integration_data[:, 2]]
+            max_cells = np.maximum(mapped_cell_0, mapped_cell_1)
+            self.subdomains[i].parent_to_submesh[integration_data[:, 0]] = max_cells
+            self.subdomains[i].parent_to_submesh[integration_data[:, 2]] = max_cells
+            self.subdomains[i].padded = True
 
+mesh, mt, ct = generate_mesh()
 
-def K_S_fun(T, K_S_0, E_K_S):
-    k_B = 8.6173303e-5
-    return K_S_0 * ufl.exp(-E_K_S / k_B / T)
+left_domain = VolumeSubdomain(3)
+mid_domain = VolumeSubdomain(4)
+right_domain = VolumeSubdomain(5)
+list_of_subdomains = [left_domain, mid_domain, right_domain]
 
 
 gdim = mesh.geometry.dim
 tdim = mesh.topology.dim
 fdim = tdim - 1
-
-num_facets_local = (
-    mesh.topology.index_map(fdim).size_local + mesh.topology.index_map(fdim).num_ghosts
+num_cells_local = (
+    mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts
 )
 
+
+
 for subdomain in list_of_subdomains:
-    subdomain.submesh, subdomain.submesh_to_mesh, subdomain.v_map = (
-        dolfinx.mesh.create_submesh(mesh, tdim, ct.find(subdomain.id))[0:3]
-    )
+    subdomain.create_subdomain(mesh, ct)
+    subdomain.transfer_meshtag(mt)
 
-    subdomain.parent_to_submesh = np.full(num_facets_local, -1, dtype=np.int32)
-    subdomain.parent_to_submesh[subdomain.submesh_to_mesh] = np.arange(
-        len(subdomain.submesh_to_mesh), dtype=np.int32
-    )
 
-    # We need to modify the cell maps, as for `dS` integrals of interfaces between submeshes, there is no entity to map to.
-    # We use the entity on the same side to fix this (as all restrictions are one-sided)
 
-    # Transfer meshtags to submesh
-    subdomain.ft, subdomain.facet_to_parent = transfer_meshtags_to_submesh(
-        mesh, mt, subdomain.submesh, subdomain.v_map, subdomain.submesh_to_mesh
-    )
 
 
-# Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
-# TODO ask Jorgen what this is for
-mesh.topology.create_connectivity(fdim, tdim)
-f_to_c = mesh.topology.connectivity(fdim, tdim)
-for interface in list_of_interfaces:
-    for facet in mt.find(interface):
-        cells = f_to_c.links(facet)
-        assert len(cells) == 2
-        for domain in list_of_interfaces[interface]:
-            map = domain.parent_to_submesh[cells]
-            domain.parent_to_submesh[cells] = max(map)
-
-# ._cpp_object needed on dolfinx 0.8.0
-entity_maps = {
-    subdomain.submesh: subdomain.parent_to_submesh for subdomain in list_of_subdomains
-}
-
-
-def define_function_spaces(subdomain: F.VolumeSubdomain):
-    # get number of species defined in the subdomain
-    all_species = [
-        species for species in list_of_species if subdomain in species.subdomains
-    ]
-
-    # instead of using the set function we use a list to keep the order
-    unique_species = []
-    for species in all_species:
-        if species not in unique_species:
-            unique_species.append(species)
-    nb_species = len(unique_species)
+i0 = Interface(mesh, mt, 5, (left_domain, mid_domain))
+i1 = Interface(mesh, mt, 6, (mid_domain, right_domain))
+interfaces = [i0, i1]
 
-    degree = 1
+def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
     element_CG = basix.ufl.element(
         basix.ElementFamily.P,
-        subdomain.submesh.basix_cell(),
+        submesh.basix_cell(),
         degree,
         basix.LagrangeVariant.equispaced,
     )
-    element = basix.ufl.mixed_element([element_CG] * nb_species)
-    V = dolfinx.fem.functionspace(subdomain.submesh, element)
-    u = dolfinx.fem.Function(V)
-    u_n = dolfinx.fem.Function(V)
-
+    element = basix.ufl.mixed_element([element_CG, element_CG])
+    V = dolfinx.fem.functionspace(submesh, element)
+    u = dolfinx.fem.Function(V, name=f"u_{value}")
     us = list(ufl.split(u))
-    u_ns = list(ufl.split(u_n))
     vs = list(ufl.TestFunctions(V))
-    for i, species in enumerate(unique_species):
-        species.subdomain_to_solution[subdomain] = us[i]
-        species.subdomain_to_prev_solution[subdomain] = u_ns[i]
-        species.subdomain_to_test_function[subdomain] = vs[i]
-    subdomain.u = u
-
-
-def define_formulation(subdomain: F.VolumeSubdomain):
-    form = 0
-    # add diffusion and time derivative for each species
-    for spe in list_of_species:
-        u = spe.subdomain_to_solution[subdomain]
-        u_n = spe.subdomain_to_prev_solution[subdomain]
-        v = spe.subdomain_to_test_function[subdomain]
-        dx = subdomain.dx
-
-        D = subdomain.material.get_diffusion_coefficient(mesh, T, spe)
-        if spe.mobile:
-            form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
-
-    for reaction in list_of_reactions:
-        if reaction.volume != subdomain:
-            continue
-        for species in reaction.reactant + reaction.product:
-            if isinstance(species, F.Species):
-                # TODO remove
-                # temporarily overide the solution and test function to the one of the subdomain
-                species.solution = species.subdomain_to_solution[subdomain]
-                species.test_function = species.subdomain_to_test_function[subdomain]
-
-        for reactant in reaction.reactant:
-            if isinstance(reactant, F.Species):
-                form += (
-                    reaction.reaction_term(T)
-                    * reactant.subdomain_to_test_function[subdomain]
-                    * dx
-                )
-
-        # product
-        if isinstance(reaction.product, list):
-            products = reaction.product
-        else:
-            products = [reaction.product]
-        for product in products:
-            form += (
-                -reaction.reaction_term(T)
-                * product.subdomain_to_test_function[subdomain]
-                * dx
-            )
-
-    subdomain.F = form
-
-
-for subdomain in list_of_subdomains:
-    define_function_spaces(subdomain)
     ct_r = dolfinx.mesh.meshtags(
         mesh,
         mesh.topology.dim,
-        subdomain.submesh_to_mesh,
-        np.full_like(subdomain.submesh_to_mesh, 1, dtype=np.int32),
+        submesh_to_mesh,
+        np.full_like(submesh_to_mesh, 1, dtype=np.int32),
     )
-    subdomain.dx = ufl.Measure("dx", domain=mesh, subdomain_data=ct_r, subdomain_id=1)
-    define_formulation(subdomain)
-    subdomain.u.name = f"u_{subdomain.id}"
+    val = dolfinx.fem.Constant(submesh, value)
+    dx_r = ufl.Measure("dx", domain=mesh, subdomain_data=ct_r, subdomain_id=1)
+    F = ufl.inner(ufl.grad(us[0]), ufl.grad(vs[0])) * dx_r - val * vs[0] * dx_r
+    k = 2
+    p = 0.1
+    n = 0.5
+    F += k * us[0] * (n - us[1]) * vs[1] * dx_r - p * us[1] * vs[1] * dx_r
+    return u, vs, F
 
-# boundary conditions
-bcs = []
-for boundary_condition in list_of_bcs:
-    volume_subdomain = surface_to_volume[boundary_condition.subdomain]
-    bc = dolfinx.fem.Function(volume_subdomain.u.function_space)
-    bc.x.array[:] = boundary_condition.value
-    volume_subdomain.submesh.topology.create_connectivity(
-        volume_subdomain.submesh.topology.dim - 1,
-        volume_subdomain.submesh.topology.dim,
-    )
-    bc = dolfinx.fem.dirichletbc(
-        bc,
-        dolfinx.fem.locate_dofs_topological(
-            volume_subdomain.u.function_space.sub(0),
-            fdim,
-            volume_subdomain.ft.find(boundary_condition.subdomain.id),
-        ),
+
+# for each subdomain, define the interior equation
+for subdomain in list_of_subdomains:
+    degree = 1
+    subdomain.u, subdomain.vs, subdomain.F = define_interior_eq(
+        mesh, degree, subdomain.submesh, subdomain.submesh_to_mesh, 0.0
     )
-    bcs.append(bc)
+    subdomain.u.name = f"u_{subdomain.id}"
 
-# Add coupling term to the interface
-# Get interface markers on submesh b
-for interface in list_of_interfaces:
 
-    dInterface = ufl.Measure(
-        "dS", domain=mesh, subdomain_data=mt, subdomain_id=interface
+def compute_mapped_interior_facet_data(interface: Interface):
+    """
+    Compute integration data for interface integrals.
+    We define the first domain on an interface as the "+" restriction,
+    meaning that we must sort all integration entities in this order
+    
+    Parameters
+        interface: Interface between two subdomains
+    Returns
+        integration_data: Integration data for interior facets
+    """
+    assert (not interface.subdomains[0].padded) and (not interface.subdomains[1].padded)
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    integration_data = compute_integration_domains(
+        dolfinx.fem.IntegralType.interior_facet, mesh.topology, interface.mt.find(interface.id), interface.mt.dim)
+
+    ordered_integration_data = integration_data.reshape(-1, 4).copy()
+
+    mapped_cell_0 = interface.subdomains[0].parent_to_submesh[integration_data[0::4]]
+    mapped_cell_1 = interface.subdomains[0].parent_to_submesh[integration_data[2::4]]
+
+    switch = mapped_cell_1 > mapped_cell_0
+    # Order restriction on one side        
+    if True in switch:
+        ordered_integration_data[switch, [0, 1, 2, 3]] = ordered_integration_data[
+            switch, [2, 3, 0, 1]
+        ]
+    
+    # Check that other restriction lies in other interface
+    domain1_cell = interface.subdomains[1].parent_to_submesh[ordered_integration_data[:, 2]]
+    assert (domain1_cell >=0).all()
+
+    return (interface.id, ordered_integration_data.reshape(-1))
+
+
+
+integral_data = [compute_mapped_interior_facet_data(interface) for interface in interfaces]
+[interface.pad_parent_maps() for interface in interfaces]
+dInterface = ufl.Measure(
+        "dS", domain=mesh, subdomain_data=integral_data
     )
-    b_res = "+"
-    t_res = "-"
-
-    # look at the first facet on interface
-    # and get the two cells that are connected to it
-    # and get the material properties of these cells
-    first_facet_interface = mt.find(interface)[0]
-    c_plus, c_minus = (
-        f_to_c.links(first_facet_interface)[0],
-        f_to_c.links(first_facet_interface)[1],
-    )
-    id_minus, id_plus = ct.values[c_minus], ct.values[c_plus]
 
-    for subdomain in list_of_interfaces[interface]:
-        if subdomain.id == id_plus:
-            subdomain_1 = subdomain
-        if subdomain.id == id_minus:
-            subdomain_2 = subdomain
+def mixed_term(u, v, n):
+    return ufl.dot(ufl.grad(u), n) * v
 
-    v_b = H.subdomain_to_test_function[subdomain_1](b_res)
-    v_t = H.subdomain_to_test_function[subdomain_2](t_res)
+n = ufl.FacetNormal(mesh)
+cr = ufl.Circumradius(mesh)
 
-    u_b = H.subdomain_to_solution[subdomain_1](b_res)
-    u_t = H.subdomain_to_solution[subdomain_2](t_res)
+gamma = 10.0
 
-    def mixed_term(u, v, n):
-        return ufl.dot(ufl.grad(u), n) * v
 
-    n = ufl.FacetNormal(mesh)
+entity_maps = {sd.submesh: sd.parent_to_submesh for sd in list_of_subdomains}
+for interface in interfaces:
+    subdomain_1, subdomain_2 = interface.subdomains
+    b_res, t_res = interface.restriction
     n_b = n(b_res)
     n_t = n(t_res)
-    cr = ufl.Circumradius(mesh)
     h_b = 2 * cr(b_res)
     h_t = 2 * cr(t_res)
-    gamma = 400.0  # this needs to be "sufficiently large"
 
-    K_b = K_S_fun(T(b_res), subdomain_1.material.K_S_0, subdomain_1.material.E_K_S)
-    K_t = K_S_fun(T(t_res), subdomain_2.material.K_S_0, subdomain_2.material.E_K_S)
+
+    v_b = subdomain_1.vs[0](b_res)
+    v_t = subdomain_2.vs[0](t_res)
+
+    u_bs = list(ufl.split(subdomain_1.u))
+    u_ts = list(ufl.split(subdomain_2.u))
+    u_b = u_bs[0](b_res)
+    u_t = u_ts[0](t_res)
+    # fabricate K
+    W_0 = dolfinx.fem.functionspace(subdomain_1.submesh, ("DG", 0))
+    K_0 = dolfinx.fem.Function(W_0, name=f"K_{subdomain_1.id}")
+    K_0.x.array[:] = 2
+    W_1 = dolfinx.fem.functionspace(subdomain_2.submesh, ("DG", 0))
+    K_1 = dolfinx.fem.Function(W_1, name=f"K_{subdomain_2.id}")
+    K_1.x.array[:] = 4
+
+    K_b = K_0(b_res)
+    K_t = K_1(t_res)
 
     F_0 = (
-        -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
-        - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
+        -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(interface.id)
+        - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
     )
 
     F_1 = (
-        +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
-        - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
+        +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(interface.id)
+        - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
     )
-    F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
-    F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
+    F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface(interface.id)
+    F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface(interface.id)
 
     subdomain_1.F += F_0
     subdomain_2.F += F_1
 
+
 J = []
 forms = []
-for subdomain1 in list_of_subdomains:
+for i,subdomain1 in enumerate(list_of_subdomains):
     jac = []
     form = subdomain1.F
-    for subdomain2 in list_of_subdomains:
+    for j, subdomain2 in enumerate(list_of_subdomains):
         jac.append(
             dolfinx.fem.form(
                 ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
             )
         )
+        sp = dolfinx.fem.create_sparsity_pattern(jac[-1])
+
     J.append(jac)
     forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
+# boundary conditions
+b_bc = dolfinx.fem.Function(left_domain.u.function_space)
+b_bc.x.array[:] = 0.2
+left_domain.submesh.topology.create_connectivity(
+    left_domain.submesh.topology.dim - 1, left_domain.submesh.topology.dim
+)
+bc_b = dolfinx.fem.dirichletbc(
+    b_bc,
+    dolfinx.fem.locate_dofs_topological(
+        left_domain.u.function_space.sub(0), fdim, left_domain.ft.find(1)
+    ),
+)
+
+t_bc = dolfinx.fem.Function(right_domain.u.function_space)
+t_bc.x.array[:] = 0.05
+right_domain.submesh.topology.create_connectivity(
+    right_domain.submesh.topology.dim - 1, right_domain.submesh.topology.dim
+)
+bc_t = dolfinx.fem.dirichletbc(
+    t_bc,
+    dolfinx.fem.locate_dofs_topological(
+        right_domain.u.function_space.sub(0), fdim, right_domain.ft.find(2)
+    ),
+)
+bcs = [bc_b, bc_t]
+
 
 solver = NewtonSolver(
     forms,
@@ -397,43 +557,22 @@ def mixed_term(u, v, n):
 
 
 # derived quantities
-entity_maps[mesh] = bottom_domain.submesh_to_mesh
+V = dolfinx.fem.functionspace(mesh, ("CG", 1))
+T = dolfinx.fem.Function(V)
+T.interpolate(lambda x: 200 + x[1])
 
-ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
-ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
-dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
-dx = ufl.Measure("dx", domain=mesh)
 
-n_b = ufl.FacetNormal(bottom_domain.submesh)
-n_t = ufl.FacetNormal(top_domain.submesh)
+T_b = dolfinx.fem.Function(right_domain.u.sub(0).collapse().function_space)
+T_b.interpolate(T)
 
-form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
-print(dolfinx.fem.assemble_scalar(form))
+ds_b = ufl.Measure("ds", domain=right_domain.submesh)
+dx_b = ufl.Measure("dx", domain=right_domain.submesh)
+dx = ufl.Measure("dx", domain=mesh)
 
-form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b)
-print(dolfinx.fem.assemble_scalar(form))
+n_b = ufl.FacetNormal(left_domain.submesh)
 
-form = dolfinx.fem.form(T * dx_b, entity_maps={mesh: bottom_domain.submesh_to_mesh})
+form = dolfinx.fem.form(left_domain.u.sub(0) * dx_b, entity_maps=entity_maps)
 print(dolfinx.fem.assemble_scalar(form))
 
-id_interface = 5
-form = dolfinx.fem.form(
-    ufl.dot(
-        D_fun(T, bottom_domain.material.D_0, bottom_domain.material.E_D)
-        * ufl.grad(bottom_domain.u.sub(0)),
-        n_b,
-    )
-    * ds_b(id_interface),
-    entity_maps={mesh: bottom_domain.submesh_to_mesh},
-)
-print(dolfinx.fem.assemble_scalar(form))
-form = dolfinx.fem.form(
-    ufl.dot(
-        D_fun(T, top_domain.material.D_0, top_domain.material.E_D)
-        * ufl.grad(top_domain.u.sub(0)),
-        n_t,
-    )
-    * ds_t(id_interface),
-    entity_maps={mesh: top_domain.submesh_to_mesh},
-)
-print(dolfinx.fem.assemble_scalar(form))
+form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(left_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
+print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file
diff --git a/festim/__init__.py b/festim/__init__.py
index d32342265..492623f4d 100644
--- a/festim/__init__.py
+++ b/festim/__init__.py
@@ -46,6 +46,7 @@
 from .subdomain.surface_subdomain_1d import SurfaceSubdomain1D
 from .subdomain.volume_subdomain import VolumeSubdomain, find_volume_from_id
 from .subdomain.volume_subdomain_1d import VolumeSubdomain1D
+from .subdomain.interface import Interface
 
 from .species import Species, Trap, ImplicitSpecies, find_species_from_name
 
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 8c1d854f1..36de48a79 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -822,50 +822,11 @@ def create_initial_conditions(self):
             )
 
     def create_submeshes(self):
-        mesh = self.mesh.mesh
-        ct = self.volume_meshtags
-        mt = self.facet_meshtags
-
-        gdim = mesh.geometry.dim
-        tdim = mesh.topology.dim
-        fdim = tdim - 1
-
-        num_facets_local = (
-            mesh.topology.index_map(fdim).size_local
-            + mesh.topology.index_map(fdim).num_ghosts
-        )
 
         for subdomain in self.volume_subdomains:
-            subdomain.submesh, subdomain.submesh_to_mesh, subdomain.v_map = (
-                dolfinx.mesh.create_submesh(mesh, tdim, ct.find(subdomain.id))[0:3]
-            )
-
-            subdomain.parent_to_submesh = np.full(num_facets_local, -1, dtype=np.int32)
-            subdomain.parent_to_submesh[subdomain.submesh_to_mesh] = np.arange(
-                len(subdomain.submesh_to_mesh), dtype=np.int32
-            )
-
-            # We need to modify the cell maps, as for `dS` integrals of interfaces between submeshes, there is no entity to map to.
-            # We use the entity on the same side to fix this (as all restrictions are one-sided)
-
-            # Transfer meshtags to submesh
-            subdomain.ft, subdomain.facet_to_parent = transfer_meshtags_to_submesh(
-                mesh, mt, subdomain.submesh, subdomain.v_map, subdomain.submesh_to_mesh
-            )
-
-        # Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
-        # TODO ask Jorgen what this is for
-        mesh.topology.create_connectivity(fdim, tdim)
-        f_to_c = mesh.topology.connectivity(fdim, tdim)
-        for interface in self.interfaces:
-            for facet in mt.find(interface):
-                cells = f_to_c.links(facet)
-                assert len(cells) == 2
-                for domain in self.interfaces[interface]:
-                    map = domain.parent_to_submesh[cells]
-                    domain.parent_to_submesh[cells] = max(map)
+            subdomain.create_subdomain(self.mesh.mesh, self.volume_meshtags)
+            subdomain.transfer_meshtag(self.mesh.mesh, self.facet_meshtags)
 
-        self.f_to_c = f_to_c
 
     def define_function_spaces(self, subdomain: F.VolumeSubdomain):
         # get number of species defined in the subdomain
@@ -972,24 +933,42 @@ def create_formulation(self):
         f_to_c = mesh.topology.connectivity(mesh.topology.dim - 1, mesh.topology.dim)
 
         for interface in self.interfaces:
+            interface.mesh = mesh
+            interface.mt = mt
 
-            dInterface = ufl.Measure(
-                "dS", domain=mesh, subdomain_data=mt, subdomain_id=interface
+        integral_data = [interface.compute_mapped_interior_facet_data(mesh) for interface in self.interfaces]
+        [interface.pad_parent_maps() for interface in self.interfaces]
+        dInterface = ufl.Measure(
+                "dS", domain=mesh, subdomain_data=integral_data
             )
-            b_res = "+"
-            t_res = "-"
+        def mixed_term(u, v, n):
+            return ufl.dot(ufl.grad(u), n) * v
+        n = ufl.FacetNormal(mesh)
+        cr = ufl.Circumradius(mesh)
+
+        gamma = 10.0
+
+        entity_maps = {sd.submesh: sd.parent_to_submesh for sd in self.volume_subdomains}
+        for interface in self.interfaces:
+
+            subdomain_1, subdomain_2 = interface.subdomains
+            b_res, t_res = interface.restriction
+            n_b = n(b_res)
+            n_t = n(t_res)
+            h_b = 2 * cr(b_res)
+            h_t = 2 * cr(t_res)
 
             # look at the first facet on interface
             # and get the two cells that are connected to it
             # and get the material properties of these cells
-            first_facet_interface = mt.find(interface)[0]
+            first_facet_interface = mt.find(interface.id)[0]
             c_plus, c_minus = (
                 f_to_c.links(first_facet_interface)[0],
                 f_to_c.links(first_facet_interface)[1],
             )
             id_minus, id_plus = ct.values[c_minus], ct.values[c_plus]
 
-            for subdomain in self.interfaces[interface]:
+            for subdomain in interface.subdomains:
                 if subdomain.id == id_plus:
                     subdomain_1 = subdomain
                 if subdomain.id == id_minus:
@@ -1006,18 +985,6 @@ def create_formulation(self):
             u_b = H.subdomain_to_solution[subdomain_1](b_res)
             u_t = H.subdomain_to_solution[subdomain_2](t_res)
 
-            def mixed_term(u, v, n):
-                return ufl.dot(ufl.grad(u), n) * v
-
-            n = ufl.FacetNormal(mesh)
-            n_b = n(b_res)
-            n_t = n(t_res)  # this doesn't seem to be used
-            cr = ufl.Circumradius(mesh)
-            h_b = 2 * cr(b_res)
-            h_t = 2 * cr(t_res)
-            # TODO make this a user parameter
-            gamma = 10.0  # this needs to be "sufficiently large"
-
             K_b = K_S_fun(
                 self.temperature_fenics(b_res),
                 subdomain_1.material.K_S_0,
@@ -1030,33 +997,35 @@ def mixed_term(u, v, n):
             )
 
             F_0 = (
-                -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
-                - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface
+                -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(interface.id)
+                - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
             )
 
             F_1 = (
-                +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
-                - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface
+                +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(interface.id)
+                - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
             )
-            F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface
-            F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface
+            F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface(interface.id)
+            F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface(interface.id)
 
             subdomain_1.F += F_0
             subdomain_2.F += F_1
 
         J = []
         forms = []
-        for subdomain1 in self.volume_subdomains:
+        for i, subdomain1 in enumerate(self.volume_subdomains):
             jac = []
             form = subdomain1.F
-            for subdomain2 in self.volume_subdomains:
+            for j, subdomain2 in enumerate(self.volume_subdomains):
                 jac.append(
                     dolfinx.fem.form(
-                        ufl.derivative(form, subdomain2.u), entity_maps=self.entity_maps
+                        ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
                     )
                 )
+                sp = dolfinx.fem.create_sparsity_pattern(jac[-1])  # NOTE this variable isn't used anywhere
             J.append(jac)
-            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=self.entity_maps))
+            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
+        
         self.forms = forms
         self.J = J
 
diff --git a/festim/mesh/mesh.py b/festim/mesh/mesh.py
index e99d1094a..aec3b17d5 100644
--- a/festim/mesh/mesh.py
+++ b/festim/mesh/mesh.py
@@ -94,7 +94,8 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
 
         # tag interfaces
         interfaces = interfaces or {}  # if interfaces is None, set it to empty dict
-        for interface_id, (domain_1, domain_2) in interfaces.items():
+        for interface in interfaces:
+            (domain_1, domain_2) = interface.subdomains
             all_1_facets = dolfinx.mesh.compute_incident_entities(
                 self.mesh.topology, volume_meshtags.find(domain_1.id), self.vdim, self.fdim
             )
@@ -102,7 +103,7 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
                 self.mesh.topology, volume_meshtags.find(domain_2.id), self.vdim, self.fdim
             )
             interface_entities = np.intersect1d(all_1_facets, all_2_facets)
-            tags_facets[interface_entities] = interface_id
+            tags_facets[interface_entities] = interface.id
 
         facet_meshtags = meshtags(self.mesh, self.fdim, mesh_facet_indices, tags_facets)
 
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
new file mode 100644
index 000000000..ab18034a1
--- /dev/null
+++ b/festim/subdomain/interface.py
@@ -0,0 +1,65 @@
+import festim as F
+import dolfinx
+from dolfinx.cpp.fem import compute_integration_domains
+import numpy as np
+
+
+class Interface():
+    id: int
+    subdomains: tuple[F.VolumeSubdomain, F.VolumeSubdomain]
+    parent_mesh: dolfinx.mesh.Mesh
+    restriction: list[str, str] = ("+", "-")
+    padded: bool
+
+    def __init__(self, parent_mesh, mt, id, subdomains):
+        self.id = id
+        self.subdomains = tuple(subdomains)
+        self.mt = mt
+        self.parent_mesh = parent_mesh
+
+    def pad_parent_maps(self):
+        """Workaround to make sparsity-pattern work without skips
+        """ 
+
+        integration_data = compute_integration_domains(
+                dolfinx.fem.IntegralType.interior_facet, self.parent_mesh.topology, self.mt.find(self.id), self.mt.dim).reshape(-1, 4)
+        for i in range(2):
+            # We pad the parent to submesh map to make sure that sparsity pattern is correct
+            mapped_cell_0 = self.subdomains[i].parent_to_submesh[integration_data[:, 0]]
+            mapped_cell_1 = self.subdomains[i].parent_to_submesh[integration_data[:, 2]]
+            max_cells = np.maximum(mapped_cell_0, mapped_cell_1)
+            self.subdomains[i].parent_to_submesh[integration_data[:, 0]] = max_cells
+            self.subdomains[i].parent_to_submesh[integration_data[:, 2]] = max_cells
+            self.subdomains[i].padded = True
+    
+    def compute_mapped_interior_facet_data(self, mesh):
+        """
+        Compute integration data for interface integrals.
+        We define the first domain on an interface as the "+" restriction,
+        meaning that we must sort all integration entities in this order
+        
+        Returns
+            integration_data: Integration data for interior facets
+        """
+        assert (not self.subdomains[0].padded) and (not self.subdomains[1].padded)
+        mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+        integration_data = compute_integration_domains(
+            dolfinx.fem.IntegralType.interior_facet, mesh.topology, self.mt.find(self.id), self.mt.dim)
+
+        ordered_integration_data = integration_data.reshape(-1, 4).copy()
+
+        mapped_cell_0 = self.subdomains[0].parent_to_submesh[integration_data[0::4]]
+        mapped_cell_1 = self.subdomains[0].parent_to_submesh[integration_data[2::4]]
+
+        switch = mapped_cell_1 > mapped_cell_0
+        # Order restriction on one side        
+        if True in switch:
+            ordered_integration_data[switch, [0, 1, 2, 3]] = ordered_integration_data[
+                switch, [2, 3, 0, 1]
+            ]
+        
+        # Check that other restriction lies in other interface
+        domain1_cell = self.subdomains[1].parent_to_submesh[ordered_integration_data[:, 2]]
+        assert (domain1_cell >=0).all()
+
+        return (self.id, ordered_integration_data.reshape(-1))
\ No newline at end of file
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index 3c476ceb8..32a39cfdd 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -1,6 +1,7 @@
 from dolfinx.mesh import locate_entities
-
+import dolfinx
 import numpy as np
+from festim.helpers_discontinuity import transfer_meshtags_to_submesh
 
 
 class VolumeSubdomain:
@@ -10,6 +11,15 @@ class VolumeSubdomain:
     Args:
         id (int): the id of the volume subdomain
     """
+    id: int
+    submesh: dolfinx.mesh.Mesh
+    submesh_to_mesh: np.ndarray
+    parent_mesh: dolfinx.mesh.Mesh
+    parent_to_submesh: np.ndarray
+    v_map: np.ndarray
+    facet_to_parent: np.ndarray
+    ft: dolfinx.mesh.MeshTags
+    padded:bool
 
     def __init__(self, id, material):
         self.id = id
@@ -31,6 +41,29 @@ def locate_subdomain_entities(self, mesh, vdim):
         entities = locate_entities(mesh, vdim, lambda x: np.full(x.shape[1], True))
         return entities
 
+    def create_subdomain(self, mesh, marker):
+        assert marker.dim == mesh.topology.dim
+        self.parent_mesh = mesh
+        self.submesh, self.submesh_to_mesh, self.v_map = (
+            dolfinx.mesh.create_submesh(mesh, marker.dim, marker.find(self.id))[0:3]
+        )
+        num_cells_local = (
+            mesh.topology.index_map(marker.dim).size_local
+            + mesh.topology.index_map(marker.dim).num_ghosts
+        )
+        self.parent_to_submesh = np.full(num_cells_local, -1, dtype=np.int32)
+        self.parent_to_submesh[self.submesh_to_mesh] = np.arange(
+            len(self.submesh_to_mesh), dtype=np.int32
+        )
+        self.padded=False
+
+    def transfer_meshtag(self, mesh, tag):
+        # Transfer meshtags to submesh
+        assert self.submesh is not None, "Need to call create_subdomain first"
+        self.ft, self.facet_to_parent = transfer_meshtags_to_submesh(
+            mesh, tag, self.submesh, self.v_map, self.submesh_to_mesh
+        )
+
 
 def find_volume_from_id(id: int, volumes: list):
     """Returns the correct volume subdomain object from a list of volume ids
diff --git a/test_implementation.py b/test_implementation.py
index 4e7df1395..c190b777d 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -84,7 +84,7 @@ def half(x):
 my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
 # we should be able to automate this
-my_model.interfaces = {5: [bottom_domain, top_domain]}
+my_model.interfaces = [F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain))]
 my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
diff --git a/test_implementation_2.py b/test_implementation_2.py
index 28ab2d256..2aa4ddd2e 100644
--- a/test_implementation_2.py
+++ b/test_implementation_2.py
@@ -84,7 +84,7 @@ def half(x):
 my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
 # we should be able to automate this
-my_model.interfaces = {5: [bottom_domain, top_domain]}
+my_model.interfaces = [F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain))]
 my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
diff --git a/test_implementation_3.py b/test_implementation_3.py
index 568bd5cbc..fb9e9a512 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -7,7 +7,7 @@
 interface_2 = 0.7
 
 # for some reason if the mesh isn't fine enough then I have a random SEGV error
-N = 1500
+N = 2000
 vertices = np.concatenate(
     [
         np.linspace(0, interface_1, num=N),
@@ -37,7 +37,10 @@
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
 # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
-my_model.interfaces = {6: [left_domain, middle_domain], 7: [middle_domain, right_domain]}
+my_model.interfaces = [
+    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)),
+    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)),
+]
 my_model.subdomains = [left_domain, middle_domain, right_domain, left_surface, right_surface]
 my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
 
diff --git a/test_implementation_transient.py b/test_implementation_transient.py
index 54589e940..b481edcab 100644
--- a/test_implementation_transient.py
+++ b/test_implementation_transient.py
@@ -36,10 +36,12 @@
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
 # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
-my_model.interfaces = {
-    6: [left_domain, middle_domain],
-    7: [middle_domain, right_domain],
-}
+# TODO remove mesh and meshtags from these arguments
+my_model.interfaces = [
+    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)),
+    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)),
+]
+
 my_model.subdomains = [
     left_domain,
     middle_domain,

From c2096009e11a85b1b22bdf1f9a24d38cc9d00ef5 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 12:05:47 +0000
Subject: [PATCH 048/134] removed unused variable

---
 discontinuity_generic.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 1277be7be..49f2a45a4 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -497,7 +497,6 @@ def mixed_term(u, v, n):
                 ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
             )
         )
-        sp = dolfinx.fem.create_sparsity_pattern(jac[-1])
 
     J.append(jac)
     forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))

From 97067b66fe98557bfbc49e60baad4b6dce2660c0 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 12:06:43 +0000
Subject: [PATCH 049/134] removed unused variables

---
 festim/hydrogen_transport_problem.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 36de48a79..93b4be455 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -1013,16 +1013,15 @@ def mixed_term(u, v, n):
 
         J = []
         forms = []
-        for i, subdomain1 in enumerate(self.volume_subdomains):
+        for subdomain1 in self.volume_subdomains:
             jac = []
             form = subdomain1.F
-            for j, subdomain2 in enumerate(self.volume_subdomains):
+            for subdomain2 in self.volume_subdomains:
                 jac.append(
                     dolfinx.fem.form(
                         ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
                     )
                 )
-                sp = dolfinx.fem.create_sparsity_pattern(jac[-1])  # NOTE this variable isn't used anywhere
             J.append(jac)
             forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
         

From ad3fd379d9fd2aadcc0b0ffd45039047daf45593 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 12:40:19 +0000
Subject: [PATCH 050/134] get_solubility_coefficient + black

---
 festim/hydrogen_transport_problem.py | 86 +++++++++++++++++-----------
 festim/material.py                   | 56 ++++++++++++++++++
 2 files changed, 108 insertions(+), 34 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 93b4be455..2656dbfb0 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -715,10 +715,6 @@ def post_processing(self):
                 export.write(float(self.t))
 
 
-def K_S_fun(T, K_S_0, E_K_S):
-    return K_S_0 * ufl.exp(-E_K_S / F.k_B / T)
-
-
 class HTransportProblemDiscontinuous(HydrogenTransportProblem):
 
     def __init__(
@@ -827,7 +823,6 @@ def create_submeshes(self):
             subdomain.create_subdomain(self.mesh.mesh, self.volume_meshtags)
             subdomain.transfer_meshtag(self.mesh.mesh, self.facet_meshtags)
 
-
     def define_function_spaces(self, subdomain: F.VolumeSubdomain):
         # get number of species defined in the subdomain
         all_species = [
@@ -860,9 +855,15 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_solution[subdomain] = us[i]
             species.subdomain_to_prev_solution[subdomain] = u_ns[i]
             species.subdomain_to_test_function[subdomain] = vs[i]
-            species.subdomain_to_post_processing_solution[subdomain] = u.sub(i).collapse()
-            species.subdomain_to_collapsed_function_space[subdomain] = (V.sub(i).collapse())
-            species.subdomain_to_post_processing_solution[subdomain].name = f"{species.name}_{subdomain.id}"
+            species.subdomain_to_post_processing_solution[subdomain] = u.sub(
+                i
+            ).collapse()
+            species.subdomain_to_collapsed_function_space[subdomain] = V.sub(
+                i
+            ).collapse()
+            species.subdomain_to_post_processing_solution[subdomain].name = (
+                f"{species.name}_{subdomain.id}"
+            )
         subdomain.u = u
         subdomain.u_n = u_n
 
@@ -936,19 +937,24 @@ def create_formulation(self):
             interface.mesh = mesh
             interface.mt = mt
 
-        integral_data = [interface.compute_mapped_interior_facet_data(mesh) for interface in self.interfaces]
+        integral_data = [
+            interface.compute_mapped_interior_facet_data(mesh)
+            for interface in self.interfaces
+        ]
         [interface.pad_parent_maps() for interface in self.interfaces]
-        dInterface = ufl.Measure(
-                "dS", domain=mesh, subdomain_data=integral_data
-            )
+        dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)
+
         def mixed_term(u, v, n):
             return ufl.dot(ufl.grad(u), n) * v
+
         n = ufl.FacetNormal(mesh)
         cr = ufl.Circumradius(mesh)
 
         gamma = 10.0
 
-        entity_maps = {sd.submesh: sd.parent_to_submesh for sd in self.volume_subdomains}
+        entity_maps = {
+            sd.submesh: sd.parent_to_submesh for sd in self.volume_subdomains
+        }
         for interface in self.interfaces:
 
             subdomain_1, subdomain_2 = interface.subdomains
@@ -985,28 +991,40 @@ def mixed_term(u, v, n):
             u_b = H.subdomain_to_solution[subdomain_1](b_res)
             u_t = H.subdomain_to_solution[subdomain_2](t_res)
 
-            K_b = K_S_fun(
-                self.temperature_fenics(b_res),
-                subdomain_1.material.K_S_0,
-                subdomain_1.material.E_K_S,
+            K_b = subdomain_1.material.get_solubility_coefficient(
+                self.mesh.mesh, self.temperature_fenics(b_res), H
             )
-            K_t = K_S_fun(
-                self.temperature_fenics(t_res),
-                subdomain_2.material.K_S_0,
-                subdomain_2.material.E_K_S,
+            K_t = subdomain_2.material.get_solubility_coefficient(
+                self.mesh.mesh, self.temperature_fenics(t_res), H
             )
 
-            F_0 = (
-                -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(interface.id)
-                - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
+            F_0 = -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(
+                interface.id
+            ) - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(
+                interface.id
             )
 
-            F_1 = (
-                +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(interface.id)
-                - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
+            F_1 = +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(
+                interface.id
+            ) - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(
+                interface.id
+            )
+            F_0 += (
+                2
+                * gamma
+                / (h_b + h_t)
+                * (u_b / K_b - u_t / K_t)
+                * v_b
+                * dInterface(interface.id)
+            )
+            F_1 += (
+                -2
+                * gamma
+                / (h_b + h_t)
+                * (u_b / K_b - u_t / K_t)
+                * v_t
+                * dInterface(interface.id)
             )
-            F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface(interface.id)
-            F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface(interface.id)
 
             subdomain_1.F += F_0
             subdomain_2.F += F_1
@@ -1024,7 +1042,7 @@ def mixed_term(u, v, n):
                 )
             J.append(jac)
             forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
-        
+
         self.forms = forms
         self.J = J
 
@@ -1059,7 +1077,9 @@ def initialise_exports(self):
             if isinstance(export, F.VTXExport):
                 species = export.field[0]
                 # override post_processing_solution attribute of species
-                species.post_processing_solution = species.subdomain_to_post_processing_solution[export.subdomain]
+                species.post_processing_solution = (
+                    species.subdomain_to_post_processing_solution[export.subdomain]
+                )
                 export.define_writer(MPI.COMM_WORLD)
             else:
                 raise NotImplementedError("Export type not implemented")
@@ -1073,9 +1093,7 @@ def post_processing(self):
                     subdomain
                 ]
                 u = subdomain.u
-                v0_to_V = species.subdomain_to_collapsed_function_space[
-                    subdomain
-                ][1]
+                v0_to_V = species.subdomain_to_collapsed_function_space[subdomain][1]
                 collapsed_function.x.array[:] = u.x.array[v0_to_V]
 
         for export in self.exports:
diff --git a/festim/material.py b/festim/material.py
index dded153a4..56f86e0b7 100644
--- a/festim/material.py
+++ b/festim/material.py
@@ -168,3 +168,59 @@ def get_diffusion_coefficient(self, mesh, temperature, species=None):
 
         else:
             raise ValueError("D_0 and E_D must be either floats or dicts")
+
+    def get_solubility_coefficient(self, mesh, temperature, species=None):
+        """Defines the diffusion coefficient
+        Args:
+
+            mesh (dolfinx.mesh.Mesh): the domain mesh
+            temperature (dolfinx.fem.Constant): the temperature
+            species (festim.Species, optional): the species we want the diffusion
+                coefficient of. Only needed if K_S_0 and E_K_S are dicts.
+        Returns:
+            ufl.algebra.Product: the solubility coefficient
+        """
+        # TODO use get_D_0 and get_E_D to refactore this method, something like:
+        # K_S_0 = self.get_K_S_0(species=species)
+        # E_K_S = self.get_E_K_S(species=species)
+
+        # K_S_0 = F.as_fenics_constant(K_S_0, mesh)
+        # E_K_S = F.as_fenics_constant(E_K_S, mesh)
+
+        # return K_S_0 * ufl.exp(-E_K_S / F.k_B / temperature)
+
+        if isinstance(self.K_S_0, (float, int)) and isinstance(
+            self.E_K_S, (float, int)
+        ):
+            K_S_0 = F.as_fenics_constant(self.K_S_0, mesh)
+            E_K_S = F.as_fenics_constant(self.E_K_S, mesh)
+
+            return K_S_0 * ufl.exp(-E_K_S / F.k_B / temperature)
+
+        elif isinstance(self.K_S_0, dict) and isinstance(self.E_K_S, dict):
+            # check D_0 and E_D have the same keys
+            # this check should go in a setter
+            if list(self.K_S_0.keys()) != list(self.E_K_S.keys()):
+                raise ValueError("K_S_0 and E_K_S have different keys")
+
+            if species is None:
+                raise ValueError(
+                    "species must be provided if K_S_0 and E_K_S are dicts"
+                )
+
+            if species in self.K_S_0:
+                K_S_0 = F.as_fenics_constant(self.K_S_0[species], mesh)
+            elif species.name in self.K_S_0:
+                K_S_0 = F.as_fenics_constant(self.K_S_0[species.name], mesh)
+            else:
+                raise ValueError(f"{species} is not in K_S_0 keys")
+
+            if species in self.E_K_S:
+                E_K_S = F.as_fenics_constant(self.E_K_S[species], mesh)
+            elif species.name in self.E_K_S:
+                E_K_S = F.as_fenics_constant(self.E_K_S[species.name], mesh)
+
+            return K_S_0 * ufl.exp(-E_K_S / F.k_B / temperature)
+
+        else:
+            raise ValueError("K_S_0 and E_K_S must be either floats or dicts")

From dfe204c7e1974eb352eb31228ac52cb6a1d40970 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 12:40:32 +0000
Subject: [PATCH 051/134] black formatted

---
 festim/exports/vtx.py                |  4 +++-
 festim/mesh/mesh.py                  | 10 +++++++--
 festim/mesh/mesh_1d.py               |  4 +++-
 festim/subdomain/interface.py        | 33 ++++++++++++++++++----------
 festim/subdomain/volume_subdomain.py | 11 +++++-----
 5 files changed, 41 insertions(+), 21 deletions(-)

diff --git a/festim/exports/vtx.py b/festim/exports/vtx.py
index 00d75eec0..c38c42d45 100644
--- a/festim/exports/vtx.py
+++ b/festim/exports/vtx.py
@@ -63,7 +63,9 @@ class VTXExport(VTXExportBase):
         ...    my_export.write(t)
     """
 
-    def __init__(self, filename: str, field, subdomain: F.VolumeSubdomain=None) -> None:
+    def __init__(
+        self, filename: str, field, subdomain: F.VolumeSubdomain = None
+    ) -> None:
         self.field = field
         self.subdomain = subdomain
         super().__init__(filename)
diff --git a/festim/mesh/mesh.py b/festim/mesh/mesh.py
index aec3b17d5..03fb8a4aa 100644
--- a/festim/mesh/mesh.py
+++ b/festim/mesh/mesh.py
@@ -97,10 +97,16 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
         for interface in interfaces:
             (domain_1, domain_2) = interface.subdomains
             all_1_facets = dolfinx.mesh.compute_incident_entities(
-                self.mesh.topology, volume_meshtags.find(domain_1.id), self.vdim, self.fdim
+                self.mesh.topology,
+                volume_meshtags.find(domain_1.id),
+                self.vdim,
+                self.fdim,
             )
             all_2_facets = dolfinx.mesh.compute_incident_entities(
-                self.mesh.topology, volume_meshtags.find(domain_2.id), self.vdim, self.fdim
+                self.mesh.topology,
+                volume_meshtags.find(domain_2.id),
+                self.vdim,
+                self.fdim,
             )
             interface_entities = np.intersect1d(all_1_facets, all_2_facets)
             tags_facets[interface_entities] = interface.id
diff --git a/festim/mesh/mesh_1d.py b/festim/mesh/mesh_1d.py
index 9e0c56802..37c98d126 100644
--- a/festim/mesh/mesh_1d.py
+++ b/festim/mesh/mesh_1d.py
@@ -75,4 +75,6 @@ def check_borders(self, volume_subdomains):
     def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None):
         # check if all borders are defined
         self.check_borders(volume_subdomains)
-        return super().define_meshtags(surface_subdomains, volume_subdomains, interfaces)
+        return super().define_meshtags(
+            surface_subdomains, volume_subdomains, interfaces
+        )
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index ab18034a1..86b20cdf4 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -4,7 +4,7 @@
 import numpy as np
 
 
-class Interface():
+class Interface:
     id: int
     subdomains: tuple[F.VolumeSubdomain, F.VolumeSubdomain]
     parent_mesh: dolfinx.mesh.Mesh
@@ -18,11 +18,14 @@ def __init__(self, parent_mesh, mt, id, subdomains):
         self.parent_mesh = parent_mesh
 
     def pad_parent_maps(self):
-        """Workaround to make sparsity-pattern work without skips
-        """ 
+        """Workaround to make sparsity-pattern work without skips"""
 
         integration_data = compute_integration_domains(
-                dolfinx.fem.IntegralType.interior_facet, self.parent_mesh.topology, self.mt.find(self.id), self.mt.dim).reshape(-1, 4)
+            dolfinx.fem.IntegralType.interior_facet,
+            self.parent_mesh.topology,
+            self.mt.find(self.id),
+            self.mt.dim,
+        ).reshape(-1, 4)
         for i in range(2):
             # We pad the parent to submesh map to make sure that sparsity pattern is correct
             mapped_cell_0 = self.subdomains[i].parent_to_submesh[integration_data[:, 0]]
@@ -31,20 +34,24 @@ def pad_parent_maps(self):
             self.subdomains[i].parent_to_submesh[integration_data[:, 0]] = max_cells
             self.subdomains[i].parent_to_submesh[integration_data[:, 2]] = max_cells
             self.subdomains[i].padded = True
-    
+
     def compute_mapped_interior_facet_data(self, mesh):
         """
         Compute integration data for interface integrals.
         We define the first domain on an interface as the "+" restriction,
         meaning that we must sort all integration entities in this order
-        
+
         Returns
             integration_data: Integration data for interior facets
         """
         assert (not self.subdomains[0].padded) and (not self.subdomains[1].padded)
         mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
         integration_data = compute_integration_domains(
-            dolfinx.fem.IntegralType.interior_facet, mesh.topology, self.mt.find(self.id), self.mt.dim)
+            dolfinx.fem.IntegralType.interior_facet,
+            mesh.topology,
+            self.mt.find(self.id),
+            self.mt.dim,
+        )
 
         ordered_integration_data = integration_data.reshape(-1, 4).copy()
 
@@ -52,14 +59,16 @@ def compute_mapped_interior_facet_data(self, mesh):
         mapped_cell_1 = self.subdomains[0].parent_to_submesh[integration_data[2::4]]
 
         switch = mapped_cell_1 > mapped_cell_0
-        # Order restriction on one side        
+        # Order restriction on one side
         if True in switch:
             ordered_integration_data[switch, [0, 1, 2, 3]] = ordered_integration_data[
                 switch, [2, 3, 0, 1]
             ]
-        
+
         # Check that other restriction lies in other interface
-        domain1_cell = self.subdomains[1].parent_to_submesh[ordered_integration_data[:, 2]]
-        assert (domain1_cell >=0).all()
+        domain1_cell = self.subdomains[1].parent_to_submesh[
+            ordered_integration_data[:, 2]
+        ]
+        assert (domain1_cell >= 0).all()
 
-        return (self.id, ordered_integration_data.reshape(-1))
\ No newline at end of file
+        return (self.id, ordered_integration_data.reshape(-1))
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index 32a39cfdd..84ebd6803 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -11,6 +11,7 @@ class VolumeSubdomain:
     Args:
         id (int): the id of the volume subdomain
     """
+
     id: int
     submesh: dolfinx.mesh.Mesh
     submesh_to_mesh: np.ndarray
@@ -19,7 +20,7 @@ class VolumeSubdomain:
     v_map: np.ndarray
     facet_to_parent: np.ndarray
     ft: dolfinx.mesh.MeshTags
-    padded:bool
+    padded: bool
 
     def __init__(self, id, material):
         self.id = id
@@ -44,9 +45,9 @@ def locate_subdomain_entities(self, mesh, vdim):
     def create_subdomain(self, mesh, marker):
         assert marker.dim == mesh.topology.dim
         self.parent_mesh = mesh
-        self.submesh, self.submesh_to_mesh, self.v_map = (
-            dolfinx.mesh.create_submesh(mesh, marker.dim, marker.find(self.id))[0:3]
-        )
+        self.submesh, self.submesh_to_mesh, self.v_map = dolfinx.mesh.create_submesh(
+            mesh, marker.dim, marker.find(self.id)
+        )[0:3]
         num_cells_local = (
             mesh.topology.index_map(marker.dim).size_local
             + mesh.topology.index_map(marker.dim).num_ghosts
@@ -55,7 +56,7 @@ def create_subdomain(self, mesh, marker):
         self.parent_to_submesh[self.submesh_to_mesh] = np.arange(
             len(self.submesh_to_mesh), dtype=np.int32
         )
-        self.padded=False
+        self.padded = False
 
     def transfer_meshtag(self, mesh, tag):
         # Transfer meshtags to submesh

From 1f714fabb5006c8d4325385e842655fe664a6379 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 13:06:01 +0000
Subject: [PATCH 052/134] removed unused import

---
 festim/hydrogen_transport_problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 2656dbfb0..bbb2f7b3e 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -6,7 +6,7 @@
 import numpy as np
 import tqdm.autonotebook
 import festim as F
-from festim.helpers_discontinuity import NewtonSolver, transfer_meshtags_to_submesh
+from festim.helpers_discontinuity import NewtonSolver
 
 
 class HydrogenTransportProblem(F.ProblemBase):

From eaf13a3bd1545d5639443a65a6c23090fde8c24a Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 14:23:00 +0000
Subject: [PATCH 053/134] dirichlet BC more generic

---
 festim/hydrogen_transport_problem.py | 7 ++++++-
 festim/species.py                    | 7 +++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index bbb2f7b3e..5c3ac9c58 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -763,6 +763,7 @@ def initialise(self):
             )
         self.define_temperature()
 
+        # TODO see if we can remove this
         self.entity_maps = {
             subdomain.submesh: subdomain.parent_to_submesh
             for subdomain in self.volume_subdomains
@@ -795,6 +796,9 @@ def initialise(self):
     def create_dirichletbc_form(self, bc):
         fdim = self.mesh.mesh.topology.dim - 1
         volume_subdomain = self.surface_to_volume[bc.subdomain]
+        # TODO check that we shouldn't use function_space.sub(i).collapse()
+        sub_function_space = bc.species.subdomain_to_function_space[volume_subdomain]
+
         bc_function = dolfinx.fem.Function(volume_subdomain.u.function_space)
         bc_function.x.array[:] = bc.value
         volume_subdomain.submesh.topology.create_connectivity(
@@ -804,7 +808,7 @@ def create_dirichletbc_form(self, bc):
         form = dolfinx.fem.dirichletbc(
             bc_function,
             dolfinx.fem.locate_dofs_topological(
-                volume_subdomain.u.function_space.sub(0),
+                sub_function_space,
                 fdim,
                 volume_subdomain.ft.find(bc.subdomain.id),
             ),
@@ -855,6 +859,7 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_solution[subdomain] = us[i]
             species.subdomain_to_prev_solution[subdomain] = u_ns[i]
             species.subdomain_to_test_function[subdomain] = vs[i]
+            species.subdomain_to_function_space[subdomain] = V.sub(i)
             species.subdomain_to_post_processing_solution[subdomain] = u.sub(
                 i
             ).collapse()
diff --git a/festim/species.py b/festim/species.py
index 89ef4b5e4..f58230af6 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -35,6 +35,12 @@ class Species:
             previous solutions
         subdomain_to_test_function (dict): a dictionary mapping subdomains to
             test functions
+        subdomain_to_post_processing_solution (dict): a dictionary mapping
+            subdomains to post processing solutions
+        subdomain_to_collapsed_function_space (dict): a dictionary mapping
+            subdomains to collapsed function spaces
+        subdomain_to_function_space (dict): a dictionary mapping subdomains to
+            function spaces
 
     Usage:
         >>> from festim import Species, HTransportProblem
@@ -64,6 +70,7 @@ def __init__(
         self.subdomain_to_test_function = {}
         self.subdomain_to_post_processing_solution = {}
         self.subdomain_to_collapsed_function_space = {}
+        self.subdomain_to_function_space = {}
 
     def __repr__(self) -> str:
         return f"Species({self.name})"

From 9c8d0295cd92504d0ec8dd02ddab0030469fecc5 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 14:42:01 +0000
Subject: [PATCH 054/134] use sub functionspace instead of full functionspace
 for dirichletBC

---
 festim/hydrogen_transport_problem.py | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 5c3ac9c58..25de677cd 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -796,23 +796,21 @@ def initialise(self):
     def create_dirichletbc_form(self, bc):
         fdim = self.mesh.mesh.topology.dim - 1
         volume_subdomain = self.surface_to_volume[bc.subdomain]
-        # TODO check that we shouldn't use function_space.sub(i).collapse()
-        sub_function_space = bc.species.subdomain_to_function_space[volume_subdomain]
+        sub_V = bc.species.subdomain_to_function_space[volume_subdomain]
+        collapsed_V, _ = sub_V.collapse()
 
-        bc_function = dolfinx.fem.Function(volume_subdomain.u.function_space)
+        bc_function = dolfinx.fem.Function(collapsed_V)
         bc_function.x.array[:] = bc.value
         volume_subdomain.submesh.topology.create_connectivity(
             volume_subdomain.submesh.topology.dim - 1,
             volume_subdomain.submesh.topology.dim,
         )
-        form = dolfinx.fem.dirichletbc(
-            bc_function,
-            dolfinx.fem.locate_dofs_topological(
-                sub_function_space,
-                fdim,
-                volume_subdomain.ft.find(bc.subdomain.id),
-            ),
+        bc_dofs = dolfinx.fem.locate_dofs_topological(
+            (sub_V, collapsed_V),
+            fdim,
+            volume_subdomain.ft.find(bc.subdomain.id),
         )
+        form = dolfinx.fem.dirichletbc(bc_function, bc_dofs, sub_V)
         return form
 
     def create_initial_conditions(self):

From 10aab1896d7d35a288fc74b1c7c47f0b650b6ca8 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 23 Aug 2024 14:50:39 +0000
Subject: [PATCH 055/134] removed entity maps as attribute

---
 festim/hydrogen_transport_problem.py |  6 ------
 test_implementation.py               | 15 +++++++++++----
 test_implementation_2.py             | 28 +++++++++++++++++++++++++---
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 25de677cd..94f3e7cd5 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -763,12 +763,6 @@ def initialise(self):
             )
         self.define_temperature()
 
-        # TODO see if we can remove this
-        self.entity_maps = {
-            subdomain.submesh: subdomain.parent_to_submesh
-            for subdomain in self.volume_subdomains
-        }
-
         self.create_source_values_fenics()
         self.create_flux_values_fenics()
         self.create_initial_conditions()
diff --git a/test_implementation.py b/test_implementation.py
index c190b777d..6f3d171b7 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -84,7 +84,11 @@ def half(x):
 my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
 # we should be able to automate this
-my_model.interfaces = [F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain))]
+my_model.interfaces = [
+    F.Interface(
+        my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
+    )
+]
 my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
@@ -129,9 +133,11 @@ def half(x):
 my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
 
 my_model.exports = [
-    F.VTXExport(f"c_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+    F.VTXExport(f"c_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    for subdomain in my_model.volume_subdomains
 ] + [
-    F.VTXExport(f"c_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+    F.VTXExport(f"c_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain)
+    for subdomain in my_model.volume_subdomains
 ]
 
 my_model.initialise()
@@ -141,7 +147,8 @@ def half(x):
 
 
 # derived quantities
-my_model.entity_maps[mesh] = bottom_domain.submesh_to_mesh
+entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
+entity_maps[mesh] = bottom_domain.submesh_to_mesh
 
 ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
 ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
diff --git a/test_implementation_2.py b/test_implementation_2.py
index 2aa4ddd2e..1c76b226a 100644
--- a/test_implementation_2.py
+++ b/test_implementation_2.py
@@ -84,7 +84,11 @@ def half(x):
 my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
 # we should be able to automate this
-my_model.interfaces = [F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain))]
+my_model.interfaces = [
+    F.Interface(
+        my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
+    )
+]
 my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
@@ -107,7 +111,8 @@ def half(x):
 my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
 
 my_model.exports = [
-    F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+    F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    for subdomain in my_model.volume_subdomains
 ]
 
 my_model.initialise()
@@ -116,16 +121,33 @@ def half(x):
 # -------------------- post processing --------------------
 
 # derived quantities
-my_model.entity_maps[mesh] = bottom_domain.submesh_to_mesh
+entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
 
 ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
 ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
+ds = ufl.Measure("ds", domain=mesh, subdomain_data=my_model.facet_meshtags)
 
+n = ufl.FacetNormal(mesh)
 n_b = ufl.FacetNormal(bottom_domain.submesh)
 n_t = ufl.FacetNormal(top_domain.submesh)
 
+D = bottom_domain.material.get_diffusion_coefficient(
+    my_model.mesh.mesh, my_model.temperature_fenics, H
+)
+form = dolfinx.fem.form(
+    ufl.dot(
+        D * ufl.grad(bottom_domain.u.sub(0)),
+        n,
+    )
+    * ds(1),
+    entity_maps=entity_maps,
+)
+print(dolfinx.fem.assemble_scalar(form))
+
+entity_maps[mesh] = bottom_domain.submesh_to_mesh
+
 form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
 print(dolfinx.fem.assemble_scalar(form))
 

From 1c83c38aa4de66b0efabec23a000e1fcb7b9206e Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 29 Aug 2024 18:48:10 +0000
Subject: [PATCH 056/134] adapted progress bar

---
 festim/hydrogen_transport_problem.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 94f3e7cd5..9da998d53 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -1101,9 +1101,10 @@ def post_processing(self):
 
     def iterate(self):
         """Iterates the model for a given time step"""
-        self.progress.update(
-            min(self.dt.value, abs(self.settings.final_time - self.t.value))
-        )
+        if self.show_progress_bar:
+            self.progress_bar.update(
+                min(self.dt.value, abs(self.settings.final_time - self.t.value))
+            )
         self.t.value += self.dt.value
 
         self.update_time_dependent_values()
@@ -1125,14 +1126,16 @@ def iterate(self):
     def run(self):
         if self.settings.transient:
             # Solve transient
-            self.progress = tqdm.autonotebook.tqdm(
-                desc=f"Solving {self.__class__.__name__}",
-                total=self.settings.final_time,
-                unit_scale=True,
-            )
+            if self.show_progress_bar:
+                self.progress_bar = tqdm.autonotebook.tqdm(
+                    desc=f"Solving {self.__class__.__name__}",
+                    total=self.settings.final_time,
+                    unit_scale=True,
+                )
             while self.t.value < self.settings.final_time:
                 self.iterate()
-            self.progress.refresh()  # refresh progress bar to show 100%
+            if self.show_progress_bar:
+                self.progress_bar.refresh()  # refresh progress bar to show 100%
         else:
             # Solve steady-state
             self.solver.solve(1e-5)

From 959e4a9eef5ac6cb4e716b51a3b2bd3447382055 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 29 Aug 2024 21:23:58 +0000
Subject: [PATCH 057/134] penalty term is now exposed

---
 festim/hydrogen_transport_problem.py | 3 +--
 festim/subdomain/interface.py        | 3 ++-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 9da998d53..198a9848c 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -947,12 +947,11 @@ def mixed_term(u, v, n):
         n = ufl.FacetNormal(mesh)
         cr = ufl.Circumradius(mesh)
 
-        gamma = 10.0
-
         entity_maps = {
             sd.submesh: sd.parent_to_submesh for sd in self.volume_subdomains
         }
         for interface in self.interfaces:
+            gamma = interface.penalty_term
 
             subdomain_1, subdomain_2 = interface.subdomains
             b_res, t_res = interface.restriction
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index 86b20cdf4..951778f53 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -11,11 +11,12 @@ class Interface:
     restriction: list[str, str] = ("+", "-")
     padded: bool
 
-    def __init__(self, parent_mesh, mt, id, subdomains):
+    def __init__(self, parent_mesh, mt, id, subdomains, penalty_term=10.0):
         self.id = id
         self.subdomains = tuple(subdomains)
         self.mt = mt
         self.parent_mesh = parent_mesh
+        self.penalty_term = penalty_term
 
     def pad_parent_maps(self):
         """Workaround to make sparsity-pattern work without skips"""

From 8adc073d093c60814490a1b61b702a3057f9bef1 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 29 Aug 2024 21:35:49 +0000
Subject: [PATCH 058/134] tolerance of solver is exposed

---
 festim/hydrogen_transport_problem.py |  4 ++--
 test_implementation.py               |  2 +-
 test_implementation_2.py             |  2 +-
 test_implementation_3.py             | 33 +++++++++++++++++++++-------
 test_implementation_transient.py     | 10 ++++++---
 test_one_material.py                 | 11 ++++++----
 6 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 198a9848c..948639829 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -1109,7 +1109,7 @@ def iterate(self):
         self.update_time_dependent_values()
 
         # solve main problem
-        self.solver.solve(1e-5)
+        self.solver.solve(self.settings.rtol)
 
         # post processing
         self.post_processing()
@@ -1137,5 +1137,5 @@ def run(self):
                 self.progress_bar.refresh()  # refresh progress bar to show 100%
         else:
             # Solve steady-state
-            self.solver.solve(1e-5)
+            self.solver.solve(self.settings.rtol)
             self.post_processing()
diff --git a/test_implementation.py b/test_implementation.py
index 6f3d171b7..be53239cd 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -130,7 +130,7 @@ def half(x):
 
 my_model.temperature = lambda x: 400 + 10 * x[1] + 100 * x[0]
 
-my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
     F.VTXExport(f"c_{subdomain.id}.bp", field=H, subdomain=subdomain)
diff --git a/test_implementation_2.py b/test_implementation_2.py
index 1c76b226a..41f49fed7 100644
--- a/test_implementation_2.py
+++ b/test_implementation_2.py
@@ -108,7 +108,7 @@ def half(x):
 
 my_model.temperature = 500.0  # lambda x: 300 + 10 * x[1] + 100 * x[0]
 
-my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
     F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
diff --git a/test_implementation_3.py b/test_implementation_3.py
index fb9e9a512..19750a89c 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -29,19 +29,35 @@
 material_right.K_S_0 = 6.0
 material_right.E_K_S = 0
 
-left_domain = F.VolumeSubdomain1D(3, borders=[vertices[0], interface_1], material=material_left)
-middle_domain = F.VolumeSubdomain1D(4, borders=[interface_1, interface_2], material=material_mid)
-right_domain = F.VolumeSubdomain1D(5, borders=[interface_2, vertices[-1]], material=material_right)
+left_domain = F.VolumeSubdomain1D(
+    3, borders=[vertices[0], interface_1], material=material_left
+)
+middle_domain = F.VolumeSubdomain1D(
+    4, borders=[interface_1, interface_2], material=material_mid
+)
+right_domain = F.VolumeSubdomain1D(
+    5, borders=[interface_2, vertices[-1]], material=material_right
+)
 
 left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
 # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
 my_model.interfaces = [
-    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)),
-    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)),
+    F.Interface(
+        my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)
+    ),
+    F.Interface(
+        my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)
+    ),
+]
+my_model.subdomains = [
+    left_domain,
+    middle_domain,
+    right_domain,
+    left_surface,
+    right_surface,
 ]
-my_model.subdomains = [left_domain, middle_domain, right_domain, left_surface, right_surface]
 my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
 
 H = F.Species("H", mobile=True)
@@ -75,10 +91,11 @@
 
 my_model.temperature = lambda x: 300 + 100 * x[0]
 
-my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    for subdomain in my_model.volume_subdomains
 ]
 
 my_model.initialise()
diff --git a/test_implementation_transient.py b/test_implementation_transient.py
index b481edcab..328fdd876 100644
--- a/test_implementation_transient.py
+++ b/test_implementation_transient.py
@@ -38,8 +38,12 @@
 # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
 # TODO remove mesh and meshtags from these arguments
 my_model.interfaces = [
-    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)),
-    F.Interface(my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)),
+    F.Interface(
+        my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)
+    ),
+    F.Interface(
+        my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)
+    ),
 ]
 
 my_model.subdomains = [
@@ -82,7 +86,7 @@
 
 my_model.temperature = lambda x: 300 + 100 * x[0]
 
-my_model.settings = F.Settings(atol=None, rtol=None, transient=True, final_time=100)
+my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=True, final_time=100)
 my_model.settings.stepsize = 1
 
 my_model.exports = [
diff --git a/test_one_material.py b/test_one_material.py
index da5391639..5660d235c 100644
--- a/test_one_material.py
+++ b/test_one_material.py
@@ -13,12 +13,14 @@
 material.K_S_0 = 2.0
 material.E_K_S = 0
 
-subdomain = F.VolumeSubdomain1D(1, borders=[vertices[0], vertices[-1]], material=material)
+subdomain = F.VolumeSubdomain1D(
+    1, borders=[vertices[0], vertices[-1]], material=material
+)
 
 left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
-my_model.subdomains = [subdomain,left_surface, right_surface]
+my_model.subdomains = [subdomain, left_surface, right_surface]
 my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
 
 H = F.Species("H", mobile=True)
@@ -52,10 +54,11 @@
 
 my_model.temperature = lambda x: 300 + 100 * x[0]
 
-my_model.settings = F.Settings(atol=None, rtol=None, transient=False)
+my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain) for subdomain in my_model.volume_subdomains
+    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    for subdomain in my_model.volume_subdomains
 ]
 
 my_model.initialise()

From 6b53b01cf3774c65dc582be9cf35697b6f142f36 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 29 Aug 2024 21:39:30 +0000
Subject: [PATCH 059/134] exposed petsc options

---
 festim/hydrogen_transport_problem.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 948639829..ec797335a 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -732,6 +732,7 @@ def __init__(
         traps=None,
         interfaces=None,
         surface_to_volume=None,
+        petsc_options=None,
     ):
         super().__init__(
             mesh,
@@ -748,6 +749,12 @@ def __init__(
         )
         self.interfaces = interfaces or {}
         self.surface_to_volume = surface_to_volume or {}
+        default_petsc_options = {
+            "ksp_type": "preonly",
+            "pc_type": "lu",
+            "pc_factor_mat_solver_type": "mumps",
+        }
+        self.petsc_options = petsc_options or default_petsc_options
 
     def initialise(self):
         self.define_meshtags_and_measures()
@@ -1049,11 +1056,7 @@ def create_solver(self):
             [subdomain.u for subdomain in self.volume_subdomains],
             bcs=self.bc_forms,
             max_iterations=10,
-            petsc_options={
-                "ksp_type": "preonly",
-                "pc_type": "lu",
-                "pc_factor_mat_solver_type": "mumps",
-            },
+            petsc_options=self.petsc_options,
         )
 
     def create_flux_values_fenics(self):

From 4005db5742b52fb860e13a03f7db7b4ae3edd8be Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 4 Oct 2024 13:47:00 +0000
Subject: [PATCH 060/134] Improved NewtonSolver

Co-authored-by: jorgensd <dokken@simula.no>
---
 festim/helpers_discontinuity.py | 48 +++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/festim/helpers_discontinuity.py b/festim/helpers_discontinuity.py
index 2ae19ccc0..bcc9f4bec 100644
--- a/festim/helpers_discontinuity.py
+++ b/festim/helpers_discontinuity.py
@@ -73,8 +73,50 @@ def solve(self, tol=1e-6, beta=1.0):
             # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
             with self.b.localForm() as b_local:
                 b_local.set(0.0)
+
+            constants_L = [
+                form and dolfinx.cpp.fem.pack_constants(form._cpp_object)
+                for form in self.F
+            ]
+            coeffs_L = [
+                dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self.F
+            ]
+
+            constants_a = [
+                [
+                    (
+                        dolfinx.cpp.fem.pack_constants(form._cpp_object)
+                        if form is not None
+                        else np.array([], dtype=PETSc.ScalarType)
+                    )
+                    for form in forms
+                ]
+                for forms in self.J
+            ]
+
+            coeffs_a = [
+                [
+                    (
+                        {}
+                        if form is None
+                        else dolfinx.cpp.fem.pack_coefficients(form._cpp_object)
+                    )
+                    for form in forms
+                ]
+                for forms in self.J
+            ]
+
             dolfinx.fem.petsc.assemble_vector_block(
-                self.b, self.F, self.J, bcs=self.bcs, x0=self.x, scale=-1.0
+                self.b,
+                self.F,
+                self.J,
+                bcs=self.bcs,
+                x0=self.x,
+                scale=-1.0,
+                coeffs_a=coeffs_a,
+                constants_a=constants_a,
+                coeffs_L=coeffs_L,
+                constants_L=constants_L,
             )
             self.b.ghostUpdate(
                 PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
@@ -82,7 +124,9 @@ def solve(self, tol=1e-6, beta=1.0):
 
             # Assemble Jacobian
             self.A.zeroEntries()
-            dolfinx.fem.petsc.assemble_matrix_block(self.A, self.J, bcs=self.bcs)
+            dolfinx.fem.petsc.assemble_matrix_block(
+                self.A, self.J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a
+            )
             self.A.assemble()
 
             self._solver.solve(self.b, self.dx)

From 4f91d914b0d26d995394e3a27d3e3a1770af5f9f Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 4 Oct 2024 13:49:58 +0000
Subject: [PATCH 061/134] new solver

---
 discontinuity_generic.py | 97 +++++++++++++++++++++++++++++-----------
 1 file changed, 70 insertions(+), 27 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 49f2a45a4..c8e3442ad 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -74,9 +74,50 @@ def solve(self, tol=1e-6, beta=1.0):
             # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
             with self.b.localForm() as b_local:
                 b_local.set(0.0)
-            
+
+            constants_L = [
+                form and dolfinx.cpp.fem.pack_constants(form._cpp_object)
+                for form in self.F
+            ]
+            coeffs_L = [
+                dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self.F
+            ]
+
+            constants_a = [
+                [
+                    (
+                        dolfinx.cpp.fem.pack_constants(form._cpp_object)
+                        if form is not None
+                        else np.array([], dtype=PETSc.ScalarType)
+                    )
+                    for form in forms
+                ]
+                for forms in self.J
+            ]
+
+            coeffs_a = [
+                [
+                    (
+                        {}
+                        if form is None
+                        else dolfinx.cpp.fem.pack_coefficients(form._cpp_object)
+                    )
+                    for form in forms
+                ]
+                for forms in self.J
+            ]
+
             dolfinx.fem.petsc.assemble_vector_block(
-                self.b, self.F, self.J, bcs=self.bcs, x0=self.x, scale=-1.0
+                self.b,
+                self.F,
+                self.J,
+                bcs=self.bcs,
+                x0=self.x,
+                scale=-1.0,
+                coeffs_a=coeffs_a,
+                constants_a=constants_a,
+                coeffs_L=coeffs_L,
+                constants_L=constants_L,
             )
             self.b.ghostUpdate(
                 PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
@@ -84,7 +125,9 @@ def solve(self, tol=1e-6, beta=1.0):
 
             # Assemble Jacobian
             self.A.zeroEntries()
-            dolfinx.fem.petsc.assemble_matrix_block(self.A, self.J, bcs=self.bcs)
+            dolfinx.fem.petsc.assemble_matrix_block(
+                self.A, self.J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a
+            )
             self.A.assemble()
 
             self._solver.solve(self.b, self.dx)
@@ -542,36 +585,36 @@ def mixed_term(u, v, n):
 )
 solver.solve(1e-5)
 
-for subdomain in list_of_subdomains:
-    u_sub_0 = subdomain.u.sub(0).collapse()
-    u_sub_0.name = "u_sub_0"
+# for subdomain in list_of_subdomains:
+#     u_sub_0 = subdomain.u.sub(0).collapse()
+#     u_sub_0.name = "u_sub_0"
 
-    u_sub_1 = subdomain.u.sub(1).collapse()
-    u_sub_1.name = "u_sub_1"
-    bp = dolfinx.io.VTXWriter(
-        mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
-    )
-    bp.write(0)
-    bp.close()
+#     u_sub_1 = subdomain.u.sub(1).collapse()
+#     u_sub_1.name = "u_sub_1"
+#     bp = dolfinx.io.VTXWriter(
+#         mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
+#     )
+#     bp.write(0)
+#     bp.close()
 
 
-# derived quantities
-V = dolfinx.fem.functionspace(mesh, ("CG", 1))
-T = dolfinx.fem.Function(V)
-T.interpolate(lambda x: 200 + x[1])
+# # derived quantities
+# V = dolfinx.fem.functionspace(mesh, ("CG", 1))
+# T = dolfinx.fem.Function(V)
+# T.interpolate(lambda x: 200 + x[1])
 
 
-T_b = dolfinx.fem.Function(right_domain.u.sub(0).collapse().function_space)
-T_b.interpolate(T)
+# T_b = dolfinx.fem.Function(right_domain.u.sub(0).collapse().function_space)
+# T_b.interpolate(T)
 
-ds_b = ufl.Measure("ds", domain=right_domain.submesh)
-dx_b = ufl.Measure("dx", domain=right_domain.submesh)
-dx = ufl.Measure("dx", domain=mesh)
+# ds_b = ufl.Measure("ds", domain=right_domain.submesh)
+# dx_b = ufl.Measure("dx", domain=right_domain.submesh)
+# dx = ufl.Measure("dx", domain=mesh)
 
-n_b = ufl.FacetNormal(left_domain.submesh)
+# n_b = ufl.FacetNormal(left_domain.submesh)
 
-form = dolfinx.fem.form(left_domain.u.sub(0) * dx_b, entity_maps=entity_maps)
-print(dolfinx.fem.assemble_scalar(form))
+# form = dolfinx.fem.form(left_domain.u.sub(0) * dx_b, entity_maps=entity_maps)
+# print(dolfinx.fem.assemble_scalar(form))
 
-form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(left_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
-print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file
+# form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(left_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
+# print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file

From aab875091386ae61f12670965d1b47c3b4a632fe Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 4 Oct 2024 14:10:28 +0000
Subject: [PATCH 062/134] scale is now alpha + topology._cpp_object

---
 festim/helpers_discontinuity.py | 2 +-
 festim/subdomain/interface.py   | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/festim/helpers_discontinuity.py b/festim/helpers_discontinuity.py
index bcc9f4bec..824860f48 100644
--- a/festim/helpers_discontinuity.py
+++ b/festim/helpers_discontinuity.py
@@ -112,7 +112,7 @@ def solve(self, tol=1e-6, beta=1.0):
                 self.J,
                 bcs=self.bcs,
                 x0=self.x,
-                scale=-1.0,
+                alpha=-1.0,
                 coeffs_a=coeffs_a,
                 constants_a=constants_a,
                 coeffs_L=coeffs_L,
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index 951778f53..0c0c30844 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -23,7 +23,7 @@ def pad_parent_maps(self):
 
         integration_data = compute_integration_domains(
             dolfinx.fem.IntegralType.interior_facet,
-            self.parent_mesh.topology,
+            self.parent_mesh.topology._cpp_object,
             self.mt.find(self.id),
             self.mt.dim,
         ).reshape(-1, 4)
@@ -49,7 +49,7 @@ def compute_mapped_interior_facet_data(self, mesh):
         mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
         integration_data = compute_integration_domains(
             dolfinx.fem.IntegralType.interior_facet,
-            mesh.topology,
+            mesh.topology._cpp_object,
             self.mt.find(self.id),
             self.mt.dim,
         )

From 1ca73b36814ebe44eef61e1bc8687f43690e02e0 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 4 Oct 2024 18:53:39 +0000
Subject: [PATCH 063/134] support for sources

---
 festim/hydrogen_transport_problem.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index ec797335a..6270b6ac1 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -800,8 +800,13 @@ def create_dirichletbc_form(self, bc):
         sub_V = bc.species.subdomain_to_function_space[volume_subdomain]
         collapsed_V, _ = sub_V.collapse()
 
-        bc_function = dolfinx.fem.Function(collapsed_V)
-        bc_function.x.array[:] = bc.value
+        bc.create_value(
+            mesh=self.mesh.mesh,
+            temperature=self.temperature_fenics,
+            function_space=collapsed_V,
+            t=self.t,
+        )
+
         volume_subdomain.submesh.topology.create_connectivity(
             volume_subdomain.submesh.topology.dim - 1,
             volume_subdomain.submesh.topology.dim,
@@ -811,7 +816,7 @@ def create_dirichletbc_form(self, bc):
             fdim,
             volume_subdomain.ft.find(bc.subdomain.id),
         )
-        form = dolfinx.fem.dirichletbc(bc_function, bc_dofs, sub_V)
+        form = dolfinx.fem.dirichletbc(bc.value_fenics, bc_dofs, sub_V)
         return form
 
     def create_initial_conditions(self):
@@ -927,6 +932,11 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
                     v = bc.species.subdomain_to_test_function[subdomain]
                     form -= bc.value_fenics * v * self.ds(bc.subdomain.id)
 
+        for source in self.sources:
+            v = source.species.subdomain_to_test_function[subdomain]
+            if source.volume == subdomain:
+                form -= source.value_fenics * v * dx
+
         subdomain.F = form
 
     def create_formulation(self):

From d3b0c071c4194596bfd97e012b17ce33f8d91972 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 4 Oct 2024 18:53:54 +0000
Subject: [PATCH 064/134] added MMS test in 2D

---
 test/system_tests/test_multi_material.py | 156 +++++++++++++++++++++++
 1 file changed, 156 insertions(+)
 create mode 100644 test/system_tests/test_multi_material.py

diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
new file mode 100644
index 000000000..0690ff9df
--- /dev/null
+++ b/test/system_tests/test_multi_material.py
@@ -0,0 +1,156 @@
+import festim as F
+
+from mpi4py import MPI
+import dolfinx
+import dolfinx.fem.petsc
+import numpy as np
+import festim as F
+import ufl
+from .tools import error_L2
+
+
+def generate_mesh(n=20):
+    def bottom_boundary(x):
+        return np.isclose(x[1], 0.0)
+
+    def top_boundary(x):
+        return np.isclose(x[1], 1.0)
+
+    def half(x):
+        return x[1] <= 0.5 + 1e-14
+
+    mesh = dolfinx.mesh.create_unit_square(
+        MPI.COMM_WORLD, n, n, dolfinx.mesh.CellType.triangle
+    )
+
+    # Split domain in half and set an interface tag of 5
+    gdim = mesh.geometry.dim
+    tdim = mesh.topology.dim
+    fdim = tdim - 1
+    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    num_facets_local = (
+        mesh.topology.index_map(fdim).size_local
+        + mesh.topology.index_map(fdim).num_ghosts
+    )
+    facets = np.arange(num_facets_local, dtype=np.int32)
+    values = np.full_like(facets, 0, dtype=np.int32)
+    values[top_facets] = 1
+    values[bottom_facets] = 2
+
+    bottom_cells = dolfinx.mesh.locate_entities(mesh, tdim, half)
+    num_cells_local = (
+        mesh.topology.index_map(tdim).size_local
+        + mesh.topology.index_map(tdim).num_ghosts
+    )
+    cells = np.full(num_cells_local, 4, dtype=np.int32)
+    cells[bottom_cells] = 3
+    ct = dolfinx.mesh.meshtags(
+        mesh, tdim, np.arange(num_cells_local, dtype=np.int32), cells
+    )
+    all_b_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(3), tdim, fdim
+    )
+    all_t_facets = dolfinx.mesh.compute_incident_entities(
+        mesh.topology, ct.find(4), tdim, fdim
+    )
+    interface = np.intersect1d(all_b_facets, all_t_facets)
+    values[interface] = 5
+
+    mt = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, values)
+    return mesh, mt, ct
+
+
+def test_2_materials_2d():
+    """
+    MMS case for a 2D problem with 2 materials
+    adapted from https://festim-vv-report.readthedocs.io/en/v1.0/verification/mms/discontinuity.html
+    """
+    K_S_top = 3.0
+    K_S_bot = 6.0
+    D_top = 2.0
+    D_bot = 5.0
+    c_exact_top_ufl = (
+        lambda x: 1 + ufl.sin(ufl.pi * (2 * x[0] + 0.5)) + ufl.cos(2 * ufl.pi * x[1])
+    )
+    c_exact_bot_ufl = lambda x: K_S_bot / K_S_top * c_exact_top_ufl(x)
+
+    c_exact_top_np = (
+        lambda x: 1 + np.sin(np.pi * (2 * x[0] + 0.5)) + np.cos(2 * np.pi * x[1])
+    )
+    c_exact_bot_np = lambda x: K_S_bot / K_S_top * c_exact_top_np(x)
+
+    mesh, mt, ct = generate_mesh(100)
+
+    my_model = F.HTransportProblemDiscontinuous()
+    my_model.mesh = F.Mesh(mesh)
+    my_model.volume_meshtags = ct
+    my_model.facet_meshtags = mt
+
+    material_bottom = F.Material(D_0=D_bot, E_D=0, K_S_0=K_S_bot, E_K_S=0)
+    material_top = F.Material(D_0=D_top, E_D=0, K_S_0=K_S_top, E_K_S=0)
+
+    top_domain = F.VolumeSubdomain(4, material=material_top)
+    bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
+
+    top_surface = F.SurfaceSubdomain(id=1)
+    bottom_surface = F.SurfaceSubdomain(id=2)
+    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+
+    my_model.interfaces = [
+        F.Interface(
+            my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
+        )
+    ]
+    my_model.surface_to_volume = {
+        top_surface: top_domain,
+        bottom_surface: bottom_domain,
+    }
+
+    H = F.Species("H", mobile=True)
+
+    my_model.species = [H]
+
+    for species in my_model.species:
+        species.subdomains = [bottom_domain, top_domain]
+
+    my_model.boundary_conditions = [
+        F.FixedConcentrationBC(top_surface, value=c_exact_top_ufl, species=H),
+        F.FixedConcentrationBC(bottom_surface, value=c_exact_bot_ufl, species=H),
+    ]
+
+    source_top_val = (
+        lambda x: 8
+        * ufl.pi**2
+        * (ufl.cos(2 * ufl.pi * x[0]) + ufl.cos(2 * ufl.pi * x[1]))
+    )
+    source_bottom_val = (
+        lambda x: 40
+        * ufl.pi**2
+        * (ufl.cos(2 * ufl.pi * x[0]) + ufl.cos(2 * ufl.pi * x[1]))
+    )
+    my_model.sources = [
+        F.ParticleSource(volume=top_domain, species=H, value=source_top_val),
+        F.ParticleSource(volume=bottom_domain, species=H, value=source_bottom_val),
+    ]
+
+    my_model.temperature = 500.0  # lambda x: 300 + 10 * x[1] + 100 * x[0]
+
+    my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
+
+    my_model.exports = [
+        F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        for subdomain in my_model.volume_subdomains
+    ]
+
+    my_model.initialise()
+    my_model.run()
+
+    c_top_computed = H.subdomain_to_post_processing_solution[top_domain]
+    c_bot_computed = H.subdomain_to_post_processing_solution[bottom_domain]
+
+    L2_error_mobile = error_L2(c_top_computed, c_exact_top_np)
+    L2_error_trapped = error_L2(c_bot_computed, c_exact_bot_np)
+
+    assert L2_error_mobile < 1e-3
+    assert L2_error_trapped < 1e-3

From b8006ba8217c9aa574d2c5011e981465ca7ce804 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 4 Oct 2024 18:58:50 +0000
Subject: [PATCH 065/134] black

---
 discontinuity_generic.py | 100 ++++++++++++++++++++++++---------------
 1 file changed, 62 insertions(+), 38 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index c8e3442ad..213d62052 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -8,6 +8,7 @@
 import basix
 import dolfinx.fem.petsc
 
+
 class NewtonSolver:
     max_iterations: int
     bcs: list[dolfinx.fem.DirichletBC]
@@ -234,7 +235,7 @@ def right_boundary(x):
 
     def left_domain(x):
         return x[0] <= 0.5 + 1e-14
-    
+
     def right_domain(x):
         return x[0] >= 0.7 - 1e-14
 
@@ -315,16 +316,17 @@ class VolumeSubdomain:
     v_map: np.ndarray
     facet_to_parent: np.ndarray
     ft: dolfinx.mesh.MeshTags
-    padded:bool
+    padded: bool
+
     def __init__(self, id):
         self.id = id
 
     def create_subdomain(self, mesh, marker):
         assert marker.dim == mesh.topology.dim
         self.parent_mesh = mesh
-        self.submesh, self.submesh_to_mesh, self.v_map = (
-            dolfinx.mesh.create_submesh(mesh, marker.dim, marker.find(self.id))[0:3]
-        )
+        self.submesh, self.submesh_to_mesh, self.v_map = dolfinx.mesh.create_submesh(
+            mesh, marker.dim, marker.find(self.id)
+        )[0:3]
         num_cells_local = (
             mesh.topology.index_map(marker.dim).size_local
             + mesh.topology.index_map(marker.dim).num_ghosts
@@ -333,7 +335,7 @@ def create_subdomain(self, mesh, marker):
         self.parent_to_submesh[self.submesh_to_mesh] = np.arange(
             len(self.submesh_to_mesh), dtype=np.int32
         )
-        self.padded=False
+        self.padded = False
 
     def transfer_meshtag(self, tag):
         # Transfer meshtags to submesh
@@ -342,23 +344,29 @@ def transfer_meshtag(self, tag):
             mesh, tag, self.submesh, self.v_map, self.submesh_to_mesh
         )
 
-class Interface():
+
+class Interface:
     id: int
     subdomains: tuple[VolumeSubdomain, VolumeSubdomain]
     parent_mesh: dolfinx.mesh.Mesh
     restriction: [str, str] = ("+", "-")
     padded: bool
-    def __init__(self, parent_mesh, mt,  id, subdomains):
+
+    def __init__(self, parent_mesh, mt, id, subdomains):
         self.id = id
         self.subdomains = tuple(subdomains)
         self.mt = mt
         self.parent_mesh = parent_mesh
+
     def pad_parent_maps(self):
-        """Workaround to make sparsity-pattern work without skips
-        """ 
+        """Workaround to make sparsity-pattern work without skips"""
 
         integration_data = compute_integration_domains(
-                dolfinx.fem.IntegralType.interior_facet, self.parent_mesh.topology, self.mt.find(self.id), self.mt.dim).reshape(-1, 4)
+            dolfinx.fem.IntegralType.interior_facet,
+            self.parent_mesh.topology,
+            self.mt.find(self.id),
+            self.mt.dim,
+        ).reshape(-1, 4)
         for i in range(2):
             # We pad the parent to submesh map to make sure that sparsity pattern is correct
             mapped_cell_0 = self.subdomains[i].parent_to_submesh[integration_data[:, 0]]
@@ -368,6 +376,7 @@ def pad_parent_maps(self):
             self.subdomains[i].parent_to_submesh[integration_data[:, 2]] = max_cells
             self.subdomains[i].padded = True
 
+
 mesh, mt, ct = generate_mesh()
 
 left_domain = VolumeSubdomain(3)
@@ -384,19 +393,16 @@ def pad_parent_maps(self):
 )
 
 
-
 for subdomain in list_of_subdomains:
     subdomain.create_subdomain(mesh, ct)
     subdomain.transfer_meshtag(mt)
 
 
-
-
-
 i0 = Interface(mesh, mt, 5, (left_domain, mid_domain))
 i1 = Interface(mesh, mt, 6, (mid_domain, right_domain))
 interfaces = [i0, i1]
 
+
 def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
     element_CG = basix.ufl.element(
         basix.ElementFamily.P,
@@ -439,7 +445,7 @@ def compute_mapped_interior_facet_data(interface: Interface):
     Compute integration data for interface integrals.
     We define the first domain on an interface as the "+" restriction,
     meaning that we must sort all integration entities in this order
-    
+
     Parameters
         interface: Interface between two subdomains
     Returns
@@ -448,7 +454,11 @@ def compute_mapped_interior_facet_data(interface: Interface):
     assert (not interface.subdomains[0].padded) and (not interface.subdomains[1].padded)
     mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
     integration_data = compute_integration_domains(
-        dolfinx.fem.IntegralType.interior_facet, mesh.topology, interface.mt.find(interface.id), interface.mt.dim)
+        dolfinx.fem.IntegralType.interior_facet,
+        mesh.topology,
+        interface.mt.find(interface.id),
+        interface.mt.dim,
+    )
 
     ordered_integration_data = integration_data.reshape(-1, 4).copy()
 
@@ -456,29 +466,32 @@ def compute_mapped_interior_facet_data(interface: Interface):
     mapped_cell_1 = interface.subdomains[0].parent_to_submesh[integration_data[2::4]]
 
     switch = mapped_cell_1 > mapped_cell_0
-    # Order restriction on one side        
+    # Order restriction on one side
     if True in switch:
         ordered_integration_data[switch, [0, 1, 2, 3]] = ordered_integration_data[
             switch, [2, 3, 0, 1]
         ]
-    
+
     # Check that other restriction lies in other interface
-    domain1_cell = interface.subdomains[1].parent_to_submesh[ordered_integration_data[:, 2]]
-    assert (domain1_cell >=0).all()
+    domain1_cell = interface.subdomains[1].parent_to_submesh[
+        ordered_integration_data[:, 2]
+    ]
+    assert (domain1_cell >= 0).all()
 
     return (interface.id, ordered_integration_data.reshape(-1))
 
 
-
-integral_data = [compute_mapped_interior_facet_data(interface) for interface in interfaces]
+integral_data = [
+    compute_mapped_interior_facet_data(interface) for interface in interfaces
+]
 [interface.pad_parent_maps() for interface in interfaces]
-dInterface = ufl.Measure(
-        "dS", domain=mesh, subdomain_data=integral_data
-    )
+dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)
+
 
 def mixed_term(u, v, n):
     return ufl.dot(ufl.grad(u), n) * v
 
+
 n = ufl.FacetNormal(mesh)
 cr = ufl.Circumradius(mesh)
 
@@ -494,7 +507,6 @@ def mixed_term(u, v, n):
     h_b = 2 * cr(b_res)
     h_t = 2 * cr(t_res)
 
-
     v_b = subdomain_1.vs[0](b_res)
     v_t = subdomain_2.vs[0](t_res)
 
@@ -513,17 +525,29 @@ def mixed_term(u, v, n):
     K_b = K_0(b_res)
     K_t = K_1(t_res)
 
-    F_0 = (
-        -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(interface.id)
-        - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
+    F_0 = -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(
+        interface.id
+    ) - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
+
+    F_1 = +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(
+        interface.id
+    ) - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
+    F_0 += (
+        2
+        * gamma
+        / (h_b + h_t)
+        * (u_b / K_b - u_t / K_t)
+        * v_b
+        * dInterface(interface.id)
     )
-
-    F_1 = (
-        +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(interface.id)
-        - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
+    F_1 += (
+        -2
+        * gamma
+        / (h_b + h_t)
+        * (u_b / K_b - u_t / K_t)
+        * v_t
+        * dInterface(interface.id)
     )
-    F_0 += 2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_b * dInterface(interface.id)
-    F_1 += -2 * gamma / (h_b + h_t) * (u_b / K_b - u_t / K_t) * v_t * dInterface(interface.id)
 
     subdomain_1.F += F_0
     subdomain_2.F += F_1
@@ -531,7 +555,7 @@ def mixed_term(u, v, n):
 
 J = []
 forms = []
-for i,subdomain1 in enumerate(list_of_subdomains):
+for i, subdomain1 in enumerate(list_of_subdomains):
     jac = []
     form = subdomain1.F
     for j, subdomain2 in enumerate(list_of_subdomains):
@@ -617,4 +641,4 @@ def mixed_term(u, v, n):
 # print(dolfinx.fem.assemble_scalar(form))
 
 # form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(left_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
-# print(dolfinx.fem.assemble_scalar(form))
\ No newline at end of file
+# print(dolfinx.fem.assemble_scalar(form))

From e2ee473755fc5a834eaa5b3f8f45d242bbe57cf4 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 10 Oct 2024 09:15:52 -0400
Subject: [PATCH 066/134] updated dolfinx version in CI + removed dev container

---
 .devcontainer/devcontainer.json | 3 ---
 .github/workflows/ci_conda.yml  | 2 +-
 .github/workflows/ci_docker.yml | 2 +-
 3 files changed, 2 insertions(+), 5 deletions(-)
 delete mode 100644 .devcontainer/devcontainer.json

diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
deleted file mode 100644
index dea4d4dd9..000000000
--- a/.devcontainer/devcontainer.json
+++ /dev/null
@@ -1,3 +0,0 @@
-{
-    "image": "dolfinx/dolfinx:nightly"
-}
\ No newline at end of file
diff --git a/.github/workflows/ci_conda.yml b/.github/workflows/ci_conda.yml
index 0af9bae05..9c3d0055d 100644
--- a/.github/workflows/ci_conda.yml
+++ b/.github/workflows/ci_conda.yml
@@ -20,7 +20,7 @@ jobs:
     - name: Create Conda environment
       shell: bash -l {0}
       run: |
-        conda install -c conda-forge fenics-dolfinx=0.8.0
+        conda install -c conda-forge fenics-dolfinx=0.9.0
 
     - name: Install local package and dependencies
       shell: bash -l {0}
diff --git a/.github/workflows/ci_docker.yml b/.github/workflows/ci_docker.yml
index 7074e6e8f..e104da607 100644
--- a/.github/workflows/ci_docker.yml
+++ b/.github/workflows/ci_docker.yml
@@ -6,7 +6,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        container_version: [v0.8.0, nightly]
+        container_version: [v0.9.0, nightly]
     container: dolfinx/dolfinx:${{ matrix.container_version }}
     steps:
         - name: Checkout code

From 45f4db765bf4f8d4e2e5143c726e49f95e554bd9 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 10 Oct 2024 11:25:48 -0400
Subject: [PATCH 067/134] fixed fenics script

---
 discontinuity_generic.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 213d62052..892773ffc 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -114,7 +114,7 @@ def solve(self, tol=1e-6, beta=1.0):
                 self.J,
                 bcs=self.bcs,
                 x0=self.x,
-                scale=-1.0,
+                alpha=-1.0,
                 coeffs_a=coeffs_a,
                 constants_a=constants_a,
                 coeffs_L=coeffs_L,
@@ -363,7 +363,7 @@ def pad_parent_maps(self):
 
         integration_data = compute_integration_domains(
             dolfinx.fem.IntegralType.interior_facet,
-            self.parent_mesh.topology,
+            self.parent_mesh.topology._cpp_object,
             self.mt.find(self.id),
             self.mt.dim,
         ).reshape(-1, 4)
@@ -455,7 +455,7 @@ def compute_mapped_interior_facet_data(interface: Interface):
     mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
     integration_data = compute_integration_domains(
         dolfinx.fem.IntegralType.interior_facet,
-        mesh.topology,
+        mesh.topology._cpp_object,
         interface.mt.find(interface.id),
         interface.mt.dim,
     )

From e64979245a376dbd2ac2b10e12ac7ca37c42bc18 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 10 Oct 2024 11:28:25 -0400
Subject: [PATCH 068/134] added failing test

---
 test/system_tests/test_multi_material.py | 63 ++++++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index 0690ff9df..d3ec085c6 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -154,3 +154,66 @@ def test_2_materials_2d():
 
     assert L2_error_mobile < 1e-3
     assert L2_error_trapped < 1e-3
+
+
+# TODO make this a MMS case
+def test_1_material_disontinuous_version():
+    my_model = F.HTransportProblemDiscontinuous()
+
+    N = 1500
+    vertices = np.linspace(0, 1, num=N)
+    my_model.mesh = F.Mesh1D(vertices)
+
+    material = F.Material(D_0=2.0, E_D=0)
+
+    material.K_S_0 = 2.0
+    material.E_K_S = 0
+
+    subdomain = F.VolumeSubdomain1D(
+        1, borders=[vertices[0], vertices[-1]], material=material
+    )
+
+    left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
+    right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
+
+    my_model.subdomains = [subdomain, left_surface, right_surface]
+    my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
+
+    H = F.Species("H", mobile=True)
+    trapped_H = F.Species("H_trapped", mobile=False)
+    empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
+
+    my_model.species = [H, trapped_H]
+
+    for species in my_model.species:
+        species.subdomains = my_model.volume_subdomains
+
+    my_model.reactions = [
+        F.Reaction(
+            reactant=[H, empty_trap],
+            product=[trapped_H],
+            k_0=2,
+            E_k=0,
+            p_0=0.1,
+            E_p=0,
+            volume=domain,
+        )
+        for domain in my_model.volume_subdomains
+    ]
+
+    my_model.boundary_conditions = [
+        F.FixedConcentrationBC(left_surface, value=0.05, species=H),
+        F.FixedConcentrationBC(right_surface, value=0.2, species=H),
+    ]
+
+    my_model.temperature = lambda x: 300 + 100 * x[0]
+
+    my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
+
+    my_model.exports = [
+        F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        for subdomain in my_model.volume_subdomains
+    ]
+
+    my_model.initialise()
+    my_model.run()

From 51a6961a871c3b9091a19fd75d06cdf0e1d6326f Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 10 Oct 2024 12:52:56 -0400
Subject: [PATCH 069/134] check for dofs functionspace

---
 festim/hydrogen_transport_problem.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 6270b6ac1..5ce7ba1f4 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -811,8 +811,16 @@ def create_dirichletbc_form(self, bc):
             volume_subdomain.submesh.topology.dim - 1,
             volume_subdomain.submesh.topology.dim,
         )
+
+        # mapping between sub_function space and collapsed is only needed if
+        # value_fenics is a function of the collapsed space
+        if isinstance(bc.value_fenics, fem.Function):
+            function_space_dofs = (sub_V, collapsed_V)
+        else:
+            function_space_dofs = sub_V
+
         bc_dofs = dolfinx.fem.locate_dofs_topological(
-            (sub_V, collapsed_V),
+            function_space_dofs,
             fdim,
             volume_subdomain.ft.find(bc.subdomain.id),
         )

From 0be3749926bd0c816e6646122eaf0ab73c8c58e1 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Thu, 10 Oct 2024 13:00:08 -0400
Subject: [PATCH 070/134] added more tests

---
 test/system_tests/test_multi_material.py | 170 ++++++++++++++++++++++-
 1 file changed, 168 insertions(+), 2 deletions(-)

diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index d3ec085c6..0a4529e20 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -61,7 +61,7 @@ def half(x):
     return mesh, mt, ct
 
 
-def test_2_materials_2d():
+def test_2_materials_2d_mms():
     """
     MMS case for a 2D problem with 2 materials
     adapted from https://festim-vv-report.readthedocs.io/en/v1.0/verification/mms/discontinuity.html
@@ -157,7 +157,7 @@ def test_2_materials_2d():
 
 
 # TODO make this a MMS case
-def test_1_material_disontinuous_version():
+def test_1_material_discontinuous_version():
     my_model = F.HTransportProblemDiscontinuous()
 
     N = 1500
@@ -217,3 +217,169 @@ def test_1_material_disontinuous_version():
 
     my_model.initialise()
     my_model.run()
+
+
+def test_3_materials_transient():
+    my_model = F.HTransportProblemDiscontinuous()
+
+    interface_1 = 0.5
+    interface_2 = 0.7
+
+    # for some reason if the mesh isn't fine enough then I have a random SEGV error
+    N = 1500
+    vertices = np.concatenate(
+        [
+            np.linspace(0, interface_1, num=N),
+            np.linspace(interface_1, interface_2, num=N),
+            np.linspace(interface_2, 1, num=N),
+        ]
+    )
+
+    my_model.mesh = F.Mesh1D(vertices)
+
+    material_left = F.Material(D_0=2.0, E_D=0, K_S_0=2.0, E_K_S=0)
+    material_mid = F.Material(D_0=2.0, E_D=0, K_S_0=4.0, E_K_S=0)
+    material_right = F.Material(D_0=2.0, E_D=0, K_S_0=6.0, E_K_S=0)
+
+    left_domain = F.VolumeSubdomain1D(
+        3, borders=[vertices[0], interface_1], material=material_left
+    )
+    middle_domain = F.VolumeSubdomain1D(
+        4, borders=[interface_1, interface_2], material=material_mid
+    )
+    right_domain = F.VolumeSubdomain1D(
+        5, borders=[interface_2, vertices[-1]], material=material_right
+    )
+
+    left_surface = F.SurfaceSubdomain1D(id=1, x=vertices[0])
+    right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
+
+    # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
+    # TODO remove mesh and meshtags from these arguments
+    my_model.interfaces = [
+        F.Interface(
+            my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)
+        ),
+        F.Interface(
+            my_model.mesh.mesh,
+            my_model.facet_meshtags,
+            7,
+            (middle_domain, right_domain),
+        ),
+    ]
+
+    my_model.subdomains = [
+        left_domain,
+        middle_domain,
+        right_domain,
+        left_surface,
+        right_surface,
+    ]
+    my_model.surface_to_volume = {
+        right_surface: right_domain,
+        left_surface: left_domain,
+    }
+
+    H = F.Species("H", mobile=True)
+    trapped_H = F.Species("H_trapped", mobile=False)
+    empty_trap = F.ImplicitSpecies(n=0.5, others=[trapped_H])
+
+    my_model.species = [H, trapped_H]
+
+    for species in my_model.species:
+        species.subdomains = my_model.volume_subdomains
+
+    my_model.reactions = [
+        F.Reaction(
+            reactant=[H, empty_trap],
+            product=[trapped_H],
+            k_0=2,
+            E_k=0,
+            p_0=0.1,
+            E_p=0,
+            volume=domain,
+        )
+        for domain in [left_domain, middle_domain, right_domain]
+    ]
+
+    my_model.boundary_conditions = [
+        F.DirichletBC(left_surface, value=0.05, species=H),
+        F.DirichletBC(right_surface, value=0.2, species=H),
+    ]
+
+    my_model.temperature = lambda x: 300 + 100 * x[0]
+
+    my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=True, final_time=100)
+    my_model.settings.stepsize = 1
+
+    my_model.exports = [
+        F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        for subdomain in my_model.volume_subdomains
+    ] + [
+        F.VTXExport(
+            filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain
+        )
+        for subdomain in my_model.volume_subdomains
+    ]
+
+    my_model.initialise()
+    my_model.run()
+
+
+def test_2_mats_particle_flux_bc():
+    mesh, mt, ct = generate_mesh()
+
+    my_model = F.HTransportProblemDiscontinuous()
+    my_model.mesh = F.Mesh(mesh)
+    my_model.volume_meshtags = ct
+    my_model.facet_meshtags = mt
+
+    material_bottom = F.Material(D_0=1, E_D=0.1)
+    material_top = F.Material(D_0=1, E_D=0.1)
+
+    material_bottom.K_S_0 = 2.0
+    material_bottom.E_K_S = 0.1
+    material_top.K_S_0 = 4.0
+    material_top.E_K_S = 0.1
+
+    top_domain = F.VolumeSubdomain(4, material=material_top)
+    bottom_domain = F.VolumeSubdomain(3, material=material_bottom)
+
+    top_surface = F.SurfaceSubdomain(id=1)
+    bottom_surface = F.SurfaceSubdomain(id=2)
+    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+
+    # we should be able to automate this
+    my_model.interfaces = [
+        F.Interface(
+            my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
+        )
+    ]
+    my_model.surface_to_volume = {
+        top_surface: top_domain,
+        bottom_surface: bottom_domain,
+    }
+
+    H = F.Species("H", mobile=True)
+
+    my_model.species = [H]
+
+    for species in my_model.species:
+        species.subdomains = [bottom_domain, top_domain]
+
+    my_model.boundary_conditions = [
+        F.DirichletBC(top_surface, value=0.05, species=H),
+        F.ParticleFluxBC(bottom_surface, value=lambda x: 1.0 + x[0], species=H),
+    ]
+
+    my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
+
+    my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
+
+    my_model.exports = [
+        F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        for subdomain in my_model.volume_subdomains
+    ]
+
+    my_model.initialise()
+    my_model.run()

From 38057f7e3e1434dc937ea8daed7411ff5617889c Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 11 Oct 2024 14:28:28 -0400
Subject: [PATCH 071/134] removed duplicate of subdomain_1 subdomain_2

---
 festim/hydrogen_transport_problem.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 5ce7ba1f4..2178e00d8 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -985,22 +985,6 @@ def mixed_term(u, v, n):
             h_b = 2 * cr(b_res)
             h_t = 2 * cr(t_res)
 
-            # look at the first facet on interface
-            # and get the two cells that are connected to it
-            # and get the material properties of these cells
-            first_facet_interface = mt.find(interface.id)[0]
-            c_plus, c_minus = (
-                f_to_c.links(first_facet_interface)[0],
-                f_to_c.links(first_facet_interface)[1],
-            )
-            id_minus, id_plus = ct.values[c_minus], ct.values[c_plus]
-
-            for subdomain in interface.subdomains:
-                if subdomain.id == id_plus:
-                    subdomain_1 = subdomain
-                if subdomain.id == id_minus:
-                    subdomain_2 = subdomain
-
             all_mobile_species = [spe for spe in self.species if spe.mobile]
             if len(all_mobile_species) > 1:
                 raise NotImplementedError("Multiple mobile species not implemented")

From 77b3f6e9d0f5ef904700c28fa39edb62346a4c81 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 11 Oct 2024 14:38:19 -0400
Subject: [PATCH 072/134] docstrings + removed unused variables

---
 festim/hydrogen_transport_problem.py | 45 +++++++++++++++++++++++-----
 1 file changed, 37 insertions(+), 8 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 2178e00d8..7d5eeb0fb 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -840,6 +840,17 @@ def create_submeshes(self):
             subdomain.transfer_meshtag(self.mesh.mesh, self.facet_meshtags)
 
     def define_function_spaces(self, subdomain: F.VolumeSubdomain):
+        """
+        Creates appropriate function space and functions for a given subdomain (submesh)
+        based on the number of species existing in this subdomain. Then stores the functionspace,
+        the current solution (``u``) and the previous solution (``u_n``) functions. It also populates the
+        correspondance dicts attributes of the species (eg. ``species.subdomain_to_solution``,
+        ``species.subdomain_to_test_function``, etc) for easy access to the right subfunctions,
+        sub-testfunctions etc.
+
+        Args:
+            subdomain (F.VolumeSubdomain): a subdomain of the geometry
+        """
         # get number of species defined in the subdomain
         all_species = [
             species for species in self.species if subdomain in species.subdomains
@@ -864,6 +875,11 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
         u = dolfinx.fem.Function(V)
         u_n = dolfinx.fem.Function(V)
 
+        # store attributes in the subdomain object
+        subdomain.u = u
+        subdomain.u_n = u_n
+
+        # split the functions and assign the subfunctions to the species
         us = list(ufl.split(u))
         u_ns = list(ufl.split(u_n))
         vs = list(ufl.TestFunctions(V))
@@ -881,10 +897,14 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_post_processing_solution[subdomain].name = (
                 f"{species.name}_{subdomain.id}"
             )
-        subdomain.u = u
-        subdomain.u_n = u_n
 
     def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
+        """
+        Creates the variational formulation for each subdomain and stores it in ``subdomain.F``
+
+        Args:
+            subdomain (F.VolumeSubdomain): a subdomain of the geometry
+        """
         form = 0
         # add diffusion and time derivative for each species
         for spe in self.species:
@@ -904,6 +924,7 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
             if spe.mobile:
                 form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
 
+        # add reaction terms
         for reaction in self.reactions:
             if reaction.volume != subdomain:
                 continue
@@ -913,6 +934,7 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
                     # temporarily overide the solution to the one of the subdomain
                     species.solution = species.subdomain_to_solution[subdomain]
 
+            # reactant
             for reactant in reaction.reactant:
                 if isinstance(reactant, F.Species):
                     form += (
@@ -936,24 +958,30 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         # add fluxes
         for bc in self.boundary_conditions:
             if isinstance(bc, F.ParticleFluxBC):
+                # check that the bc is applied on a surface
+                # belonging to this subdomain
                 if subdomain == self.surface_to_volume[bc.subdomain]:
                     v = bc.species.subdomain_to_test_function[subdomain]
                     form -= bc.value_fenics * v * self.ds(bc.subdomain.id)
 
+        # add volumetric sources
         for source in self.sources:
             v = source.species.subdomain_to_test_function[subdomain]
             if source.volume == subdomain:
                 form -= source.value_fenics * v * dx
 
+        # store the form in the subdomain object
         subdomain.F = form
 
     def create_formulation(self):
-        # Add coupling term to the interface
-        # Get interface markers on submesh b
+        """
+        Takes all the formulations for each subdomain and adds the interface conditions.
+
+        Finally compute the jacobian matrix and store it in the ``J`` attribute,
+        adds the ``entity_maps`` to the forms and store them in the ``forms`` attribute
+        """
         mesh = self.mesh.mesh
-        ct = self.volume_meshtags
         mt = self.facet_meshtags
-        f_to_c = mesh.topology.connectivity(mesh.topology.dim - 1, mesh.topology.dim)
 
         for interface in self.interfaces:
             interface.mesh = mesh
@@ -1062,9 +1090,8 @@ def create_solver(self):
         )
 
     def create_flux_values_fenics(self):
-        """For each particle flux create the value_fenics"""
+        """For each particle flux create the ``value_fenics`` attribute"""
         for bc in self.boundary_conditions:
-            # create value_fenics for all F.ParticleFluxBC objects
             if isinstance(bc, F.ParticleFluxBC):
                 volume_subdomain = self.surface_to_volume[bc.subdomain]
                 bc.create_value_fenics(
@@ -1086,6 +1113,8 @@ def initialise_exports(self):
                 raise NotImplementedError("Export type not implemented")
 
     def post_processing(self):
+        # update post-processing solutions (for each species in each subdomain)
+        # with new solution
         for subdomain in self.volume_subdomains:
             for species in self.species:
                 if subdomain not in species.subdomains:

From c92121abc5f1d92fe50b9c9f54c85325c7784359 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 11 Oct 2024 15:17:32 -0400
Subject: [PATCH 073/134] more docstrings

---
 festim/hydrogen_transport_problem.py | 61 +++++++++++++++++++++-------
 festim/subdomain/interface.py        | 19 ++++++++-
 festim/subdomain/volume_subdomain.py | 19 +++++++--
 3 files changed, 81 insertions(+), 18 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 7d5eeb0fb..89f4e8889 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -716,6 +716,9 @@ def post_processing(self):
 
 
 class HTransportProblemDiscontinuous(HydrogenTransportProblem):
+    interfaces: list[F.Interface]
+    petsc_options: dict
+    surface_to_volume: dict
 
     def __init__(
         self,
@@ -730,10 +733,30 @@ def __init__(
         settings=None,
         exports=None,
         traps=None,
-        interfaces=None,
-        surface_to_volume=None,
-        petsc_options=None,
+        interfaces: list[F.Interface] = None,
+        surface_to_volume: dict = None,
+        petsc_options: dict = None,
     ):
+        """Class for a multi-material hydrogen transport problem
+        For other arguments see ``festim.HydrogenTransportProblem``.
+
+        Args:
+            interfaces (list, optional): list of interfaces (``festim.Interface``
+                objects). Defaults to None.
+            surface_to_volume (dict, optional): correspondance dictionary linking
+                each ``festim.SurfaceSubdomain`` objects to a ``festim.VolumeSubdomain``
+                object). Defaults to None.
+            petsc_options (dict, optional): petsc options to be passed to the
+                ``festim.NewtonSolver`` object. If None, the default options are:
+                ```
+                default_petsc_options = {
+                    "ksp_type": "preonly",
+                    "pc_type": "lu",
+                    "pc_factor_mat_solver_type": "mumps",
+                }
+                ```
+                Defaults to None.
+        """
         super().__init__(
             mesh,
             subdomains,
@@ -747,7 +770,7 @@ def __init__(
             exports,
             traps,
         )
-        self.interfaces = interfaces or {}
+        self.interfaces = interfaces or []
         self.surface_to_volume = surface_to_volume or {}
         default_petsc_options = {
             "ksp_type": "preonly",
@@ -758,7 +781,12 @@ def __init__(
 
     def initialise(self):
         self.define_meshtags_and_measures()
-        self.create_submeshes()
+
+        # create submeshes and transfer meshtags to subdomains
+        for subdomain in self.volume_subdomains:
+            subdomain.create_subdomain(self.mesh.mesh, self.volume_meshtags)
+            subdomain.transfer_meshtag(self.mesh.mesh, self.facet_meshtags)
+
         self.create_species_from_traps()
 
         self.t = fem.Constant(self.mesh.mesh, 0.0)
@@ -768,8 +796,8 @@ def initialise(self):
             self.dt = F.as_fenics_constant(
                 self.settings.stepsize.initial_value, self.mesh.mesh
             )
-        self.define_temperature()
 
+        self.define_temperature()
         self.create_source_values_fenics()
         self.create_flux_values_fenics()
         self.create_initial_conditions()
@@ -789,12 +817,23 @@ def initialise(self):
             subdomain.u.name = f"u_{subdomain.id}"
 
         self.define_boundary_conditions()
-
         self.create_formulation()
         self.create_solver()
         self.initialise_exports()
 
-    def create_dirichletbc_form(self, bc):
+    def create_dirichletbc_form(self, bc: F.FixedConcentrationBC):
+        """
+        Creates the ``value_fenics`` attribute for a given
+        ``festim.FixedConcentrationBC`` and returns the appropriate
+        ``dolfinx.fem.DirichletBC`` object.
+
+        Args:
+            bc (festim.FixedConcentrationBC): the dirichlet BC
+
+        Returns:
+            dolfinx.fem.DirichletBC: the appropriate dolfinx representation
+                generated from ``dolfinx.fem.dirichletbc()``
+        """
         fdim = self.mesh.mesh.topology.dim - 1
         volume_subdomain = self.surface_to_volume[bc.subdomain]
         sub_V = bc.species.subdomain_to_function_space[volume_subdomain]
@@ -833,12 +872,6 @@ def create_initial_conditions(self):
                 "initial conditions not yet implemented for discontinuous"
             )
 
-    def create_submeshes(self):
-
-        for subdomain in self.volume_subdomains:
-            subdomain.create_subdomain(self.mesh.mesh, self.volume_meshtags)
-            subdomain.transfer_meshtag(self.mesh.mesh, self.facet_meshtags)
-
     def define_function_spaces(self, subdomain: F.VolumeSubdomain):
         """
         Creates appropriate function space and functions for a given subdomain (submesh)
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index 0c0c30844..600d541d8 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -11,7 +11,24 @@ class Interface:
     restriction: list[str, str] = ("+", "-")
     padded: bool
 
-    def __init__(self, parent_mesh, mt, id, subdomains, penalty_term=10.0):
+    def __init__(
+        self,
+        parent_mesh: dolfinx.mesh.Mesh,
+        mt: dolfinx.mesh.MeshTags,
+        id: int,
+        subdomains: list[F.VolumeSubdomain],
+        penalty_term: float = 10.0,
+    ):
+        """Class representing an interface between two subdomains.
+
+        Args:
+            parent_mesh (dolfinx.mesh.Mesh): the parent mesh
+            mt (dolfinx.mesh.MeshTags): the facet meshtags (on the parent mesh)
+            id (int): the tag of the interface subdomain in the parent meshtags
+            subdomains (list[F.VolumeSubdomain]): the subdomains sharing this interface
+            penalty_term (float, optional): Penalty term in the Nietsche DG formulation.
+                Needs to be "sufficiently large". Defaults to 10.0.
+        """
         self.id = id
         self.subdomains = tuple(subdomains)
         self.mt = mt
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index 84ebd6803..60041f5e1 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -10,6 +10,7 @@ class VolumeSubdomain:
 
     Args:
         id (int): the id of the volume subdomain
+        material (festim.Material): the material assigned to the subdomain
     """
 
     id: int
@@ -42,9 +43,21 @@ def locate_subdomain_entities(self, mesh, vdim):
         entities = locate_entities(mesh, vdim, lambda x: np.full(x.shape[1], True))
         return entities
 
-    def create_subdomain(self, mesh, marker):
+    def create_subdomain(self, mesh: dolfinx.mesh.Mesh, marker: dolfinx.mesh.MeshTags):
+        """
+        Creates the following attributes: ``.parent_mesh``, ``.submesh``, ``.submesh_to_mesh``,
+        ``.v_map``, ``padded``, and the entity map ``parent_to_submesh``.
+
+        Only used in ``festim.HTransportProblemDiscontinuous``
+
+        Args:
+            mesh (dolfinx.mesh.Mesh): the parent mesh
+            marker (dolfinx.mesh.MeshTags): the parent volume markers
+        """
         assert marker.dim == mesh.topology.dim
-        self.parent_mesh = mesh
+        self.parent_mesh = (
+            mesh  # NOTE: it doesn't seem like we use this attribute anywhere
+        )
         self.submesh, self.submesh_to_mesh, self.v_map = dolfinx.mesh.create_submesh(
             mesh, marker.dim, marker.find(self.id)
         )[0:3]
@@ -58,7 +71,7 @@ def create_subdomain(self, mesh, marker):
         )
         self.padded = False
 
-    def transfer_meshtag(self, mesh, tag):
+    def transfer_meshtag(self, mesh: dolfinx.mesh.Mesh, tag: dolfinx.mesh.MeshTags):
         # Transfer meshtags to submesh
         assert self.submesh is not None, "Need to call create_subdomain first"
         self.ft, self.facet_to_parent = transfer_meshtags_to_submesh(

From 1282892dd9367a05a577d781a9b28c0ce3ef77c7 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 11 Oct 2024 15:21:49 -0400
Subject: [PATCH 074/134] removed mt and parent mesh from constructor of
 interface

---
 festim/hydrogen_transport_problem.py     |  4 ++++
 festim/subdomain/interface.py            |  7 +------
 test/system_tests/test_multi_material.py | 23 ++++-------------------
 3 files changed, 9 insertions(+), 25 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 89f4e8889..4c141a443 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -787,6 +787,10 @@ def initialise(self):
             subdomain.create_subdomain(self.mesh.mesh, self.volume_meshtags)
             subdomain.transfer_meshtag(self.mesh.mesh, self.facet_meshtags)
 
+        for interface in self.interfaces:
+            interface.mt = self.volume_meshtags
+            interface.parent_mesh = self.mesh.mesh
+
         self.create_species_from_traps()
 
         self.t = fem.Constant(self.mesh.mesh, 0.0)
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index 600d541d8..78bd2b2be 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -8,13 +8,12 @@ class Interface:
     id: int
     subdomains: tuple[F.VolumeSubdomain, F.VolumeSubdomain]
     parent_mesh: dolfinx.mesh.Mesh
+    mt: dolfinx.mesh.MeshTags
     restriction: list[str, str] = ("+", "-")
     padded: bool
 
     def __init__(
         self,
-        parent_mesh: dolfinx.mesh.Mesh,
-        mt: dolfinx.mesh.MeshTags,
         id: int,
         subdomains: list[F.VolumeSubdomain],
         penalty_term: float = 10.0,
@@ -22,8 +21,6 @@ def __init__(
         """Class representing an interface between two subdomains.
 
         Args:
-            parent_mesh (dolfinx.mesh.Mesh): the parent mesh
-            mt (dolfinx.mesh.MeshTags): the facet meshtags (on the parent mesh)
             id (int): the tag of the interface subdomain in the parent meshtags
             subdomains (list[F.VolumeSubdomain]): the subdomains sharing this interface
             penalty_term (float, optional): Penalty term in the Nietsche DG formulation.
@@ -31,8 +28,6 @@ def __init__(
         """
         self.id = id
         self.subdomains = tuple(subdomains)
-        self.mt = mt
-        self.parent_mesh = parent_mesh
         self.penalty_term = penalty_term
 
     def pad_parent_maps(self):
diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index 0a4529e20..f5d5711ad 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -97,11 +97,7 @@ def test_2_materials_2d_mms():
     bottom_surface = F.SurfaceSubdomain(id=2)
     my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
-    my_model.interfaces = [
-        F.Interface(
-            my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
-        )
-    ]
+    my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
     my_model.surface_to_volume = {
         top_surface: top_domain,
         bottom_surface: bottom_domain,
@@ -257,15 +253,8 @@ def test_3_materials_transient():
     # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
     # TODO remove mesh and meshtags from these arguments
     my_model.interfaces = [
-        F.Interface(
-            my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)
-        ),
-        F.Interface(
-            my_model.mesh.mesh,
-            my_model.facet_meshtags,
-            7,
-            (middle_domain, right_domain),
-        ),
+        F.Interface(6, (left_domain, middle_domain)),
+        F.Interface(7, (middle_domain, right_domain)),
     ]
 
     my_model.subdomains = [
@@ -350,11 +339,7 @@ def test_2_mats_particle_flux_bc():
     my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
     # we should be able to automate this
-    my_model.interfaces = [
-        F.Interface(
-            my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
-        )
-    ]
+    my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
     my_model.surface_to_volume = {
         top_surface: top_domain,
         bottom_surface: bottom_domain,

From 01948a4bf8a43e269cdb0ca8d775aa0e9fc443be Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 11 Oct 2024 15:50:20 -0400
Subject: [PATCH 075/134] more type hinting

---
 festim/hydrogen_transport_problem.py |  2 +-
 festim/species.py                    | 14 ++++++++++----
 festim/subdomain/volume_subdomain.py |  2 ++
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 4c141a443..01fa9e39f 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -173,7 +173,7 @@ def multispecies(self):
         return len(self.species) > 1
 
     @property
-    def species(self):
+    def species(self) -> list[F.Species]:
         return self._species
 
     @species.setter
diff --git a/festim/species.py b/festim/species.py
index f58230af6..b474d9aff 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Union
 import festim as F
 
 
@@ -52,9 +52,15 @@ class Species:
 
     """
 
-    def __init__(
-        self, name: str = None, mobile=True, subdomains: F.VolumeSubdomain = None
-    ) -> None:
+    subdomains: Union[List[F.VolumeSubdomain], F.VolumeSubdomain]
+    subdomain_to_solution: dict
+    subdomain_to_prev_solution: dict
+    subdomain_to_test_function: dict
+    subdomain_to_post_processing_solution: dict
+    subdomain_to_collapsed_function_space: dict
+    subdomain_to_function_space: dict
+
+    def __init__(self, name: str = None, mobile=True, subdomains=None) -> None:
         self.name = name
         self.mobile = mobile
         self.solution = None
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index 60041f5e1..4d8b055c9 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -22,6 +22,8 @@ class VolumeSubdomain:
     facet_to_parent: np.ndarray
     ft: dolfinx.mesh.MeshTags
     padded: bool
+    u: dolfinx.fem.Function
+    u_n: dolfinx.fem.Function
 
     def __init__(self, id, material):
         self.id = id

From e39946dd41aeba8d3f69257beca8ae089764e4c6 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 11 Oct 2024 15:50:34 -0400
Subject: [PATCH 076/134] type check for species.subdomains

---
 festim/hydrogen_transport_problem.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 01fa9e39f..e909c97ea 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -780,6 +780,12 @@ def __init__(
         self.petsc_options = petsc_options or default_petsc_options
 
     def initialise(self):
+        # check that all species have a list of F.VolumeSubdomain as this is
+        # different from F.HydrogenTransportProblem
+        for spe in self.species:
+            if not isinstance(spe.subdomains, list):
+                raise TypeError("subdomains attribute should be list")
+
         self.define_meshtags_and_measures()
 
         # create submeshes and transfer meshtags to subdomains

From 60b9b2e52dd0ced4bfd34f35a0aac6adcf17d233 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Sun, 13 Oct 2024 05:48:00 +0800
Subject: [PATCH 077/134] fixed script

---
 discontinuity_generic.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 892773ffc..5917e9f51 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -581,17 +581,17 @@ def mixed_term(u, v, n):
     ),
 )
 
-t_bc = dolfinx.fem.Function(right_domain.u.function_space)
-t_bc.x.array[:] = 0.05
+t_bc = dolfinx.fem.Constant(right_domain.submesh, 0.05)
+sub_V = right_domain.u.function_space.sub(0)
 right_domain.submesh.topology.create_connectivity(
     right_domain.submesh.topology.dim - 1, right_domain.submesh.topology.dim
 )
-bc_t = dolfinx.fem.dirichletbc(
-    t_bc,
-    dolfinx.fem.locate_dofs_topological(
-        right_domain.u.function_space.sub(0), fdim, right_domain.ft.find(2)
-    ),
+bc_dofs = dolfinx.fem.locate_dofs_topological(
+    sub_V,
+    fdim,
+    right_domain.ft.find(2),
 )
+bc_t = dolfinx.fem.dirichletbc(t_bc, bc_dofs, sub_V)
 bcs = [bc_b, bc_t]
 
 

From 6a1456b23f0842a501af51200fd52565a1c11ad8 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 09:22:22 -0400
Subject: [PATCH 078/134] fixed the examples

---
 test_implementation.py           | 10 +++-------
 test_implementation_2.py         |  6 +-----
 test_implementation_3.py         |  8 ++------
 test_implementation_transient.py |  8 ++------
 4 files changed, 8 insertions(+), 24 deletions(-)

diff --git a/test_implementation.py b/test_implementation.py
index be53239cd..27f80eb53 100644
--- a/test_implementation.py
+++ b/test_implementation.py
@@ -84,11 +84,7 @@ def half(x):
 my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
 # we should be able to automate this
-my_model.interfaces = [
-    F.Interface(
-        my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
-    )
-]
+my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
 my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
@@ -123,8 +119,8 @@ def half(x):
 ]
 
 my_model.boundary_conditions = [
-    F.DirichletBC(top_surface, value=0.05, species=H),
-    F.DirichletBC(bottom_surface, value=0.2, species=H),
+    F.FixedConcentrationBC(top_surface, value=0.05, species=H),
+    F.FixedConcentrationBC(bottom_surface, value=0.2, species=H),
 ]
 
 
diff --git a/test_implementation_2.py b/test_implementation_2.py
index 41f49fed7..e26c8e89a 100644
--- a/test_implementation_2.py
+++ b/test_implementation_2.py
@@ -84,11 +84,7 @@ def half(x):
 my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
 # we should be able to automate this
-my_model.interfaces = [
-    F.Interface(
-        my_model.mesh.mesh, my_model.facet_meshtags, 5, (bottom_domain, top_domain)
-    )
-]
+my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
 my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
diff --git a/test_implementation_3.py b/test_implementation_3.py
index 19750a89c..bd5e7f80b 100644
--- a/test_implementation_3.py
+++ b/test_implementation_3.py
@@ -44,12 +44,8 @@
 
 # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
 my_model.interfaces = [
-    F.Interface(
-        my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)
-    ),
-    F.Interface(
-        my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)
-    ),
+    F.Interface(6, (left_domain, middle_domain)),
+    F.Interface(7, (middle_domain, right_domain)),
 ]
 my_model.subdomains = [
     left_domain,
diff --git a/test_implementation_transient.py b/test_implementation_transient.py
index 328fdd876..342ba6b50 100644
--- a/test_implementation_transient.py
+++ b/test_implementation_transient.py
@@ -38,12 +38,8 @@
 # the ids here are arbitrary in 1D, you can put anything as long as it's not the same as the surfaces
 # TODO remove mesh and meshtags from these arguments
 my_model.interfaces = [
-    F.Interface(
-        my_model.mesh.mesh, my_model.facet_meshtags, 6, (left_domain, middle_domain)
-    ),
-    F.Interface(
-        my_model.mesh.mesh, my_model.facet_meshtags, 7, (middle_domain, right_domain)
-    ),
+    F.Interface(6, (left_domain, middle_domain)),
+    F.Interface(7, (middle_domain, right_domain)),
 ]
 
 my_model.subdomains = [

From 56940d03537ad7eb076c01c5d084e6618fd489dc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 12:06:42 -0400
Subject: [PATCH 079/134] UPdate python version

---
 setup.cfg | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index b3595d9a2..c1a292ef3 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -13,11 +13,11 @@ license_file = LICENSE
 classifiers =
     Natural Language :: English
     Topic :: Scientific/Engineering
-    Programming Language :: Python :: 3
-    Programming Language :: Python :: 3.6
-    Programming Language :: Python :: 3.7
-    Programming Language :: Python :: 3.8
     Programming Language :: Python :: 3.9
+    Programming Language :: Python :: 3.10
+    Programming Language :: Python :: 3.11
+    Programming Language :: Python :: 3.12
+    Programming Language :: Python :: 3.13
     License :: OSI Approved :: Apache Software License
     Operating System :: OS Independent
 project_urls =
@@ -26,7 +26,7 @@ project_urls =
 
 [options]
 packages = find:
-python_requires= >=3.6
+python_requires= >=3.9
 install_requires =
     tqdm
 

From 2f16118a907bb2da914dcb4f1c0368aef66eacfe Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:07:49 -0400
Subject: [PATCH 080/134] removed discontinuity-generic since we have lots of
 examples

---
 discontinuity_generic.py | 644 ---------------------------------------
 1 file changed, 644 deletions(-)
 delete mode 100644 discontinuity_generic.py

diff --git a/discontinuity_generic.py b/discontinuity_generic.py
deleted file mode 100644
index 5917e9f51..000000000
--- a/discontinuity_generic.py
+++ /dev/null
@@ -1,644 +0,0 @@
-from mpi4py import MPI
-from petsc4py import PETSc
-import dolfinx
-from dolfinx.cpp.fem import compute_integration_domains
-
-import ufl
-import numpy as np
-import basix
-import dolfinx.fem.petsc
-
-
-class NewtonSolver:
-    max_iterations: int
-    bcs: list[dolfinx.fem.DirichletBC]
-    A: PETSc.Mat
-    b: PETSc.Vec
-    J: dolfinx.fem.Form
-    b: dolfinx.fem.Form
-    dx: PETSc.Vec
-
-    def __init__(
-        self,
-        F: list[dolfinx.fem.form],
-        J: list[list[dolfinx.fem.form]],
-        w: list[dolfinx.fem.Function],
-        bcs: list[dolfinx.fem.DirichletBC] | None = None,
-        max_iterations: int = 5,
-        petsc_options: dict[str, str | float | int | None] = None,
-        problem_prefix="newton",
-    ):
-        self.max_iterations = max_iterations
-        self.bcs = [] if bcs is None else bcs
-        self.b = dolfinx.fem.petsc.create_vector_block(F)
-        self.F = F
-        self.J = J
-        self.A = dolfinx.fem.petsc.create_matrix_block(J)
-        self.dx = self.A.createVecLeft()
-        self.w = w
-        self.x = dolfinx.fem.petsc.create_vector_block(F)
-
-        # Set PETSc options
-        opts = PETSc.Options()
-        if petsc_options is not None:
-            for k, v in petsc_options.items():
-                opts[k] = v
-
-        # Define KSP solver
-        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
-        self._solver.setOperators(self.A)
-        self._solver.setFromOptions()
-
-        # Set matrix and vector PETSc options
-        self.A.setFromOptions()
-        self.b.setFromOptions()
-
-    def solve(self, tol=1e-6, beta=1.0):
-        i = 0
-
-        while i < self.max_iterations:
-            dolfinx.cpp.la.petsc.scatter_local_vectors(
-                self.x,
-                [si.x.petsc_vec.array_r for si in self.w],
-                [
-                    (
-                        si.function_space.dofmap.index_map,
-                        si.function_space.dofmap.index_map_bs,
-                    )
-                    for si in self.w
-                ],
-            )
-            self.x.ghostUpdate(
-                addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
-            )
-
-            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
-            with self.b.localForm() as b_local:
-                b_local.set(0.0)
-
-            constants_L = [
-                form and dolfinx.cpp.fem.pack_constants(form._cpp_object)
-                for form in self.F
-            ]
-            coeffs_L = [
-                dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self.F
-            ]
-
-            constants_a = [
-                [
-                    (
-                        dolfinx.cpp.fem.pack_constants(form._cpp_object)
-                        if form is not None
-                        else np.array([], dtype=PETSc.ScalarType)
-                    )
-                    for form in forms
-                ]
-                for forms in self.J
-            ]
-
-            coeffs_a = [
-                [
-                    (
-                        {}
-                        if form is None
-                        else dolfinx.cpp.fem.pack_coefficients(form._cpp_object)
-                    )
-                    for form in forms
-                ]
-                for forms in self.J
-            ]
-
-            dolfinx.fem.petsc.assemble_vector_block(
-                self.b,
-                self.F,
-                self.J,
-                bcs=self.bcs,
-                x0=self.x,
-                alpha=-1.0,
-                coeffs_a=coeffs_a,
-                constants_a=constants_a,
-                coeffs_L=coeffs_L,
-                constants_L=constants_L,
-            )
-            self.b.ghostUpdate(
-                PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
-            )
-
-            # Assemble Jacobian
-            self.A.zeroEntries()
-            dolfinx.fem.petsc.assemble_matrix_block(
-                self.A, self.J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a
-            )
-            self.A.assemble()
-
-            self._solver.solve(self.b, self.dx)
-            # self._solver.view()
-            assert (
-                self._solver.getConvergedReason() > 0
-            ), "Linear solver did not converge"
-            offset_start = 0
-            for s in self.w:
-                num_sub_dofs = (
-                    s.function_space.dofmap.index_map.size_local
-                    * s.function_space.dofmap.index_map_bs
-                )
-                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
-                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
-                )
-                s.x.petsc_vec.ghostUpdate(
-                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
-                )
-                offset_start += num_sub_dofs
-            # Compute norm of update
-
-            correction_norm = self.dx.norm(0)
-            print(f"Iteration {i}: Correction norm {correction_norm}")
-            if correction_norm < tol:
-                break
-            i += 1
-
-    def __del__(self):
-        self.A.destroy()
-        self.b.destroy()
-        self.dx.destroy()
-        self._solver.destroy()
-        self.x.destroy()
-
-
-def transfer_meshtags_to_submesh(
-    mesh, entity_tag, submesh, sub_vertex_to_parent, sub_cell_to_parent
-):
-    """
-    Transfer a meshtag from a parent mesh to a sub-mesh.
-    """
-
-    tdim = mesh.topology.dim
-    cell_imap = mesh.topology.index_map(tdim)
-    num_cells = cell_imap.size_local + cell_imap.num_ghosts
-    mesh_to_submesh = np.full(num_cells, -1)
-    mesh_to_submesh[sub_cell_to_parent] = np.arange(
-        len(sub_cell_to_parent), dtype=np.int32
-    )
-    sub_vertex_to_parent = np.asarray(sub_vertex_to_parent)
-
-    submesh.topology.create_connectivity(entity_tag.dim, 0)
-
-    num_child_entities = (
-        submesh.topology.index_map(entity_tag.dim).size_local
-        + submesh.topology.index_map(entity_tag.dim).num_ghosts
-    )
-    submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)
-
-    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
-    c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)
-
-    child_markers = np.full(num_child_entities, 0, dtype=np.int32)
-
-    mesh.topology.create_connectivity(entity_tag.dim, 0)
-    mesh.topology.create_connectivity(entity_tag.dim, mesh.topology.dim)
-    p_f_to_v = mesh.topology.connectivity(entity_tag.dim, 0)
-    p_f_to_c = mesh.topology.connectivity(entity_tag.dim, mesh.topology.dim)
-    sub_to_parent_entity_map = np.full(num_child_entities, -1, dtype=np.int32)
-    for facet, value in zip(entity_tag.indices, entity_tag.values):
-        facet_found = False
-        for cell in p_f_to_c.links(facet):
-            if facet_found:
-                break
-            if (child_cell := mesh_to_submesh[cell]) != -1:
-                for child_facet in c_c_to_e.links(child_cell):
-                    child_vertices = c_e_to_v.links(child_facet)
-                    child_vertices_as_parent = sub_vertex_to_parent[child_vertices]
-                    is_facet = np.isin(
-                        child_vertices_as_parent, p_f_to_v.links(facet)
-                    ).all()
-                    if is_facet:
-                        child_markers[child_facet] = value
-                        facet_found = True
-                        sub_to_parent_entity_map[child_facet] = facet
-    tags = dolfinx.mesh.meshtags(
-        submesh,
-        entity_tag.dim,
-        np.arange(num_child_entities, dtype=np.int32),
-        child_markers,
-    )
-    tags.name = entity_tag.name
-    return tags, sub_to_parent_entity_map
-
-
-# ---------------- Generate a mesh ----------------
-def generate_mesh():
-    def left_boundary(x):
-        return np.isclose(x[0], 0.0)
-
-    def right_boundary(x):
-        return np.isclose(x[0], 1.0)
-
-    def left_domain(x):
-        return x[0] <= 0.5 + 1e-14
-
-    def right_domain(x):
-        return x[0] >= 0.7 - 1e-14
-
-    interface_1 = 0.5
-    interface_2 = 0.7
-
-    # with N = 2000 I've never had the error but with N = 10 I had it
-    N = 3
-    vertices = np.concatenate(
-        [
-            np.linspace(0, interface_1, num=N),
-            np.linspace(interface_1, interface_2, num=N),
-            np.linspace(interface_2, 1, num=N),
-        ]
-    )
-
-    vertices = np.sort(np.unique(vertices)).astype(float)
-    degree = 1
-    domain = ufl.Mesh(
-        basix.ufl.element(basix.ElementFamily.P, "interval", degree, shape=(1,))
-    )
-    mesh_points = np.reshape(vertices, (len(vertices), 1))
-    indexes = np.arange(vertices.shape[0])
-    cells = np.stack((indexes[:-1], indexes[1:]), axis=-1)
-
-    mesh = dolfinx.mesh.create_mesh(MPI.COMM_WORLD, cells, mesh_points, domain)
-
-    # Split domain in half and set an interface tag of 5
-    tdim = mesh.topology.dim
-    fdim = tdim - 1
-    left_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, left_boundary)
-    right_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, right_boundary)
-    num_facets_local = (
-        mesh.topology.index_map(fdim).size_local
-        + mesh.topology.index_map(fdim).num_ghosts
-    )
-    facets = np.arange(num_facets_local, dtype=np.int32)
-    values = np.full_like(facets, 0, dtype=np.int32)
-    values[left_facets] = 1
-    values[right_facets] = 2
-
-    left_cells = dolfinx.mesh.locate_entities(mesh, tdim, left_domain)
-    right_cells = dolfinx.mesh.locate_entities(mesh, tdim, right_domain)
-    num_cells_local = (
-        mesh.topology.index_map(tdim).size_local
-        + mesh.topology.index_map(tdim).num_ghosts
-    )
-    cells = np.full(num_cells_local, 4, dtype=np.int32)
-    cells[left_cells] = 3
-    cells[right_cells] = 5
-    ct = dolfinx.mesh.meshtags(
-        mesh, tdim, np.arange(num_cells_local, dtype=np.int32), cells
-    )
-    all_l_facets = dolfinx.mesh.compute_incident_entities(
-        mesh.topology, ct.find(3), tdim, fdim
-    )
-    all_m_facets = dolfinx.mesh.compute_incident_entities(
-        mesh.topology, ct.find(4), tdim, fdim
-    )
-    all_r_facets = dolfinx.mesh.compute_incident_entities(
-        mesh.topology, ct.find(5), tdim, fdim
-    )
-    interface1 = np.intersect1d(all_l_facets, all_m_facets)
-    interface2 = np.intersect1d(all_m_facets, all_r_facets)
-    values[interface1] = 5
-    values[interface2] = 6
-
-    mt = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, values)
-    return mesh, mt, ct
-
-
-class VolumeSubdomain:
-    id: int
-    submesh: dolfinx.mesh.Mesh
-    submesh_to_mesh: np.ndarray
-    parent_mesh: dolfinx.mesh.Mesh
-    parent_to_submesh: np.ndarray
-    v_map: np.ndarray
-    facet_to_parent: np.ndarray
-    ft: dolfinx.mesh.MeshTags
-    padded: bool
-
-    def __init__(self, id):
-        self.id = id
-
-    def create_subdomain(self, mesh, marker):
-        assert marker.dim == mesh.topology.dim
-        self.parent_mesh = mesh
-        self.submesh, self.submesh_to_mesh, self.v_map = dolfinx.mesh.create_submesh(
-            mesh, marker.dim, marker.find(self.id)
-        )[0:3]
-        num_cells_local = (
-            mesh.topology.index_map(marker.dim).size_local
-            + mesh.topology.index_map(marker.dim).num_ghosts
-        )
-        self.parent_to_submesh = np.full(num_cells_local, -1, dtype=np.int32)
-        self.parent_to_submesh[self.submesh_to_mesh] = np.arange(
-            len(self.submesh_to_mesh), dtype=np.int32
-        )
-        self.padded = False
-
-    def transfer_meshtag(self, tag):
-        # Transfer meshtags to submesh
-        assert self.submesh is not None, "Need to call create_subdomain first"
-        self.ft, self.facet_to_parent = transfer_meshtags_to_submesh(
-            mesh, tag, self.submesh, self.v_map, self.submesh_to_mesh
-        )
-
-
-class Interface:
-    id: int
-    subdomains: tuple[VolumeSubdomain, VolumeSubdomain]
-    parent_mesh: dolfinx.mesh.Mesh
-    restriction: [str, str] = ("+", "-")
-    padded: bool
-
-    def __init__(self, parent_mesh, mt, id, subdomains):
-        self.id = id
-        self.subdomains = tuple(subdomains)
-        self.mt = mt
-        self.parent_mesh = parent_mesh
-
-    def pad_parent_maps(self):
-        """Workaround to make sparsity-pattern work without skips"""
-
-        integration_data = compute_integration_domains(
-            dolfinx.fem.IntegralType.interior_facet,
-            self.parent_mesh.topology._cpp_object,
-            self.mt.find(self.id),
-            self.mt.dim,
-        ).reshape(-1, 4)
-        for i in range(2):
-            # We pad the parent to submesh map to make sure that sparsity pattern is correct
-            mapped_cell_0 = self.subdomains[i].parent_to_submesh[integration_data[:, 0]]
-            mapped_cell_1 = self.subdomains[i].parent_to_submesh[integration_data[:, 2]]
-            max_cells = np.maximum(mapped_cell_0, mapped_cell_1)
-            self.subdomains[i].parent_to_submesh[integration_data[:, 0]] = max_cells
-            self.subdomains[i].parent_to_submesh[integration_data[:, 2]] = max_cells
-            self.subdomains[i].padded = True
-
-
-mesh, mt, ct = generate_mesh()
-
-left_domain = VolumeSubdomain(3)
-mid_domain = VolumeSubdomain(4)
-right_domain = VolumeSubdomain(5)
-list_of_subdomains = [left_domain, mid_domain, right_domain]
-
-
-gdim = mesh.geometry.dim
-tdim = mesh.topology.dim
-fdim = tdim - 1
-num_cells_local = (
-    mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts
-)
-
-
-for subdomain in list_of_subdomains:
-    subdomain.create_subdomain(mesh, ct)
-    subdomain.transfer_meshtag(mt)
-
-
-i0 = Interface(mesh, mt, 5, (left_domain, mid_domain))
-i1 = Interface(mesh, mt, 6, (mid_domain, right_domain))
-interfaces = [i0, i1]
-
-
-def define_interior_eq(mesh, degree, submesh, submesh_to_mesh, value):
-    element_CG = basix.ufl.element(
-        basix.ElementFamily.P,
-        submesh.basix_cell(),
-        degree,
-        basix.LagrangeVariant.equispaced,
-    )
-    element = basix.ufl.mixed_element([element_CG, element_CG])
-    V = dolfinx.fem.functionspace(submesh, element)
-    u = dolfinx.fem.Function(V, name=f"u_{value}")
-    us = list(ufl.split(u))
-    vs = list(ufl.TestFunctions(V))
-    ct_r = dolfinx.mesh.meshtags(
-        mesh,
-        mesh.topology.dim,
-        submesh_to_mesh,
-        np.full_like(submesh_to_mesh, 1, dtype=np.int32),
-    )
-    val = dolfinx.fem.Constant(submesh, value)
-    dx_r = ufl.Measure("dx", domain=mesh, subdomain_data=ct_r, subdomain_id=1)
-    F = ufl.inner(ufl.grad(us[0]), ufl.grad(vs[0])) * dx_r - val * vs[0] * dx_r
-    k = 2
-    p = 0.1
-    n = 0.5
-    F += k * us[0] * (n - us[1]) * vs[1] * dx_r - p * us[1] * vs[1] * dx_r
-    return u, vs, F
-
-
-# for each subdomain, define the interior equation
-for subdomain in list_of_subdomains:
-    degree = 1
-    subdomain.u, subdomain.vs, subdomain.F = define_interior_eq(
-        mesh, degree, subdomain.submesh, subdomain.submesh_to_mesh, 0.0
-    )
-    subdomain.u.name = f"u_{subdomain.id}"
-
-
-def compute_mapped_interior_facet_data(interface: Interface):
-    """
-    Compute integration data for interface integrals.
-    We define the first domain on an interface as the "+" restriction,
-    meaning that we must sort all integration entities in this order
-
-    Parameters
-        interface: Interface between two subdomains
-    Returns
-        integration_data: Integration data for interior facets
-    """
-    assert (not interface.subdomains[0].padded) and (not interface.subdomains[1].padded)
-    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
-    integration_data = compute_integration_domains(
-        dolfinx.fem.IntegralType.interior_facet,
-        mesh.topology._cpp_object,
-        interface.mt.find(interface.id),
-        interface.mt.dim,
-    )
-
-    ordered_integration_data = integration_data.reshape(-1, 4).copy()
-
-    mapped_cell_0 = interface.subdomains[0].parent_to_submesh[integration_data[0::4]]
-    mapped_cell_1 = interface.subdomains[0].parent_to_submesh[integration_data[2::4]]
-
-    switch = mapped_cell_1 > mapped_cell_0
-    # Order restriction on one side
-    if True in switch:
-        ordered_integration_data[switch, [0, 1, 2, 3]] = ordered_integration_data[
-            switch, [2, 3, 0, 1]
-        ]
-
-    # Check that other restriction lies in other interface
-    domain1_cell = interface.subdomains[1].parent_to_submesh[
-        ordered_integration_data[:, 2]
-    ]
-    assert (domain1_cell >= 0).all()
-
-    return (interface.id, ordered_integration_data.reshape(-1))
-
-
-integral_data = [
-    compute_mapped_interior_facet_data(interface) for interface in interfaces
-]
-[interface.pad_parent_maps() for interface in interfaces]
-dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)
-
-
-def mixed_term(u, v, n):
-    return ufl.dot(ufl.grad(u), n) * v
-
-
-n = ufl.FacetNormal(mesh)
-cr = ufl.Circumradius(mesh)
-
-gamma = 10.0
-
-
-entity_maps = {sd.submesh: sd.parent_to_submesh for sd in list_of_subdomains}
-for interface in interfaces:
-    subdomain_1, subdomain_2 = interface.subdomains
-    b_res, t_res = interface.restriction
-    n_b = n(b_res)
-    n_t = n(t_res)
-    h_b = 2 * cr(b_res)
-    h_t = 2 * cr(t_res)
-
-    v_b = subdomain_1.vs[0](b_res)
-    v_t = subdomain_2.vs[0](t_res)
-
-    u_bs = list(ufl.split(subdomain_1.u))
-    u_ts = list(ufl.split(subdomain_2.u))
-    u_b = u_bs[0](b_res)
-    u_t = u_ts[0](t_res)
-    # fabricate K
-    W_0 = dolfinx.fem.functionspace(subdomain_1.submesh, ("DG", 0))
-    K_0 = dolfinx.fem.Function(W_0, name=f"K_{subdomain_1.id}")
-    K_0.x.array[:] = 2
-    W_1 = dolfinx.fem.functionspace(subdomain_2.submesh, ("DG", 0))
-    K_1 = dolfinx.fem.Function(W_1, name=f"K_{subdomain_2.id}")
-    K_1.x.array[:] = 4
-
-    K_b = K_0(b_res)
-    K_t = K_1(t_res)
-
-    F_0 = -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(
-        interface.id
-    ) - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
-
-    F_1 = +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(
-        interface.id
-    ) - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(interface.id)
-    F_0 += (
-        2
-        * gamma
-        / (h_b + h_t)
-        * (u_b / K_b - u_t / K_t)
-        * v_b
-        * dInterface(interface.id)
-    )
-    F_1 += (
-        -2
-        * gamma
-        / (h_b + h_t)
-        * (u_b / K_b - u_t / K_t)
-        * v_t
-        * dInterface(interface.id)
-    )
-
-    subdomain_1.F += F_0
-    subdomain_2.F += F_1
-
-
-J = []
-forms = []
-for i, subdomain1 in enumerate(list_of_subdomains):
-    jac = []
-    form = subdomain1.F
-    for j, subdomain2 in enumerate(list_of_subdomains):
-        jac.append(
-            dolfinx.fem.form(
-                ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
-            )
-        )
-
-    J.append(jac)
-    forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
-
-# boundary conditions
-b_bc = dolfinx.fem.Function(left_domain.u.function_space)
-b_bc.x.array[:] = 0.2
-left_domain.submesh.topology.create_connectivity(
-    left_domain.submesh.topology.dim - 1, left_domain.submesh.topology.dim
-)
-bc_b = dolfinx.fem.dirichletbc(
-    b_bc,
-    dolfinx.fem.locate_dofs_topological(
-        left_domain.u.function_space.sub(0), fdim, left_domain.ft.find(1)
-    ),
-)
-
-t_bc = dolfinx.fem.Constant(right_domain.submesh, 0.05)
-sub_V = right_domain.u.function_space.sub(0)
-right_domain.submesh.topology.create_connectivity(
-    right_domain.submesh.topology.dim - 1, right_domain.submesh.topology.dim
-)
-bc_dofs = dolfinx.fem.locate_dofs_topological(
-    sub_V,
-    fdim,
-    right_domain.ft.find(2),
-)
-bc_t = dolfinx.fem.dirichletbc(t_bc, bc_dofs, sub_V)
-bcs = [bc_b, bc_t]
-
-
-solver = NewtonSolver(
-    forms,
-    J,
-    [subdomain.u for subdomain in list_of_subdomains],
-    bcs=bcs,
-    max_iterations=10,
-    petsc_options={
-        "ksp_type": "preonly",
-        "pc_type": "lu",
-        "pc_factor_mat_solver_type": "mumps",
-    },
-)
-solver.solve(1e-5)
-
-# for subdomain in list_of_subdomains:
-#     u_sub_0 = subdomain.u.sub(0).collapse()
-#     u_sub_0.name = "u_sub_0"
-
-#     u_sub_1 = subdomain.u.sub(1).collapse()
-#     u_sub_1.name = "u_sub_1"
-#     bp = dolfinx.io.VTXWriter(
-#         mesh.comm, f"u_{subdomain.id}.bp", [u_sub_0, u_sub_1], engine="BP4"
-#     )
-#     bp.write(0)
-#     bp.close()
-
-
-# # derived quantities
-# V = dolfinx.fem.functionspace(mesh, ("CG", 1))
-# T = dolfinx.fem.Function(V)
-# T.interpolate(lambda x: 200 + x[1])
-
-
-# T_b = dolfinx.fem.Function(right_domain.u.sub(0).collapse().function_space)
-# T_b.interpolate(T)
-
-# ds_b = ufl.Measure("ds", domain=right_domain.submesh)
-# dx_b = ufl.Measure("dx", domain=right_domain.submesh)
-# dx = ufl.Measure("dx", domain=mesh)
-
-# n_b = ufl.FacetNormal(left_domain.submesh)
-
-# form = dolfinx.fem.form(left_domain.u.sub(0) * dx_b, entity_maps=entity_maps)
-# print(dolfinx.fem.assemble_scalar(form))
-
-# form = dolfinx.fem.form(T_b * ufl.dot(ufl.grad(left_domain.u.sub(0)), n_b) * ds_b, entity_maps=entity_maps)
-# print(dolfinx.fem.assemble_scalar(form))

From 67d6e4ad7654853dc60a87ca75cc3c47d87c4799 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:10:37 -0400
Subject: [PATCH 081/134] moved the examples to examples folder

---
 test_implementation_3.py => examples/multi_material_1d.py         | 0
 test_implementation_2.py => examples/multi_material_2d.py         | 0
 test_implementation.py => examples/multi_material_2d_bis.py       | 0
 .../multi_material_transient.py                                   | 0
 .../multi_material_with_one_material.py                           | 0
 5 files changed, 0 insertions(+), 0 deletions(-)
 rename test_implementation_3.py => examples/multi_material_1d.py (100%)
 rename test_implementation_2.py => examples/multi_material_2d.py (100%)
 rename test_implementation.py => examples/multi_material_2d_bis.py (100%)
 rename test_implementation_transient.py => examples/multi_material_transient.py (100%)
 rename test_one_material.py => examples/multi_material_with_one_material.py (100%)

diff --git a/test_implementation_3.py b/examples/multi_material_1d.py
similarity index 100%
rename from test_implementation_3.py
rename to examples/multi_material_1d.py
diff --git a/test_implementation_2.py b/examples/multi_material_2d.py
similarity index 100%
rename from test_implementation_2.py
rename to examples/multi_material_2d.py
diff --git a/test_implementation.py b/examples/multi_material_2d_bis.py
similarity index 100%
rename from test_implementation.py
rename to examples/multi_material_2d_bis.py
diff --git a/test_implementation_transient.py b/examples/multi_material_transient.py
similarity index 100%
rename from test_implementation_transient.py
rename to examples/multi_material_transient.py
diff --git a/test_one_material.py b/examples/multi_material_with_one_material.py
similarity index 100%
rename from test_one_material.py
rename to examples/multi_material_with_one_material.py

From 4fbbab47e3ad00cd9579af2889e6b49eda57aa8d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Delaporte-Mathurin?=
 <40028739+RemDelaporteMathurin@users.noreply.github.com>
Date: Wed, 23 Oct 2024 12:12:27 -0400
Subject: [PATCH 082/134] Fix docstrings

Co-authored-by: James Dark <65899899+jhdark@users.noreply.github.com>
---
 festim/material.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/festim/material.py b/festim/material.py
index 56f86e0b7..13daef98a 100644
--- a/festim/material.py
+++ b/festim/material.py
@@ -170,7 +170,7 @@ def get_diffusion_coefficient(self, mesh, temperature, species=None):
             raise ValueError("D_0 and E_D must be either floats or dicts")
 
     def get_solubility_coefficient(self, mesh, temperature, species=None):
-        """Defines the diffusion coefficient
+        """Defines the solubility coefficient
         Args:
 
             mesh (dolfinx.mesh.Mesh): the domain mesh

From ad6c4b0851cc5732936ed7bc15bd004ca6659d94 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Delaporte-Mathurin?=
 <40028739+RemDelaporteMathurin@users.noreply.github.com>
Date: Wed, 23 Oct 2024 12:12:44 -0400
Subject: [PATCH 083/134] Fixed docstrings again

Co-authored-by: James Dark <65899899+jhdark@users.noreply.github.com>
---
 festim/material.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/festim/material.py b/festim/material.py
index 13daef98a..a698b0588 100644
--- a/festim/material.py
+++ b/festim/material.py
@@ -175,7 +175,7 @@ def get_solubility_coefficient(self, mesh, temperature, species=None):
 
             mesh (dolfinx.mesh.Mesh): the domain mesh
             temperature (dolfinx.fem.Constant): the temperature
-            species (festim.Species, optional): the species we want the diffusion
+            species (festim.Species, optional): the species we want the solubility
                 coefficient of. Only needed if K_S_0 and E_K_S are dicts.
         Returns:
             ufl.algebra.Product: the solubility coefficient

From a6e1de6d2cf59c550bb350f90548ef441f5b0cc6 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:17:22 -0400
Subject: [PATCH 084/134] domain 0 and 1 instead of 1 and 2

---
 festim/mesh/mesh.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/festim/mesh/mesh.py b/festim/mesh/mesh.py
index 03fb8a4aa..74fe486f5 100644
--- a/festim/mesh/mesh.py
+++ b/festim/mesh/mesh.py
@@ -95,20 +95,20 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
         # tag interfaces
         interfaces = interfaces or {}  # if interfaces is None, set it to empty dict
         for interface in interfaces:
-            (domain_1, domain_2) = interface.subdomains
-            all_1_facets = dolfinx.mesh.compute_incident_entities(
+            (domain_0, domain_1) = interface.subdomains
+            all_0_facets = dolfinx.mesh.compute_incident_entities(
                 self.mesh.topology,
-                volume_meshtags.find(domain_1.id),
+                volume_meshtags.find(domain_0.id),
                 self.vdim,
                 self.fdim,
             )
-            all_2_facets = dolfinx.mesh.compute_incident_entities(
+            all_1_facets = dolfinx.mesh.compute_incident_entities(
                 self.mesh.topology,
-                volume_meshtags.find(domain_2.id),
+                volume_meshtags.find(domain_1.id),
                 self.vdim,
                 self.fdim,
             )
-            interface_entities = np.intersect1d(all_1_facets, all_2_facets)
+            interface_entities = np.intersect1d(all_0_facets, all_1_facets)
             tags_facets[interface_entities] = interface.id
 
         facet_meshtags = meshtags(self.mesh, self.fdim, mesh_facet_indices, tags_facets)

From 5df377bd0b8ad22ba82aa28071025f72cabe0bbb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 12:24:19 -0400
Subject: [PATCH 085/134] - Add in/outfiles to git ignore - Change "CG" ->
 "Lagrange" - Use `scifem` NewtonSolver` - Update Python version support - Get
 rid of scary behavior in XDMFFile exporter

---
 .gitignore                           |   6 +
 discontinuity_generic.py             |  23 ++--
 festim/exports/xdmf.py               |   7 ++
 festim/helpers_discontinuity.py      | 164 +--------------------------
 festim/hydrogen_transport_problem.py |  43 ++++---
 festim/subdomain/volume_subdomain.py |   3 +-
 setup.cfg                            |   1 +
 test/test_average_surface.py         |   2 +-
 test/test_average_volume.py          |   5 +-
 test/test_maximum_surface.py         |   2 +-
 test/test_maximum_volume.py          |   5 +-
 test/test_minimum_surface.py         |   2 +-
 test/test_minimum_volume.py          |   5 +-
 test/test_total_surface.py           |   2 +-
 14 files changed, 73 insertions(+), 197 deletions(-)

diff --git a/.gitignore b/.gitignore
index 1d5907d26..26288e881 100644
--- a/.gitignore
+++ b/.gitignore
@@ -105,3 +105,9 @@ venv.bak/
 
 # version file
 **_version.py
+
+# Input/Output files
+*.bp
+*.xdmf
+*.h5
+*.txt
\ No newline at end of file
diff --git a/discontinuity_generic.py b/discontinuity_generic.py
index 5917e9f51..4766fa37b 100644
--- a/discontinuity_generic.py
+++ b/discontinuity_generic.py
@@ -143,7 +143,8 @@ def solve(self, tol=1e-6, beta=1.0):
                     * s.function_space.dofmap.index_map_bs
                 )
                 s.x.petsc_vec.array_w[:num_sub_dofs] -= (
-                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
+                    beta *
+                    self.dx.array_r[offset_start: offset_start + num_sub_dofs]
                 )
                 s.x.petsc_vec.ghostUpdate(
                     addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
@@ -189,7 +190,8 @@ def transfer_meshtags_to_submesh(
     )
     submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)
 
-    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
+    c_c_to_e = submesh.topology.connectivity(
+        submesh.topology.dim, entity_tag.dim)
     c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)
 
     child_markers = np.full(num_child_entities, 0, dtype=np.int32)
@@ -255,7 +257,8 @@ def right_domain(x):
     vertices = np.sort(np.unique(vertices)).astype(float)
     degree = 1
     domain = ufl.Mesh(
-        basix.ufl.element(basix.ElementFamily.P, "interval", degree, shape=(1,))
+        basix.ufl.element(basix.ElementFamily.P,
+                          "interval", degree, shape=(1,))
     )
     mesh_points = np.reshape(vertices, (len(vertices), 1))
     indexes = np.arange(vertices.shape[0])
@@ -266,8 +269,10 @@ def right_domain(x):
     # Split domain in half and set an interface tag of 5
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    left_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, left_boundary)
-    right_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, right_boundary)
+    left_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, left_boundary)
+    right_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, right_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
@@ -389,7 +394,8 @@ def pad_parent_maps(self):
 tdim = mesh.topology.dim
 fdim = tdim - 1
 num_cells_local = (
-    mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts
+    mesh.topology.index_map(tdim).size_local +
+    mesh.topology.index_map(tdim).num_ghosts
 )
 
 
@@ -451,7 +457,8 @@ def compute_mapped_interior_facet_data(interface: Interface):
     Returns
         integration_data: Integration data for interior facets
     """
-    assert (not interface.subdomains[0].padded) and (not interface.subdomains[1].padded)
+    assert (not interface.subdomains[0].padded) and (
+        not interface.subdomains[1].padded)
     mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
     integration_data = compute_integration_domains(
         dolfinx.fem.IntegralType.interior_facet,
@@ -623,7 +630,7 @@ def mixed_term(u, v, n):
 
 
 # # derived quantities
-# V = dolfinx.fem.functionspace(mesh, ("CG", 1))
+# V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
 # T = dolfinx.fem.Function(V)
 # T.interpolate(lambda x: 200 + x[1])
 
diff --git a/festim/exports/xdmf.py b/festim/exports/xdmf.py
index 83c8f6dce..593123c52 100644
--- a/festim/exports/xdmf.py
+++ b/festim/exports/xdmf.py
@@ -15,10 +15,12 @@ class XDMFExport:
         writer (dolfinx.io.XDMFFile): the XDMF writer
         field (festim.Species, list of festim.Species): the field index to export
     """
+    _mesh_written: bool
 
     def __init__(self, filename: str, field) -> None:
         self.filename = filename
         self.field = field
+        self._mesh_written = False
 
     @property
     def filename(self):
@@ -70,5 +72,10 @@ def write(self, t: float):
         Args:
             t (float): the time of export
         """
+        if not self._mesh_written:
+            self.writer.write_mesh(
+                self.field[0].post_processing_solution.function_space.mesh)
+            self._mesh_written = True
+
         for field in self.field:
             self.writer.write_function(field.post_processing_solution, t)
diff --git a/festim/helpers_discontinuity.py b/festim/helpers_discontinuity.py
index 824860f48..0bc9ea60c 100644
--- a/festim/helpers_discontinuity.py
+++ b/festim/helpers_discontinuity.py
@@ -1,166 +1,5 @@
-from mpi4py import MPI
 import dolfinx
-import dolfinx.fem.petsc
-import ufl
 import numpy as np
-from petsc4py import PETSc
-import basix
-
-
-class NewtonSolver:
-    max_iterations: int
-    bcs: list[dolfinx.fem.DirichletBC]
-    A: PETSc.Mat
-    b: PETSc.Vec
-    J: dolfinx.fem.Form
-    b: dolfinx.fem.Form
-    dx: PETSc.Vec
-
-    def __init__(
-        self,
-        F: list[dolfinx.fem.form],
-        J: list[list[dolfinx.fem.form]],
-        w: list[dolfinx.fem.Function],
-        bcs: list[dolfinx.fem.DirichletBC] | None = None,
-        max_iterations: int = 5,
-        petsc_options: dict[str, str | float | int | None] = None,
-        problem_prefix="newton",
-    ):
-        self.max_iterations = max_iterations
-        self.bcs = [] if bcs is None else bcs
-        self.b = dolfinx.fem.petsc.create_vector_block(F)
-        self.F = F
-        self.J = J
-        self.A = dolfinx.fem.petsc.create_matrix_block(J)
-        self.dx = self.A.createVecLeft()
-        self.w = w
-        self.x = dolfinx.fem.petsc.create_vector_block(F)
-
-        # Set PETSc options
-        opts = PETSc.Options()
-        if petsc_options is not None:
-            for k, v in petsc_options.items():
-                opts[k] = v
-
-        # Define KSP solver
-        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
-        self._solver.setOperators(self.A)
-        self._solver.setFromOptions()
-
-        # Set matrix and vector PETSc options
-        self.A.setFromOptions()
-        self.b.setFromOptions()
-
-    def solve(self, tol=1e-6, beta=1.0):
-        i = 0
-
-        while i < self.max_iterations:
-            dolfinx.cpp.la.petsc.scatter_local_vectors(
-                self.x,
-                [si.x.petsc_vec.array_r for si in self.w],
-                [
-                    (
-                        si.function_space.dofmap.index_map,
-                        si.function_space.dofmap.index_map_bs,
-                    )
-                    for si in self.w
-                ],
-            )
-            self.x.ghostUpdate(
-                addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
-            )
-
-            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
-            with self.b.localForm() as b_local:
-                b_local.set(0.0)
-
-            constants_L = [
-                form and dolfinx.cpp.fem.pack_constants(form._cpp_object)
-                for form in self.F
-            ]
-            coeffs_L = [
-                dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self.F
-            ]
-
-            constants_a = [
-                [
-                    (
-                        dolfinx.cpp.fem.pack_constants(form._cpp_object)
-                        if form is not None
-                        else np.array([], dtype=PETSc.ScalarType)
-                    )
-                    for form in forms
-                ]
-                for forms in self.J
-            ]
-
-            coeffs_a = [
-                [
-                    (
-                        {}
-                        if form is None
-                        else dolfinx.cpp.fem.pack_coefficients(form._cpp_object)
-                    )
-                    for form in forms
-                ]
-                for forms in self.J
-            ]
-
-            dolfinx.fem.petsc.assemble_vector_block(
-                self.b,
-                self.F,
-                self.J,
-                bcs=self.bcs,
-                x0=self.x,
-                alpha=-1.0,
-                coeffs_a=coeffs_a,
-                constants_a=constants_a,
-                coeffs_L=coeffs_L,
-                constants_L=constants_L,
-            )
-            self.b.ghostUpdate(
-                PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
-            )
-
-            # Assemble Jacobian
-            self.A.zeroEntries()
-            dolfinx.fem.petsc.assemble_matrix_block(
-                self.A, self.J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a
-            )
-            self.A.assemble()
-
-            self._solver.solve(self.b, self.dx)
-            # self._solver.view()
-            assert (
-                self._solver.getConvergedReason() > 0
-            ), "Linear solver did not converge"
-            offset_start = 0
-            for s in self.w:
-                num_sub_dofs = (
-                    s.function_space.dofmap.index_map.size_local
-                    * s.function_space.dofmap.index_map_bs
-                )
-                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
-                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
-                )
-                s.x.petsc_vec.ghostUpdate(
-                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
-                )
-                offset_start += num_sub_dofs
-            # Compute norm of update
-
-            correction_norm = self.dx.norm(0)
-            print(f"Iteration {i}: Correction norm {correction_norm}")
-            if correction_norm < tol:
-                break
-            i += 1
-
-    def __del__(self):
-        self.A.destroy()
-        self.b.destroy()
-        self.dx.destroy()
-        self._solver.destroy()
-        self.x.destroy()
 
 
 def transfer_meshtags_to_submesh(
@@ -187,7 +26,8 @@ def transfer_meshtags_to_submesh(
     )
     submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)
 
-    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
+    c_c_to_e = submesh.topology.connectivity(
+        submesh.topology.dim, entity_tag.dim)
     c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)
 
     child_markers = np.full(num_child_entities, 0, dtype=np.int32)
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index e909c97ea..fe34557d2 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -6,7 +6,7 @@
 import numpy as np
 import tqdm.autonotebook
 import festim as F
-from festim.helpers_discontinuity import NewtonSolver
+from scifem import NewtonSolver
 
 
 class HydrogenTransportProblem(F.ProblemBase):
@@ -182,7 +182,8 @@ def species(self, value):
         for spe in value:
             if not isinstance(spe, F.Species):
                 raise TypeError(
-                    f"elements of species must be of type festim.Species not {type(spe)}"
+                    f"elements of species must be of type festim.Species not {
+                        type(spe)}"
                 )
         self._species = value
 
@@ -270,11 +271,13 @@ def define_temperature(self):
             if "t" in arguments and "x" not in arguments:
                 if not isinstance(self.temperature(t=float(self.t)), (float, int)):
                     raise ValueError(
-                        f"self.temperature should return a float or an int, not {type(self.temperature(t=float(self.t)))} "
+                        f"self.temperature should return a float or an int, not {
+                            type(self.temperature(t=float(self.t)))} "
                     )
                 # only t is an argument
                 self.temperature_fenics = F.as_fenics_constant(
-                    mesh=self.mesh.mesh, value=self.temperature(t=float(self.t))
+                    mesh=self.mesh.mesh, value=self.temperature(
+                        t=float(self.t))
                 )
             else:
                 x = ufl.SpatialCoordinate(self.mesh.mesh)
@@ -288,7 +291,8 @@ def define_temperature(self):
                 function_space_temperature = fem.functionspace(
                     self.mesh.mesh, element_temperature
                 )
-                self.temperature_fenics = fem.Function(function_space_temperature)
+                self.temperature_fenics = fem.Function(
+                    function_space_temperature)
                 kwargs = {}
                 if "t" in arguments:
                     kwargs["t"] = self.t
@@ -316,12 +320,11 @@ def initialise_exports(self):
                             field, self.species
                         )
             elif isinstance(export.field, str):
-                export.field = F.find_species_from_name(export.field, self.species)
+                export.field = F.find_species_from_name(
+                    export.field, self.species)
 
             if isinstance(export, (F.VTXExport, F.XDMFExport)):
                 export.define_writer(MPI.COMM_WORLD)
-                if isinstance(export, F.XDMFExport):
-                    export.writer.write_mesh(self.mesh.mesh)
 
         # compute diffusivity function for surface fluxes
 
@@ -360,7 +363,8 @@ def define_D_global(self, species):
         D_0 = fem.Function(self.V_DG_0)
         E_D = fem.Function(self.V_DG_0)
         for vol in self.volume_subdomains:
-            cell_indices = vol.locate_subdomain_entities(self.mesh.mesh, self.mesh.vdim)
+            cell_indices = vol.locate_subdomain_entities(
+                self.mesh.mesh, self.mesh.vdim)
 
             # replace values of D_0 and E_D by values from the material
             D_0.x.array[cell_indices] = vol.material.get_D_0(species=species)
@@ -370,9 +374,11 @@ def define_D_global(self, species):
         D = fem.Function(self.V_DG_1)
 
         expr = D_0 * ufl.exp(
-            -E_D / F.as_fenics_constant(F.k_B, self.mesh.mesh) / self.temperature_fenics
+            -E_D / F.as_fenics_constant(F.k_B,
+                                        self.mesh.mesh) / self.temperature_fenics
         )
-        D_expr = fem.Expression(expr, self.V_DG_1.element.interpolation_points())
+        D_expr = fem.Expression(
+            expr, self.V_DG_1.element.interpolation_points())
         D.interpolate(D_expr)
         return D, D_expr
 
@@ -565,7 +571,8 @@ def create_initial_conditions(self):
 
             # assign to previous solution of species
             if not self.multispecies:
-                condition.species.prev_solution.interpolate(condition.expr_fenics)
+                condition.species.prev_solution.interpolate(
+                    condition.expr_fenics)
             else:
                 idx = self.species.index(condition.species)
                 self.u_n.sub(idx).interpolate(condition.expr_fenics)
@@ -591,7 +598,8 @@ def create_formulation(self):
                     )
 
                 if self.settings.transient:
-                    self.formulation += ((u - u_n) / self.dt) * v * self.dx(vol.id)
+                    self.formulation += ((u - u_n) / self.dt) * \
+                        v * self.dx(vol.id)
 
         for reaction in self.reactions:
             for reactant in reaction.reactant:
@@ -1035,7 +1043,8 @@ def create_formulation(self):
             for interface in self.interfaces
         ]
         [interface.pad_parent_maps() for interface in self.interfaces]
-        dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)
+        dInterface = ufl.Measure(
+            "dS", domain=mesh, subdomain_data=integral_data)
 
         def mixed_term(u, v, n):
             return ufl.dot(ufl.grad(u), n) * v
@@ -1058,7 +1067,8 @@ def mixed_term(u, v, n):
 
             all_mobile_species = [spe for spe in self.species if spe.mobile]
             if len(all_mobile_species) > 1:
-                raise NotImplementedError("Multiple mobile species not implemented")
+                raise NotImplementedError(
+                    "Multiple mobile species not implemented")
             H = all_mobile_species[0]
 
             v_b = H.subdomain_to_test_function[subdomain_1](b_res)
@@ -1117,7 +1127,8 @@ def mixed_term(u, v, n):
                     )
                 )
             J.append(jac)
-            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
+            forms.append(dolfinx.fem.form(
+                subdomain1.F, entity_maps=entity_maps))
 
         self.forms = forms
         self.J = J
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index 4d8b055c9..c727fda0c 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -42,7 +42,8 @@ def locate_subdomain_entities(self, mesh, vdim):
         """
         # By default, all entities are included
         # return array like x full of True
-        entities = locate_entities(mesh, vdim, lambda x: np.full(x.shape[1], True))
+        entities = locate_entities(
+            mesh, vdim, lambda x: np.full(x.shape[1], True))
         return entities
 
     def create_subdomain(self, mesh: dolfinx.mesh.Mesh, marker: dolfinx.mesh.MeshTags):
diff --git a/setup.cfg b/setup.cfg
index c1a292ef3..38feea1a9 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -29,6 +29,7 @@ packages = find:
 python_requires= >=3.9
 install_requires =
     tqdm
+    scifem>=0.2.8
 
 [options.extras_require]
 test = 
diff --git a/test/test_average_surface.py b/test/test_average_surface.py
index 3c573f4cb..142bf7cfc 100644
--- a/test/test_average_surface.py
+++ b/test/test_average_surface.py
@@ -23,7 +23,7 @@ def test_average_surface_compute_1D():
     ds = ufl.Measure("ds", domain=my_mesh.mesh, subdomain_data=facet_meshtags)
 
     # give function to species
-    V = fem.functionspace(my_mesh.mesh, ("CG", 1))
+    V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))
     c = fem.Function(V)
     c.interpolate(lambda x: x[0] * 0.5 + 2)
 
diff --git a/test/test_average_volume.py b/test/test_average_volume.py
index 72ad79ec0..ca95dc17b 100644
--- a/test/test_average_volume.py
+++ b/test/test_average_volume.py
@@ -14,14 +14,15 @@ def test_average_volume_compute_1D():
     # BUILD
     L = 6.0
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
-    dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_mat)
+    dummy_volume = F.VolumeSubdomain1D(
+        id=1, borders=[0, L], material=dummy_mat)
     temp, cell_meshtags = my_mesh.define_meshtags(
         surface_subdomains=[], volume_subdomains=[dummy_volume]
     )
     dx = ufl.Measure("dx", domain=my_mesh.mesh, subdomain_data=cell_meshtags)
 
     # give function to species
-    V = fem.functionspace(my_mesh.mesh, ("CG", 1))
+    V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))
     c = fem.Function(V)
     c.interpolate(lambda x: x[0] * 0.5 + 2)
 
diff --git a/test/test_maximum_surface.py b/test/test_maximum_surface.py
index 81ad631fc..7ba03ca61 100644
--- a/test/test_maximum_surface.py
+++ b/test/test_maximum_surface.py
@@ -14,7 +14,7 @@ def test_maximum_surface_compute_1D():
     dummy_surface.locate_boundary_facet_indices(mesh=my_mesh.mesh)
 
     # give function to species
-    V = fem.functionspace(my_mesh.mesh, ("CG", 1))
+    V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))
     c = fem.Function(V)
     c.interpolate(lambda x: (x[0] - 3) ** 2)
 
diff --git a/test/test_maximum_volume.py b/test/test_maximum_volume.py
index bd460ac27..b6f459082 100644
--- a/test/test_maximum_volume.py
+++ b/test/test_maximum_volume.py
@@ -12,11 +12,12 @@ def test_maximum_volume_compute_1D():
     L = 6
     dummy_material = F.Material(D_0=1.5, E_D=1, name="dummy")
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
-    dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_material)
+    dummy_volume = F.VolumeSubdomain1D(
+        id=1, borders=[0, L], material=dummy_material)
     dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh, vdim=0)
 
     # give function to species
-    V = fem.functionspace(my_mesh.mesh, ("CG", 1))
+    V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))
     c = fem.Function(V)
     c.interpolate(lambda x: (x[0] - 1) ** 2 + 2)
 
diff --git a/test/test_minimum_surface.py b/test/test_minimum_surface.py
index d64b8bc25..eb3706a40 100644
--- a/test/test_minimum_surface.py
+++ b/test/test_minimum_surface.py
@@ -18,7 +18,7 @@ def test_minimum_surface_export_compute_1D():
     dummy_surface.locate_boundary_facet_indices(mesh=my_mesh.mesh)
 
     # give function to species
-    V = fem.functionspace(my_mesh.mesh, ("CG", 1))
+    V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))
     c = fem.Function(V)
     c.interpolate(lambda x: (x[0] - 2) ** 2)
 
diff --git a/test/test_minimum_volume.py b/test/test_minimum_volume.py
index 03f9e2d26..6cd5445c0 100644
--- a/test/test_minimum_volume.py
+++ b/test/test_minimum_volume.py
@@ -12,11 +12,12 @@ def test_minimum_volume_compute_1D():
     L = 6
     dummy_material = F.Material(D_0=1.5, E_D=1, name="dummy")
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
-    dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_material)
+    dummy_volume = F.VolumeSubdomain1D(
+        id=1, borders=[0, L], material=dummy_material)
     dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh, vdim=0)
 
     # give function to species
-    V = fem.functionspace(my_mesh.mesh, ("CG", 1))
+    V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))
     c = fem.Function(V)
     c.interpolate(lambda x: (x[0] - 2) ** 2 + 0.5)
 
diff --git a/test/test_total_surface.py b/test/test_total_surface.py
index 49007344d..8aef39755 100644
--- a/test/test_total_surface.py
+++ b/test/test_total_surface.py
@@ -22,7 +22,7 @@ def test_total_surface_compute_1D():
     ds = ufl.Measure("ds", domain=my_mesh.mesh, subdomain_data=facet_meshtags)
 
     # give function to species
-    V = fem.functionspace(my_mesh.mesh, ("CG", 1))
+    V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))
     c = fem.Function(V)
     c.interpolate(lambda x: x[0] * 0.5 + 1)
 

From 124fb2958884535455a339d2305ab5aa34c00225 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:36:49 -0400
Subject: [PATCH 086/134] removed top and bottom subscripts

---
 festim/hydrogen_transport_problem.py | 44 ++++++++++++++--------------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index e909c97ea..f63c32e58 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -1049,46 +1049,46 @@ def mixed_term(u, v, n):
         for interface in self.interfaces:
             gamma = interface.penalty_term
 
-            subdomain_1, subdomain_2 = interface.subdomains
-            b_res, t_res = interface.restriction
-            n_b = n(b_res)
-            n_t = n(t_res)
-            h_b = 2 * cr(b_res)
-            h_t = 2 * cr(t_res)
+            subdomain_0, subdomain_1 = interface.subdomains
+            res = interface.restriction
+            n_0 = n(res[0])
+            n_1 = n(res[1])
+            h_0 = 2 * cr(res[0])
+            h_1 = 2 * cr(res[1])
 
             all_mobile_species = [spe for spe in self.species if spe.mobile]
             if len(all_mobile_species) > 1:
                 raise NotImplementedError("Multiple mobile species not implemented")
             H = all_mobile_species[0]
 
-            v_b = H.subdomain_to_test_function[subdomain_1](b_res)
-            v_t = H.subdomain_to_test_function[subdomain_2](t_res)
+            v_b = H.subdomain_to_test_function[subdomain_0](res[0])
+            v_t = H.subdomain_to_test_function[subdomain_1](res[1])
 
-            u_b = H.subdomain_to_solution[subdomain_1](b_res)
-            u_t = H.subdomain_to_solution[subdomain_2](t_res)
+            u_b = H.subdomain_to_solution[subdomain_0](res[0])
+            u_t = H.subdomain_to_solution[subdomain_1](res[1])
 
-            K_b = subdomain_1.material.get_solubility_coefficient(
-                self.mesh.mesh, self.temperature_fenics(b_res), H
+            K_b = subdomain_0.material.get_solubility_coefficient(
+                self.mesh.mesh, self.temperature_fenics(res[0]), H
             )
-            K_t = subdomain_2.material.get_solubility_coefficient(
-                self.mesh.mesh, self.temperature_fenics(t_res), H
+            K_t = subdomain_1.material.get_solubility_coefficient(
+                self.mesh.mesh, self.temperature_fenics(res[1]), H
             )
 
-            F_0 = -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface(
+            F_0 = -0.5 * mixed_term((u_b + u_t), v_b, n_0) * dInterface(
                 interface.id
-            ) - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_b) * dInterface(
+            ) - 0.5 * mixed_term(v_b, (u_b / K_b - u_t / K_t), n_0) * dInterface(
                 interface.id
             )
 
-            F_1 = +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface(
+            F_1 = +0.5 * mixed_term((u_b + u_t), v_t, n_0) * dInterface(
                 interface.id
-            ) - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_b) * dInterface(
+            ) - 0.5 * mixed_term(v_t, (u_b / K_b - u_t / K_t), n_0) * dInterface(
                 interface.id
             )
             F_0 += (
                 2
                 * gamma
-                / (h_b + h_t)
+                / (h_0 + h_1)
                 * (u_b / K_b - u_t / K_t)
                 * v_b
                 * dInterface(interface.id)
@@ -1096,14 +1096,14 @@ def mixed_term(u, v, n):
             F_1 += (
                 -2
                 * gamma
-                / (h_b + h_t)
+                / (h_0 + h_1)
                 * (u_b / K_b - u_t / K_t)
                 * v_t
                 * dInterface(interface.id)
             )
 
-            subdomain_1.F += F_0
-            subdomain_2.F += F_1
+            subdomain_0.F += F_0
+            subdomain_1.F += F_1
 
         J = []
         forms = []

From eab5fa536f54a822e871c7ed6a50c12007ca51ec Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:44:41 -0400
Subject: [PATCH 087/134] type hinting for as_fenics_constant

---
 festim/helpers.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/festim/helpers.py b/festim/helpers.py
index 9b17fb92b..f0e741ffc 100644
--- a/festim/helpers.py
+++ b/festim/helpers.py
@@ -1,21 +1,24 @@
 from dolfinx import fem
+import dolfinx
 
 
-def as_fenics_constant(value, mesh):
+def as_fenics_constant(
+    value: float | int | fem.Constant, mesh: dolfinx.mesh.Mesh
+) -> fem.Constant:
     """Converts a value to a dolfinx.Constant
 
     Args:
-        value (float, int or dolfinx.Constant): the value to convert
-        mesh (dolfinx.mesh.Mesh): the mesh of the domiain
+        value: the value to convert
+        mesh: the mesh of the domiain
 
     Returns:
-        dolfinx.Constant: the converted value
+        The converted value
 
     Raises:
         TypeError: if the value is not a float, an int or a dolfinx.Constant
     """
     if isinstance(value, (float, int)):
-        return fem.Constant(mesh, float(value))
+        return fem.Constant(mesh, dolfinx.default_scalar_type(value))
     elif isinstance(value, fem.Constant):
         return value
     else:

From d46fc1a60a6fc831bc8b2b661022289f1534fc60 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:46:53 -0400
Subject: [PATCH 088/134] removed unused var

---
 festim/hydrogen_transport_problem.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index f63c32e58..206de4290 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -1052,7 +1052,6 @@ def mixed_term(u, v, n):
             subdomain_0, subdomain_1 = interface.subdomains
             res = interface.restriction
             n_0 = n(res[0])
-            n_1 = n(res[1])
             h_0 = 2 * cr(res[0])
             h_1 = 2 * cr(res[1])
 

From 8ed278ceb7a3d05b496ece593b21610ca7766ff3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 12:47:43 -0400
Subject: [PATCH 089/134] Rewrite docs

---
 festim/hydrogen_transport_problem.py | 100 ++++++++++++++++-----------
 1 file changed, 60 insertions(+), 40 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index fe34557d2..4f2ce2a10 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -7,6 +7,8 @@
 import tqdm.autonotebook
 import festim as F
 from scifem import NewtonSolver
+from typing import Callable
+import numpy.typing as npt
 
 
 class HydrogenTransportProblem(F.ProblemBase):
@@ -14,12 +16,11 @@ class HydrogenTransportProblem(F.ProblemBase):
     Hydrogen Transport Problem.
 
     Args:
-        mesh (festim.Mesh): the mesh of the model
-        subdomains (list of festim.Subdomain): the subdomains of the model
-        species (list of festim.Species): the species of the model
-        reactions (list of festim.Reaction): the reactions of the model
-        temperature (float, int, fem.Constant, fem.Function or callable): the
-            temperature of the model (K)
+        mesh: The mesh
+        subdomains: List containing the subdomains
+        species: List containing the species
+        reactions: List containing the reactions
+        temperature: The temperature of the model (K)
         sources (list of festim.Source): the hydrogen sources of the model
         initial_conditions (list of festim.InitialCondition): the initial conditions
             of the model
@@ -30,12 +31,11 @@ class HydrogenTransportProblem(F.ProblemBase):
         traps (list of F.Trap): the traps of the model
 
     Attributes:
-        mesh (festim.Mesh): the mesh of the model
-        subdomains (list of festim.Subdomain): the subdomains of the model
-        species (list of festim.Species): the species of the model
-        reactions (list of festim.Reaction): the reactions of the model
-        temperature (float, int, fem.Constant, fem.Function or callable): the
-            temperature of the model (K)
+        mesh : The mesh of the model
+        subdomains: The subdomains of the model
+        species: The species of the model
+        reactions: the reactions of the model
+        temperature: The temperature of the model (K)
         sources (list of festim.Source): the hydrogen sources of the model
         initial_conditions (list of festim.InitialCondition): the initial conditions
             of the model
@@ -71,35 +71,55 @@ class HydrogenTransportProblem(F.ProblemBase):
             of the model
 
 
-    Usage:
-        >>> import festim as F
-        >>> my_model = F.HydrogenTransportProblem()
-        >>> my_model.mesh = F.Mesh(...)
-        >>> my_model.subdomains = [F.Subdomain(...)]
-        >>> my_model.species = [F.Species(name="H"), F.Species(name="Trap")]
-        >>> my_model.temperature = 500
-        >>> my_model.sources = [F.ParticleSource(...)]
-        >>> my_model.boundary_conditions = [F.BoundaryCondition(...)]
-        >>> my_model.initialise()
+    Examples:
+        Can be used as either
+
+        .. highlight:: python
+        .. code-block:: python
+
+            import festim as F
+            my_model = F.HydrogenTransportProblem()
+            my_model.mesh = F.Mesh(...)
+            my_model.subdomains = [F.Subdomain(...)]
+            my_model.species = [F.Species(name="H"), F.Species(name="Trap")]
+            my_model.temperature = 500
+            my_model.sources = [F.ParticleSource(...)]
+            my_model.boundary_conditions = [F.BoundaryCondition(...)]
+            my_model.initialise()
 
         or
 
-        >>> my_model = F.HydrogenTransportProblem(
-        ...     mesh=F.Mesh(...),
-        ...     subdomains=[F.Subdomain(...)],
-        ...     species=[F.Species(name="H"), F.Species(name="Trap")],
-        ... )
-        >>> my_model.initialise()
+        .. highlight:: python
+        .. code-block:: python
+
+            my_model = F.HydrogenTransportProblem(
+                mesh=F.Mesh(...),
+                subdomains=[F.Subdomain(...)],
+                species=[F.Species(name="H"), F.Species(name="Trap")],
+            )
+            my_model.initialise()
 
     """
 
     def __init__(
         self,
-        mesh=None,
-        subdomains=None,
-        species=None,
-        reactions=None,
-        temperature=None,
+        mesh: F.Mesh | None = None,
+        subdomains: list[F.VolumeSubdomain | F.SurfaceSubdomain] | None = None,
+        species: list[F.Species] | None = None,
+        reactions: list[F.Reaction] | None = None,
+        temperature: float
+        | int
+        | fem.Constant
+        | fem.Function
+        | Callable[
+            [npt.NDArray[dolfinx.default_scalar_type]],
+            npt.NDArray[dolfinx.default_scalar_type],
+        ]
+        | Callable[
+            [npt.NDArray[dolfinx.default_scalar_type], fem.Constant],
+            npt.NDArray[dolfinx.default_scalar_type],
+        ]
+        | None = None,
         sources=None,
         initial_conditions=None,
         boundary_conditions=None,
@@ -137,7 +157,7 @@ def temperature(self, value):
             self._temperature = value
         else:
             raise TypeError(
-                f"Value must be a float, int, fem.Constant, fem.Function, or callable"
+                "Value must be a float, int, fem.Constant, fem.Function, or callable"
             )
 
     @property
@@ -153,7 +173,7 @@ def temperature_fenics(self, value):
             value,
             (fem.Constant, fem.Function),
         ):
-            raise TypeError(f"Value must be a fem.Constant or fem.Function")
+            raise TypeError("Value must be a fem.Constant or fem.Function")
         self._temperature_fenics = value
 
     @property
@@ -198,7 +218,7 @@ def facet_meshtags(self, value):
         elif isinstance(value, dolfinx.mesh.MeshTags):
             self._facet_meshtags = value
         else:
-            raise TypeError(f"value must be of type dolfinx.mesh.MeshTags")
+            raise TypeError("value must be of type dolfinx.mesh.MeshTags")
 
     @property
     def volume_meshtags(self):
@@ -211,7 +231,7 @@ def volume_meshtags(self, value):
         elif isinstance(value, dolfinx.mesh.MeshTags):
             self._volume_meshtags = value
         else:
-            raise TypeError(f"value must be of type dolfinx.mesh.MeshTags")
+            raise TypeError("value must be of type dolfinx.mesh.MeshTags")
 
     def initialise(self):
         self.create_species_from_traps()
@@ -945,9 +965,9 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_collapsed_function_space[subdomain] = V.sub(
                 i
             ).collapse()
-            species.subdomain_to_post_processing_solution[subdomain].name = (
-                f"{species.name}_{subdomain.id}"
-            )
+            species.subdomain_to_post_processing_solution[
+                subdomain
+            ].name = f"{species.name}_{subdomain.id}"
 
     def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         """

From 8b2db98d78075a8965f0401ec398cff3a959afb6 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:48:24 -0400
Subject: [PATCH 090/134] format

---
 festim/hydrogen_transport_problem.py | 33 ++++++++++------------------
 1 file changed, 11 insertions(+), 22 deletions(-)

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index fe34557d2..a13d9cc86 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -276,8 +276,7 @@ def define_temperature(self):
                     )
                 # only t is an argument
                 self.temperature_fenics = F.as_fenics_constant(
-                    mesh=self.mesh.mesh, value=self.temperature(
-                        t=float(self.t))
+                    mesh=self.mesh.mesh, value=self.temperature(t=float(self.t))
                 )
             else:
                 x = ufl.SpatialCoordinate(self.mesh.mesh)
@@ -291,8 +290,7 @@ def define_temperature(self):
                 function_space_temperature = fem.functionspace(
                     self.mesh.mesh, element_temperature
                 )
-                self.temperature_fenics = fem.Function(
-                    function_space_temperature)
+                self.temperature_fenics = fem.Function(function_space_temperature)
                 kwargs = {}
                 if "t" in arguments:
                     kwargs["t"] = self.t
@@ -320,8 +318,7 @@ def initialise_exports(self):
                             field, self.species
                         )
             elif isinstance(export.field, str):
-                export.field = F.find_species_from_name(
-                    export.field, self.species)
+                export.field = F.find_species_from_name(export.field, self.species)
 
             if isinstance(export, (F.VTXExport, F.XDMFExport)):
                 export.define_writer(MPI.COMM_WORLD)
@@ -363,8 +360,7 @@ def define_D_global(self, species):
         D_0 = fem.Function(self.V_DG_0)
         E_D = fem.Function(self.V_DG_0)
         for vol in self.volume_subdomains:
-            cell_indices = vol.locate_subdomain_entities(
-                self.mesh.mesh, self.mesh.vdim)
+            cell_indices = vol.locate_subdomain_entities(self.mesh.mesh, self.mesh.vdim)
 
             # replace values of D_0 and E_D by values from the material
             D_0.x.array[cell_indices] = vol.material.get_D_0(species=species)
@@ -374,11 +370,9 @@ def define_D_global(self, species):
         D = fem.Function(self.V_DG_1)
 
         expr = D_0 * ufl.exp(
-            -E_D / F.as_fenics_constant(F.k_B,
-                                        self.mesh.mesh) / self.temperature_fenics
+            -E_D / F.as_fenics_constant(F.k_B, self.mesh.mesh) / self.temperature_fenics
         )
-        D_expr = fem.Expression(
-            expr, self.V_DG_1.element.interpolation_points())
+        D_expr = fem.Expression(expr, self.V_DG_1.element.interpolation_points())
         D.interpolate(D_expr)
         return D, D_expr
 
@@ -571,8 +565,7 @@ def create_initial_conditions(self):
 
             # assign to previous solution of species
             if not self.multispecies:
-                condition.species.prev_solution.interpolate(
-                    condition.expr_fenics)
+                condition.species.prev_solution.interpolate(condition.expr_fenics)
             else:
                 idx = self.species.index(condition.species)
                 self.u_n.sub(idx).interpolate(condition.expr_fenics)
@@ -598,8 +591,7 @@ def create_formulation(self):
                     )
 
                 if self.settings.transient:
-                    self.formulation += ((u - u_n) / self.dt) * \
-                        v * self.dx(vol.id)
+                    self.formulation += ((u - u_n) / self.dt) * v * self.dx(vol.id)
 
         for reaction in self.reactions:
             for reactant in reaction.reactant:
@@ -1043,8 +1035,7 @@ def create_formulation(self):
             for interface in self.interfaces
         ]
         [interface.pad_parent_maps() for interface in self.interfaces]
-        dInterface = ufl.Measure(
-            "dS", domain=mesh, subdomain_data=integral_data)
+        dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)
 
         def mixed_term(u, v, n):
             return ufl.dot(ufl.grad(u), n) * v
@@ -1067,8 +1058,7 @@ def mixed_term(u, v, n):
 
             all_mobile_species = [spe for spe in self.species if spe.mobile]
             if len(all_mobile_species) > 1:
-                raise NotImplementedError(
-                    "Multiple mobile species not implemented")
+                raise NotImplementedError("Multiple mobile species not implemented")
             H = all_mobile_species[0]
 
             v_b = H.subdomain_to_test_function[subdomain_1](b_res)
@@ -1127,8 +1117,7 @@ def mixed_term(u, v, n):
                     )
                 )
             J.append(jac)
-            forms.append(dolfinx.fem.form(
-                subdomain1.F, entity_maps=entity_maps))
+            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
         self.forms = forms
         self.J = J

From 4d6954ee7cf1bcf35b9ca53406989aa89c144393 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 12:51:33 -0400
Subject: [PATCH 091/134] added API

---
 docs/source/api/index.rst | 13 +++++++++++++
 docs/source/index.rst     |  1 +
 2 files changed, 14 insertions(+)
 create mode 100644 docs/source/api/index.rst

diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst
new file mode 100644
index 000000000..23396ca21
--- /dev/null
+++ b/docs/source/api/index.rst
@@ -0,0 +1,13 @@
+FESTIM
+======
+
+.. automodule:: festim
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :inherited-members:
+    :exclude-members: __weakref__
+    :private-members:
+    :special-members:
+    :inherited-members:
+    :exclude-members: __weakref__
\ No newline at end of file
diff --git a/docs/source/index.rst b/docs/source/index.rst
index df0b5649e..7a103b06f 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -34,6 +34,7 @@ Contents
    userguide/index
    devguide/index
    publications
+   api/index
 
 
 Indices and tables

From ec2ceb4b268481618b7ee3975fbe7d2cb0019249 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 17:08:28 -0400
Subject: [PATCH 092/134] fixed docstrings

---
 festim/exports/volume_quantity.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/festim/exports/volume_quantity.py b/festim/exports/volume_quantity.py
index 6d0f1d0a9..60118e9b8 100644
--- a/festim/exports/volume_quantity.py
+++ b/festim/exports/volume_quantity.py
@@ -8,7 +8,7 @@ class VolumeQuantity:
 
     Args:
         field (festim.Species): species for which the volume quantity is computed
-        surface (festim.VolumeSubdomain): volume subdomain
+        volume (festim.VolumeSubdomain): volume subdomain
         filename (str, optional): name of the file to which the volume quantity is exported
 
     Attributes:

From 586d7c12e7cdea09466a1e08e4aede2481742a78 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 17:47:21 -0400
Subject: [PATCH 093/134] fixed dolfinx version in benchmark

---
 .github/workflows/benchmark.yml | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
index af332b990..9fb3b48c2 100644
--- a/.github/workflows/benchmark.yml
+++ b/.github/workflows/benchmark.yml
@@ -4,16 +4,16 @@ on: [pull_request, push]
 jobs:
   benchmark:
     runs-on: ubuntu-latest
-    container: dolfinx/dolfinx:v0.7.3
+    container: dolfinx/dolfinx:v0.9.0
 
     steps:
-    - name: Checkout code
-      uses: actions/checkout@v2
+      - name: Checkout code
+        uses: actions/checkout@v2
 
-    - name: Install local package and dependencies
-      run: |
-        pip install .[test]
+      - name: Install local package and dependencies
+        run: |
+          pip install .[test]
 
-    - name: Run benchmark
-      run: |
-        python3 test/benchmark.py
+      - name: Run benchmark
+        run: |
+          python3 test/benchmark.py

From 9fcd1e5c32f1740402d4f8ad75f9e64d5bfd2d1c Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 17:50:15 -0400
Subject: [PATCH 094/134] black

---
 festim/exports/xdmf.py               |  4 +++-
 festim/helpers_discontinuity.py      |  3 +--
 festim/hydrogen_transport_problem.py | 34 +++++++++++++++-------------
 festim/subdomain/volume_subdomain.py |  3 +--
 4 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/festim/exports/xdmf.py b/festim/exports/xdmf.py
index 593123c52..23e2ed1af 100644
--- a/festim/exports/xdmf.py
+++ b/festim/exports/xdmf.py
@@ -15,6 +15,7 @@ class XDMFExport:
         writer (dolfinx.io.XDMFFile): the XDMF writer
         field (festim.Species, list of festim.Species): the field index to export
     """
+
     _mesh_written: bool
 
     def __init__(self, filename: str, field) -> None:
@@ -74,7 +75,8 @@ def write(self, t: float):
         """
         if not self._mesh_written:
             self.writer.write_mesh(
-                self.field[0].post_processing_solution.function_space.mesh)
+                self.field[0].post_processing_solution.function_space.mesh
+            )
             self._mesh_written = True
 
         for field in self.field:
diff --git a/festim/helpers_discontinuity.py b/festim/helpers_discontinuity.py
index 0bc9ea60c..2c225a522 100644
--- a/festim/helpers_discontinuity.py
+++ b/festim/helpers_discontinuity.py
@@ -26,8 +26,7 @@ def transfer_meshtags_to_submesh(
     )
     submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)
 
-    c_c_to_e = submesh.topology.connectivity(
-        submesh.topology.dim, entity_tag.dim)
+    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
     c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)
 
     child_markers = np.full(num_child_entities, 0, dtype=np.int32)
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 882596616..79d8df573 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -107,19 +107,21 @@ def __init__(
         subdomains: list[F.VolumeSubdomain | F.SurfaceSubdomain] | None = None,
         species: list[F.Species] | None = None,
         reactions: list[F.Reaction] | None = None,
-        temperature: float
-        | int
-        | fem.Constant
-        | fem.Function
-        | Callable[
-            [npt.NDArray[dolfinx.default_scalar_type]],
-            npt.NDArray[dolfinx.default_scalar_type],
-        ]
-        | Callable[
-            [npt.NDArray[dolfinx.default_scalar_type], fem.Constant],
-            npt.NDArray[dolfinx.default_scalar_type],
-        ]
-        | None = None,
+        temperature: (
+            float
+            | int
+            | fem.Constant
+            | fem.Function
+            | Callable[
+                [npt.NDArray[dolfinx.default_scalar_type]],
+                npt.NDArray[dolfinx.default_scalar_type],
+            ]
+            | Callable[
+                [npt.NDArray[dolfinx.default_scalar_type], fem.Constant],
+                npt.NDArray[dolfinx.default_scalar_type],
+            ]
+            | None
+        ) = None,
         sources=None,
         initial_conditions=None,
         boundary_conditions=None,
@@ -957,9 +959,9 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_collapsed_function_space[subdomain] = V.sub(
                 i
             ).collapse()
-            species.subdomain_to_post_processing_solution[
-                subdomain
-            ].name = f"{species.name}_{subdomain.id}"
+            species.subdomain_to_post_processing_solution[subdomain].name = (
+                f"{species.name}_{subdomain.id}"
+            )
 
     def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         """
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index c727fda0c..4d8b055c9 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -42,8 +42,7 @@ def locate_subdomain_entities(self, mesh, vdim):
         """
         # By default, all entities are included
         # return array like x full of True
-        entities = locate_entities(
-            mesh, vdim, lambda x: np.full(x.shape[1], True))
+        entities = locate_entities(mesh, vdim, lambda x: np.full(x.shape[1], True))
         return entities
 
     def create_subdomain(self, mesh: dolfinx.mesh.Mesh, marker: dolfinx.mesh.MeshTags):

From 9ece9f41b548748d5b779f4f530646c2e2d2f86b Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 17:51:42 -0400
Subject: [PATCH 095/134] black 2

---
 test/test_average_volume.py | 3 +--
 test/test_maximum_volume.py | 3 +--
 test/test_minimum_volume.py | 3 +--
 3 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/test/test_average_volume.py b/test/test_average_volume.py
index ca95dc17b..aac2832b3 100644
--- a/test/test_average_volume.py
+++ b/test/test_average_volume.py
@@ -14,8 +14,7 @@ def test_average_volume_compute_1D():
     # BUILD
     L = 6.0
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
-    dummy_volume = F.VolumeSubdomain1D(
-        id=1, borders=[0, L], material=dummy_mat)
+    dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_mat)
     temp, cell_meshtags = my_mesh.define_meshtags(
         surface_subdomains=[], volume_subdomains=[dummy_volume]
     )
diff --git a/test/test_maximum_volume.py b/test/test_maximum_volume.py
index b6f459082..3a275a358 100644
--- a/test/test_maximum_volume.py
+++ b/test/test_maximum_volume.py
@@ -12,8 +12,7 @@ def test_maximum_volume_compute_1D():
     L = 6
     dummy_material = F.Material(D_0=1.5, E_D=1, name="dummy")
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
-    dummy_volume = F.VolumeSubdomain1D(
-        id=1, borders=[0, L], material=dummy_material)
+    dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_material)
     dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh, vdim=0)
 
     # give function to species
diff --git a/test/test_minimum_volume.py b/test/test_minimum_volume.py
index 6cd5445c0..71f11a387 100644
--- a/test/test_minimum_volume.py
+++ b/test/test_minimum_volume.py
@@ -12,8 +12,7 @@ def test_minimum_volume_compute_1D():
     L = 6
     dummy_material = F.Material(D_0=1.5, E_D=1, name="dummy")
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
-    dummy_volume = F.VolumeSubdomain1D(
-        id=1, borders=[0, L], material=dummy_material)
+    dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_material)
     dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh, vdim=0)
 
     # give function to species

From fd34d03166bf03714904b373070f910286ea965c Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 23 Oct 2024 17:58:46 -0400
Subject: [PATCH 096/134] fixed benchmark for dolfinx 0.9

---
 test/benchmark.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/test/benchmark.py b/test/benchmark.py
index 18b5b0c37..5bbe89457 100644
--- a/test/benchmark.py
+++ b/test/benchmark.py
@@ -23,7 +23,6 @@
     FacetNormal,
     Cell,
     Mesh,
-    VectorElement,
     Measure,
 )
 import basix
@@ -38,8 +37,7 @@ def fenics_test_permeation_problem(mesh_size=1001):
     L = 3e-04
     indices = np.linspace(0, L, num=mesh_size)
     gdim, shape, degree = 1, "interval", 1
-    cell = Cell(shape, geometric_dimension=gdim)
-    domain = Mesh(VectorElement("Lagrange", cell, degree))
+    domain = Mesh(basix.ufl.element("Lagrange", shape, degree, shape=(gdim,)))
     mesh_points = np.reshape(indices, (len(indices), 1))
     indexes = np.arange(mesh_points.shape[0])
     cells = np.stack((indexes[:-1], indexes[1:]), axis=-1)

From 5dc7432dd0c4c79350fba1ce3e326d4464432a52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 18:17:22 -0400
Subject: [PATCH 097/134] Refactor writers

---
 convergence_rates.py                         |  20 +--
 examples/multi_material_1d.py                |   6 +-
 examples/multi_material_2d.py                |  20 ++-
 examples/multi_material_2d_bis.py            |  23 ++-
 examples/multi_material_transient.py         |  12 +-
 examples/multi_material_with_one_material.py |   6 +-
 examples/tds_example.py                      |  10 +-
 examples/trapping_example.py                 |   8 +-
 festim/__init__.py                           |  15 +-
 festim/exports/vtx.py                        | 154 ++++++++++---------
 festim/exports/xdmf.py                       |  75 +++++----
 festim/heat_transfer_problem.py              |  30 ++--
 festim/hydrogen_transport_problem.py         |  74 +++++----
 festim/problem.py                            |   8 +
 festim/reaction.py                           |  28 ++--
 festim/species.py                            | 126 ++-------------
 festim/trap.py                               | 107 +++++++++++++
 test/system_tests/test_multi_material.py     |  52 ++++---
 test/test_heat_transfer_problem.py           |  64 ++++----
 test/test_permeation_problem.py              |  14 +-
 test/test_vtx.py                             |  67 +++-----
 test/test_xdmf.py                            |  34 +---
 22 files changed, 501 insertions(+), 452 deletions(-)
 create mode 100644 festim/trap.py

diff --git a/convergence_rates.py b/convergence_rates.py
index 5281ef028..65d4c3c91 100644
--- a/convergence_rates.py
+++ b/convergence_rates.py
@@ -58,19 +58,20 @@ def error_L2(u_computed, u_exact, degree_raise=3):
 
 
 def run(N):
-    exact_solution = lambda x, t: 2 * x[0] ** 2 + 20 * t
+    def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
 
-    density = lambda T: 0.2 * T + 2
-    heat_capacity = lambda T: 0.2 * T + 3
-    thermal_conductivity = lambda T: 0.1 * T + 4
+    def density(T): return 0.2 * T + 2
+    def heat_capacity(T): return 0.2 * T + 3
+    def thermal_conductivity(T): return 0.1 * T + 4
 
     mms_source_from_sp = source_from_exact_solution(
         exact_solution,
         density=lambda x, t: density(exact_solution(x, t)),
         heat_capacity=lambda x, t: heat_capacity(exact_solution(x, t)),
-        thermal_conductivity=lambda x, t: thermal_conductivity(exact_solution(x, t)),
+        thermal_conductivity=lambda x, t: thermal_conductivity(
+            exact_solution(x, t)),
     )
-    mms_source = lambda x, t: mms_source_from_sp((x[0], None, None), t)
+    def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
 
     my_problem = F.HeatTransferProblem()
 
@@ -88,7 +89,8 @@ def run(N):
     ]
 
     # NOTE: it's good to check that without the IC the solution is not the exact one
-    my_problem.initial_condition = F.InitialTemperature(lambda x: exact_solution(x, 0))
+    my_problem.initial_condition = F.InitialTemperature(
+        lambda x: exact_solution(x, 0))
 
     my_problem.boundary_conditions = [
         F.FixedTemperatureBC(subdomain=left, value=exact_solution),
@@ -110,7 +112,7 @@ def run(N):
     my_problem.settings.stepsize = F.Stepsize(0.05)
 
     my_problem.exports = [
-        F.VTXExportForTemperature(filename="test_transient_heat_transfer.bp")
+        F.VTXSpeciesExport(filename="test_transient_heat_transfer.bp")
     ]
 
     my_problem.initialise()
@@ -121,7 +123,7 @@ def run(N):
         my_problem.t.value
     )  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
-    exact_solution_end = lambda x: exact_solution(x, final_time_sim)
+    def exact_solution_end(x): return exact_solution(x, final_time_sim)
     L2_error = error_L2(computed_solution, exact_solution_end)
     return L2_error
 
diff --git a/examples/multi_material_1d.py b/examples/multi_material_1d.py
index bd5e7f80b..2041f0b20 100644
--- a/examples/multi_material_1d.py
+++ b/examples/multi_material_1d.py
@@ -54,7 +54,8 @@
     left_surface,
     right_surface,
 ]
-my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
+my_model.surface_to_volume = {
+    right_surface: right_domain, left_surface: left_domain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -90,7 +91,8 @@
 my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(
+        filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
diff --git a/examples/multi_material_2d.py b/examples/multi_material_2d.py
index e26c8e89a..fc59277ba 100644
--- a/examples/multi_material_2d.py
+++ b/examples/multi_material_2d.py
@@ -27,8 +27,10 @@ def half(x):
     gdim = mesh.geometry.dim
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
-    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    top_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, bottom_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
@@ -85,7 +87,8 @@ def half(x):
 
 # we should be able to automate this
 my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
-my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+my_model.surface_to_volume = {
+    top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
 
@@ -107,7 +110,7 @@ def half(x):
 my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
@@ -117,10 +120,13 @@ def half(x):
 # -------------------- post processing --------------------
 
 # derived quantities
-entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
+entity_maps = {
+    sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
 
-ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
-ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh,
+                   subdomain_data=bottom_domain.ft)
+ds_t = ufl.Measure("ds", domain=top_domain.submesh,
+                   subdomain_data=top_domain.ft)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
 ds = ufl.Measure("ds", domain=mesh, subdomain_data=my_model.facet_meshtags)
diff --git a/examples/multi_material_2d_bis.py b/examples/multi_material_2d_bis.py
index 27f80eb53..cdd0d0c23 100644
--- a/examples/multi_material_2d_bis.py
+++ b/examples/multi_material_2d_bis.py
@@ -27,8 +27,10 @@ def half(x):
     gdim = mesh.geometry.dim
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
-    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    top_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, bottom_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
@@ -85,7 +87,8 @@ def half(x):
 
 # we should be able to automate this
 my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
-my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+my_model.surface_to_volume = {
+    top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -129,10 +132,11 @@ def half(x):
 my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXExport(f"c_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(f"c_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ] + [
-    F.VTXExport(f"c_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain)
+    F.VTXSpeciesExport(f"c_t_{subdomain.id}.bp",
+                       field=trapped_H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
@@ -143,11 +147,14 @@ def half(x):
 
 
 # derived quantities
-entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
+entity_maps = {
+    sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
 entity_maps[mesh] = bottom_domain.submesh_to_mesh
 
-ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
-ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh,
+                   subdomain_data=bottom_domain.ft)
+ds_t = ufl.Measure("ds", domain=top_domain.submesh,
+                   subdomain_data=top_domain.ft)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
 
diff --git a/examples/multi_material_transient.py b/examples/multi_material_transient.py
index 342ba6b50..d00b1a2b9 100644
--- a/examples/multi_material_transient.py
+++ b/examples/multi_material_transient.py
@@ -49,7 +49,8 @@
     left_surface,
     right_surface,
 ]
-my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
+my_model.surface_to_volume = {
+    right_surface: right_domain, left_surface: left_domain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -82,14 +83,17 @@
 
 my_model.temperature = lambda x: 300 + 100 * x[0]
 
-my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=True, final_time=100)
+my_model.settings = F.Settings(
+    atol=None, rtol=1e-5, transient=True, final_time=100)
 my_model.settings.stepsize = 1
 
 my_model.exports = [
-    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(
+        filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ] + [
-    F.VTXExport(filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain)
+    F.VTXSpeciesExport(
+        filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
diff --git a/examples/multi_material_with_one_material.py b/examples/multi_material_with_one_material.py
index 5660d235c..281791ca1 100644
--- a/examples/multi_material_with_one_material.py
+++ b/examples/multi_material_with_one_material.py
@@ -21,7 +21,8 @@
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
 my_model.subdomains = [subdomain, left_surface, right_surface]
-my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
+my_model.surface_to_volume = {
+    right_surface: subdomain, left_surface: subdomain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -57,7 +58,8 @@
 my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(
+        filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
diff --git a/examples/tds_example.py b/examples/tds_example.py
index d0cf87013..6f1f8fda8 100644
--- a/examples/tds_example.py
+++ b/examples/tds_example.py
@@ -97,7 +97,8 @@ def left_conc_value(t):
 
 
 my_model.boundary_conditions = [
-    F.DirichletBC(subdomain=left_surface, value=left_conc_value, species=mobile_H),
+    F.DirichletBC(subdomain=left_surface,
+                  value=left_conc_value, species=mobile_H),
     F.DirichletBC(subdomain=right_surface, value=0, species=mobile_H),
 ]
 
@@ -107,8 +108,8 @@ def left_conc_value(t):
 right_flux = F.SurfaceFlux(field=mobile_H, surface=right_surface)
 
 my_model.exports = [
-    # F.VTXExport("mobile_concentration_h.bp", field=mobile_H),
-    # F.VTXExport("trapped_concentration_h.bp", field=trapped_H1),
+    # F.VTXSpeciesExport("mobile_concentration_h.bp", field=mobile_H),
+    # F.VTXSpeciesExport("trapped_concentration_h.bp", field=trapped_H1),
     F.XDMFExport("mobile_concentration_h.xdmf", field=mobile_H),
     F.XDMFExport("trapped_concentration_h1.xdmf", field=trapped_H1),
     F.XDMFExport("trapped_concentration_h2.xdmf", field=trapped_H2),
@@ -118,7 +119,8 @@ def left_conc_value(t):
 
 # -------- Settings --------- #
 
-my_model.settings = F.Settings(atol=1e10, rtol=1e-10, max_iterations=30, final_time=500)
+my_model.settings = F.Settings(
+    atol=1e10, rtol=1e-10, max_iterations=30, final_time=500)
 
 my_model.settings.stepsize = F.Stepsize(initial_value=0.5)
 
diff --git a/examples/trapping_example.py b/examples/trapping_example.py
index 2ed82a7d7..4d980a24b 100644
--- a/examples/trapping_example.py
+++ b/examples/trapping_example.py
@@ -1,3 +1,4 @@
+from petsc4py import PETSc
 import festim as F
 import numpy as np
 
@@ -43,8 +44,10 @@
     F.DirichletBC(subdomain=left_surface, value=1e12, species=mobile_H),
 ]
 my_model.exports = [
-    F.VTXExport("mobile_concentration_h.bp", field=mobile_H),  # produces 0 in file
-    F.VTXExport("trapped_concentration_h.bp", field=trapped_H),  # produces 0 in file
+    F.VTXSpeciesExport("mobile_concentration_h.bp",
+                       field=mobile_H),  # produces 0 in file
+    F.VTXSpeciesExport("trapped_concentration_h.bp",
+                       field=trapped_H),  # produces 0 in file
     F.XDMFExport("mobile_concentration_h.xdmf", field=mobile_H),
     F.XDMFExport("trapped_concentration_h.xdmf", field=trapped_H),
 ]
@@ -60,7 +63,6 @@
 
 my_model.initialise()
 
-from petsc4py import PETSc
 
 my_model.solver.convergence_criterion = "incremental"
 ksp = my_model.solver.krylov_solver
diff --git a/festim/__init__.py b/festim/__init__.py
index 492623f4d..f67fb779f 100644
--- a/festim/__init__.py
+++ b/festim/__init__.py
@@ -1,11 +1,4 @@
-try:
-    # Python 3.8+
-    from importlib import metadata
-except ImportError:
-    try:
-        import importlib_metadata as metadata
-    except ImportError:
-        __version__ = "unknown"
+from importlib import metadata
 
 try:
     __version__ = metadata.version("FESTIM")
@@ -48,8 +41,8 @@
 from .subdomain.volume_subdomain_1d import VolumeSubdomain1D
 from .subdomain.interface import Interface
 
-from .species import Species, Trap, ImplicitSpecies, find_species_from_name
-
+from .species import Species, ImplicitSpecies, find_species_from_name
+from .trap import Trap
 from .stepsize import Stepsize
 
 from .exports.surface_quantity import SurfaceQuantity
@@ -63,7 +56,7 @@
 from .exports.minimum_surface import MinimumSurface
 from .exports.average_surface import AverageSurface
 from .exports.surface_flux import SurfaceFlux
-from .exports.vtx import VTXExport, VTXExportForTemperature
+from .exports.vtx import VTXSpeciesExport, VTXTemperatureExport
 from .exports.xdmf import XDMFExport
 
 from .reaction import Reaction
diff --git a/festim/exports/vtx.py b/festim/exports/vtx.py
index c38c42d45..dafbbe2fb 100644
--- a/festim/exports/vtx.py
+++ b/festim/exports/vtx.py
@@ -1,104 +1,116 @@
-from dolfinx.io import VTXWriter
-import mpi4py
+from dolfinx.fem import Function as _Function
+from pathlib import Path
+from festim.species import Species as _Species
+from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
+import warnings
 
-import festim as F
 
+class ExportBaseClass():
+    _filename: Path | str
 
-class VTXExportBase:
-    def __init__(self, filename: str) -> None:
-        self.filename = filename
+    def __init__(self, filename: str | Path, ext: str) -> None:
+        name = Path(filename)
+        if name.suffix != ext:
+            warnings.warn(
+                f"Filename {filename} does not have {ext} extension, adding it.")
+            name = name.with_suffix(ext)
+
+        self._filename = Path(filename)
 
     @property
     def filename(self):
         return self._filename
 
-    @filename.setter
-    def filename(self, value):
-        if not isinstance(value, str):
-            raise TypeError("filename must be of type str")
-        if not value.endswith(".bp"):
-            raise ValueError("filename must end with .bp")
-        self._filename = value
-
-    def define_writer(self, comm: mpi4py.MPI.Intracomm) -> None:
-        """Define the writer
-
-        Args:
-            comm (mpi4py.MPI.Intracomm): the MPI communicator
-        """
-        self.writer = VTXWriter(
-            comm,
-            self.filename,
-            self.functions,
-            "BP4",
-        )
-
-    def write(self, t: float):
-        """Write functions to VTX file
 
-        Args:
-            t (float): the time of export
-        """
-        self.writer.write(t)
+class VTXTemperatureExport(ExportBaseClass):
+    def __init__(
+            self, filename: str | Path):
+        super().__init__(filename, ".bp")
 
 
-class VTXExport(VTXExportBase):
+class VTXSpeciesExport(ExportBaseClass):
     """Export functions to VTX file
 
     Args:
-        filename (str): the name of the output file
-        field (int): the field index to export
-
-    Attributes:
-        filename (str): the name of the output file
-        writer (dolfinx.io.VTXWriter): the VTX writer
-        field (festim.Species, list of festim.Species): the field index to export
-
-    Usage:
-        >>> u = dolfinx.fem.Function(V)
-        >>> my_export = festim.VTXExport("my_export.bp")
-        >>> my_export.define_writer(mesh.comm, [u])
-        >>> for t in range(10):
-        ...    u.interpolate(lambda x: t * (x[0] ** 2 + x[1] ** 2 + x[2] ** 2))
-        ...    my_export.write(t)
+        filename: The name of the output file
+        field: Set of species to export
+        subdomain: A field can be defined on multiple domains.
+            This arguments specifies what subdomains we export on.
+            If `None` we export on all domains.
+
     """
 
+    field: list[_Species]
+    _subdomain: _VolumeSubdomain
+
     def __init__(
-        self, filename: str, field, subdomain: F.VolumeSubdomain = None
+        self, filename: str | Path,
+        field: _Species | list[_Species],
+        subdomain: _VolumeSubdomain = None
     ) -> None:
+        super().__init__(filename, ".bp")
         self.field = field
-        self.subdomain = subdomain
-        super().__init__(filename)
+        self._subdomain = subdomain
 
     @property
-    def field(self):
+    def field(self) -> list[_Species]:
         return self._field
 
     @field.setter
-    def field(self, value):
-        # check that field is festim.Species or list of festim.Species
-        if not isinstance(value, (F.Species, str)) and not isinstance(value, list):
-            raise TypeError(
-                "field must be of type festim.Species or str or a list of festim.Species or str"
-            )
+    def field(self, value: _Species | list[_Species]):
+        """
+        Update the field to export.
+
+        Note:
+            This also creates a new writer with the updated field.
+
+        Args:
+            value: The species to export
+
+        Raises:
+            TypeError: If input field is not a Species or a list of Species
+        """
         # check that all elements of list are festim.Species
         if isinstance(value, list):
             for element in value:
-                if not isinstance(element, (F.Species, str)):
+                if not isinstance(element, (_Species, str)):
                     raise TypeError(
-                        "field must be of type festim.Species or str or a list of festim.Species or str"
+                        "field must be of type festim.Species or a list of festim.Species or str"
                     )
-        # if field is festim.Species, convert to list
-        if not isinstance(value, list):
-            value = [value]
-
-        self._field = value
+            val = value
+        elif isinstance(value, _Species):
+            val = [value]
+        else:
+            raise TypeError(
+                "field must be of type festim.Species or a list of festim.Species or str",
+                f"got {type(value)}."
+            )
+        self._field = val
 
     @property
-    def functions(self):
-        return [field.post_processing_solution for field in self.field]
+    def subdomain(self) -> _VolumeSubdomain:
+        return self._subdomain
 
+    def get_functions(self) -> list[_Function]:
+        """
+        Returns list of species for a given subdomain.
+        If using legacy mode, return the whole species.
+        """
 
-class VTXExportForTemperature(VTXExportBase):
-    def __init__(self, filename: str) -> None:
-        super().__init__(filename)
+        legacy_output: bool = False
+        for field in self._field:
+            if field.legacy:
+                legacy_output = True
+                break
+        if legacy_output:
+            return [field.post_processing_solution for field in self._field]
+        else:
+            if self._subdomain is None:
+                raise ValueError("Subdomain must be specified")
+            else:
+                outfiles = []
+                for field in self._field:
+                    if self._subdomain in field.subdomains:
+                        outfiles.append(
+                            field.subdomain_to_post_processing_solution[self._subdomain])
+                return outfiles
diff --git a/festim/exports/xdmf.py b/festim/exports/xdmf.py
index 23e2ed1af..a3fa25873 100644
--- a/festim/exports/xdmf.py
+++ b/festim/exports/xdmf.py
@@ -1,63 +1,54 @@
 import mpi4py
 from dolfinx.io import XDMFFile
-import festim as F
+from pathlib import Path
+from festim.species import Species as _Species
+from .vtx import ExportBaseClass
 
 
-class XDMFExport:
+class XDMFExport(ExportBaseClass):
     """Export functions to XDMFfile
 
     Args:
-        filename (str): the name of the output file
-        field (int or festim.Species): the field index to export
+        filename: The name of the output file
+        field: The field(s) to export
 
     Attributes:
-        filename (str): the name of the output file
-        writer (dolfinx.io.XDMFFile): the XDMF writer
-        field (festim.Species, list of festim.Species): the field index to export
+        _writer (dolfinx.io.XDMFFile): the XDMF writer
+        _field (festim.Species, list of festim.Species): the field index to export
     """
 
     _mesh_written: bool
+    _filename: Path
+    _writer: XDMFFile | None
 
-    def __init__(self, filename: str, field) -> None:
-        self.filename = filename
+    def __init__(self, filename: str | Path, field: list[_Species] | _Species) -> None:
+        # Initializes the writer
+        self._writer = None
+        super().__init__(filename, ".xdmf")
         self.field = field
         self._mesh_written = False
 
     @property
-    def filename(self):
-        return self._filename
-
-    @filename.setter
-    def filename(self, value):
-        if not isinstance(value, str):
-            raise TypeError("filename must be of type str")
-        if not value.endswith(".xdmf"):
-            raise ValueError("filename must end with .xdmf")
-        self._filename = value
-
-    @property
-    def field(self):
+    def field(self) -> list[_Species]:
         return self._field
 
     @field.setter
-    def field(self, value):
+    def field(self, value: _Species | list[_Species]):
         # check that field is festim.Species or list of festim.Species
-        if not isinstance(value, (F.Species, str)) and not isinstance(value, list):
-            raise TypeError(
-                "field must be of type festim.Species or str or a list of festim.Species or str"
-            )
-        # check that all elements of list are festim.Species
         if isinstance(value, list):
             for element in value:
-                if not isinstance(element, (F.Species, str)):
-                    raise TypeError(
-                        "field must be of type festim.Species or str or a list of festim.Species or str"
-                    )
-        # if field is festim.Species, convert to list
-        if not isinstance(value, list):
-            value = [value]
-
-        self._field = value
+                if not isinstance(element, _Species):
+                    raise TypeError(f"Each element in the list must be a species, got {type(element)}."
+                                    )
+            val = value
+        elif isinstance(value, _Species):
+            val = [value]
+        else:
+            raise TypeError(
+                f"field must be of type festim.Species or a list of festim.Species, got {
+                    type(value)}."
+            )
+        self._field = val
 
     def define_writer(self, comm: mpi4py.MPI.Intracomm) -> None:
         """Define the writer
@@ -65,7 +56,7 @@ def define_writer(self, comm: mpi4py.MPI.Intracomm) -> None:
         Args:
             comm (mpi4py.MPI.Intracomm): the MPI communicator
         """
-        self.writer = XDMFFile(comm, self.filename, "w")
+        self._writer = XDMFFile(comm, self.filename, "w")
 
     def write(self, t: float):
         """Write functions to VTX file
@@ -74,10 +65,14 @@ def write(self, t: float):
             t (float): the time of export
         """
         if not self._mesh_written:
-            self.writer.write_mesh(
+            self._writer.write_mesh(
                 self.field[0].post_processing_solution.function_space.mesh
             )
             self._mesh_written = True
 
         for field in self.field:
-            self.writer.write_function(field.post_processing_solution, t)
+            self._writer.write_function(field.post_processing_solution, t)
+
+    def __del__(self):
+        if self._writer is not None:
+            self._writer.close()
diff --git a/festim/heat_transfer_problem.py b/festim/heat_transfer_problem.py
index 6f3ee5c77..27c6b21df 100644
--- a/festim/heat_transfer_problem.py
+++ b/festim/heat_transfer_problem.py
@@ -1,8 +1,7 @@
 from dolfinx import fem
-from dolfinx.nls.petsc import NewtonSolver
+from dolfinx.io import VTXWriter
 import basix
 import ufl
-from mpi4py import MPI
 import festim as F
 
 
@@ -27,6 +26,7 @@ def __init__(
         )
 
         self.initial_condition = initial_condition
+        self._vtxfile: VTXWriter | None = None
 
     @property
     def sources(self):
@@ -35,7 +35,8 @@ def sources(self):
     @sources.setter
     def sources(self, value):
         if not all(isinstance(source, F.HeatSource) for source in value):
-            raise TypeError("sources must be a list of festim.HeatSource objects")
+            raise TypeError(
+                "sources must be a list of festim.HeatSource objects")
         self._sources = value
 
     @property
@@ -198,14 +199,16 @@ def create_formulation(self):
         # add sources
         for source in self.sources:
             self.formulation -= (
-                source.value_fenics * self.test_function * self.dx(source.volume.id)
+                source.value_fenics * self.test_function *
+                self.dx(source.volume.id)
             )
 
         # add fluxes
         for bc in self.boundary_conditions:
             if isinstance(bc, F.HeatFluxBC):
                 self.formulation -= (
-                    bc.value_fenics * self.test_function * self.ds(bc.subdomain.id)
+                    bc.value_fenics * self.test_function *
+                    self.ds(bc.subdomain.id)
                 )
 
     def initialise_exports(self):
@@ -217,9 +220,10 @@ def initialise_exports(self):
                 raise NotImplementedError(
                     "XDMF export is not implemented yet for heat transfer problems"
                 )
-            if isinstance(export, (F.VTXExportForTemperature, F.XDMFExport)):
-                export.functions = [self.u]
-                export.define_writer(MPI.COMM_WORLD)
+            if isinstance(export, F.VTXTemperatureExport):
+                self._vtxfile = VTXWriter(self.u.function_space.mesh.comm,
+                                          export.filename, [self.u],
+                                          engine="BP5")
 
     def post_processing(self):
         """Post processes the model"""
@@ -240,6 +244,12 @@ def post_processing(self):
                 # if filename given write export data to file
                 if export.filename is not None:
                     export.write(t=float(self.t))
-
-            if isinstance(export, (F.VTXExportForTemperature, F.XDMFExport)):
+            if isinstance(export, F.XDMFExport):
                 export.write(float(self.t))
+
+        if self._vtxfile is not None:
+            self._vtxfile.write(float(self.t))
+
+    def __del__(self):
+        if self._vtxfile is not None:
+            self._vtxfile.close()
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 79d8df573..f0ecaad48 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -20,29 +20,26 @@ class HydrogenTransportProblem(F.ProblemBase):
         subdomains: List containing the subdomains
         species: List containing the species
         reactions: List containing the reactions
-        temperature: The temperature of the model (K)
-        sources (list of festim.Source): the hydrogen sources of the model
-        initial_conditions (list of festim.InitialCondition): the initial conditions
-            of the model
-        boundary_conditions (list of festim.BoundaryCondition): the boundary
-            conditions of the model
+        temperature: The temperature or a function describing the temperature as
+            a model of either space or space and time. Unit (K)
+        sources: The hydrogen sources
+        initial_conditions: The initial conditions
+        boundary_conditions: The boundary conditions
         solver_parameters (dict): the solver parameters of the model
         exports (list of festim.Export): the exports of the model
         traps (list of F.Trap): the traps of the model
 
     Attributes:
-        mesh : The mesh of the model
-        subdomains: The subdomains of the model
-        species: The species of the model
-        reactions: the reactions of the model
-        temperature: The temperature of the model (K)
-        sources (list of festim.Source): the hydrogen sources of the model
-        initial_conditions (list of festim.InitialCondition): the initial conditions
-            of the model
-        boundary_conditions (list of festim.BoundaryCondition): the boundary
-            conditions of the model
-        solver_parameters (dict): the solver parameters of the model
-        exports (list of festim.Export): the exports of the model
+        mesh : The mesh
+        subdomains: The subdomains
+        species: The species
+        reactions: the reaction
+        temperature: The temperature in unit `K`
+        sources: The hydrogen sources
+        initial_conditions: The initial conditions
+        boundary_conditions: List of Dirichlet boundary conditions
+        solver_parameters (dict): the solver parameters
+        exports (list of festim.Export): the export
         traps (list of F.Trap): the traps of the model
         dx (dolfinx.fem.dx): the volume measure of the model
         ds (dolfinx.fem.ds): the surface measure of the model
@@ -125,6 +122,7 @@ def __init__(
         sources=None,
         initial_conditions=None,
         boundary_conditions=None,
+
         settings=None,
         exports=None,
         traps=None,
@@ -144,6 +142,7 @@ def __init__(
         self.initial_conditions = initial_conditions or []
         self.traps = traps or []
         self.temperature_fenics = None
+        self._vtxfiles: list[dolfinx.io.VTXWriter] = []
 
     @property
     def temperature(self):
@@ -342,9 +341,13 @@ def initialise_exports(self):
             elif isinstance(export.field, str):
                 export.field = F.find_species_from_name(export.field, self.species)
 
-            if isinstance(export, (F.VTXExport, F.XDMFExport)):
+            # Initialize XDMFFile for writer
+            if isinstance(export, F.XDMFExport):
                 export.define_writer(MPI.COMM_WORLD)
-
+            if isinstance(export, F.VTXSpeciesExport):
+                functions = export.get_functions()
+                self._vtxfiles.append(dolfinx.io.VTXWriter(functions[0].function_space.mesh.comm,
+                                                           export.filename, functions, engine="BP5"))
         # compute diffusivity function for surface fluxes
 
         spe_to_D_global = {}  # links species to global D function
@@ -733,7 +736,7 @@ def post_processing(self):
                 # if filename given write export data to file
                 if export.filename is not None:
                     export.write(t=float(self.t))
-            if isinstance(export, (F.VTXExport, F.XDMFExport)):
+            if isinstance(export, F.XDMFExport):
                 export.write(float(self.t))
 
 
@@ -800,6 +803,7 @@ def __init__(
             "pc_factor_mat_solver_type": "mumps",
         }
         self.petsc_options = petsc_options or default_petsc_options
+        self._vtxfiles: list[dolfinx.io.VTXWriter] = []
 
     def initialise(self):
         # check that all species have a list of F.VolumeSubdomain as this is
@@ -1167,15 +1171,13 @@ def create_flux_values_fenics(self):
 
     def initialise_exports(self):
         for export in self.exports:
-            if isinstance(export, F.VTXExport):
-                species = export.field[0]
-                # override post_processing_solution attribute of species
-                species.post_processing_solution = (
-                    species.subdomain_to_post_processing_solution[export.subdomain]
-                )
-                export.define_writer(MPI.COMM_WORLD)
+            if isinstance(export, F.VTXSpeciesExport):
+                functions = export.get_functions()
+                self._vtxfiles.append(dolfinx.io.VTXWriter(functions[0].function_space.mesh.comm,
+                                                           export.filename, functions, engine="BP5"))
             else:
-                raise NotImplementedError("Export type not implemented")
+                raise NotImplementedError(
+                    f"Export type {type(export)} not implemented")
 
     def post_processing(self):
         # update post-processing solutions (for each species in each subdomain)
@@ -1191,11 +1193,13 @@ def post_processing(self):
                 v0_to_V = species.subdomain_to_collapsed_function_space[subdomain][1]
                 collapsed_function.x.array[:] = u.x.array[v0_to_V]
 
+        for vtxfile in self._vtxfiles:
+            vtxfile.write(float(self.t))
+
         for export in self.exports:
-            if isinstance(export, F.VTXExport):
-                export.write(float(self.t))
-            else:
-                raise NotImplementedError("Export type not implemented")
+            if not isinstance(export, F.VTXSpeciesExport):
+                raise NotImplementedError(
+                    f"Export type {type(export)} not implemented")
 
     def iterate(self):
         """Iterates the model for a given time step"""
@@ -1238,3 +1242,7 @@ def run(self):
             # Solve steady-state
             self.solver.solve(self.settings.rtol)
             self.post_processing()
+
+    def __del__(self):
+        for vtxfile in self._vtxfiles:
+            vtxfile.close()
diff --git a/festim/problem.py b/festim/problem.py
index 6773f2673..fb14f43e2 100644
--- a/festim/problem.py
+++ b/festim/problem.py
@@ -5,6 +5,10 @@
 import numpy as np
 import tqdm.autonotebook
 import festim as F
+from festim.mesh.mesh import Mesh as _Mesh
+from festim.source import SourceBase as _SourceBase
+from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
+from typing import Any
 
 
 class ProblemBase:
@@ -15,6 +19,10 @@ class ProblemBase:
         simulation
     progress_bar (tqdm.autonotebook.tqdm) the progress bar
     """
+    mesh: _Mesh
+    sources: list[_SourceBase]
+    exports: list[Any]
+    subdomains: list[_VolumeSubdomain]
 
     def __init__(
         self,
diff --git a/festim/reaction.py b/festim/reaction.py
index 08b25ce8e..43ad7afab 100644
--- a/festim/reaction.py
+++ b/festim/reaction.py
@@ -1,7 +1,9 @@
-import festim as F
 from typing import Union, Optional, List
 
 from ufl import exp
+from festim.species import Species as _Species, ImplicitSpecies as _ImplicitSpecies
+from festim.subdomain.volume_subdomain_1d import VolumeSubdomain1D as VS1D
+from festim import k_B as _k_B
 
 
 class Reaction:
@@ -45,13 +47,12 @@ class Reaction:
 
     def __init__(
         self,
-        reactant: Union[
-            F.Species, F.ImplicitSpecies, List[Union[F.Species, F.ImplicitSpecies]]
-        ],
+        reactant:
+            _Species | _ImplicitSpecies | list[_Species | _ImplicitSpecies],
         k_0: float,
         E_k: float,
-        volume: F.VolumeSubdomain1D,
-        product: Optional[Union[F.Species, List[F.Species]]] = [],
+        volume: VS1D,
+        product: Optional[Union[_Species, List[_Species]]] = [],
         p_0: float = None,
         E_p: float = None,
     ) -> None:
@@ -76,9 +77,10 @@ def reactant(self, value):
                 f"reactant must be an entry of one or more species objects, not an empty list."
             )
         for i in value:
-            if not isinstance(i, (F.Species, F.ImplicitSpecies)):
+            if not isinstance(i, (_Species, _ImplicitSpecies)):
                 raise TypeError(
-                    f"reactant must be an F.Species or F.ImplicitSpecies, not {type(i)}"
+                    f"reactant must be an F.Species or F.ImplicitSpecies, not {
+                        type(i)}"
                 )
         self._reactant = value
 
@@ -109,11 +111,13 @@ def reaction_term(self, temperature):
         if self.product == []:
             if self.p_0 is not None:
                 raise ValueError(
-                    f"p_0 must be None, not {self.p_0} when no products are present."
+                    f"p_0 must be None, not {
+                        self.p_0} when no products are present."
                 )
             if self.E_p is not None:
                 raise ValueError(
-                    f"E_p must be None, not {self.E_p} when no products are present."
+                    f"E_p must be None, not {
+                        self.E_p} when no products are present."
                 )
         else:
             if self.p_0 == None:
@@ -125,10 +129,10 @@ def reaction_term(self, temperature):
                     f"E_p cannot be None when reaction products are present."
                 )
 
-        k = self.k_0 * exp(-self.E_k / (F.k_B * temperature))
+        k = self.k_0 * exp(-self.E_k / (_k_B * temperature))
 
         if self.p_0 and self.E_p:
-            p = self.p_0 * exp(-self.E_p / (F.k_B * temperature))
+            p = self.p_0 * exp(-self.E_p / (_k_B * temperature))
         elif self.p_0:
             p = self.p_0
         else:
diff --git a/festim/species.py b/festim/species.py
index b474d9aff..ce031b728 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -1,5 +1,4 @@
-from typing import List, Union
-import festim as F
+from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
 
 
 class Species:
@@ -52,7 +51,7 @@ class Species:
 
     """
 
-    subdomains: Union[List[F.VolumeSubdomain], F.VolumeSubdomain]
+    subdomains: list[_VolumeSubdomain] | _VolumeSubdomain
     subdomain_to_solution: dict
     subdomain_to_prev_solution: dict
     subdomain_to_test_function: dict
@@ -88,109 +87,15 @@ def __str__(self) -> str:
     def concentration(self):
         return self.solution
 
-
-class Trap(Species):
-    """Trap species class for H transport simulation.
-
-    This class only works for 1 mobile species and 1 trapping level and is
-    for convenience, for more details see notes.
-
-    Args:
-        name (str, optional): a name given to the trap. Defaults to None.
-        mobile_species (F.Species): the mobile species to be trapped
-        k_0 (float): the trapping rate constant pre-exponential factor (m3 s-1)
-        E_k (float): the trapping rate constant activation energy (eV)
-        p_0 (float): the detrapping rate constant pre-exponential factor (s-1)
-        E_p (float): the detrapping rate constant activation energy (eV)
-        volume (F.VolumeSubdomain1D): The volume subdomain where the trap is.
-
-
-    Attributes:
-        name (str, optional): a name given to the trap. Defaults to None.
-        mobile_species (F.Species): the mobile species to be trapped
-        k_0 (float): the trapping rate constant pre-exponential factor (m3 s-1)
-        E_k (float): the trapping rate constant activation energy (eV)
-        p_0 (float): the detrapping rate constant pre-exponential factor (s-1)
-        E_p (float): the detrapping rate constant activation energy (eV)
-        volume (F.VolumeSubdomain1D): The volume subdomain where the trap is.
-        trapped_concentration (F.Species): The immobile trapped concentration
-        trap_reaction (F.Reaction): The reaction for trapping the mobile conc.
-
-    Usage:
-        >>> import festim as F
-        >>> trap = F.Trap(name="Trap", species=H, k_0=1.0, E_k=0.2, p_0=0.1, E_p=0.3, volume=my_vol)
-        >>> trap.name
-        'Trap'
-        >>> my_model = F.HydorgenTransportProblem()
-        >>> my_model.traps = [trap]
-
-    Notes:
-        This convenience class replaces the need to specify an implicit species and
-        the associated reaction, thus:
-
-        cm = F.Species("mobile")
-        my_trap = F.Trap(
-            name="trapped",
-            mobile_species=cm,
-            k_0=1,
-            E_k=1,
-            p_0=1,
-            E_p=1,
-            n=1,
-            volume=my_vol,
-        )
-        my_model.species = [cm]
-        my_model.traps = [my_trap]
-
-        is equivalent to:
-
-        cm = F.Species("mobile")
-        ct = F.Species("trapped")
-        trap_sites = F.ImplicitSpecies(n=1, others=[ct])
-        trap_reaction = F.Reaction(
-            reactant=[cm, trap_sites],
-            product=ct,
-            k_0=1,
-            E_k=1,
-            p_0=1,
-            E_p=1,
-            volume=my_vol,
-        )
-        my_model.species = [cm, ct]
-        my_model.reactions = [trap_reaction]
-
-
-    """
-
-    def __init__(
-        self, name: str, mobile_species, k_0, E_k, p_0, E_p, n, volume
-    ) -> None:
-        super().__init__(name)
-        self.mobile_species = mobile_species
-        self.k_0 = k_0
-        self.E_k = E_k
-        self.p_0 = p_0
-        self.E_p = E_p
-        self.n = n
-        self.volume = volume
-
-        self.trapped_concentration = None
-        self.reaction = None
-
-    def create_species_and_reaction(self):
-        """create the immobile trapped species object and the reaction for trapping"""
-        self.trapped_concentration = F.Species(name=self.name, mobile=False)
-        trap_site = F.ImplicitSpecies(n=self.n, others=[self.trapped_concentration])
-
-        self.reaction = F.Reaction(
-            reactant=[self.mobile_species, trap_site],
-            product=self.trapped_concentration,
-            k_0=self.k_0,
-            E_k=self.E_k,
-            p_0=self.p_0,
-            E_p=self.E_p,
-            volume=self.volume,
-        )
+    @property
+    def legacy(self) -> bool:
+        """
+        Check if we are using FESTIM 1.0 implementation or FESTIM 2.0
+        """
+        if not self.subdomain_to_solution:
+            return True
+        else:
+            return False
 
 
 class ImplicitSpecies:
@@ -199,14 +104,14 @@ class ImplicitSpecies:
 
     Args:
         n (float): the total concentration of the species
-        others (List[Species]): the list of species from which the implicit
+        others (list[Species]): the list of species from which the implicit
             species concentration is computed (c = n - others)
         name (str, optional): a name given to the species. Defaults to None.
 
     Attributes:
         name (str): a name given to the species.
         n (float): the total concentration of the species
-        others (List[Species]): the list of species from which the implicit
+        others (list[Species]): the list of species from which the implicit
             species concentration is computed (c = n - others)
         concentration (form): the concentration of the species
 
@@ -215,7 +120,7 @@ class ImplicitSpecies:
     def __init__(
         self,
         n: float,
-        others: List[Species] = None,
+        others: list[Species] = None,
         name: str = None,
     ) -> None:
         self.name = name
@@ -234,7 +139,8 @@ def concentration(self):
             for other in self.others:
                 if other.solution is None:
                     raise ValueError(
-                        f"Cannot compute concentration of {self.name} because {other.name} has no solution"
+                        f"Cannot compute concentration of {
+                            self.name} because {other.name} has no solution"
                     )
         return self.n - sum([other.solution for other in self.others])
 
diff --git a/festim/trap.py b/festim/trap.py
new file mode 100644
index 000000000..8e30edd01
--- /dev/null
+++ b/festim/trap.py
@@ -0,0 +1,107 @@
+from festim.species import Species as _Species, ImplicitSpecies as _ImplicitSpecies
+from festim.reaction import Reaction as _Reaction
+
+
+class Trap(_Species):
+    """Trap species class for H transport simulation.
+
+    This class only works for 1 mobile species and 1 trapping level and is
+    for convenience, for more details see notes.
+
+    Args:
+        name (str, optional): a name given to the trap. Defaults to None.
+        mobile_species (_Species): the mobile species to be trapped
+        k_0 (float): the trapping rate constant pre-exponential factor (m3 s-1)
+        E_k (float): the trapping rate constant activation energy (eV)
+        p_0 (float): the detrapping rate constant pre-exponential factor (s-1)
+        E_p (float): the detrapping rate constant activation energy (eV)
+        volume (F.VolumeSubdomain1D): The volume subdomain where the trap is.
+
+
+    Attributes:
+        name (str, optional): a name given to the trap. Defaults to None.
+        mobile_species (_Species): the mobile species to be trapped
+        k_0 (float): the trapping rate constant pre-exponential factor (m3 s-1)
+        E_k (float): the trapping rate constant activation energy (eV)
+        p_0 (float): the detrapping rate constant pre-exponential factor (s-1)
+        E_p (float): the detrapping rate constant activation energy (eV)
+        volume (F.VolumeSubdomain1D): The volume subdomain where the trap is.
+        trapped_concentration (_Species): The immobile trapped concentration
+        trap_reaction (_Reaction): The reaction for trapping the mobile conc.
+
+    Usage:
+        >>> import festim as F
+        >>> trap = F.Trap(name="Trap", species=H, k_0=1.0, E_k=0.2, p_0=0.1, E_p=0.3, volume=my_vol)
+        >>> trap.name
+        'Trap'
+        >>> my_model = F.HydorgenTransportProblem()
+        >>> my_model.traps = [trap]
+
+    Notes:
+        This convenience class replaces the need to specify an implicit species and
+        the associated reaction, thus:
+
+        cm = _Species("mobile")
+        my_trap = F.Trap(
+            name="trapped",
+            mobile_species=cm,
+            k_0=1,
+            E_k=1,
+            p_0=1,
+            E_p=1,
+            n=1,
+            volume=my_vol,
+        )
+        my_model.species = [cm]
+        my_model.traps = [my_trap]
+
+        is equivalent to:
+
+        cm = _Species("mobile")
+        ct = _Species("trapped")
+        trap_sites = F.ImplicitSpecies(n=1, others=[ct])
+        trap_reaction = _Reaction(
+            reactant=[cm, trap_sites],
+            product=ct,
+            k_0=1,
+            E_k=1,
+            p_0=1,
+            E_p=1,
+            volume=my_vol,
+        )
+        my_model.species = [cm, ct]
+        my_model.reactions = [trap_reaction]
+
+
+    """
+
+    def __init__(
+        self, name: str, mobile_species, k_0, E_k, p_0, E_p, n, volume
+    ) -> None:
+        super().__init__(name)
+        self.mobile_species = mobile_species
+        self.k_0 = k_0
+        self.E_k = E_k
+        self.p_0 = p_0
+        self.E_p = E_p
+        self.n = n
+        self.volume = volume
+
+        self.trapped_concentration = None
+        self.reaction = None
+
+    def create_species_and_reaction(self):
+        """create the immobile trapped species object and the reaction for trapping"""
+        self.trapped_concentration = _Species(name=self.name, mobile=False)
+        trap_site = _ImplicitSpecies(
+            n=self.n, others=[self.trapped_concentration])
+
+        self.reaction = _Reaction(
+            reactant=[self.mobile_species, trap_site],
+            product=self.trapped_concentration,
+            k_0=self.k_0,
+            E_k=self.E_k,
+            p_0=self.p_0,
+            E_p=self.E_p,
+            volume=self.volume,
+        )
diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index f5d5711ad..9df1b9799 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -27,8 +27,10 @@ def half(x):
     gdim = mesh.geometry.dim
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
-    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
+    top_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(
+        mesh, fdim, bottom_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
@@ -71,14 +73,16 @@ def test_2_materials_2d_mms():
     D_top = 2.0
     D_bot = 5.0
     c_exact_top_ufl = (
-        lambda x: 1 + ufl.sin(ufl.pi * (2 * x[0] + 0.5)) + ufl.cos(2 * ufl.pi * x[1])
+        lambda x: 1 +
+        ufl.sin(ufl.pi * (2 * x[0] + 0.5)) + ufl.cos(2 * ufl.pi * x[1])
     )
-    c_exact_bot_ufl = lambda x: K_S_bot / K_S_top * c_exact_top_ufl(x)
+    def c_exact_bot_ufl(x): return K_S_bot / K_S_top * c_exact_top_ufl(x)
 
     c_exact_top_np = (
-        lambda x: 1 + np.sin(np.pi * (2 * x[0] + 0.5)) + np.cos(2 * np.pi * x[1])
+        lambda x: 1 +
+        np.sin(np.pi * (2 * x[0] + 0.5)) + np.cos(2 * np.pi * x[1])
     )
-    c_exact_bot_np = lambda x: K_S_bot / K_S_top * c_exact_top_np(x)
+    def c_exact_bot_np(x): return K_S_bot / K_S_top * c_exact_top_np(x)
 
     mesh, mt, ct = generate_mesh(100)
 
@@ -95,7 +99,8 @@ def test_2_materials_2d_mms():
 
     top_surface = F.SurfaceSubdomain(id=1)
     bottom_surface = F.SurfaceSubdomain(id=2)
-    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+    my_model.subdomains = [bottom_domain,
+                           top_domain, top_surface, bottom_surface]
 
     my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
     my_model.surface_to_volume = {
@@ -112,7 +117,8 @@ def test_2_materials_2d_mms():
 
     my_model.boundary_conditions = [
         F.FixedConcentrationBC(top_surface, value=c_exact_top_ufl, species=H),
-        F.FixedConcentrationBC(bottom_surface, value=c_exact_bot_ufl, species=H),
+        F.FixedConcentrationBC(
+            bottom_surface, value=c_exact_bot_ufl, species=H),
     ]
 
     source_top_val = (
@@ -127,15 +133,16 @@ def test_2_materials_2d_mms():
     )
     my_model.sources = [
         F.ParticleSource(volume=top_domain, species=H, value=source_top_val),
-        F.ParticleSource(volume=bottom_domain, species=H, value=source_bottom_val),
+        F.ParticleSource(volume=bottom_domain, species=H,
+                         value=source_bottom_val),
     ]
 
     my_model.temperature = 500.0  # lambda x: 300 + 10 * x[1] + 100 * x[0]
 
     my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
-
     my_model.exports = [
-        F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        F.VTXSpeciesExport(f"u_{subdomain.id}.bp",
+                           field=H, subdomain=subdomain)
         for subdomain in my_model.volume_subdomains
     ]
 
@@ -173,7 +180,8 @@ def test_1_material_discontinuous_version():
     right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
     my_model.subdomains = [subdomain, left_surface, right_surface]
-    my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
+    my_model.surface_to_volume = {
+        right_surface: subdomain, left_surface: subdomain}
 
     H = F.Species("H", mobile=True)
     trapped_H = F.Species("H_trapped", mobile=False)
@@ -207,7 +215,8 @@ def test_1_material_discontinuous_version():
     my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
     my_model.exports = [
-        F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        F.VTXSpeciesExport(
+            filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
         for subdomain in my_model.volume_subdomains
     ]
 
@@ -298,14 +307,16 @@ def test_3_materials_transient():
 
     my_model.temperature = lambda x: 300 + 100 * x[0]
 
-    my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=True, final_time=100)
+    my_model.settings = F.Settings(
+        atol=None, rtol=1e-5, transient=True, final_time=100)
     my_model.settings.stepsize = 1
 
     my_model.exports = [
-        F.VTXExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        F.VTXSpeciesExport(
+            filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
         for subdomain in my_model.volume_subdomains
     ] + [
-        F.VTXExport(
+        F.VTXSpeciesExport(
             filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain
         )
         for subdomain in my_model.volume_subdomains
@@ -336,7 +347,8 @@ def test_2_mats_particle_flux_bc():
 
     top_surface = F.SurfaceSubdomain(id=1)
     bottom_surface = F.SurfaceSubdomain(id=2)
-    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+    my_model.subdomains = [bottom_domain,
+                           top_domain, top_surface, bottom_surface]
 
     # we should be able to automate this
     my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
@@ -354,7 +366,8 @@ def test_2_mats_particle_flux_bc():
 
     my_model.boundary_conditions = [
         F.DirichletBC(top_surface, value=0.05, species=H),
-        F.ParticleFluxBC(bottom_surface, value=lambda x: 1.0 + x[0], species=H),
+        F.ParticleFluxBC(
+            bottom_surface, value=lambda x: 1.0 + x[0], species=H),
     ]
 
     my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
@@ -362,7 +375,8 @@ def test_2_mats_particle_flux_bc():
     my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
     my_model.exports = [
-        F.VTXExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+        F.VTXSpeciesExport(f"u_{subdomain.id}.bp",
+                           field=H, subdomain=subdomain)
         for subdomain in my_model.volume_subdomains
     ]
 
diff --git a/test/test_heat_transfer_problem.py b/test/test_heat_transfer_problem.py
index 30027dd4a..76007be48 100644
--- a/test/test_heat_transfer_problem.py
+++ b/test/test_heat_transfer_problem.py
@@ -64,7 +64,7 @@ def error_L2(u_computed, u_exact, degree_raise=3):
 
 def test_MMS_1():
     thermal_conductivity = 4.0
-    exact_solution = lambda x: 2 * x[0] ** 2
+    def exact_solution(x): return 2 * x[0] ** 2
     mms_source = -4 * thermal_conductivity
 
     my_problem = F.HeatTransferProblem()
@@ -106,9 +106,10 @@ def test_MMS_1():
 
 def test_MMS_T_dependent_thermal_cond():
     """MMS test with space T dependent thermal cond"""
-    thermal_conductivity = lambda T: 3 * T + 2
-    exact_solution = lambda x: 2 * x[0] ** 2 + 1
-    mms_source = lambda x: -(72 * x[0] ** 2 + 20)  # TODO would be nice to automate
+    def thermal_conductivity(T): return 3 * T + 2
+    def exact_solution(x): return 2 * x[0] ** 2 + 1
+    # TODO would be nice to automate
+    def mms_source(x): return -(72 * x[0] ** 2 + 20)
 
     my_problem = F.HeatTransferProblem()
 
@@ -154,7 +155,7 @@ def test_heat_transfer_transient():
     density = 2
     heat_capacity = 3
     thermal_conductivity = 4
-    exact_solution = lambda x, t: 2 * x[0] ** 2 + 20 * t
+    def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
     dTdt = 20
     mms_source = density * heat_capacity * dTdt - thermal_conductivity * 4
 
@@ -174,7 +175,8 @@ def test_heat_transfer_transient():
         F.VolumeSubdomain1D(id=1, borders=[2, 3], material=mat),
     ]
     # NOTE: it's good to check that without the IC the solution is not the exact one
-    my_problem.initial_condition = F.InitialTemperature(lambda x: exact_solution(x, 0))
+    my_problem.initial_condition = F.InitialTemperature(
+        lambda x: exact_solution(x, 0))
 
     my_problem.boundary_conditions = [
         F.FixedTemperatureBC(subdomain=left, value=exact_solution),
@@ -196,7 +198,7 @@ def test_heat_transfer_transient():
     my_problem.settings.stepsize = F.Stepsize(0.1)
 
     my_problem.exports = [
-        F.VTXExportForTemperature(filename="test_transient_heat_transfer.bp")
+        F.VTXTemperatureExport(filename="test_transient_heat_transfer.bp")
     ]
 
     my_problem.initialise()
@@ -206,7 +208,7 @@ def test_heat_transfer_transient():
     final_time_sim = (
         my_problem.t.value
     )  # we use the exact final time of the simulation which may differ from the one specified in the settings
-    exact_solution_end = lambda x: exact_solution(x, final_time_sim)
+    def exact_solution_end(x): return exact_solution(x, final_time_sim)
     L2_error = error_L2(computed_solution, exact_solution_end)
     assert L2_error < 1e-7
 
@@ -234,7 +236,8 @@ def test_MES():
     ]
 
     my_problem.sources = [
-        F.HeatSource(value=8 * mat.thermal_conductivity, volume=volume_subdomain)
+        F.HeatSource(value=8 * mat.thermal_conductivity,
+                     volume=volume_subdomain)
     ]
 
     my_problem.settings = F.Settings(
@@ -247,26 +250,27 @@ def test_MES():
     my_problem.run()
 
     computed_solution = my_problem.u
-    analytical_solution = lambda x: 4 * x[0] * (1 - x[0])
+    def analytical_solution(x): return 4 * x[0] * (1 - x[0])
     L2_error = error_L2(computed_solution, analytical_solution)
     assert L2_error < 1e-7
 
 
 # TODO populate this in other tests
 def test_sympify():
-    exact_solution = lambda x, t: 2 * x[0] ** 2 + 20 * t
+    def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
 
-    density = lambda T: 0.2 * T + 2
-    heat_capacity = lambda T: 0.2 * T + 3
-    thermal_conductivity = lambda T: 0.1 * T + 4
+    def density(T): return 0.2 * T + 2
+    def heat_capacity(T): return 0.2 * T + 3
+    def thermal_conductivity(T): return 0.1 * T + 4
 
     mms_source_from_sp = source_from_exact_solution(
         exact_solution,
         density=lambda x, t: density(exact_solution(x, t)),
         heat_capacity=lambda x, t: heat_capacity(exact_solution(x, t)),
-        thermal_conductivity=lambda x, t: thermal_conductivity(exact_solution(x, t)),
+        thermal_conductivity=lambda x, t: thermal_conductivity(
+            exact_solution(x, t)),
     )
-    mms_source = lambda x, t: mms_source_from_sp((x[0], None, None), t)
+    def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
 
     my_problem = F.HeatTransferProblem()
 
@@ -285,7 +289,8 @@ def test_sympify():
     ]
 
     # NOTE: it's good to check that without the IC the solution is not the exact one
-    my_problem.initial_condition = F.InitialTemperature(lambda x: exact_solution(x, 0))
+    my_problem.initial_condition = F.InitialTemperature(
+        lambda x: exact_solution(x, 0))
 
     my_problem.boundary_conditions = [
         F.FixedTemperatureBC(subdomain=left, value=exact_solution),
@@ -307,7 +312,7 @@ def test_sympify():
     my_problem.settings.stepsize = F.Stepsize(0.05)
 
     my_problem.exports = [
-        F.VTXExportForTemperature(filename="test_transient_heat_transfer.bp")
+        F.VTXTemperatureExport(filename="test_transient_heat_transfer.bp")
     ]
 
     my_problem.initialise()
@@ -318,7 +323,7 @@ def test_sympify():
         my_problem.t.value
     )  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
-    exact_solution_end = lambda x: exact_solution(x, final_time_sim)
+    def exact_solution_end(x): return exact_solution(x, final_time_sim)
     L2_error = error_L2(computed_solution, exact_solution_end)
     assert L2_error < 1e-7
 
@@ -338,7 +343,8 @@ def test_sources():
     # Test that setting invalid sources raises a TypeError
     with pytest.raises(TypeError, match="festim.HeatSource objects"):
         spe = F.Species("H")
-        htp.sources = [F.ParticleSource(1, vol, spe), F.ParticleSource(1, vol, spe)]
+        htp.sources = [F.ParticleSource(
+            1, vol, spe), F.ParticleSource(1, vol, spe)]
 
 
 def test_boundary_conditions():
@@ -391,7 +397,8 @@ def test_meshtags_from_xdmf(tmp_path, mesh):
     facet_tags = np.array(facet_tags).flatten()
     facet_indices = np.array(facet_indices).flatten()
 
-    facet_meshtags = dolfinx.mesh.meshtags(mesh, fdim, facet_indices, facet_tags)
+    facet_meshtags = dolfinx.mesh.meshtags(
+        mesh, fdim, facet_indices, facet_tags)
 
     # create volume meshtags
     num_cells = mesh.topology.index_map(vdim).size_local
@@ -413,7 +420,8 @@ def test_meshtags_from_xdmf(tmp_path, mesh):
     tags_volumes[volume_indices_left] = 2
     tags_volumes[volume_indices_right] = 3
 
-    volume_meshtags = dolfinx.mesh.meshtags(mesh, vdim, mesh_cell_indices, tags_volumes)
+    volume_meshtags = dolfinx.mesh.meshtags(
+        mesh, vdim, mesh_cell_indices, tags_volumes)
 
     # write files
     surface_file_path = os.path.join(tmp_path, "facets_file.xdmf")
@@ -503,7 +511,8 @@ def test_adaptive_timestepping_grows():
     my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 4], material=dummy_mat)
     my_model = F.HeatTransferProblem(
         mesh=test_mesh,
-        settings=F.Settings(atol=1e-10, rtol=1e-10, transient=True, final_time=10),
+        settings=F.Settings(atol=1e-10, rtol=1e-10,
+                            transient=True, final_time=10),
         subdomains=[my_vol],
     )
 
@@ -537,7 +546,8 @@ def test_adaptive_timestepping_shrinks():
     my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 4], material=dummy_mat)
     my_model = F.HeatTransferProblem(
         mesh=test_mesh,
-        settings=F.Settings(atol=1e-10, rtol=1e-10, transient=True, final_time=10),
+        settings=F.Settings(atol=1e-10, rtol=1e-10,
+                            transient=True, final_time=10),
         subdomains=[my_vol],
     )
 
@@ -585,7 +595,8 @@ def test_update_time_dependent_values_HeatFluxBC(bc_value, expected_values):
     my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 4], material=dummy_mat)
     surface = F.SurfaceSubdomain1D(id=2, x=0)
 
-    my_model = F.HeatTransferProblem(mesh=test_mesh, subdomains=[my_vol, surface])
+    my_model = F.HeatTransferProblem(
+        mesh=test_mesh, subdomains=[my_vol, surface])
     my_model.t = fem.Constant(my_model.mesh.mesh, 0.0)
     dt = fem.Constant(test_mesh.mesh, 1.0)
 
@@ -608,5 +619,6 @@ def test_update_time_dependent_values_HeatFluxBC(bc_value, expected_values):
 
         # TEST
         if isinstance(my_model.boundary_conditions[0].value_fenics, fem.Constant):
-            computed_value = float(my_model.boundary_conditions[0].value_fenics)
+            computed_value = float(
+                my_model.boundary_conditions[0].value_fenics)
             assert np.isclose(computed_value, expected_values[i])
diff --git a/test/test_permeation_problem.py b/test/test_permeation_problem.py
index 624ce635e..f7b0190aa 100644
--- a/test/test_permeation_problem.py
+++ b/test/test_permeation_problem.py
@@ -24,7 +24,8 @@ def relative_error_computed_to_analytical(
     computed_flux = computed_flux[indices]
 
     # evaulate relative error compared to analytical solution
-    relative_error = np.abs((computed_flux - analytical_flux) / analytical_flux)
+    relative_error = np.abs(
+        (computed_flux - analytical_flux) / analytical_flux)
     error = relative_error.mean()
 
     return error
@@ -128,12 +129,15 @@ def test_permeation_problem_multi_volume(tmp_path):
     my_model.mesh = my_mesh
 
     my_mat = F.Material(D_0=1.9e-7, E_D=0.2, name="my_mat")
-    my_subdomain_1 = F.VolumeSubdomain1D(id=1, borders=[0, L / 4], material=my_mat)
-    my_subdomain_2 = F.VolumeSubdomain1D(id=2, borders=[L / 4, L / 2], material=my_mat)
+    my_subdomain_1 = F.VolumeSubdomain1D(
+        id=1, borders=[0, L / 4], material=my_mat)
+    my_subdomain_2 = F.VolumeSubdomain1D(
+        id=2, borders=[L / 4, L / 2], material=my_mat)
     my_subdomain_3 = F.VolumeSubdomain1D(
         id=3, borders=[L / 2, 3 * L / 4], material=my_mat
     )
-    my_subdomain_4 = F.VolumeSubdomain1D(id=4, borders=[3 * L / 4, L], material=my_mat)
+    my_subdomain_4 = F.VolumeSubdomain1D(
+        id=4, borders=[3 * L / 4, L], material=my_mat)
     left_surface = F.SurfaceSubdomain1D(id=1, x=0)
     right_surface = F.SurfaceSubdomain1D(id=2, x=L)
     my_model.subdomains = [
@@ -163,7 +167,7 @@ def test_permeation_problem_multi_volume(tmp_path):
         surface=right_surface,
     )
     my_model.exports = [
-        F.VTXExport(
+        F.VTXSpeciesExport(
             os.path.join(tmp_path, "mobile_concentration_h.bp"), field=mobile_H
         ),
         outgassing_flux,
diff --git a/test/test_vtx.py b/test/test_vtx.py
index 0a3f00253..44cabc859 100644
--- a/test/test_vtx.py
+++ b/test/test_vtx.py
@@ -15,13 +15,11 @@ def test_vtx_export_one_function(tmpdir):
     sp = F.Species("H")
     sp.post_processing_solution = u
     filename = str(tmpdir.join("my_export.bp"))
-    my_export = F.VTXExport(filename, field=sp)
-    my_export.define_writer(mesh.comm)
+    my_export = F.VTXSpeciesExport(filename, field=sp)
 
-    for t in range(10):
-        u.interpolate(lambda x: t * (x[0] ** 2 + x[1] ** 2 + x[2] ** 2))
-
-        my_export.write(t)
+    functions = my_export.get_functions()
+    assert len(functions) == 1
+    assert functions[0] == u
 
 
 def test_vtx_export_two_functions(tmpdir):
@@ -34,15 +32,19 @@ def test_vtx_export_two_functions(tmpdir):
     sp1.post_processing_solution = u
     sp2.post_processing_solution = v
     filename = str(tmpdir.join("my_export.bp"))
-    my_export = F.VTXExport(filename, field=[sp1, sp2])
+    my_export = F.VTXSpeciesExport(filename, field=[sp1, sp2])
 
-    my_export.define_writer(mesh.comm)
+    functions = my_export.get_functions()
+    assert len(functions) == 2
+    assert functions[0] == u
+    assert functions[1] == v
 
-    for t in range(10):
-        u.interpolate(lambda x: t * (x[0] ** 2 + x[1] ** 2 + x[2] ** 2))
-        v.interpolate(lambda x: t * (x[0] ** 2 + x[1] ** 2 + x[2] ** 2))
 
-        my_export.write(t)
+@pytest.mark.skip(reason="Not implemented")
+def test_vtx_export_subdomain():
+    """Test that given multiple subdomains in problem, 
+    only correct functions are extracted from species"""
+    pass
 
 
 def test_vtx_integration_with_h_transport_problem(tmpdir):
@@ -58,23 +60,22 @@ def test_vtx_integration_with_h_transport_problem(tmpdir):
     my_model.temperature = 500
 
     filename = str(tmpdir.join("my_export.bp"))
-    my_export = F.VTXExport(filename, field=my_model.species[0])
+    my_export = F.VTXSpeciesExport(filename, field=my_model.species[0])
     my_model.exports = [my_export]
     my_model.settings = F.Settings(atol=1, rtol=0.1)
     my_model.settings.stepsize = F.Stepsize(initial_value=1)
 
     my_model.initialise()
-
-    for t in range(10):
-        my_export.write(t)
+    assert len(my_export.get_functions()) == 1
+    assert len(my_model._vtxfiles) == 1
 
 
 def test_field_attribute_is_always_list():
     """Test that the field attribute is always a list"""
-    my_export = F.VTXExport("my_export.bp", field=F.Species("H"))
+    my_export = F.VTXSpeciesExport("my_export.bp", field=F.Species("H"))
     assert isinstance(my_export.field, list)
 
-    my_export = F.VTXExport("my_export.bp", field=[F.Species("H")])
+    my_export = F.VTXSpeciesExport("my_export.bp", field=[F.Species("H")])
     assert isinstance(my_export.field, list)
 
 
@@ -82,36 +83,10 @@ def test_field_attribute_is_always_list():
 def test_field_attribute_raises_error_when_invalid_type(field):
     """Test that the field attribute raises an error if the type is not festim.Species or list"""
     with pytest.raises(TypeError):
-        F.VTXExport("my_export.bp", field=field)
-
-
-def test_filename_raises_error_with_wrong_extension():
-    """Test that the filename attribute raises an error if the extension is not .bp"""
-    with pytest.raises(ValueError):
-        F.VTXExport("my_export.txt", field=[F.Species("H")])
+        F.VTXSpeciesExport("my_export.bp", field=field)
 
 
 def test_filename_raises_error_when_wrong_type():
     """Test that the filename attribute raises an error if the extension is not .bp"""
     with pytest.raises(TypeError):
-        F.VTXExport(1, field=[F.Species("H")])
-
-
-def test_vtx_field_as_string_found_in_species(tmpdir):
-    """Test that the field attribute can be a string and is found in the species list"""
-    my_model = F.HydrogenTransportProblem()
-    my_model.mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
-    my_mat = F.Material(D_0=1, E_D=0, name="mat")
-    my_model.subdomains = [
-        F.VolumeSubdomain1D(1, borders=[0.0, 4.0], material=my_mat),
-    ]
-    my_model.species = [F.Species("H")]
-    my_model.temperature = 500
-
-    filename = str(tmpdir.join("my_export.bp"))
-    my_export = F.VTXExport(filename, field="H")
-    my_model.exports = [my_export]
-    my_model.settings = F.Settings(atol=1, rtol=0.1)
-    my_model.settings.stepsize = F.Stepsize(initial_value=1)
-
-    my_model.initialise()
+        F.VTXSpeciesExport(1, field=[F.Species("H")])
diff --git a/test/test_xdmf.py b/test/test_xdmf.py
index 8f6afc768..23a6b2581 100644
--- a/test/test_xdmf.py
+++ b/test/test_xdmf.py
@@ -4,6 +4,7 @@
 import os
 import numpy as np
 import pytest
+from pathlib import Path
 
 
 def test_init():
@@ -11,7 +12,7 @@ def test_init():
     species = F.Species("H")
     my_export = F.XDMFExport(filename="my_export.xdmf", field=species)
 
-    assert my_export.filename == "my_export.xdmf"
+    assert my_export.filename == Path("my_export.xdmf")
     assert my_export.field == [species]
 
 
@@ -20,9 +21,7 @@ def test_write(tmp_path):
     species = F.Species("H")
     filename = os.path.join(tmp_path, "test.xdmf")
     my_export = F.XDMFExport(filename=filename, field=species)
-
     my_export.define_writer(MPI.COMM_WORLD)
-
     domain = mesh.create_unit_cube(MPI.COMM_WORLD, 10, 10, 10)
     V = fem.functionspace(domain, ("Lagrange", 1))
     u = fem.Function(V)
@@ -45,7 +44,8 @@ def test_integration_with_HTransportProblem(tmp_path):
     my_model.temperature = 500.0
     my_model.species = [F.Species("H")]
     filename = os.path.join(tmp_path, "test.xdmf")
-    my_model.exports = [F.XDMFExport(filename=filename, field=my_model.species)]
+    my_model.exports = [F.XDMFExport(
+        filename=filename, field=my_model.species)]
 
     my_model.settings = F.Settings(atol=1, rtol=0.1, final_time=1)
     my_model.settings.stepsize = F.Stepsize(initial_value=0.5)
@@ -72,33 +72,7 @@ def test_field_attribute_raises_error_when_invalid_type(field):
         F.XDMFExport("my_export.xdmf", field=field)
 
 
-def test_filename_raises_error_with_wrong_extension():
-    """Test that the filename attribute raises an error if the extension is not .xdmf"""
-    with pytest.raises(ValueError):
-        F.XDMFExport("my_export.txt", field=[F.Species("H")])
-
-
 def test_filename_raises_error_when_wrong_type():
     """Test that the filename attribute raises an error if the file is not str"""
     with pytest.raises(TypeError):
         F.XDMFExport(1, field=[F.Species("H")])
-
-
-def test_xdmf_field_as_string_found_in_species(tmpdir):
-    """Test that the field attribute can be a string and is found in the species list"""
-    my_model = F.HydrogenTransportProblem()
-    my_model.mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
-    my_mat = F.Material(D_0=1, E_D=0, name="mat")
-    my_model.subdomains = [
-        F.VolumeSubdomain1D(1, borders=[0.0, 4.0], material=my_mat),
-    ]
-    my_model.species = [F.Species("H")]
-    my_model.temperature = 500
-
-    filename = str(tmpdir.join("my_export.xdmf"))
-    my_export = F.XDMFExport(filename, field="H")
-    my_model.exports = [my_export]
-    my_model.settings = F.Settings(atol=1, rtol=0.1)
-    my_model.settings.stepsize = F.Stepsize(initial_value=1)
-
-    my_model.initialise()

From b278110f85eb716b005f5da6b18738dc51e91e22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 18:25:06 -0400
Subject: [PATCH 098/134] Black formatting / ruff

---
 convergence_rates.py                         | 28 ++++---
 examples/multi_material_1d.py                |  6 +-
 examples/multi_material_2d.py                | 18 ++--
 examples/multi_material_2d_bis.py            | 21 ++---
 examples/multi_material_transient.py         | 12 ++-
 examples/multi_material_with_one_material.py |  6 +-
 examples/tds_example.py                      |  6 +-
 examples/trapping_example.py                 | 10 ++-
 festim/exports/vtx.py                        | 18 ++--
 festim/exports/xdmf.py                       |  5 +-
 festim/heat_transfer_problem.py              | 18 ++--
 festim/hydrogen_transport_problem.py         | 27 ++++--
 festim/problem.py                            |  1 +
 festim/reaction.py                           |  3 +-
 festim/trap.py                               |  3 +-
 test/system_tests/test_multi_material.py     | 53 ++++++------
 test/test_heat_transfer_problem.py           | 87 ++++++++++++--------
 test/test_permeation_problem.py              | 12 +--
 test/test_vtx.py                             |  2 +-
 test/test_xdmf.py                            |  3 +-
 20 files changed, 172 insertions(+), 167 deletions(-)

diff --git a/convergence_rates.py b/convergence_rates.py
index 65d4c3c91..037270442 100644
--- a/convergence_rates.py
+++ b/convergence_rates.py
@@ -58,20 +58,27 @@ def error_L2(u_computed, u_exact, degree_raise=3):
 
 
 def run(N):
-    def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
+    def exact_solution(x, t):
+        return 2 * x[0] ** 2 + 20 * t
 
-    def density(T): return 0.2 * T + 2
-    def heat_capacity(T): return 0.2 * T + 3
-    def thermal_conductivity(T): return 0.1 * T + 4
+    def density(T):
+        return 0.2 * T + 2
+
+    def heat_capacity(T):
+        return 0.2 * T + 3
+
+    def thermal_conductivity(T):
+        return 0.1 * T + 4
 
     mms_source_from_sp = source_from_exact_solution(
         exact_solution,
         density=lambda x, t: density(exact_solution(x, t)),
         heat_capacity=lambda x, t: heat_capacity(exact_solution(x, t)),
-        thermal_conductivity=lambda x, t: thermal_conductivity(
-            exact_solution(x, t)),
+        thermal_conductivity=lambda x, t: thermal_conductivity(exact_solution(x, t)),
     )
-    def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
+
+    def mms_source(x, t):
+        return mms_source_from_sp((x[0], None, None), t)
 
     my_problem = F.HeatTransferProblem()
 
@@ -89,8 +96,7 @@ def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
     ]
 
     # NOTE: it's good to check that without the IC the solution is not the exact one
-    my_problem.initial_condition = F.InitialTemperature(
-        lambda x: exact_solution(x, 0))
+    my_problem.initial_condition = F.InitialTemperature(lambda x: exact_solution(x, 0))
 
     my_problem.boundary_conditions = [
         F.FixedTemperatureBC(subdomain=left, value=exact_solution),
@@ -123,7 +129,9 @@ def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
         my_problem.t.value
     )  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
-    def exact_solution_end(x): return exact_solution(x, final_time_sim)
+    def exact_solution_end(x):
+        return exact_solution(x, final_time_sim)
+
     L2_error = error_L2(computed_solution, exact_solution_end)
     return L2_error
 
diff --git a/examples/multi_material_1d.py b/examples/multi_material_1d.py
index 2041f0b20..76221480e 100644
--- a/examples/multi_material_1d.py
+++ b/examples/multi_material_1d.py
@@ -54,8 +54,7 @@
     left_surface,
     right_surface,
 ]
-my_model.surface_to_volume = {
-    right_surface: right_domain, left_surface: left_domain}
+my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -91,8 +90,7 @@
 my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXSpeciesExport(
-        filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
diff --git a/examples/multi_material_2d.py b/examples/multi_material_2d.py
index fc59277ba..bc55a8d95 100644
--- a/examples/multi_material_2d.py
+++ b/examples/multi_material_2d.py
@@ -27,10 +27,8 @@ def half(x):
     gdim = mesh.geometry.dim
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    top_facets = dolfinx.mesh.locate_entities_boundary(
-        mesh, fdim, top_boundary)
-    bottom_facets = dolfinx.mesh.locate_entities_boundary(
-        mesh, fdim, bottom_boundary)
+    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
@@ -87,8 +85,7 @@ def half(x):
 
 # we should be able to automate this
 my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
-my_model.surface_to_volume = {
-    top_surface: top_domain, bottom_surface: bottom_domain}
+my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
 
@@ -120,13 +117,10 @@ def half(x):
 # -------------------- post processing --------------------
 
 # derived quantities
-entity_maps = {
-    sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
+entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
 
-ds_b = ufl.Measure("ds", domain=bottom_domain.submesh,
-                   subdomain_data=bottom_domain.ft)
-ds_t = ufl.Measure("ds", domain=top_domain.submesh,
-                   subdomain_data=top_domain.ft)
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
+ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
 ds = ufl.Measure("ds", domain=mesh, subdomain_data=my_model.facet_meshtags)
diff --git a/examples/multi_material_2d_bis.py b/examples/multi_material_2d_bis.py
index cdd0d0c23..a1f022a78 100644
--- a/examples/multi_material_2d_bis.py
+++ b/examples/multi_material_2d_bis.py
@@ -27,10 +27,8 @@ def half(x):
     gdim = mesh.geometry.dim
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    top_facets = dolfinx.mesh.locate_entities_boundary(
-        mesh, fdim, top_boundary)
-    bottom_facets = dolfinx.mesh.locate_entities_boundary(
-        mesh, fdim, bottom_boundary)
+    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
@@ -87,8 +85,7 @@ def half(x):
 
 # we should be able to automate this
 my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
-my_model.surface_to_volume = {
-    top_surface: top_domain, bottom_surface: bottom_domain}
+my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -135,8 +132,7 @@ def half(x):
     F.VTXSpeciesExport(f"c_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ] + [
-    F.VTXSpeciesExport(f"c_t_{subdomain.id}.bp",
-                       field=trapped_H, subdomain=subdomain)
+    F.VTXSpeciesExport(f"c_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
@@ -147,14 +143,11 @@ def half(x):
 
 
 # derived quantities
-entity_maps = {
-    sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
+entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
 entity_maps[mesh] = bottom_domain.submesh_to_mesh
 
-ds_b = ufl.Measure("ds", domain=bottom_domain.submesh,
-                   subdomain_data=bottom_domain.ft)
-ds_t = ufl.Measure("ds", domain=top_domain.submesh,
-                   subdomain_data=top_domain.ft)
+ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
+ds_t = ufl.Measure("ds", domain=top_domain.submesh, subdomain_data=top_domain.ft)
 dx_b = ufl.Measure("dx", domain=bottom_domain.submesh)
 dx = ufl.Measure("dx", domain=mesh)
 
diff --git a/examples/multi_material_transient.py b/examples/multi_material_transient.py
index d00b1a2b9..4469625b9 100644
--- a/examples/multi_material_transient.py
+++ b/examples/multi_material_transient.py
@@ -49,8 +49,7 @@
     left_surface,
     right_surface,
 ]
-my_model.surface_to_volume = {
-    right_surface: right_domain, left_surface: left_domain}
+my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -83,17 +82,16 @@
 
 my_model.temperature = lambda x: 300 + 100 * x[0]
 
-my_model.settings = F.Settings(
-    atol=None, rtol=1e-5, transient=True, final_time=100)
+my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=True, final_time=100)
 my_model.settings.stepsize = 1
 
 my_model.exports = [
-    F.VTXSpeciesExport(
-        filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ] + [
     F.VTXSpeciesExport(
-        filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain)
+        filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain
+    )
     for subdomain in my_model.volume_subdomains
 ]
 
diff --git a/examples/multi_material_with_one_material.py b/examples/multi_material_with_one_material.py
index 281791ca1..7bf502bf1 100644
--- a/examples/multi_material_with_one_material.py
+++ b/examples/multi_material_with_one_material.py
@@ -21,8 +21,7 @@
 right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
 my_model.subdomains = [subdomain, left_surface, right_surface]
-my_model.surface_to_volume = {
-    right_surface: subdomain, left_surface: subdomain}
+my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
@@ -58,8 +57,7 @@
 my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
 my_model.exports = [
-    F.VTXSpeciesExport(
-        filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+    F.VTXSpeciesExport(filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
     for subdomain in my_model.volume_subdomains
 ]
 
diff --git a/examples/tds_example.py b/examples/tds_example.py
index 6f1f8fda8..fa60689c7 100644
--- a/examples/tds_example.py
+++ b/examples/tds_example.py
@@ -97,8 +97,7 @@ def left_conc_value(t):
 
 
 my_model.boundary_conditions = [
-    F.DirichletBC(subdomain=left_surface,
-                  value=left_conc_value, species=mobile_H),
+    F.DirichletBC(subdomain=left_surface, value=left_conc_value, species=mobile_H),
     F.DirichletBC(subdomain=right_surface, value=0, species=mobile_H),
 ]
 
@@ -119,8 +118,7 @@ def left_conc_value(t):
 
 # -------- Settings --------- #
 
-my_model.settings = F.Settings(
-    atol=1e10, rtol=1e-10, max_iterations=30, final_time=500)
+my_model.settings = F.Settings(atol=1e10, rtol=1e-10, max_iterations=30, final_time=500)
 
 my_model.settings.stepsize = F.Stepsize(initial_value=0.5)
 
diff --git a/examples/trapping_example.py b/examples/trapping_example.py
index 4d980a24b..adc1eae85 100644
--- a/examples/trapping_example.py
+++ b/examples/trapping_example.py
@@ -44,10 +44,12 @@
     F.DirichletBC(subdomain=left_surface, value=1e12, species=mobile_H),
 ]
 my_model.exports = [
-    F.VTXSpeciesExport("mobile_concentration_h.bp",
-                       field=mobile_H),  # produces 0 in file
-    F.VTXSpeciesExport("trapped_concentration_h.bp",
-                       field=trapped_H),  # produces 0 in file
+    F.VTXSpeciesExport(
+        "mobile_concentration_h.bp", field=mobile_H
+    ),  # produces 0 in file
+    F.VTXSpeciesExport(
+        "trapped_concentration_h.bp", field=trapped_H
+    ),  # produces 0 in file
     F.XDMFExport("mobile_concentration_h.xdmf", field=mobile_H),
     F.XDMFExport("trapped_concentration_h.xdmf", field=trapped_H),
 ]
diff --git a/festim/exports/vtx.py b/festim/exports/vtx.py
index dafbbe2fb..df3fd2df6 100644
--- a/festim/exports/vtx.py
+++ b/festim/exports/vtx.py
@@ -5,14 +5,15 @@
 import warnings
 
 
-class ExportBaseClass():
+class ExportBaseClass:
     _filename: Path | str
 
     def __init__(self, filename: str | Path, ext: str) -> None:
         name = Path(filename)
         if name.suffix != ext:
             warnings.warn(
-                f"Filename {filename} does not have {ext} extension, adding it.")
+                f"Filename {filename} does not have {ext} extension, adding it."
+            )
             name = name.with_suffix(ext)
 
         self._filename = Path(filename)
@@ -23,8 +24,7 @@ def filename(self):
 
 
 class VTXTemperatureExport(ExportBaseClass):
-    def __init__(
-            self, filename: str | Path):
+    def __init__(self, filename: str | Path):
         super().__init__(filename, ".bp")
 
 
@@ -44,9 +44,10 @@ class VTXSpeciesExport(ExportBaseClass):
     _subdomain: _VolumeSubdomain
 
     def __init__(
-        self, filename: str | Path,
+        self,
+        filename: str | Path,
         field: _Species | list[_Species],
-        subdomain: _VolumeSubdomain = None
+        subdomain: _VolumeSubdomain = None,
     ) -> None:
         super().__init__(filename, ".bp")
         self.field = field
@@ -83,7 +84,7 @@ def field(self, value: _Species | list[_Species]):
         else:
             raise TypeError(
                 "field must be of type festim.Species or a list of festim.Species or str",
-                f"got {type(value)}."
+                f"got {type(value)}.",
             )
         self._field = val
 
@@ -112,5 +113,6 @@ def get_functions(self) -> list[_Function]:
                 for field in self._field:
                     if self._subdomain in field.subdomains:
                         outfiles.append(
-                            field.subdomain_to_post_processing_solution[self._subdomain])
+                            field.subdomain_to_post_processing_solution[self._subdomain]
+                        )
                 return outfiles
diff --git a/festim/exports/xdmf.py b/festim/exports/xdmf.py
index a3fa25873..e2956caf2 100644
--- a/festim/exports/xdmf.py
+++ b/festim/exports/xdmf.py
@@ -38,8 +38,9 @@ def field(self, value: _Species | list[_Species]):
         if isinstance(value, list):
             for element in value:
                 if not isinstance(element, _Species):
-                    raise TypeError(f"Each element in the list must be a species, got {type(element)}."
-                                    )
+                    raise TypeError(
+                        f"Each element in the list must be a species, got {type(element)}."
+                    )
             val = value
         elif isinstance(value, _Species):
             val = [value]
diff --git a/festim/heat_transfer_problem.py b/festim/heat_transfer_problem.py
index 27c6b21df..9bad85e08 100644
--- a/festim/heat_transfer_problem.py
+++ b/festim/heat_transfer_problem.py
@@ -35,8 +35,7 @@ def sources(self):
     @sources.setter
     def sources(self, value):
         if not all(isinstance(source, F.HeatSource) for source in value):
-            raise TypeError(
-                "sources must be a list of festim.HeatSource objects")
+            raise TypeError("sources must be a list of festim.HeatSource objects")
         self._sources = value
 
     @property
@@ -199,16 +198,14 @@ def create_formulation(self):
         # add sources
         for source in self.sources:
             self.formulation -= (
-                source.value_fenics * self.test_function *
-                self.dx(source.volume.id)
+                source.value_fenics * self.test_function * self.dx(source.volume.id)
             )
 
         # add fluxes
         for bc in self.boundary_conditions:
             if isinstance(bc, F.HeatFluxBC):
                 self.formulation -= (
-                    bc.value_fenics * self.test_function *
-                    self.ds(bc.subdomain.id)
+                    bc.value_fenics * self.test_function * self.ds(bc.subdomain.id)
                 )
 
     def initialise_exports(self):
@@ -221,9 +218,12 @@ def initialise_exports(self):
                     "XDMF export is not implemented yet for heat transfer problems"
                 )
             if isinstance(export, F.VTXTemperatureExport):
-                self._vtxfile = VTXWriter(self.u.function_space.mesh.comm,
-                                          export.filename, [self.u],
-                                          engine="BP5")
+                self._vtxfile = VTXWriter(
+                    self.u.function_space.mesh.comm,
+                    export.filename,
+                    [self.u],
+                    engine="BP5",
+                )
 
     def post_processing(self):
         """Post processes the model"""
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index f0ecaad48..7a3be07d4 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -122,7 +122,6 @@ def __init__(
         sources=None,
         initial_conditions=None,
         boundary_conditions=None,
-
         settings=None,
         exports=None,
         traps=None,
@@ -346,8 +345,14 @@ def initialise_exports(self):
                 export.define_writer(MPI.COMM_WORLD)
             if isinstance(export, F.VTXSpeciesExport):
                 functions = export.get_functions()
-                self._vtxfiles.append(dolfinx.io.VTXWriter(functions[0].function_space.mesh.comm,
-                                                           export.filename, functions, engine="BP5"))
+                self._vtxfiles.append(
+                    dolfinx.io.VTXWriter(
+                        functions[0].function_space.mesh.comm,
+                        export.filename,
+                        functions,
+                        engine="BP5",
+                    )
+                )
         # compute diffusivity function for surface fluxes
 
         spe_to_D_global = {}  # links species to global D function
@@ -1173,11 +1178,16 @@ def initialise_exports(self):
         for export in self.exports:
             if isinstance(export, F.VTXSpeciesExport):
                 functions = export.get_functions()
-                self._vtxfiles.append(dolfinx.io.VTXWriter(functions[0].function_space.mesh.comm,
-                                                           export.filename, functions, engine="BP5"))
+                self._vtxfiles.append(
+                    dolfinx.io.VTXWriter(
+                        functions[0].function_space.mesh.comm,
+                        export.filename,
+                        functions,
+                        engine="BP5",
+                    )
+                )
             else:
-                raise NotImplementedError(
-                    f"Export type {type(export)} not implemented")
+                raise NotImplementedError(f"Export type {type(export)} not implemented")
 
     def post_processing(self):
         # update post-processing solutions (for each species in each subdomain)
@@ -1198,8 +1208,7 @@ def post_processing(self):
 
         for export in self.exports:
             if not isinstance(export, F.VTXSpeciesExport):
-                raise NotImplementedError(
-                    f"Export type {type(export)} not implemented")
+                raise NotImplementedError(f"Export type {type(export)} not implemented")
 
     def iterate(self):
         """Iterates the model for a given time step"""
diff --git a/festim/problem.py b/festim/problem.py
index fb14f43e2..99d4f71bc 100644
--- a/festim/problem.py
+++ b/festim/problem.py
@@ -19,6 +19,7 @@ class ProblemBase:
         simulation
     progress_bar (tqdm.autonotebook.tqdm) the progress bar
     """
+
     mesh: _Mesh
     sources: list[_SourceBase]
     exports: list[Any]
diff --git a/festim/reaction.py b/festim/reaction.py
index 43ad7afab..730065dc0 100644
--- a/festim/reaction.py
+++ b/festim/reaction.py
@@ -47,8 +47,7 @@ class Reaction:
 
     def __init__(
         self,
-        reactant:
-            _Species | _ImplicitSpecies | list[_Species | _ImplicitSpecies],
+        reactant: _Species | _ImplicitSpecies | list[_Species | _ImplicitSpecies],
         k_0: float,
         E_k: float,
         volume: VS1D,
diff --git a/festim/trap.py b/festim/trap.py
index 8e30edd01..9bb3ca6dc 100644
--- a/festim/trap.py
+++ b/festim/trap.py
@@ -93,8 +93,7 @@ def __init__(
     def create_species_and_reaction(self):
         """create the immobile trapped species object and the reaction for trapping"""
         self.trapped_concentration = _Species(name=self.name, mobile=False)
-        trap_site = _ImplicitSpecies(
-            n=self.n, others=[self.trapped_concentration])
+        trap_site = _ImplicitSpecies(n=self.n, others=[self.trapped_concentration])
 
         self.reaction = _Reaction(
             reactant=[self.mobile_species, trap_site],
diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index 9df1b9799..02478ea9b 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -27,10 +27,8 @@ def half(x):
     gdim = mesh.geometry.dim
     tdim = mesh.topology.dim
     fdim = tdim - 1
-    top_facets = dolfinx.mesh.locate_entities_boundary(
-        mesh, fdim, top_boundary)
-    bottom_facets = dolfinx.mesh.locate_entities_boundary(
-        mesh, fdim, bottom_boundary)
+    top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
+    bottom_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary)
     num_facets_local = (
         mesh.topology.index_map(fdim).size_local
         + mesh.topology.index_map(fdim).num_ghosts
@@ -73,16 +71,18 @@ def test_2_materials_2d_mms():
     D_top = 2.0
     D_bot = 5.0
     c_exact_top_ufl = (
-        lambda x: 1 +
-        ufl.sin(ufl.pi * (2 * x[0] + 0.5)) + ufl.cos(2 * ufl.pi * x[1])
+        lambda x: 1 + ufl.sin(ufl.pi * (2 * x[0] + 0.5)) + ufl.cos(2 * ufl.pi * x[1])
     )
-    def c_exact_bot_ufl(x): return K_S_bot / K_S_top * c_exact_top_ufl(x)
+
+    def c_exact_bot_ufl(x):
+        return K_S_bot / K_S_top * c_exact_top_ufl(x)
 
     c_exact_top_np = (
-        lambda x: 1 +
-        np.sin(np.pi * (2 * x[0] + 0.5)) + np.cos(2 * np.pi * x[1])
+        lambda x: 1 + np.sin(np.pi * (2 * x[0] + 0.5)) + np.cos(2 * np.pi * x[1])
     )
-    def c_exact_bot_np(x): return K_S_bot / K_S_top * c_exact_top_np(x)
+
+    def c_exact_bot_np(x):
+        return K_S_bot / K_S_top * c_exact_top_np(x)
 
     mesh, mt, ct = generate_mesh(100)
 
@@ -99,8 +99,7 @@ def c_exact_bot_np(x): return K_S_bot / K_S_top * c_exact_top_np(x)
 
     top_surface = F.SurfaceSubdomain(id=1)
     bottom_surface = F.SurfaceSubdomain(id=2)
-    my_model.subdomains = [bottom_domain,
-                           top_domain, top_surface, bottom_surface]
+    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
     my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
     my_model.surface_to_volume = {
@@ -117,8 +116,7 @@ def c_exact_bot_np(x): return K_S_bot / K_S_top * c_exact_top_np(x)
 
     my_model.boundary_conditions = [
         F.FixedConcentrationBC(top_surface, value=c_exact_top_ufl, species=H),
-        F.FixedConcentrationBC(
-            bottom_surface, value=c_exact_bot_ufl, species=H),
+        F.FixedConcentrationBC(bottom_surface, value=c_exact_bot_ufl, species=H),
     ]
 
     source_top_val = (
@@ -133,16 +131,14 @@ def c_exact_bot_np(x): return K_S_bot / K_S_top * c_exact_top_np(x)
     )
     my_model.sources = [
         F.ParticleSource(volume=top_domain, species=H, value=source_top_val),
-        F.ParticleSource(volume=bottom_domain, species=H,
-                         value=source_bottom_val),
+        F.ParticleSource(volume=bottom_domain, species=H, value=source_bottom_val),
     ]
 
     my_model.temperature = 500.0  # lambda x: 300 + 10 * x[1] + 100 * x[0]
 
     my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
     my_model.exports = [
-        F.VTXSpeciesExport(f"u_{subdomain.id}.bp",
-                           field=H, subdomain=subdomain)
+        F.VTXSpeciesExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
         for subdomain in my_model.volume_subdomains
     ]
 
@@ -180,8 +176,7 @@ def test_1_material_discontinuous_version():
     right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
     my_model.subdomains = [subdomain, left_surface, right_surface]
-    my_model.surface_to_volume = {
-        right_surface: subdomain, left_surface: subdomain}
+    my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
 
     H = F.Species("H", mobile=True)
     trapped_H = F.Species("H_trapped", mobile=False)
@@ -216,7 +211,8 @@ def test_1_material_discontinuous_version():
 
     my_model.exports = [
         F.VTXSpeciesExport(
-            filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+            filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain
+        )
         for subdomain in my_model.volume_subdomains
     ]
 
@@ -307,13 +303,13 @@ def test_3_materials_transient():
 
     my_model.temperature = lambda x: 300 + 100 * x[0]
 
-    my_model.settings = F.Settings(
-        atol=None, rtol=1e-5, transient=True, final_time=100)
+    my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=True, final_time=100)
     my_model.settings.stepsize = 1
 
     my_model.exports = [
         F.VTXSpeciesExport(
-            filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
+            filename=f"u_{subdomain.id}.bp", field=H, subdomain=subdomain
+        )
         for subdomain in my_model.volume_subdomains
     ] + [
         F.VTXSpeciesExport(
@@ -347,8 +343,7 @@ def test_2_mats_particle_flux_bc():
 
     top_surface = F.SurfaceSubdomain(id=1)
     bottom_surface = F.SurfaceSubdomain(id=2)
-    my_model.subdomains = [bottom_domain,
-                           top_domain, top_surface, bottom_surface]
+    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
 
     # we should be able to automate this
     my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
@@ -366,8 +361,7 @@ def test_2_mats_particle_flux_bc():
 
     my_model.boundary_conditions = [
         F.DirichletBC(top_surface, value=0.05, species=H),
-        F.ParticleFluxBC(
-            bottom_surface, value=lambda x: 1.0 + x[0], species=H),
+        F.ParticleFluxBC(bottom_surface, value=lambda x: 1.0 + x[0], species=H),
     ]
 
     my_model.temperature = lambda x: 300 + 10 * x[1] + 100 * x[0]
@@ -375,8 +369,7 @@ def test_2_mats_particle_flux_bc():
     my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=False)
 
     my_model.exports = [
-        F.VTXSpeciesExport(f"u_{subdomain.id}.bp",
-                           field=H, subdomain=subdomain)
+        F.VTXSpeciesExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
         for subdomain in my_model.volume_subdomains
     ]
 
diff --git a/test/test_heat_transfer_problem.py b/test/test_heat_transfer_problem.py
index 76007be48..dfc335f46 100644
--- a/test/test_heat_transfer_problem.py
+++ b/test/test_heat_transfer_problem.py
@@ -64,7 +64,10 @@ def error_L2(u_computed, u_exact, degree_raise=3):
 
 def test_MMS_1():
     thermal_conductivity = 4.0
-    def exact_solution(x): return 2 * x[0] ** 2
+
+    def exact_solution(x):
+        return 2 * x[0] ** 2
+
     mms_source = -4 * thermal_conductivity
 
     my_problem = F.HeatTransferProblem()
@@ -106,10 +109,16 @@ def exact_solution(x): return 2 * x[0] ** 2
 
 def test_MMS_T_dependent_thermal_cond():
     """MMS test with space T dependent thermal cond"""
-    def thermal_conductivity(T): return 3 * T + 2
-    def exact_solution(x): return 2 * x[0] ** 2 + 1
+
+    def thermal_conductivity(T):
+        return 3 * T + 2
+
+    def exact_solution(x):
+        return 2 * x[0] ** 2 + 1
+
     # TODO would be nice to automate
-    def mms_source(x): return -(72 * x[0] ** 2 + 20)
+    def mms_source(x):
+        return -(72 * x[0] ** 2 + 20)
 
     my_problem = F.HeatTransferProblem()
 
@@ -155,7 +164,10 @@ def test_heat_transfer_transient():
     density = 2
     heat_capacity = 3
     thermal_conductivity = 4
-    def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
+
+    def exact_solution(x, t):
+        return 2 * x[0] ** 2 + 20 * t
+
     dTdt = 20
     mms_source = density * heat_capacity * dTdt - thermal_conductivity * 4
 
@@ -175,8 +187,7 @@ def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
         F.VolumeSubdomain1D(id=1, borders=[2, 3], material=mat),
     ]
     # NOTE: it's good to check that without the IC the solution is not the exact one
-    my_problem.initial_condition = F.InitialTemperature(
-        lambda x: exact_solution(x, 0))
+    my_problem.initial_condition = F.InitialTemperature(lambda x: exact_solution(x, 0))
 
     my_problem.boundary_conditions = [
         F.FixedTemperatureBC(subdomain=left, value=exact_solution),
@@ -208,7 +219,10 @@ def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
     final_time_sim = (
         my_problem.t.value
     )  # we use the exact final time of the simulation which may differ from the one specified in the settings
-    def exact_solution_end(x): return exact_solution(x, final_time_sim)
+
+    def exact_solution_end(x):
+        return exact_solution(x, final_time_sim)
+
     L2_error = error_L2(computed_solution, exact_solution_end)
     assert L2_error < 1e-7
 
@@ -236,8 +250,7 @@ def test_MES():
     ]
 
     my_problem.sources = [
-        F.HeatSource(value=8 * mat.thermal_conductivity,
-                     volume=volume_subdomain)
+        F.HeatSource(value=8 * mat.thermal_conductivity, volume=volume_subdomain)
     ]
 
     my_problem.settings = F.Settings(
@@ -250,27 +263,37 @@ def test_MES():
     my_problem.run()
 
     computed_solution = my_problem.u
-    def analytical_solution(x): return 4 * x[0] * (1 - x[0])
+
+    def analytical_solution(x):
+        return 4 * x[0] * (1 - x[0])
+
     L2_error = error_L2(computed_solution, analytical_solution)
     assert L2_error < 1e-7
 
 
 # TODO populate this in other tests
 def test_sympify():
-    def exact_solution(x, t): return 2 * x[0] ** 2 + 20 * t
+    def exact_solution(x, t):
+        return 2 * x[0] ** 2 + 20 * t
+
+    def density(T):
+        return 0.2 * T + 2
 
-    def density(T): return 0.2 * T + 2
-    def heat_capacity(T): return 0.2 * T + 3
-    def thermal_conductivity(T): return 0.1 * T + 4
+    def heat_capacity(T):
+        return 0.2 * T + 3
+
+    def thermal_conductivity(T):
+        return 0.1 * T + 4
 
     mms_source_from_sp = source_from_exact_solution(
         exact_solution,
         density=lambda x, t: density(exact_solution(x, t)),
         heat_capacity=lambda x, t: heat_capacity(exact_solution(x, t)),
-        thermal_conductivity=lambda x, t: thermal_conductivity(
-            exact_solution(x, t)),
+        thermal_conductivity=lambda x, t: thermal_conductivity(exact_solution(x, t)),
     )
-    def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
+
+    def mms_source(x, t):
+        return mms_source_from_sp((x[0], None, None), t)
 
     my_problem = F.HeatTransferProblem()
 
@@ -289,8 +312,7 @@ def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
     ]
 
     # NOTE: it's good to check that without the IC the solution is not the exact one
-    my_problem.initial_condition = F.InitialTemperature(
-        lambda x: exact_solution(x, 0))
+    my_problem.initial_condition = F.InitialTemperature(lambda x: exact_solution(x, 0))
 
     my_problem.boundary_conditions = [
         F.FixedTemperatureBC(subdomain=left, value=exact_solution),
@@ -323,7 +345,9 @@ def mms_source(x, t): return mms_source_from_sp((x[0], None, None), t)
         my_problem.t.value
     )  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
-    def exact_solution_end(x): return exact_solution(x, final_time_sim)
+    def exact_solution_end(x):
+        return exact_solution(x, final_time_sim)
+
     L2_error = error_L2(computed_solution, exact_solution_end)
     assert L2_error < 1e-7
 
@@ -343,8 +367,7 @@ def test_sources():
     # Test that setting invalid sources raises a TypeError
     with pytest.raises(TypeError, match="festim.HeatSource objects"):
         spe = F.Species("H")
-        htp.sources = [F.ParticleSource(
-            1, vol, spe), F.ParticleSource(1, vol, spe)]
+        htp.sources = [F.ParticleSource(1, vol, spe), F.ParticleSource(1, vol, spe)]
 
 
 def test_boundary_conditions():
@@ -397,8 +420,7 @@ def test_meshtags_from_xdmf(tmp_path, mesh):
     facet_tags = np.array(facet_tags).flatten()
     facet_indices = np.array(facet_indices).flatten()
 
-    facet_meshtags = dolfinx.mesh.meshtags(
-        mesh, fdim, facet_indices, facet_tags)
+    facet_meshtags = dolfinx.mesh.meshtags(mesh, fdim, facet_indices, facet_tags)
 
     # create volume meshtags
     num_cells = mesh.topology.index_map(vdim).size_local
@@ -420,8 +442,7 @@ def test_meshtags_from_xdmf(tmp_path, mesh):
     tags_volumes[volume_indices_left] = 2
     tags_volumes[volume_indices_right] = 3
 
-    volume_meshtags = dolfinx.mesh.meshtags(
-        mesh, vdim, mesh_cell_indices, tags_volumes)
+    volume_meshtags = dolfinx.mesh.meshtags(mesh, vdim, mesh_cell_indices, tags_volumes)
 
     # write files
     surface_file_path = os.path.join(tmp_path, "facets_file.xdmf")
@@ -511,8 +532,7 @@ def test_adaptive_timestepping_grows():
     my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 4], material=dummy_mat)
     my_model = F.HeatTransferProblem(
         mesh=test_mesh,
-        settings=F.Settings(atol=1e-10, rtol=1e-10,
-                            transient=True, final_time=10),
+        settings=F.Settings(atol=1e-10, rtol=1e-10, transient=True, final_time=10),
         subdomains=[my_vol],
     )
 
@@ -546,8 +566,7 @@ def test_adaptive_timestepping_shrinks():
     my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 4], material=dummy_mat)
     my_model = F.HeatTransferProblem(
         mesh=test_mesh,
-        settings=F.Settings(atol=1e-10, rtol=1e-10,
-                            transient=True, final_time=10),
+        settings=F.Settings(atol=1e-10, rtol=1e-10, transient=True, final_time=10),
         subdomains=[my_vol],
     )
 
@@ -595,8 +614,7 @@ def test_update_time_dependent_values_HeatFluxBC(bc_value, expected_values):
     my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 4], material=dummy_mat)
     surface = F.SurfaceSubdomain1D(id=2, x=0)
 
-    my_model = F.HeatTransferProblem(
-        mesh=test_mesh, subdomains=[my_vol, surface])
+    my_model = F.HeatTransferProblem(mesh=test_mesh, subdomains=[my_vol, surface])
     my_model.t = fem.Constant(my_model.mesh.mesh, 0.0)
     dt = fem.Constant(test_mesh.mesh, 1.0)
 
@@ -619,6 +637,5 @@ def test_update_time_dependent_values_HeatFluxBC(bc_value, expected_values):
 
         # TEST
         if isinstance(my_model.boundary_conditions[0].value_fenics, fem.Constant):
-            computed_value = float(
-                my_model.boundary_conditions[0].value_fenics)
+            computed_value = float(my_model.boundary_conditions[0].value_fenics)
             assert np.isclose(computed_value, expected_values[i])
diff --git a/test/test_permeation_problem.py b/test/test_permeation_problem.py
index f7b0190aa..5226724f0 100644
--- a/test/test_permeation_problem.py
+++ b/test/test_permeation_problem.py
@@ -24,8 +24,7 @@ def relative_error_computed_to_analytical(
     computed_flux = computed_flux[indices]
 
     # evaulate relative error compared to analytical solution
-    relative_error = np.abs(
-        (computed_flux - analytical_flux) / analytical_flux)
+    relative_error = np.abs((computed_flux - analytical_flux) / analytical_flux)
     error = relative_error.mean()
 
     return error
@@ -129,15 +128,12 @@ def test_permeation_problem_multi_volume(tmp_path):
     my_model.mesh = my_mesh
 
     my_mat = F.Material(D_0=1.9e-7, E_D=0.2, name="my_mat")
-    my_subdomain_1 = F.VolumeSubdomain1D(
-        id=1, borders=[0, L / 4], material=my_mat)
-    my_subdomain_2 = F.VolumeSubdomain1D(
-        id=2, borders=[L / 4, L / 2], material=my_mat)
+    my_subdomain_1 = F.VolumeSubdomain1D(id=1, borders=[0, L / 4], material=my_mat)
+    my_subdomain_2 = F.VolumeSubdomain1D(id=2, borders=[L / 4, L / 2], material=my_mat)
     my_subdomain_3 = F.VolumeSubdomain1D(
         id=3, borders=[L / 2, 3 * L / 4], material=my_mat
     )
-    my_subdomain_4 = F.VolumeSubdomain1D(
-        id=4, borders=[3 * L / 4, L], material=my_mat)
+    my_subdomain_4 = F.VolumeSubdomain1D(id=4, borders=[3 * L / 4, L], material=my_mat)
     left_surface = F.SurfaceSubdomain1D(id=1, x=0)
     right_surface = F.SurfaceSubdomain1D(id=2, x=L)
     my_model.subdomains = [
diff --git a/test/test_vtx.py b/test/test_vtx.py
index 44cabc859..cfd37ca79 100644
--- a/test/test_vtx.py
+++ b/test/test_vtx.py
@@ -42,7 +42,7 @@ def test_vtx_export_two_functions(tmpdir):
 
 @pytest.mark.skip(reason="Not implemented")
 def test_vtx_export_subdomain():
-    """Test that given multiple subdomains in problem, 
+    """Test that given multiple subdomains in problem,
     only correct functions are extracted from species"""
     pass
 
diff --git a/test/test_xdmf.py b/test/test_xdmf.py
index 23a6b2581..2fdf7d120 100644
--- a/test/test_xdmf.py
+++ b/test/test_xdmf.py
@@ -44,8 +44,7 @@ def test_integration_with_HTransportProblem(tmp_path):
     my_model.temperature = 500.0
     my_model.species = [F.Species("H")]
     filename = os.path.join(tmp_path, "test.xdmf")
-    my_model.exports = [F.XDMFExport(
-        filename=filename, field=my_model.species)]
+    my_model.exports = [F.XDMFExport(filename=filename, field=my_model.species)]
 
     my_model.settings = F.Settings(atol=1, rtol=0.1, final_time=1)
     my_model.settings.stepsize = F.Stepsize(initial_value=0.5)

From 28c66f1412be62e33f478e3b4d007a6faadd7975 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 18:26:47 -0400
Subject: [PATCH 099/134] Sort imports

---
 convergence_rates.py                          | 10 +--
 examples/multi_isotope_trapping_example.py    |  2 +-
 examples/multi_material_1d.py                 |  1 +
 examples/multi_material_2d.py                 |  6 +-
 examples/multi_material_2d_bis.py             |  6 +-
 examples/multi_material_transient.py          |  1 +
 examples/multi_material_with_one_material.py  |  1 +
 examples/tds_example.py                       |  2 +-
 examples/trapping_example.py                  |  4 +-
 festim/__init__.py                            | 72 ++++++++-----------
 festim/boundary_conditions/dirichlet_bc.py    |  3 +-
 festim/boundary_conditions/flux_bc.py         |  3 +-
 festim/exports/average_surface.py             |  1 -
 festim/exports/average_volume.py              |  1 -
 festim/exports/maximum_surface.py             |  3 +-
 festim/exports/maximum_volume.py              |  3 +-
 festim/exports/minimum_surface.py             |  3 +-
 festim/exports/minimum_volume.py              |  3 +-
 festim/exports/surface_flux.py                |  2 +-
 festim/exports/surface_quantity.py            |  3 +-
 festim/exports/total_surface.py               |  2 +-
 festim/exports/total_volume.py                |  2 +-
 festim/exports/volume_quantity.py             |  3 +-
 festim/exports/vtx.py                         |  5 +-
 festim/exports/xdmf.py                        |  5 +-
 festim/heat_transfer_problem.py               |  6 +-
 festim/helpers_discontinuity.py               |  3 +-
 festim/hydrogen_transport_problem.py          | 17 +++--
 festim/initial_condition.py                   |  3 +-
 festim/material.py                            |  2 +-
 festim/mesh/mesh.py                           |  7 +-
 festim/mesh/mesh_1d.py                        |  8 ++-
 festim/mesh/mesh_from_xdmf.py                 |  3 +-
 festim/problem.py                             | 11 +--
 festim/reaction.py                            | 15 ++--
 festim/source.py                              |  5 +-
 festim/stepsize.py                            |  3 -
 festim/subdomain/interface.py                 |  5 +-
 festim/subdomain/surface_subdomain.py         |  3 +-
 festim/subdomain/surface_subdomain_1d.py      |  3 +-
 festim/subdomain/volume_subdomain.py          |  5 +-
 festim/subdomain/volume_subdomain_1d.py       |  3 +-
 festim/trap.py                                |  3 +-
 pyproject.toml                                | 42 +++++++++--
 test/benchmark.py                             | 36 +++++-----
 test/system_tests/test_1_mobile_1_trap.py     | 12 ++--
 test/system_tests/test_1_mobile_species.py    | 11 +--
 test/system_tests/test_fluxes.py              |  6 +-
 test/system_tests/test_heat_transfer.py       |  6 +-
 test/system_tests/test_multi_material.py      |  7 +-
 .../test_reaction_not_in_every_volume.py      |  3 +-
 test/system_tests/tools.py                    |  6 +-
 test/test_average_surface.py                  |  7 +-
 test/test_average_volume.py                   |  7 +-
 test/test_complex_reaction.py                 |  3 +-
 test/test_derived_quantities.py               |  6 +-
 test/test_dirichlet_bc.py                     |  9 +--
 test/test_flux_bc.py                          |  9 +--
 test/test_h_transport_problem.py              | 12 ++--
 test/test_heat_transfer_problem.py            | 18 ++---
 test/test_helpers.py                          |  5 +-
 test/test_henrysbc.py                         | 10 +--
 test/test_initial_condition.py                |  6 +-
 test/test_material.py                         |  3 +-
 test/test_maximum_surface.py                  |  3 +-
 test/test_maximum_volume.py                   |  5 +-
 test/test_mesh.py                             | 11 +--
 test/test_minimum_surface.py                  |  7 +-
 test/test_minimum_volume.py                   |  5 +-
 test/test_multispecies_problem.py             |  4 +-
 test/test_permeation_problem.py               | 10 +--
 test/test_reaction.py                         |  6 +-
 test/test_settings.py                         |  3 +-
 test/test_sievertsbc.py                       | 10 +--
 test/test_source.py                           |  9 +--
 test/test_species.py                          | 12 ++--
 test/test_stepsize.py                         |  3 +-
 test/test_subdomains.py                       |  6 +-
 test/test_surface_quantity.py                 | 10 +--
 test/test_total_surface.py                    |  4 +-
 test/test_total_volume.py                     |  9 +--
 test/test_volume_quantity.py                  |  1 +
 test/test_vtx.py                              |  7 +-
 test/test_xdmf.py                             | 11 +--
 84 files changed, 358 insertions(+), 254 deletions(-)

diff --git a/convergence_rates.py b/convergence_rates.py
index 037270442..0e2e140a6 100644
--- a/convergence_rates.py
+++ b/convergence_rates.py
@@ -1,9 +1,11 @@
-import festim as F
-import numpy as np
-from dolfinx import fem
-import ufl
 import mpi4py.MPI as MPI
+
 import matplotlib.pyplot as plt
+import numpy as np
+
+import festim as F
+import ufl
+from dolfinx import fem
 
 
 def source_from_exact_solution(
diff --git a/examples/multi_isotope_trapping_example.py b/examples/multi_isotope_trapping_example.py
index 9312e62c9..209929f71 100644
--- a/examples/multi_isotope_trapping_example.py
+++ b/examples/multi_isotope_trapping_example.py
@@ -1,6 +1,6 @@
-import festim as F
 import numpy as np
 
+import festim as F
 
 my_model = F.HydrogenTransportProblem()
 
diff --git a/examples/multi_material_1d.py b/examples/multi_material_1d.py
index 76221480e..1cd22ec29 100644
--- a/examples/multi_material_1d.py
+++ b/examples/multi_material_1d.py
@@ -1,4 +1,5 @@
 import numpy as np
+
 import festim as F
 
 my_model = F.HTransportProblemDiscontinuous()
diff --git a/examples/multi_material_2d.py b/examples/multi_material_2d.py
index bc55a8d95..94fc4ab15 100644
--- a/examples/multi_material_2d.py
+++ b/examples/multi_material_2d.py
@@ -1,9 +1,9 @@
-import festim as F
-
 from mpi4py import MPI
+
+import numpy as np
+
 import dolfinx
 import dolfinx.fem.petsc
-import numpy as np
 import festim as F
 import ufl
 
diff --git a/examples/multi_material_2d_bis.py b/examples/multi_material_2d_bis.py
index a1f022a78..86429615a 100644
--- a/examples/multi_material_2d_bis.py
+++ b/examples/multi_material_2d_bis.py
@@ -1,9 +1,9 @@
-import festim as F
-
 from mpi4py import MPI
+
+import numpy as np
+
 import dolfinx
 import dolfinx.fem.petsc
-import numpy as np
 import festim as F
 import ufl
 
diff --git a/examples/multi_material_transient.py b/examples/multi_material_transient.py
index 4469625b9..423a4deb4 100644
--- a/examples/multi_material_transient.py
+++ b/examples/multi_material_transient.py
@@ -1,4 +1,5 @@
 import numpy as np
+
 import festim as F
 
 my_model = F.HTransportProblemDiscontinuous()
diff --git a/examples/multi_material_with_one_material.py b/examples/multi_material_with_one_material.py
index 7bf502bf1..497b2dce4 100644
--- a/examples/multi_material_with_one_material.py
+++ b/examples/multi_material_with_one_material.py
@@ -1,4 +1,5 @@
 import numpy as np
+
 import festim as F
 
 my_model = F.HTransportProblemDiscontinuous()
diff --git a/examples/tds_example.py b/examples/tds_example.py
index fa60689c7..7f7d7bcbb 100644
--- a/examples/tds_example.py
+++ b/examples/tds_example.py
@@ -1,6 +1,6 @@
-import festim as F
 import numpy as np
 
+import festim as F
 
 my_model = F.HydrogenTransportProblem()
 
diff --git a/examples/trapping_example.py b/examples/trapping_example.py
index adc1eae85..57d339907 100644
--- a/examples/trapping_example.py
+++ b/examples/trapping_example.py
@@ -1,7 +1,9 @@
 from petsc4py import PETSc
-import festim as F
+
 import numpy as np
 
+import festim as F
+
 L = 3e-04
 vertices = np.linspace(0, L, num=1000)
 
diff --git a/festim/__init__.py b/festim/__init__.py
index f67fb779f..991508e09 100644
--- a/festim/__init__.py
+++ b/festim/__init__.py
@@ -9,60 +9,48 @@
 R = 8.314462618  # Gas constant J.mol-1.K-1
 k_B = 8.6173303e-5  # Boltzmann constant eV.K-1
 
-from .helpers import as_fenics_constant
-from .problem import ProblemBase
-
 from .boundary_conditions.dirichlet_bc import (
-    DirichletBCBase,
     DirichletBC,
-    FixedTemperatureBC,
+    DirichletBCBase,
     FixedConcentrationBC,
+    FixedTemperatureBC,
 )
-from .boundary_conditions.sieverts_bc import SievertsBC
+from .boundary_conditions.flux_bc import FluxBCBase, HeatFluxBC, ParticleFluxBC
 from .boundary_conditions.henrys_bc import HenrysBC
-from .boundary_conditions.flux_bc import FluxBCBase, ParticleFluxBC, HeatFluxBC
-
+from .boundary_conditions.sieverts_bc import SievertsBC
+from .exports.average_surface import AverageSurface
+from .exports.average_volume import AverageVolume
+from .exports.maximum_surface import MaximumSurface
+from .exports.maximum_volume import MaximumVolume
+from .exports.minimum_surface import MinimumSurface
+from .exports.minimum_volume import MinimumVolume
+from .exports.surface_flux import SurfaceFlux
+from .exports.surface_quantity import SurfaceQuantity
+from .exports.total_surface import TotalSurface
+from .exports.total_volume import TotalVolume
+from .exports.volume_quantity import VolumeQuantity
+from .exports.vtx import VTXSpeciesExport, VTXTemperatureExport
+from .exports.xdmf import XDMFExport
+from .heat_transfer_problem import HeatTransferProblem
+from .helpers import as_fenics_constant
+from .hydrogen_transport_problem import (
+    HTransportProblemDiscontinuous,
+    HydrogenTransportProblem,
+)
+from .initial_condition import InitialCondition, InitialTemperature
 from .material import Material
-
 from .mesh.mesh import Mesh
 from .mesh.mesh_1d import Mesh1D
 from .mesh.mesh_from_xdmf import MeshFromXDMF
-
-
-from .initial_condition import InitialCondition, InitialTemperature
-
+from .problem import ProblemBase
+from .reaction import Reaction
 from .settings import Settings
-
-from .source import ParticleSource, HeatSource, SourceBase
-
+from .source import HeatSource, ParticleSource, SourceBase
+from .species import ImplicitSpecies, Species, find_species_from_name
+from .stepsize import Stepsize
+from .subdomain.interface import Interface
 from .subdomain.surface_subdomain import SurfaceSubdomain, find_surface_from_id
 from .subdomain.surface_subdomain_1d import SurfaceSubdomain1D
 from .subdomain.volume_subdomain import VolumeSubdomain, find_volume_from_id
 from .subdomain.volume_subdomain_1d import VolumeSubdomain1D
-from .subdomain.interface import Interface
-
-from .species import Species, ImplicitSpecies, find_species_from_name
 from .trap import Trap
-from .stepsize import Stepsize
-
-from .exports.surface_quantity import SurfaceQuantity
-from .exports.volume_quantity import VolumeQuantity
-from .exports.total_volume import TotalVolume
-from .exports.average_volume import AverageVolume
-from .exports.maximum_volume import MaximumVolume
-from .exports.minimum_volume import MinimumVolume
-from .exports.total_surface import TotalSurface
-from .exports.maximum_surface import MaximumSurface
-from .exports.minimum_surface import MinimumSurface
-from .exports.average_surface import AverageSurface
-from .exports.surface_flux import SurfaceFlux
-from .exports.vtx import VTXSpeciesExport, VTXTemperatureExport
-from .exports.xdmf import XDMFExport
-
-from .reaction import Reaction
-
-from .hydrogen_transport_problem import (
-    HydrogenTransportProblem,
-    HTransportProblemDiscontinuous,
-)
-from .heat_transfer_problem import HeatTransferProblem
diff --git a/festim/boundary_conditions/dirichlet_bc.py b/festim/boundary_conditions/dirichlet_bc.py
index 8281cabc3..2bae86a6f 100644
--- a/festim/boundary_conditions/dirichlet_bc.py
+++ b/festim/boundary_conditions/dirichlet_bc.py
@@ -1,7 +1,8 @@
+import numpy as np
+
 import festim as F
 import ufl
 from dolfinx import fem
-import numpy as np
 
 
 class DirichletBCBase:
diff --git a/festim/boundary_conditions/flux_bc.py b/festim/boundary_conditions/flux_bc.py
index 4699c7391..c7243b972 100644
--- a/festim/boundary_conditions/flux_bc.py
+++ b/festim/boundary_conditions/flux_bc.py
@@ -1,7 +1,8 @@
+import numpy as np
+
 import festim as F
 import ufl
 from dolfinx import fem
-import numpy as np
 
 
 class FluxBCBase:
diff --git a/festim/exports/average_surface.py b/festim/exports/average_surface.py
index ffcc14a08..51601cf25 100644
--- a/festim/exports/average_surface.py
+++ b/festim/exports/average_surface.py
@@ -1,6 +1,5 @@
 import festim as F
 from dolfinx import fem
-import numpy as np
 
 
 class AverageSurface(F.SurfaceQuantity):
diff --git a/festim/exports/average_volume.py b/festim/exports/average_volume.py
index 7295044c4..53de6f794 100644
--- a/festim/exports/average_volume.py
+++ b/festim/exports/average_volume.py
@@ -1,6 +1,5 @@
 import festim as F
 from dolfinx import fem
-import numpy as np
 
 
 class AverageVolume(F.VolumeQuantity):
diff --git a/festim/exports/maximum_surface.py b/festim/exports/maximum_surface.py
index bcca2f49e..29da50d1f 100644
--- a/festim/exports/maximum_surface.py
+++ b/festim/exports/maximum_surface.py
@@ -1,6 +1,7 @@
-import festim as F
 import numpy as np
 
+import festim as F
+
 
 class MaximumSurface(F.SurfaceQuantity):
     """Computes the maximum value of a field on a given surface
diff --git a/festim/exports/maximum_volume.py b/festim/exports/maximum_volume.py
index c582c4907..3a26cacf3 100644
--- a/festim/exports/maximum_volume.py
+++ b/festim/exports/maximum_volume.py
@@ -1,6 +1,7 @@
-import festim as F
 import numpy as np
 
+import festim as F
+
 
 class MaximumVolume(F.VolumeQuantity):
     """Computes the maximum value of a field in a given volume
diff --git a/festim/exports/minimum_surface.py b/festim/exports/minimum_surface.py
index 96dd5e6e8..120ba8991 100644
--- a/festim/exports/minimum_surface.py
+++ b/festim/exports/minimum_surface.py
@@ -1,6 +1,7 @@
-import festim as F
 import numpy as np
 
+import festim as F
+
 
 class MinimumSurface(F.SurfaceQuantity):
     """Computes the minimum value of a field on a given surface
diff --git a/festim/exports/minimum_volume.py b/festim/exports/minimum_volume.py
index 0bd95a7dc..a2e7c14b6 100644
--- a/festim/exports/minimum_volume.py
+++ b/festim/exports/minimum_volume.py
@@ -1,6 +1,7 @@
-import festim as F
 import numpy as np
 
+import festim as F
+
 
 class MinimumVolume(F.VolumeQuantity):
     """Computes the minmum value of a field in a given volume
diff --git a/festim/exports/surface_flux.py b/festim/exports/surface_flux.py
index 43b0cf6a1..aa68f7ad0 100644
--- a/festim/exports/surface_flux.py
+++ b/festim/exports/surface_flux.py
@@ -1,6 +1,6 @@
-from dolfinx import fem
 import festim as F
 import ufl
+from dolfinx import fem
 
 
 class SurfaceFlux(F.SurfaceQuantity):
diff --git a/festim/exports/surface_quantity.py b/festim/exports/surface_quantity.py
index 32646ad9d..d13f93268 100644
--- a/festim/exports/surface_quantity.py
+++ b/festim/exports/surface_quantity.py
@@ -1,6 +1,7 @@
-import festim as F
 import csv
 
+import festim as F
+
 
 class SurfaceQuantity:
     """Export SurfaceQuantity
diff --git a/festim/exports/total_surface.py b/festim/exports/total_surface.py
index 56807f3ef..1d535583e 100644
--- a/festim/exports/total_surface.py
+++ b/festim/exports/total_surface.py
@@ -1,6 +1,6 @@
 import festim as F
-from dolfinx import fem
 import ufl
+from dolfinx import fem
 
 
 class TotalSurface(F.SurfaceQuantity):
diff --git a/festim/exports/total_volume.py b/festim/exports/total_volume.py
index ac27170d6..e16a96c4a 100644
--- a/festim/exports/total_volume.py
+++ b/festim/exports/total_volume.py
@@ -1,6 +1,6 @@
 import festim as F
-from dolfinx import fem
 import ufl
+from dolfinx import fem
 
 
 class TotalVolume(F.VolumeQuantity):
diff --git a/festim/exports/volume_quantity.py b/festim/exports/volume_quantity.py
index 60118e9b8..37ccc4edb 100644
--- a/festim/exports/volume_quantity.py
+++ b/festim/exports/volume_quantity.py
@@ -1,7 +1,8 @@
-import festim as F
 import csv
 import os
 
+import festim as F
+
 
 class VolumeQuantity:
     """Export VolumeQuantity
diff --git a/festim/exports/vtx.py b/festim/exports/vtx.py
index df3fd2df6..ec5ccf648 100644
--- a/festim/exports/vtx.py
+++ b/festim/exports/vtx.py
@@ -1,8 +1,9 @@
-from dolfinx.fem import Function as _Function
+import warnings
 from pathlib import Path
+
+from dolfinx.fem import Function as _Function
 from festim.species import Species as _Species
 from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
-import warnings
 
 
 class ExportBaseClass:
diff --git a/festim/exports/xdmf.py b/festim/exports/xdmf.py
index e2956caf2..d813dbe84 100644
--- a/festim/exports/xdmf.py
+++ b/festim/exports/xdmf.py
@@ -1,7 +1,10 @@
+from pathlib import Path
+
 import mpi4py
+
 from dolfinx.io import XDMFFile
-from pathlib import Path
 from festim.species import Species as _Species
+
 from .vtx import ExportBaseClass
 
 
diff --git a/festim/heat_transfer_problem.py b/festim/heat_transfer_problem.py
index 9bad85e08..cd5b0b029 100644
--- a/festim/heat_transfer_problem.py
+++ b/festim/heat_transfer_problem.py
@@ -1,8 +1,8 @@
-from dolfinx import fem
-from dolfinx.io import VTXWriter
 import basix
-import ufl
 import festim as F
+import ufl
+from dolfinx import fem
+from dolfinx.io import VTXWriter
 
 
 class HeatTransferProblem(F.ProblemBase):
diff --git a/festim/helpers_discontinuity.py b/festim/helpers_discontinuity.py
index 2c225a522..1b3aacf84 100644
--- a/festim/helpers_discontinuity.py
+++ b/festim/helpers_discontinuity.py
@@ -1,6 +1,7 @@
-import dolfinx
 import numpy as np
 
+import dolfinx
+
 
 def transfer_meshtags_to_submesh(
     mesh, entity_tag, submesh, sub_vertex_to_parent, sub_cell_to_parent
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 7a3be07d4..178631578 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -1,14 +1,17 @@
-import dolfinx
-from dolfinx import fem
-import basix
-import ufl
+from typing import Callable
+
 from mpi4py import MPI
+
 import numpy as np
+import numpy.typing as npt
 import tqdm.autonotebook
-import festim as F
 from scifem import NewtonSolver
-from typing import Callable
-import numpy.typing as npt
+
+import basix
+import dolfinx
+import festim as F
+import ufl
+from dolfinx import fem
 
 
 class HydrogenTransportProblem(F.ProblemBase):
diff --git a/festim/initial_condition.py b/festim/initial_condition.py
index 189fe9013..6270b59e3 100644
--- a/festim/initial_condition.py
+++ b/festim/initial_condition.py
@@ -1,6 +1,7 @@
+import numpy as np
+
 import ufl
 from dolfinx import fem
-import numpy as np
 
 
 # TODO rename this to InitialConcentration and create a new base class
diff --git a/festim/material.py b/festim/material.py
index a698b0588..8cbf3cf62 100644
--- a/festim/material.py
+++ b/festim/material.py
@@ -1,5 +1,5 @@
-import ufl
 import festim as F
+import ufl
 
 
 class Material:
diff --git a/festim/mesh/mesh.py b/festim/mesh/mesh.py
index 74fe486f5..bf4d53ef6 100644
--- a/festim/mesh/mesh.py
+++ b/festim/mesh/mesh.py
@@ -1,6 +1,7 @@
-import ufl
-import dolfinx
 import numpy as np
+
+import dolfinx
+import ufl
 from dolfinx.mesh import meshtags
 
 
@@ -41,7 +42,7 @@ def mesh(self, value):
         if isinstance(value, dolfinx.mesh.Mesh):
             self._mesh = value
         else:
-            raise TypeError(f"Mesh must be of type dolfinx.mesh.Mesh")
+            raise TypeError("Mesh must be of type dolfinx.mesh.Mesh")
 
     @property
     def vdim(self):
diff --git a/festim/mesh/mesh_1d.py b/festim/mesh/mesh_1d.py
index 37c98d126..322ead5a0 100644
--- a/festim/mesh/mesh_1d.py
+++ b/festim/mesh/mesh_1d.py
@@ -1,9 +1,11 @@
-import dolfinx.mesh
 from mpi4py import MPI
-import basix.ufl
-import ufl
+
 import numpy as np
+
+import basix.ufl
+import dolfinx.mesh
 import festim as F
+import ufl
 
 
 class Mesh1D(F.Mesh):
diff --git a/festim/mesh/mesh_from_xdmf.py b/festim/mesh/mesh_from_xdmf.py
index b4f9b4092..4b744554d 100644
--- a/festim/mesh/mesh_from_xdmf.py
+++ b/festim/mesh/mesh_from_xdmf.py
@@ -1,6 +1,7 @@
-from dolfinx.io import XDMFFile
 from mpi4py import MPI
+
 import festim as F
+from dolfinx.io import XDMFFile
 
 
 class MeshFromXDMF(F.Mesh):
diff --git a/festim/problem.py b/festim/problem.py
index 99d4f71bc..350ca1285 100644
--- a/festim/problem.py
+++ b/festim/problem.py
@@ -1,14 +1,17 @@
-from dolfinx import fem
-from dolfinx.nls.petsc import NewtonSolver
-import ufl
+from typing import Any
+
 from mpi4py import MPI
+
 import numpy as np
 import tqdm.autonotebook
+
 import festim as F
+import ufl
+from dolfinx import fem
+from dolfinx.nls.petsc import NewtonSolver
 from festim.mesh.mesh import Mesh as _Mesh
 from festim.source import SourceBase as _SourceBase
 from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
-from typing import Any
 
 
 class ProblemBase:
diff --git a/festim/reaction.py b/festim/reaction.py
index 730065dc0..23cdac11b 100644
--- a/festim/reaction.py
+++ b/festim/reaction.py
@@ -1,9 +1,10 @@
-from typing import Union, Optional, List
+from typing import List, Optional, Union
 
-from ufl import exp
-from festim.species import Species as _Species, ImplicitSpecies as _ImplicitSpecies
-from festim.subdomain.volume_subdomain_1d import VolumeSubdomain1D as VS1D
 from festim import k_B as _k_B
+from festim.species import ImplicitSpecies as _ImplicitSpecies
+from festim.species import Species as _Species
+from festim.subdomain.volume_subdomain_1d import VolumeSubdomain1D as VS1D
+from ufl import exp
 
 
 class Reaction:
@@ -73,7 +74,7 @@ def reactant(self, value):
             value = [value]
         if len(value) == 0:
             raise ValueError(
-                f"reactant must be an entry of one or more species objects, not an empty list."
+                "reactant must be an entry of one or more species objects, not an empty list."
             )
         for i in value:
             if not isinstance(i, (_Species, _ImplicitSpecies)):
@@ -121,11 +122,11 @@ def reaction_term(self, temperature):
         else:
             if self.p_0 == None:
                 raise ValueError(
-                    f"p_0 cannot be None when reaction products are present."
+                    "p_0 cannot be None when reaction products are present."
                 )
             elif self.E_p == None:
                 raise ValueError(
-                    f"E_p cannot be None when reaction products are present."
+                    "E_p cannot be None when reaction products are present."
                 )
 
         k = self.k_0 * exp(-self.E_k / (_k_B * temperature))
diff --git a/festim/source.py b/festim/source.py
index 65b641c83..d53ebf315 100644
--- a/festim/source.py
+++ b/festim/source.py
@@ -1,8 +1,9 @@
+import numpy as np
+
+import dolfinx
 import festim as F
 import ufl
-import dolfinx
 from dolfinx import fem
-import numpy as np
 
 
 class SourceBase:
diff --git a/festim/stepsize.py b/festim/stepsize.py
index 5ca0ff2c3..aab36a7c2 100644
--- a/festim/stepsize.py
+++ b/festim/stepsize.py
@@ -1,6 +1,3 @@
-import festim as F
-
-
 class Stepsize:
     """
     A class for evaluating the stepsize of transient simulations.
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index 78bd2b2be..3c3a4aac3 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -1,7 +1,8 @@
-import festim as F
+import numpy as np
+
 import dolfinx
+import festim as F
 from dolfinx.cpp.fem import compute_integration_domains
-import numpy as np
 
 
 class Interface:
diff --git a/festim/subdomain/surface_subdomain.py b/festim/subdomain/surface_subdomain.py
index 578545804..b28b550c0 100644
--- a/festim/subdomain/surface_subdomain.py
+++ b/festim/subdomain/surface_subdomain.py
@@ -1,6 +1,7 @@
-import dolfinx.mesh
 import numpy as np
 
+import dolfinx.mesh
+
 
 class SurfaceSubdomain:
     """
diff --git a/festim/subdomain/surface_subdomain_1d.py b/festim/subdomain/surface_subdomain_1d.py
index 79a0a7726..e757914ff 100644
--- a/festim/subdomain/surface_subdomain_1d.py
+++ b/festim/subdomain/surface_subdomain_1d.py
@@ -1,5 +1,6 @@
-import dolfinx.mesh
 import numpy as np
+
+import dolfinx.mesh
 import festim as F
 
 
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index 4d8b055c9..fb23bf520 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -1,6 +1,7 @@
-from dolfinx.mesh import locate_entities
-import dolfinx
 import numpy as np
+
+import dolfinx
+from dolfinx.mesh import locate_entities
 from festim.helpers_discontinuity import transfer_meshtags_to_submesh
 
 
diff --git a/festim/subdomain/volume_subdomain_1d.py b/festim/subdomain/volume_subdomain_1d.py
index 630c38865..f0b4a0acd 100644
--- a/festim/subdomain/volume_subdomain_1d.py
+++ b/festim/subdomain/volume_subdomain_1d.py
@@ -1,6 +1,7 @@
-from dolfinx.mesh import locate_entities
 import numpy as np
+
 import festim as F
+from dolfinx.mesh import locate_entities
 
 
 class VolumeSubdomain1D(F.VolumeSubdomain):
diff --git a/festim/trap.py b/festim/trap.py
index 9bb3ca6dc..ac7b45831 100644
--- a/festim/trap.py
+++ b/festim/trap.py
@@ -1,5 +1,6 @@
-from festim.species import Species as _Species, ImplicitSpecies as _ImplicitSpecies
 from festim.reaction import Reaction as _Reaction
+from festim.species import ImplicitSpecies as _ImplicitSpecies
+from festim.species import Species as _Species
 
 
 class Trap(_Species):
diff --git a/pyproject.toml b/pyproject.toml
index a7ce20481..82c76bf1c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,10 +1,42 @@
 [build-system]
-requires = [
-    "setuptools >= 45",
-    "wheel",
-    "setuptools_scm[toml] >= 7.0.5"
-]
+requires = ["setuptools >= 45", "wheel", "setuptools_scm[toml] >= 7.0.5"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools_scm]
 write_to = "festim/_version.py"
+
+[tool.ruff]
+line-length = 88
+indent-width = 4
+
+
+[tool.ruff.lint]
+select = [
+    "E",      # pycodestyle
+    "W",      # pycodestyle
+    "F",      # pyflakes
+    "I",      # isort - use standalone isort
+    "RUF",    # Ruff-specific rules
+    "UP",     # pyupgrade
+    "ICN",    # flake8-import-conventions
+    "NPY",    # numpy-specific rules
+    "FLY",    # use f-string not static joins
+    "NPY201", # numpy 2.x ruleset
+]
+ignore = ["UP007", "RUF012"]
+allowed-confusables = ["σ"]
+
+[tool.ruff.lint.isort]
+known-first-party = ["festim", "basix", "dolfinx", "ffcx", "ufl"]
+known-third-party = ["gmsh", "numba", "numpy", "pytest", "pyvista", "pyamg"]
+section-order = [
+    "future",
+    "standard-library",
+    "mpi",
+    "third-party",
+    "first-party",
+    "local-folder",
+]
+
+[tool.ruff.lint.isort.sections]
+"mpi" = ["mpi4py", "petsc4py"]
diff --git a/test/benchmark.py b/test/benchmark.py
index 5bbe89457..0fbdec4a5 100644
--- a/test/benchmark.py
+++ b/test/benchmark.py
@@ -1,36 +1,38 @@
+import time
+
 from mpi4py import MPI
 from petsc4py import PETSc
-from dolfinx.io import XDMFFile
+
+import numpy as np
+import tqdm.autonotebook
+from test_permeation_problem import test_permeation_problem
+
+import basix
 from dolfinx.fem import (
     Constant,
-    dirichletbc,
     Function,
+    assemble_scalar,
+    dirichletbc,
+    form,
     functionspace,
-    locate_dofs_topological,
     locate_dofs_geometrical,
-    form,
-    assemble_scalar,
+    locate_dofs_topological,
 )
 from dolfinx.fem.petsc import (
     NonlinearProblem,
 )
+from dolfinx.io import XDMFFile
+from dolfinx.mesh import create_mesh, locate_entities, meshtags
 from dolfinx.nls.petsc import NewtonSolver
 from ufl import (
-    dot,
-    grad,
-    TestFunction,
-    exp,
     FacetNormal,
-    Cell,
-    Mesh,
     Measure,
+    Mesh,
+    TestFunction,
+    dot,
+    exp,
+    grad,
 )
-import basix
-from dolfinx.mesh import create_mesh, meshtags, locate_entities
-import numpy as np
-import tqdm.autonotebook
-import time
-from test_permeation_problem import test_permeation_problem
 
 
 def fenics_test_permeation_problem(mesh_size=1001):
diff --git a/test/system_tests/test_1_mobile_1_trap.py b/test/system_tests/test_1_mobile_1_trap.py
index c8cfb344d..98a1e5ad5 100644
--- a/test/system_tests/test_1_mobile_1_trap.py
+++ b/test/system_tests/test_1_mobile_1_trap.py
@@ -1,11 +1,13 @@
-import festim as F
+from mpi4py import MPI
+
 import numpy as np
-from dolfinx import fem
+
+import festim as F
 import ufl
-from .tools import error_L2
-from dolfinx.mesh import create_unit_square, create_unit_cube, locate_entities
-from mpi4py import MPI
+from dolfinx import fem
+from dolfinx.mesh import create_unit_cube, create_unit_square, locate_entities
 
+from .tools import error_L2
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
 test_mesh_2d = create_unit_square(MPI.COMM_WORLD, 50, 50)
diff --git a/test/system_tests/test_1_mobile_species.py b/test/system_tests/test_1_mobile_species.py
index 27e574107..64da01ac8 100644
--- a/test/system_tests/test_1_mobile_species.py
+++ b/test/system_tests/test_1_mobile_species.py
@@ -1,10 +1,13 @@
-import festim as F
+from mpi4py import MPI
+
 import numpy as np
-from dolfinx import fem
+
+import festim as F
 import ufl
+from dolfinx import fem
+from dolfinx.mesh import create_unit_cube, create_unit_square, locate_entities
+
 from .tools import error_L2
-from dolfinx.mesh import create_unit_square, create_unit_cube, locate_entities
-from mpi4py import MPI
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
 test_mesh_2d = create_unit_square(MPI.COMM_WORLD, 50, 50)
diff --git a/test/system_tests/test_fluxes.py b/test/system_tests/test_fluxes.py
index abecbf391..a01bfacb9 100644
--- a/test/system_tests/test_fluxes.py
+++ b/test/system_tests/test_fluxes.py
@@ -1,7 +1,9 @@
-import festim as F
 import numpy as np
-from dolfinx import fem
+
+import festim as F
 import ufl
+from dolfinx import fem
+
 from .tools import error_L2
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
diff --git a/test/system_tests/test_heat_transfer.py b/test/system_tests/test_heat_transfer.py
index f9cb1654e..9486b73b4 100644
--- a/test/system_tests/test_heat_transfer.py
+++ b/test/system_tests/test_heat_transfer.py
@@ -1,7 +1,9 @@
-import festim as F
 import numpy as np
-from dolfinx import fem
+
+import festim as F
 import ufl
+from dolfinx import fem
+
 from .tools import error_L2
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index 02478ea9b..35a8e6233 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -1,11 +1,12 @@
-import festim as F
-
 from mpi4py import MPI
+
+import numpy as np
+
 import dolfinx
 import dolfinx.fem.petsc
-import numpy as np
 import festim as F
 import ufl
+
 from .tools import error_L2
 
 
diff --git a/test/system_tests/test_reaction_not_in_every_volume.py b/test/system_tests/test_reaction_not_in_every_volume.py
index ea315458a..b1cdea67b 100644
--- a/test/system_tests/test_reaction_not_in_every_volume.py
+++ b/test/system_tests/test_reaction_not_in_every_volume.py
@@ -1,6 +1,7 @@
-import festim as F
 import numpy as np
 
+import festim as F
+
 
 def test_sim_reaction_not_in_every_volume():
     """Tests that a steady simulation can be run if a reaction is not defined
diff --git a/test/system_tests/tools.py b/test/system_tests/tools.py
index 171afe7d8..9d13aa670 100644
--- a/test/system_tests/tools.py
+++ b/test/system_tests/tools.py
@@ -1,7 +1,9 @@
+import mpi4py.MPI as MPI
+
+import numpy as np
+
 import ufl
 from dolfinx import fem
-import numpy as np
-import mpi4py.MPI as MPI
 
 
 def error_L2(u_computed, u_exact, degree_raise=3):
diff --git a/test/test_average_surface.py b/test/test_average_surface.py
index 142bf7cfc..131ff6d6a 100644
--- a/test/test_average_surface.py
+++ b/test/test_average_surface.py
@@ -1,9 +1,8 @@
-import festim as F
 import numpy as np
-from dolfinx import fem
+
+import festim as F
 import ufl
-from dolfinx.mesh import meshtags
-import pytest
+from dolfinx import fem
 
 
 def test_average_surface_compute_1D():
diff --git a/test/test_average_volume.py b/test/test_average_volume.py
index aac2832b3..f9370c9c5 100644
--- a/test/test_average_volume.py
+++ b/test/test_average_volume.py
@@ -1,9 +1,8 @@
-import festim as F
 import numpy as np
-from dolfinx import fem
+
+import festim as F
 import ufl
-from dolfinx.mesh import meshtags
-import pytest
+from dolfinx import fem
 
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy")
 
diff --git a/test/test_complex_reaction.py b/test/test_complex_reaction.py
index dfaaab6a7..36e9dd3d6 100644
--- a/test/test_complex_reaction.py
+++ b/test/test_complex_reaction.py
@@ -1,7 +1,8 @@
 import numpy as np
-import festim as F
 import pytest
 
+import festim as F
+
 
 def concentration_A_exact(t, c_A_0, k, p):
     """Analytical solution for the concentration of species A in a reaction A + B <-> C + D
diff --git a/test/test_derived_quantities.py b/test/test_derived_quantities.py
index 73667ca4d..edfd6ae37 100644
--- a/test/test_derived_quantities.py
+++ b/test/test_derived_quantities.py
@@ -1,7 +1,9 @@
-import festim as F
-import pytest
 import os
 
+import pytest
+
+import festim as F
+
 mobile_H = F.Species("H")
 mobile_D = F.Species("D")
 surf_1 = F.SurfaceSubdomain(id=1)
diff --git a/test/test_dirichlet_bc.py b/test/test_dirichlet_bc.py
index 8b933d8aa..76ecebb05 100644
--- a/test/test_dirichlet_bc.py
+++ b/test/test_dirichlet_bc.py
@@ -1,12 +1,13 @@
+from mpi4py import MPI
+
 import numpy as np
 import pytest
-import ufl
-from ufl.conditional import Conditional
 
-from dolfinx import fem
 import dolfinx.mesh
-from mpi4py import MPI
 import festim as F
+import ufl
+from dolfinx import fem
+from ufl.conditional import Conditional
 
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy_mat")
 
diff --git a/test/test_flux_bc.py b/test/test_flux_bc.py
index 5aa05c8bc..b6865f9ce 100644
--- a/test/test_flux_bc.py
+++ b/test/test_flux_bc.py
@@ -1,12 +1,13 @@
+from mpi4py import MPI
+
 import numpy as np
 import pytest
+
+import dolfinx.mesh
+import festim as F
 import ufl
-from ufl.conditional import Conditional
 import ufl.core
 from dolfinx import fem
-import dolfinx.mesh
-from mpi4py import MPI
-import festim as F
 
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy_mat")
 mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
diff --git a/test/test_h_transport_problem.py b/test/test_h_transport_problem.py
index bec6e2fd7..6a60bc658 100644
--- a/test/test_h_transport_problem.py
+++ b/test/test_h_transport_problem.py
@@ -1,11 +1,13 @@
-import festim as F
-import tqdm.autonotebook
 import mpi4py.MPI as MPI
-import dolfinx.mesh
-from dolfinx import fem, nls
-import ufl
+
 import numpy as np
 import pytest
+import tqdm.autonotebook
+
+import dolfinx.mesh
+import festim as F
+import ufl
+from dolfinx import fem, nls
 
 test_mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
 x = ufl.SpatialCoordinate(test_mesh.mesh)
diff --git a/test/test_heat_transfer_problem.py b/test/test_heat_transfer_problem.py
index dfc335f46..961da0a99 100644
--- a/test/test_heat_transfer_problem.py
+++ b/test/test_heat_transfer_problem.py
@@ -1,14 +1,16 @@
-import festim as F
-import numpy as np
-import dolfinx
-from dolfinx.io import XDMFFile
-from dolfinx import fem
-import ufl
+import os
+
 import mpi4py.MPI as MPI
-import tqdm.autonotebook
 
+import numpy as np
 import pytest
-import os
+import tqdm.autonotebook
+
+import dolfinx
+import festim as F
+import ufl
+from dolfinx import fem
+from dolfinx.io import XDMFFile
 
 
 def source_from_exact_solution(
diff --git a/test/test_helpers.py b/test/test_helpers.py
index 14b269bda..e9c70f138 100644
--- a/test/test_helpers.py
+++ b/test/test_helpers.py
@@ -1,8 +1,9 @@
-import festim as F
-from dolfinx import fem
 import numpy as np
 import pytest
+
+import festim as F
 import ufl
+from dolfinx import fem
 
 test_mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
 x = ufl.SpatialCoordinate(test_mesh.mesh)
diff --git a/test/test_henrysbc.py b/test/test_henrysbc.py
index cc2b0f89d..0bc4ca41b 100644
--- a/test/test_henrysbc.py
+++ b/test/test_henrysbc.py
@@ -1,10 +1,12 @@
+from mpi4py import MPI
+
+import numpy as np
+import pytest
+
+import dolfinx.mesh
 import festim as F
 import ufl
-import pytest
-import numpy as np
 from dolfinx import fem
-import dolfinx.mesh
-from mpi4py import MPI
 
 
 def henrys_law(T, H_0, E_H, pressure):
diff --git a/test/test_initial_condition.py b/test/test_initial_condition.py
index 64eabe54a..f8b6eb459 100644
--- a/test/test_initial_condition.py
+++ b/test/test_initial_condition.py
@@ -1,7 +1,9 @@
+from types import LambdaType
+
 import numpy as np
-import festim as F
 import pytest
-from types import LambdaType
+
+import festim as F
 from dolfinx import fem
 
 dummy_mat = F.Material(D_0=1, E_D=0.1, name="dummy_mat")
diff --git a/test/test_material.py b/test/test_material.py
index c897a7daa..29541d8fb 100644
--- a/test/test_material.py
+++ b/test/test_material.py
@@ -1,7 +1,8 @@
-import festim as F
 import numpy as np
 import pytest
 
+import festim as F
+
 test_mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
 
 
diff --git a/test/test_maximum_surface.py b/test/test_maximum_surface.py
index 7ba03ca61..2e5c0c649 100644
--- a/test/test_maximum_surface.py
+++ b/test/test_maximum_surface.py
@@ -1,5 +1,6 @@
-import festim as F
 import numpy as np
+
+import festim as F
 from dolfinx import fem
 
 
diff --git a/test/test_maximum_volume.py b/test/test_maximum_volume.py
index 3a275a358..6a4fff030 100644
--- a/test/test_maximum_volume.py
+++ b/test/test_maximum_volume.py
@@ -1,7 +1,6 @@
-import festim as F
 import numpy as np
-from dolfinx.mesh import meshtags
-import ufl
+
+import festim as F
 from dolfinx import fem
 
 
diff --git a/test/test_mesh.py b/test/test_mesh.py
index 3c9a3f95f..c0f9c0b43 100644
--- a/test/test_mesh.py
+++ b/test/test_mesh.py
@@ -1,9 +1,12 @@
-import festim as F
-from dolfinx import mesh as fenics_mesh
+import os
+
 from mpi4py import MPI
-import pytest
+
 import numpy as np
-import os
+import pytest
+
+import festim as F
+from dolfinx import mesh as fenics_mesh
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import meshtags
 
diff --git a/test/test_minimum_surface.py b/test/test_minimum_surface.py
index eb3706a40..9d4ec5467 100644
--- a/test/test_minimum_surface.py
+++ b/test/test_minimum_surface.py
@@ -1,10 +1,7 @@
-import festim as F
 import numpy as np
-import ufl
-from dolfinx.mesh import meshtags
+
+import festim as F
 from dolfinx import fem
-import pytest
-import os
 
 
 def test_minimum_surface_export_compute_1D():
diff --git a/test/test_minimum_volume.py b/test/test_minimum_volume.py
index 71f11a387..affa829aa 100644
--- a/test/test_minimum_volume.py
+++ b/test/test_minimum_volume.py
@@ -1,7 +1,6 @@
-import festim as F
 import numpy as np
-from dolfinx.mesh import meshtags
-import ufl
+
+import festim as F
 from dolfinx import fem
 
 
diff --git a/test/test_multispecies_problem.py b/test/test_multispecies_problem.py
index e1782d84c..552e24133 100644
--- a/test/test_multispecies_problem.py
+++ b/test/test_multispecies_problem.py
@@ -1,6 +1,8 @@
+from petsc4py import PETSc
+
 import numpy as np
+
 import festim as F
-from petsc4py import PETSc
 from dolfinx.fem import Constant
 from ufl import exp
 
diff --git a/test/test_permeation_problem.py b/test/test_permeation_problem.py
index 5226724f0..e2961cc74 100644
--- a/test/test_permeation_problem.py
+++ b/test/test_permeation_problem.py
@@ -1,10 +1,12 @@
+import os
+
 from petsc4py import PETSc
-from dolfinx.fem import Constant
-from ufl import exp
+
 import numpy as np
+
 import festim as F
-import os
-import ufl
+from dolfinx.fem import Constant
+from ufl import exp
 
 
 def relative_error_computed_to_analytical(
diff --git a/test/test_reaction.py b/test/test_reaction.py
index e1dc47617..6842f29f4 100644
--- a/test/test_reaction.py
+++ b/test/test_reaction.py
@@ -1,8 +1,10 @@
+from mpi4py import MPI
+
 import pytest
+
 import festim as F
-from dolfinx.fem import functionspace, Function
+from dolfinx.fem import Function, functionspace
 from dolfinx.mesh import create_unit_cube
-from mpi4py import MPI
 from ufl import exp
 
 my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 1], material=None)
diff --git a/test/test_settings.py b/test/test_settings.py
index 5ffeeaa60..c518cffec 100644
--- a/test/test_settings.py
+++ b/test/test_settings.py
@@ -1,7 +1,8 @@
-import festim as F
 import numpy as np
 import pytest
 
+import festim as F
+
 
 @pytest.mark.parametrize("test_type", [int, F.Stepsize, float])
 def test_stepsize_value(test_type):
diff --git a/test/test_sievertsbc.py b/test/test_sievertsbc.py
index 26a3b85cd..c25be1404 100644
--- a/test/test_sievertsbc.py
+++ b/test/test_sievertsbc.py
@@ -1,10 +1,12 @@
+from mpi4py import MPI
+
+import numpy as np
+import pytest
+
+import dolfinx.mesh
 import festim as F
 import ufl
-import pytest
-import numpy as np
 from dolfinx import fem
-import dolfinx.mesh
-from mpi4py import MPI
 
 
 def sieverts_law(T, S_0, E_S, pressure):
diff --git a/test/test_source.py b/test/test_source.py
index 242362236..986bcb833 100644
--- a/test/test_source.py
+++ b/test/test_source.py
@@ -1,10 +1,11 @@
-import festim as F
-import ufl
+from mpi4py import MPI
+
 import pytest
+
+import dolfinx.mesh
+import festim as F
 import ufl
 from dolfinx import fem
-import dolfinx.mesh
-from mpi4py import MPI
 
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy_mat")
 
diff --git a/test/test_species.py b/test/test_species.py
index d2e672923..7183587d8 100644
--- a/test/test_species.py
+++ b/test/test_species.py
@@ -1,11 +1,13 @@
-import festim as F
-import dolfinx
-import ufl
+from mpi4py import MPI
+
 import numpy as np
 import pytest
-from dolfinx.fem import functionspace, Function
+
+import dolfinx
+import festim as F
+import ufl
+from dolfinx.fem import Function, functionspace
 from dolfinx.mesh import create_unit_cube
-from mpi4py import MPI
 
 
 def test_assign_functions_to_species():
diff --git a/test/test_stepsize.py b/test/test_stepsize.py
index 56cead680..79f009407 100644
--- a/test/test_stepsize.py
+++ b/test/test_stepsize.py
@@ -1,7 +1,8 @@
-import festim as F
 import numpy as np
 import pytest
 
+import festim as F
+
 
 @pytest.mark.parametrize("growth_factor, target", [(10, 5), (1.2, 2), (1, 1)])
 def test_adaptive_stepsize_grows(growth_factor, target):
diff --git a/test/test_subdomains.py b/test/test_subdomains.py
index 2da15c72b..3b537900c 100644
--- a/test/test_subdomains.py
+++ b/test/test_subdomains.py
@@ -1,8 +1,10 @@
+from mpi4py import MPI
+
 import numpy as np
-import festim as F
 import pytest
+
 import dolfinx.mesh
-from mpi4py import MPI
+import festim as F
 
 
 def test_different_surface_ids():
diff --git a/test/test_surface_quantity.py b/test/test_surface_quantity.py
index ecdac5e00..8680f356a 100644
--- a/test/test_surface_quantity.py
+++ b/test/test_surface_quantity.py
@@ -1,10 +1,12 @@
-import festim as F
+import os
+
 import numpy as np
+import pytest
+
+import festim as F
 import ufl
-from dolfinx.mesh import meshtags
 from dolfinx import fem
-import pytest
-import os
+from dolfinx.mesh import meshtags
 
 
 def test_surface_flux_export_compute():
diff --git a/test/test_total_surface.py b/test/test_total_surface.py
index 8aef39755..18fbca27a 100644
--- a/test/test_total_surface.py
+++ b/test/test_total_surface.py
@@ -1,7 +1,7 @@
-import festim as F
 import numpy as np
+
+import festim as F
 import ufl
-from dolfinx.mesh import meshtags
 from dolfinx import fem
 
 
diff --git a/test/test_total_volume.py b/test/test_total_volume.py
index 8a67a8343..ae1908bb1 100644
--- a/test/test_total_volume.py
+++ b/test/test_total_volume.py
@@ -1,10 +1,11 @@
-import festim as F
+import os
+
 import numpy as np
-from dolfinx.mesh import meshtags
-from dolfinx import fem
 import pytest
+
+import festim as F
 import ufl
-import os
+from dolfinx import fem
 
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy")
 
diff --git a/test/test_volume_quantity.py b/test/test_volume_quantity.py
index a80a063df..d93e07c66 100644
--- a/test/test_volume_quantity.py
+++ b/test/test_volume_quantity.py
@@ -1,4 +1,5 @@
 import pytest
+
 import festim as F
 
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy")
diff --git a/test/test_vtx.py b/test/test_vtx.py
index cfd37ca79..daa597c0b 100644
--- a/test/test_vtx.py
+++ b/test/test_vtx.py
@@ -1,10 +1,11 @@
-import festim as F
-import dolfinx
 from mpi4py import MPI
-import numpy as np
 
+import numpy as np
 import pytest
 
+import dolfinx
+import festim as F
+
 mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 10, 10, 10)
 V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
 
diff --git a/test/test_xdmf.py b/test/test_xdmf.py
index 2fdf7d120..d386bd507 100644
--- a/test/test_xdmf.py
+++ b/test/test_xdmf.py
@@ -1,10 +1,13 @@
-import festim as F
-from dolfinx import fem, mesh
-import mpi4py.MPI as MPI
 import os
+from pathlib import Path
+
+import mpi4py.MPI as MPI
+
 import numpy as np
 import pytest
-from pathlib import Path
+
+import festim as F
+from dolfinx import fem, mesh
 
 
 def test_init():

From 60dcb5b87264a1accd69da5f56d6bcc336f0b44f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 20:55:29 -0400
Subject: [PATCH 100/134] Remove setup files

---
 convergence_rates.py                       |  4 +--
 examples/multi_material_1d.py              |  5 ++-
 examples/multi_material_2d.py              |  5 ++-
 examples/multi_material_2d_bis.py          |  5 ++-
 examples/multi_material_transient.py       |  5 ++-
 festim/boundary_conditions/dirichlet_bc.py |  6 +++-
 festim/exports/vtx.py                      |  4 ++-
 festim/hydrogen_transport_problem.py       |  9 ++---
 festim/problem.py                          |  4 ++-
 festim/species.py                          |  4 ++-
 pyproject.toml                             | 36 +++++++++++++++++++-
 setup.cfg                                  | 38 ----------------------
 setup.py                                   |  4 ---
 test/system_tests/test_multi_material.py   | 23 ++++++++++---
 test/test_dirichlet_bc.py                  | 13 ++++++--
 test/test_flux_bc.py                       |  8 +++--
 test/test_h_transport_problem.py           | 35 +++++++++++++++-----
 test/test_heat_transfer_problem.py         | 13 ++++----
 test/test_mesh.py                          |  5 ++-
 test/test_permeation_problem.py            | 12 +++++--
 test/test_source.py                        |  5 ++-
 21 files changed, 157 insertions(+), 86 deletions(-)
 delete mode 100644 setup.cfg
 delete mode 100644 setup.py

diff --git a/convergence_rates.py b/convergence_rates.py
index 0e2e140a6..47371c7af 100644
--- a/convergence_rates.py
+++ b/convergence_rates.py
@@ -127,9 +127,7 @@ def mms_source(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = (
-        my_problem.t.value
-    )  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = my_problem.t.value  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)
diff --git a/examples/multi_material_1d.py b/examples/multi_material_1d.py
index 1cd22ec29..da15fb043 100644
--- a/examples/multi_material_1d.py
+++ b/examples/multi_material_1d.py
@@ -55,7 +55,10 @@
     left_surface,
     right_surface,
 ]
-my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
+my_model.surface_to_volume = {
+    right_surface: right_domain,
+    left_surface: left_domain,
+}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
diff --git a/examples/multi_material_2d.py b/examples/multi_material_2d.py
index 94fc4ab15..7f1a4d030 100644
--- a/examples/multi_material_2d.py
+++ b/examples/multi_material_2d.py
@@ -85,7 +85,10 @@ def half(x):
 
 # we should be able to automate this
 my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
-my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+my_model.surface_to_volume = {
+    top_surface: top_domain,
+    bottom_surface: bottom_domain,
+}
 
 H = F.Species("H", mobile=True)
 
diff --git a/examples/multi_material_2d_bis.py b/examples/multi_material_2d_bis.py
index 86429615a..8cd105933 100644
--- a/examples/multi_material_2d_bis.py
+++ b/examples/multi_material_2d_bis.py
@@ -85,7 +85,10 @@ def half(x):
 
 # we should be able to automate this
 my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
-my_model.surface_to_volume = {top_surface: top_domain, bottom_surface: bottom_domain}
+my_model.surface_to_volume = {
+    top_surface: top_domain,
+    bottom_surface: bottom_domain,
+}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
diff --git a/examples/multi_material_transient.py b/examples/multi_material_transient.py
index 423a4deb4..8f190c6c1 100644
--- a/examples/multi_material_transient.py
+++ b/examples/multi_material_transient.py
@@ -50,7 +50,10 @@
     left_surface,
     right_surface,
 ]
-my_model.surface_to_volume = {right_surface: right_domain, left_surface: left_domain}
+my_model.surface_to_volume = {
+    right_surface: right_domain,
+    left_surface: left_domain,
+}
 
 H = F.Species("H", mobile=True)
 trapped_H = F.Species("H_trapped", mobile=False)
diff --git a/festim/boundary_conditions/dirichlet_bc.py b/festim/boundary_conditions/dirichlet_bc.py
index 2bae86a6f..084fe4190 100644
--- a/festim/boundary_conditions/dirichlet_bc.py
+++ b/festim/boundary_conditions/dirichlet_bc.py
@@ -128,7 +128,11 @@ def temperature_dependent(self):
             return False
 
     def create_value(
-        self, mesh, function_space: fem.FunctionSpace, temperature, t: fem.Constant
+        self,
+        mesh,
+        function_space: fem.FunctionSpace,
+        temperature,
+        t: fem.Constant,
     ):
         """Creates the value of the boundary condition as a fenics object and sets it to
         self.value_fenics.
diff --git a/festim/exports/vtx.py b/festim/exports/vtx.py
index ec5ccf648..53c7c6f13 100644
--- a/festim/exports/vtx.py
+++ b/festim/exports/vtx.py
@@ -3,7 +3,9 @@
 
 from dolfinx.fem import Function as _Function
 from festim.species import Species as _Species
-from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
+from festim.subdomain.volume_subdomain import (
+    VolumeSubdomain as _VolumeSubdomain,
+)
 
 
 class ExportBaseClass:
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 178631578..6c614233d 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -971,9 +971,9 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_collapsed_function_space[subdomain] = V.sub(
                 i
             ).collapse()
-            species.subdomain_to_post_processing_solution[subdomain].name = (
-                f"{species.name}_{subdomain.id}"
-            )
+            species.subdomain_to_post_processing_solution[
+                subdomain
+            ].name = f"{species.name}_{subdomain.id}"
 
     def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
         """
@@ -1147,7 +1147,8 @@ def mixed_term(u, v, n):
             for subdomain2 in self.volume_subdomains:
                 jac.append(
                     dolfinx.fem.form(
-                        ufl.derivative(form, subdomain2.u), entity_maps=entity_maps
+                        ufl.derivative(form, subdomain2.u),
+                        entity_maps=entity_maps,
                     )
                 )
             J.append(jac)
diff --git a/festim/problem.py b/festim/problem.py
index 350ca1285..f26084ef3 100644
--- a/festim/problem.py
+++ b/festim/problem.py
@@ -11,7 +11,9 @@
 from dolfinx.nls.petsc import NewtonSolver
 from festim.mesh.mesh import Mesh as _Mesh
 from festim.source import SourceBase as _SourceBase
-from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
+from festim.subdomain.volume_subdomain import (
+    VolumeSubdomain as _VolumeSubdomain,
+)
 
 
 class ProblemBase:
diff --git a/festim/species.py b/festim/species.py
index ce031b728..ca9d28921 100644
--- a/festim/species.py
+++ b/festim/species.py
@@ -1,4 +1,6 @@
-from festim.subdomain.volume_subdomain import VolumeSubdomain as _VolumeSubdomain
+from festim.subdomain.volume_subdomain import (
+    VolumeSubdomain as _VolumeSubdomain,
+)
 
 
 class Species:
diff --git a/pyproject.toml b/pyproject.toml
index 82c76bf1c..44224ca11 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,41 @@
 [build-system]
-requires = ["setuptools >= 45", "wheel", "setuptools_scm[toml] >= 7.0.5"]
+requires = ["setuptools >= 61", "wheel", "setuptools-scm[toml] >= 7.0.5"]
 build-backend = "setuptools.build_meta"
 
+[project]
+authors = [
+    { name = "Remi Delaporte-Mathurin" },
+    { email = "rdelaportemathurin@gmail.com" },
+]
+dynamic = ["version"]
+name = "FESTIM"
+description = "Finite element simulations of hydrogen transport"
+readme = "README/md"
+requires-python = ">=3.9"
+license = { file = "LICENSE" }
+dependencies = ["tqdm", "scifem>=0.2.8"]
+classifiers = [
+    "Natural Language :: English",
+    "Topic :: Scientific/Engineering",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent",
+]
+
+[project.optional-dependencies]
+test = ["pytest >= 5.4.3", "pytest-cov", "sympy"]
+lint = ["ruff"]
+docs = ["sphinx"]
+
+[project.urls]
+Homepage = "https://github.com/RemDelaporteMathurin/FESTIM"
+Issues = "https://github.com/RemDelaporteMathurin/FESTIM/issues"
+
+
 [tool.setuptools_scm]
 write_to = "festim/_version.py"
 
diff --git a/setup.cfg b/setup.cfg
deleted file mode 100644
index 38feea1a9..000000000
--- a/setup.cfg
+++ /dev/null
@@ -1,38 +0,0 @@
-
-
-[metadata]
-name = FESTIM
-author = Remi Delaporte-Mathurin
-author_email = rdelaportemathurin@gmail.com
-description = Finite element simulations of hydrogen transport
-long_description = file: README.md
-long_description_content_type = text/markdown
-url = https://github.com/RemDelaporteMathurin/FESTIM
-license = Apache-2.0
-license_file = LICENSE
-classifiers =
-    Natural Language :: English
-    Topic :: Scientific/Engineering
-    Programming Language :: Python :: 3.9
-    Programming Language :: Python :: 3.10
-    Programming Language :: Python :: 3.11
-    Programming Language :: Python :: 3.12
-    Programming Language :: Python :: 3.13
-    License :: OSI Approved :: Apache Software License
-    Operating System :: OS Independent
-project_urls =
-    Source = https://github.com/RemDelaporteMathurin/FESTIM
-    Tracker = https://github.com/RemDelaporteMathurin/FESTIM/issues
-
-[options]
-packages = find:
-python_requires= >=3.9
-install_requires =
-    tqdm
-    scifem>=0.2.8
-
-[options.extras_require]
-test = 
-    pytest >= 5.4.3
-    pytest-cov
-    sympy
diff --git a/setup.py b/setup.py
deleted file mode 100644
index 1abbd068c..000000000
--- a/setup.py
+++ /dev/null
@@ -1,4 +0,0 @@
-import setuptools
-
-if __name__ == "__main__":
-    setuptools.setup()
diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index 35a8e6233..1b2343088 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -100,7 +100,12 @@ def c_exact_bot_np(x):
 
     top_surface = F.SurfaceSubdomain(id=1)
     bottom_surface = F.SurfaceSubdomain(id=2)
-    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+    my_model.subdomains = [
+        bottom_domain,
+        top_domain,
+        top_surface,
+        bottom_surface,
+    ]
 
     my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
     my_model.surface_to_volume = {
@@ -177,7 +182,10 @@ def test_1_material_discontinuous_version():
     right_surface = F.SurfaceSubdomain1D(id=2, x=vertices[-1])
 
     my_model.subdomains = [subdomain, left_surface, right_surface]
-    my_model.surface_to_volume = {right_surface: subdomain, left_surface: subdomain}
+    my_model.surface_to_volume = {
+        right_surface: subdomain,
+        left_surface: subdomain,
+    }
 
     H = F.Species("H", mobile=True)
     trapped_H = F.Species("H_trapped", mobile=False)
@@ -314,7 +322,9 @@ def test_3_materials_transient():
         for subdomain in my_model.volume_subdomains
     ] + [
         F.VTXSpeciesExport(
-            filename=f"u_t_{subdomain.id}.bp", field=trapped_H, subdomain=subdomain
+            filename=f"u_t_{subdomain.id}.bp",
+            field=trapped_H,
+            subdomain=subdomain,
         )
         for subdomain in my_model.volume_subdomains
     ]
@@ -344,7 +354,12 @@ def test_2_mats_particle_flux_bc():
 
     top_surface = F.SurfaceSubdomain(id=1)
     bottom_surface = F.SurfaceSubdomain(id=2)
-    my_model.subdomains = [bottom_domain, top_domain, top_surface, bottom_surface]
+    my_model.subdomains = [
+        bottom_domain,
+        top_domain,
+        top_surface,
+        bottom_surface,
+    ]
 
     # we should be able to automate this
     my_model.interfaces = [F.Interface(5, (bottom_domain, top_domain))]
diff --git a/test/test_dirichlet_bc.py b/test/test_dirichlet_bc.py
index 76ecebb05..fc49ccbad 100644
--- a/test/test_dirichlet_bc.py
+++ b/test/test_dirichlet_bc.py
@@ -63,7 +63,9 @@ def test_callable_for_value():
     bc = F.DirichletBC(subdomain, value, species)
 
     my_model = F.HydrogenTransportProblem(
-        mesh=F.Mesh(mesh), subdomains=[subdomain, vol_subdomain], species=[species]
+        mesh=F.Mesh(mesh),
+        subdomains=[subdomain, vol_subdomain],
+        species=[species],
     )
 
     my_model.define_function_spaces()
@@ -101,7 +103,9 @@ def test_value_callable_x_t_T():
     bc = F.DirichletBC(subdomain, value, species)
 
     my_model = F.HydrogenTransportProblem(
-        mesh=F.Mesh(mesh), subdomains=[subdomain, vol_subdomain], species=[species]
+        mesh=F.Mesh(mesh),
+        subdomains=[subdomain, vol_subdomain],
+        species=[species],
     )
 
     my_model.define_function_spaces()
@@ -440,7 +444,10 @@ def test_bc_time_dependent_attribute(input, expected_value):
         (lambda t: 1.0 + t, False),
         (lambda x, T: 1.0 + x[0] + T, True),
         (lambda x, t, T: 1.0 + x[0] + t + T, True),
-        (lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0), False),
+        (
+            lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0),
+            False,
+        ),
     ],
 )
 def test_bc_temperature_dependent_attribute(input, expected_value):
diff --git a/test/test_flux_bc.py b/test/test_flux_bc.py
index b6865f9ce..ccd655161 100644
--- a/test/test_flux_bc.py
+++ b/test/test_flux_bc.py
@@ -174,7 +174,8 @@ def test_bc_time_dependent_attribute_raises_error_when_value_none():
     my_flux_bc = F.FluxBCBase(subdomain=surface, value=None)
 
     with pytest.raises(
-        TypeError, match="Value must be given to determine if its time dependent"
+        TypeError,
+        match="Value must be given to determine if its time dependent",
     ):
         my_flux_bc.time_dependent
 
@@ -189,7 +190,10 @@ def test_bc_time_dependent_attribute_raises_error_when_value_none():
         (lambda t: 1.0 + t, False),
         (lambda x, T: 1.0 + x[0] + T, True),
         (lambda x, t, T: 1.0 + x[0] + t + T, True),
-        (lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0), False),
+        (
+            lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0),
+            False,
+        ),
     ],
 )
 def test_bc_temperature_dependent_attribute(input, expected_value):
diff --git a/test/test_h_transport_problem.py b/test/test_h_transport_problem.py
index 6a60bc658..e085d266a 100644
--- a/test/test_h_transport_problem.py
+++ b/test/test_h_transport_problem.py
@@ -16,7 +16,8 @@
 
 # TODO test all the methods in the class
 @pytest.mark.parametrize(
-    "value", [1, fem.Constant(test_mesh.mesh, 1.0), 1.0, "coucou", lambda x: 2 * x[0]]
+    "value",
+    [1, fem.Constant(test_mesh.mesh, 1.0), 1.0, "coucou", lambda x: 2 * x[0]],
 )
 def test_temperature_setter_type(value):
     """Test that the temperature type is correctly set"""
@@ -78,7 +79,10 @@ def test_define_temperature_value_error_raised():
         (lambda t: 1.0 + t, fem.Constant),
         (lambda x: 1.0 + x[0], fem.Function),
         (lambda x, t: 1.0 + x[0] + t, fem.Function),
-        (lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0), fem.Function),
+        (
+            lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0),
+            fem.Function,
+        ),
         (lambda t: 100.0 if t < 1 else 0.0, fem.Constant),
     ],
 )
@@ -118,7 +122,8 @@ def test_define_temperature_error_if_ufl_conditional_t_only(input):
     my_model.temperature = input
 
     with pytest.raises(
-        ValueError, match="self.temperature should return a float or an int, not "
+        ValueError,
+        match="self.temperature should return a float or an int, not ",
     ):
         my_model.define_temperature()
 
@@ -399,7 +404,10 @@ def test_post_processing_update_D_global():
     # Build the model
     my_model = F.HydrogenTransportProblem(
         mesh=my_mesh,
-        subdomains=[F.VolumeSubdomain1D(id=1, borders=[0, 1], material=my_mat), surf],
+        subdomains=[
+            F.VolumeSubdomain1D(id=1, borders=[0, 1], material=my_mat),
+            surf,
+        ],
         species=[H],
         temperature=lambda t: 500 * t,
         exports=[my_export],
@@ -446,7 +454,10 @@ def test_post_processing_update_D_global_2():
     # Build the model
     my_model = F.HydrogenTransportProblem(
         mesh=my_mesh,
-        subdomains=[F.VolumeSubdomain1D(id=1, borders=[0, 1], material=my_mat), surf],
+        subdomains=[
+            F.VolumeSubdomain1D(id=1, borders=[0, 1], material=my_mat),
+            surf,
+        ],
         species=[H],
         temperature=lambda t: 500 * t,
         exports=[my_export],
@@ -563,7 +574,9 @@ def test_update_time_dependent_bcs_with_time_dependent_temperature(
     volume_subdomain = F.VolumeSubdomain1D(id=1, borders=[0, 4], material=dummy_mat)
     surface_subdomain = F.SurfaceSubdomain1D(id=1, x=4)
     my_model = F.HydrogenTransportProblem(
-        mesh=test_mesh, species=[H], subdomains=[volume_subdomain, surface_subdomain]
+        mesh=test_mesh,
+        species=[H],
+        subdomains=[volume_subdomain, surface_subdomain],
     )
     my_model.t = fem.Constant(my_model.mesh.mesh, 0.0)
     dt = fem.Constant(test_mesh.mesh, 1.0)
@@ -1000,7 +1013,10 @@ def test_define_meshtags_and_measures_with_custom_fenics_mesh():
     mesh_1D = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
     # 1D meshtags
     my_surface_meshtags = dolfinx.mesh.meshtags(
-        mesh_1D, 0, np.array([0, 10], dtype=np.int32), np.array([1, 2], dtype=np.int32)
+        mesh_1D,
+        0,
+        np.array([0, 10], dtype=np.int32),
+        np.array([1, 2], dtype=np.int32),
     )
 
     num_cells = mesh_1D.topology.index_map(1).size_local
@@ -1068,7 +1084,10 @@ def test_update_time_dependent_values_flux(bc_value, expected_values):
     surface = F.SurfaceSubdomain1D(id=2, x=0)
     H = F.Species("H")
     my_model = F.HydrogenTransportProblem(
-        mesh=test_mesh, temperature=10, subdomains=[my_vol, surface], species=[H]
+        mesh=test_mesh,
+        temperature=10,
+        subdomains=[my_vol, surface],
+        species=[H],
     )
     my_model.t = fem.Constant(my_model.mesh.mesh, 0.0)
     dt = fem.Constant(test_mesh.mesh, 1.0)
diff --git a/test/test_heat_transfer_problem.py b/test/test_heat_transfer_problem.py
index 961da0a99..b8b138e45 100644
--- a/test/test_heat_transfer_problem.py
+++ b/test/test_heat_transfer_problem.py
@@ -218,9 +218,7 @@ def exact_solution(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = (
-        my_problem.t.value
-    )  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = my_problem.t.value  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)
@@ -343,9 +341,7 @@ def mms_source(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = (
-        my_problem.t.value
-    )  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = my_problem.t.value  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)
@@ -369,7 +365,10 @@ def test_sources():
     # Test that setting invalid sources raises a TypeError
     with pytest.raises(TypeError, match="festim.HeatSource objects"):
         spe = F.Species("H")
-        htp.sources = [F.ParticleSource(1, vol, spe), F.ParticleSource(1, vol, spe)]
+        htp.sources = [
+            F.ParticleSource(1, vol, spe),
+            F.ParticleSource(1, vol, spe),
+        ]
 
 
 def test_boundary_conditions():
diff --git a/test/test_mesh.py b/test/test_mesh.py
index c0f9c0b43..74d6caeb0 100644
--- a/test/test_mesh.py
+++ b/test/test_mesh.py
@@ -16,7 +16,10 @@
 
 # 1D meshtags
 my_surface_meshtags = meshtags(
-    mesh_1D, 0, np.array([0, 10], dtype=np.int32), np.array([1, 2], dtype=np.int32)
+    mesh_1D,
+    0,
+    np.array([0, 10], dtype=np.int32),
+    np.array([1, 2], dtype=np.int32),
 )
 
 num_cells = mesh_1D.topology.index_map(1).size_local
diff --git a/test/test_permeation_problem.py b/test/test_permeation_problem.py
index e2961cc74..d6c4dcc57 100644
--- a/test/test_permeation_problem.py
+++ b/test/test_permeation_problem.py
@@ -60,7 +60,11 @@ def test_permeation_problem(mesh_size=1001):
     my_model.boundary_conditions = [
         F.DirichletBC(subdomain=right_surface, value=0, species="H"),
         F.SievertsBC(
-            subdomain=left_surface, S_0=4.02e21, E_S=1.04, pressure=100, species="H"
+            subdomain=left_surface,
+            S_0=4.02e21,
+            E_S=1.04,
+            pressure=100,
+            species="H",
         ),
     ]
     outgassing_flux = F.SurfaceFlux(
@@ -156,7 +160,11 @@ def test_permeation_problem_multi_volume(tmp_path):
     my_model.boundary_conditions = [
         F.DirichletBC(subdomain=right_surface, value=0, species="H"),
         F.SievertsBC(
-            subdomain=left_surface, S_0=4.02e21, E_S=1.04, pressure=100, species="H"
+            subdomain=left_surface,
+            S_0=4.02e21,
+            E_S=1.04,
+            pressure=100,
+            species="H",
         ),
     ]
     outgassing_flux = F.SurfaceFlux(
diff --git a/test/test_source.py b/test/test_source.py
index 986bcb833..614c5c21c 100644
--- a/test/test_source.py
+++ b/test/test_source.py
@@ -120,7 +120,10 @@ def test_source_time_dependent_attribute(input, expected_value):
         (lambda t: 1.0 + t, False),
         (lambda x, T: 1.0 + x[0] + T, True),
         (lambda x, t, T: 1.0 + x[0] + t + T, True),
-        (lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0), False),
+        (
+            lambda x, t: ufl.conditional(ufl.lt(t, 1.0), 100.0 + x[0], 0.0),
+            False,
+        ),
     ],
 )
 def test_source_temperature_dependent_attribute(input, expected_value):

From 3101c40f2a98d5060328627297e7080d383886ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 21:03:11 -0400
Subject: [PATCH 101/134] Change linting

---
 .github/workflows/code_formatting.yml | 23 +++++++++++++++++++++--
 convergence_rates.py                  |  3 ++-
 pyproject.toml                        |  2 +-
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/code_formatting.yml b/.github/workflows/code_formatting.yml
index e9a96b395..11576717f 100644
--- a/.github/workflows/code_formatting.yml
+++ b/.github/workflows/code_formatting.yml
@@ -6,5 +6,24 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: psf/black@stable
\ No newline at end of file
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install linting tools
+        run: pip install mypy ruff
+
+      - name: ruff format
+        run: |
+          ruff format --check .
+
+      - name: ruff check
+        continue-on-error: true
+        run: |
+          ruff check .
+
+      - name: mypy
+        run: |
+          python -m mypy .
diff --git a/convergence_rates.py b/convergence_rates.py
index 47371c7af..abc9d4b43 100644
--- a/convergence_rates.py
+++ b/convergence_rates.py
@@ -127,7 +127,8 @@ def mms_source(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = my_problem.t.value  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = my_problem.t.value
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)
diff --git a/pyproject.toml b/pyproject.toml
index 44224ca11..96f570659 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ classifiers = [
 
 [project.optional-dependencies]
 test = ["pytest >= 5.4.3", "pytest-cov", "sympy"]
-lint = ["ruff"]
+lint = ["ruff", "mypy"]
 docs = ["sphinx"]
 
 [project.urls]

From 729272e771988251a32f4fa6b7706b5f53f0e2f6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Wed, 23 Oct 2024 21:03:49 -0400
Subject: [PATCH 102/134] Add dependabot

---
 .github/dependabot.yml | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 .github/dependabot.yml

diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 000000000..8d8c19045
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,15 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file
+
+version: 2
+updates:
+  - package-ecosystem: "github-actions" # See documentation for possible values
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "weekly"
+  - package-ecosystem: "pip" # See documentation for possible values
+    directory: "python/"
+    schedule:
+      interval: "weekly"

From 6b174cbdadadc7262f2c205903389a2852ca5e6e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= <dokken92@gmail.com>
Date: Thu, 24 Oct 2024 09:36:27 -0400
Subject: [PATCH 103/134] Apply suggestions from code review

---
 .github/dependabot.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/dependabot.yml b/.github/dependabot.yml
index 8d8c19045..b55084063 100644
--- a/.github/dependabot.yml
+++ b/.github/dependabot.yml
@@ -10,6 +10,6 @@ updates:
     schedule:
       interval: "weekly"
   - package-ecosystem: "pip" # See documentation for possible values
-    directory: "python/"
+    directory: "./"
     schedule:
       interval: "weekly"

From ca99a72cde038d3f56d8800e6494fdcbf4b2e263 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 09:41:06 -0400
Subject: [PATCH 104/134] Formatting and importing within coer

---
 convergence_rates.py                       |   4 +-
 examples/multi_material_2d.py              |   6 +-
 examples/multi_material_2d_bis.py          |   6 +-
 festim/boundary_conditions/__init__.py     |   4 +
 festim/boundary_conditions/dirichlet_bc.py |   4 +-
 festim/boundary_conditions/flux_bc.py      |   4 +-
 festim/boundary_conditions/henrys_bc.py    |   3 +-
 festim/boundary_conditions/sieverts_bc.py  |   3 +-
 festim/exports/__init__.py                 |  22 ++++
 festim/exports/average_surface.py          |   5 +-
 festim/exports/average_volume.py           |   5 +-
 festim/exports/maximum_surface.py          |   4 +-
 festim/exports/maximum_volume.py           |   4 +-
 festim/exports/minimum_surface.py          |   4 +-
 festim/exports/minimum_volume.py           |  10 +-
 festim/exports/surface_flux.py             |   5 +-
 festim/exports/total_surface.py            |   5 +-
 festim/exports/total_volume.py             |   5 +-
 festim/exports/vtx.py                      |   1 +
 festim/exports/xdmf.py                     |   1 +
 festim/heat_transfer_problem.py            |  31 +++--
 festim/helpers_discontinuity.py            |   3 +-
 festim/hydrogen_transport_problem.py       | 127 ++++++++++++---------
 festim/initial_condition.py                |   1 -
 festim/material.py                         |   3 +-
 festim/mesh/__init__.py                    |   3 +
 festim/mesh/mesh.py                        |   3 +-
 festim/mesh/mesh_1d.py                     |   6 +-
 festim/mesh/mesh_from_xdmf.py              |   3 +-
 festim/problem.py                          |   4 +-
 festim/reaction.py                         |   3 +-
 festim/source.py                           |   6 +-
 festim/subdomain/__init__.py               |   6 +
 festim/subdomain/interface.py              |  10 +-
 festim/subdomain/surface_subdomain.py      |   3 +-
 festim/subdomain/surface_subdomain_1d.py   |   2 +-
 festim/subdomain/volume_subdomain.py       |   4 +-
 festim/subdomain/volume_subdomain_1d.py    |   6 +-
 pyproject.toml                             |  27 ++++-
 test/benchmark.py                          |   5 +-
 test/system_tests/test_1_mobile_1_trap.py  |   4 +-
 test/system_tests/test_1_mobile_species.py |   4 +-
 test/system_tests/test_fluxes.py           |   4 +-
 test/system_tests/test_heat_transfer.py    |   4 +-
 test/system_tests/test_multi_material.py   |   6 +-
 test/system_tests/tools.py                 |   1 -
 test/test_average_surface.py               |   4 +-
 test/test_average_volume.py                |   4 +-
 test/test_dirichlet_bc.py                  |  27 +++--
 test/test_flux_bc.py                       |  12 +-
 test/test_h_transport_problem.py           |  18 ++-
 test/test_heat_transfer_problem.py         |  14 ++-
 test/test_helpers.py                       |  13 ++-
 test/test_henrysbc.py                      |   6 +-
 test/test_initial_condition.py             |   2 +-
 test/test_maximum_surface.py               |   2 +-
 test/test_maximum_volume.py                |   2 +-
 test/test_mesh.py                          |   4 +-
 test/test_minimum_surface.py               |   2 +-
 test/test_minimum_volume.py                |   2 +-
 test/test_multispecies_problem.py          |   4 +-
 test/test_permeation_problem.py            |   4 +-
 test/test_reaction.py                      |   4 +-
 test/test_sievertsbc.py                    |   6 +-
 test/test_source.py                        |  10 +-
 test/test_species.py                       |   6 +-
 test/test_subdomains.py                    |   2 +-
 test/test_surface_quantity.py              |   4 +-
 test/test_total_surface.py                 |   4 +-
 test/test_total_volume.py                  |   4 +-
 test/test_vtx.py                           |   2 +-
 test/test_xdmf.py                          |   7 +-
 72 files changed, 339 insertions(+), 214 deletions(-)

diff --git a/convergence_rates.py b/convergence_rates.py
index abc9d4b43..6286c8e5e 100644
--- a/convergence_rates.py
+++ b/convergence_rates.py
@@ -2,11 +2,11 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 def source_from_exact_solution(
     exact_solution, thermal_conductivity, density, heat_capacity
diff --git a/examples/multi_material_2d.py b/examples/multi_material_2d.py
index 7f1a4d030..31d9cca7a 100644
--- a/examples/multi_material_2d.py
+++ b/examples/multi_material_2d.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
-import numpy as np
-
 import dolfinx
 import dolfinx.fem.petsc
-import festim as F
+import numpy as np
 import ufl
 
+import festim as F
+
 
 # ---------------- Generate a mesh ----------------
 def generate_mesh():
diff --git a/examples/multi_material_2d_bis.py b/examples/multi_material_2d_bis.py
index 8cd105933..e1ca09ee8 100644
--- a/examples/multi_material_2d_bis.py
+++ b/examples/multi_material_2d_bis.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
-import numpy as np
-
 import dolfinx
 import dolfinx.fem.petsc
-import festim as F
+import numpy as np
 import ufl
 
+import festim as F
+
 
 # ---------------- Generate a mesh ----------------
 def generate_mesh():
diff --git a/festim/boundary_conditions/__init__.py b/festim/boundary_conditions/__init__.py
index e69de29bb..de6d9f19f 100644
--- a/festim/boundary_conditions/__init__.py
+++ b/festim/boundary_conditions/__init__.py
@@ -0,0 +1,4 @@
+from .dirichlet_bc import FixedConcentrationBC, FixedTemperatureBC
+from .flux_bc import HeatFluxBC, ParticleFluxBC
+
+__all__ = ["FixedConcentrationBC", "FixedTemperatureBC", "ParticleFluxBC", "HeatFluxBC"]
diff --git a/festim/boundary_conditions/dirichlet_bc.py b/festim/boundary_conditions/dirichlet_bc.py
index 084fe4190..b7022c0c4 100644
--- a/festim/boundary_conditions/dirichlet_bc.py
+++ b/festim/boundary_conditions/dirichlet_bc.py
@@ -1,9 +1,9 @@
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 class DirichletBCBase:
     """
diff --git a/festim/boundary_conditions/flux_bc.py b/festim/boundary_conditions/flux_bc.py
index c7243b972..fea235571 100644
--- a/festim/boundary_conditions/flux_bc.py
+++ b/festim/boundary_conditions/flux_bc.py
@@ -1,9 +1,9 @@
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 class FluxBCBase:
     """
diff --git a/festim/boundary_conditions/henrys_bc.py b/festim/boundary_conditions/henrys_bc.py
index 16558dd10..8bc6b1362 100644
--- a/festim/boundary_conditions/henrys_bc.py
+++ b/festim/boundary_conditions/henrys_bc.py
@@ -1,6 +1,7 @@
-import festim as F
 import ufl
 
+import festim as F
+
 
 def henrys_law(T, H_0, E_H, pressure):
     """Applies the Henry's law to compute the concentration at the boundary"""
diff --git a/festim/boundary_conditions/sieverts_bc.py b/festim/boundary_conditions/sieverts_bc.py
index c70ee160c..b894588e2 100644
--- a/festim/boundary_conditions/sieverts_bc.py
+++ b/festim/boundary_conditions/sieverts_bc.py
@@ -1,6 +1,7 @@
-import festim as F
 import ufl
 
+import festim as F
+
 
 def sieverts_law(T, S_0, E_S, pressure):
     """Applies the Sieverts law to compute the concentration at the boundary"""
diff --git a/festim/exports/__init__.py b/festim/exports/__init__.py
index e69de29bb..ca4eb00a7 100644
--- a/festim/exports/__init__.py
+++ b/festim/exports/__init__.py
@@ -0,0 +1,22 @@
+__all__ = [
+    "VolumeQuantity",
+    "SurfaceQuantity",
+    "VTXSpeciesExport",
+    "VTXTemperatureExport",
+    "XDMFExport",
+    "SurfaceFlux",
+    "TotalSurface",
+    "AverageSurface",
+    "AverageVolume",
+    "TotalVolume",
+]
+
+from .average_surface import AverageSurface
+from .average_volume import AverageVolume
+from .surface_flux import SurfaceFlux
+from .surface_quantity import SurfaceQuantity
+from .total_surface import TotalSurface
+from .total_volume import TotalVolume
+from .volume_quantity import VolumeQuantity
+from .vtx import VTXSpeciesExport, VTXTemperatureExport
+from .xdmf import XDMFExport
diff --git a/festim/exports/average_surface.py b/festim/exports/average_surface.py
index 51601cf25..5dee531aa 100644
--- a/festim/exports/average_surface.py
+++ b/festim/exports/average_surface.py
@@ -1,8 +1,9 @@
-import festim as F
 from dolfinx import fem
 
+import festim.exports.surface_quantity as sq
 
-class AverageSurface(F.SurfaceQuantity):
+
+class AverageSurface(sq.SurfaceQuantity):
     """Computes the average value of a field on a given surface
 
     Args:
diff --git a/festim/exports/average_volume.py b/festim/exports/average_volume.py
index 53de6f794..e7677e341 100644
--- a/festim/exports/average_volume.py
+++ b/festim/exports/average_volume.py
@@ -1,8 +1,9 @@
-import festim as F
 from dolfinx import fem
 
+from festim.exports import VolumeQuantity
 
-class AverageVolume(F.VolumeQuantity):
+
+class AverageVolume(VolumeQuantity):
     """Computes the average value of a field in a given volume
 
     Args:
diff --git a/festim/exports/maximum_surface.py b/festim/exports/maximum_surface.py
index 29da50d1f..b398f91c0 100644
--- a/festim/exports/maximum_surface.py
+++ b/festim/exports/maximum_surface.py
@@ -1,9 +1,9 @@
 import numpy as np
 
-import festim as F
+import festim.exports.surface_quantity as sq
 
 
-class MaximumSurface(F.SurfaceQuantity):
+class MaximumSurface(sq.SurfaceQuantity):
     """Computes the maximum value of a field on a given surface
 
     Args:
diff --git a/festim/exports/maximum_volume.py b/festim/exports/maximum_volume.py
index 3a26cacf3..80534e777 100644
--- a/festim/exports/maximum_volume.py
+++ b/festim/exports/maximum_volume.py
@@ -1,9 +1,9 @@
 import numpy as np
 
-import festim as F
+from festim.exports import VolumeQuantity
 
 
-class MaximumVolume(F.VolumeQuantity):
+class MaximumVolume(VolumeQuantity):
     """Computes the maximum value of a field in a given volume
 
     Args:
diff --git a/festim/exports/minimum_surface.py b/festim/exports/minimum_surface.py
index 120ba8991..c5b53e3a7 100644
--- a/festim/exports/minimum_surface.py
+++ b/festim/exports/minimum_surface.py
@@ -1,9 +1,9 @@
 import numpy as np
 
-import festim as F
+import festim.exports.surface_quantity as sq
 
 
-class MinimumSurface(F.SurfaceQuantity):
+class MinimumSurface(sq.SurfaceQuantity):
     """Computes the minimum value of a field on a given surface
 
     Args:
diff --git a/festim/exports/minimum_volume.py b/festim/exports/minimum_volume.py
index a2e7c14b6..0a29e17f8 100644
--- a/festim/exports/minimum_volume.py
+++ b/festim/exports/minimum_volume.py
@@ -1,15 +1,15 @@
 import numpy as np
 
-import festim as F
+from festim.exports import VolumeQuantity
 
 
-class MinimumVolume(F.VolumeQuantity):
-    """Computes the minmum value of a field in a given volume
+class MinimumVolume(VolumeQuantity):
+    """Computes the minimum value of a field in a given volume
 
     Args:
-        field (festim.Species): species for which the minmum volume is computed
+        field (festim.Species): species for which the minimum volume is computed
         volume (festim.VolumeSubdomain): volume subdomain
-        filename (str, optional): name of the file to which the minmum volume is exported
+        filename (str, optional): name of the file to which the minimum volume is exported
 
     Attributes:
         see `festim.VolumeQuantity`
diff --git a/festim/exports/surface_flux.py b/festim/exports/surface_flux.py
index aa68f7ad0..7ded96409 100644
--- a/festim/exports/surface_flux.py
+++ b/festim/exports/surface_flux.py
@@ -1,9 +1,10 @@
-import festim as F
 import ufl
 from dolfinx import fem
 
+from festim.exports import SurfaceQuantity
 
-class SurfaceFlux(F.SurfaceQuantity):
+
+class SurfaceFlux(SurfaceQuantity):
     """Computes the flux of a field on a given surface
 
     Args:
diff --git a/festim/exports/total_surface.py b/festim/exports/total_surface.py
index 1d535583e..12c28c774 100644
--- a/festim/exports/total_surface.py
+++ b/festim/exports/total_surface.py
@@ -1,9 +1,10 @@
-import festim as F
 import ufl
 from dolfinx import fem
 
+from .surface_quantity import SurfaceQuantity
 
-class TotalSurface(F.SurfaceQuantity):
+
+class TotalSurface(SurfaceQuantity):
     """Computes the total value of a field on a given surface
 
     Args:
diff --git a/festim/exports/total_volume.py b/festim/exports/total_volume.py
index e16a96c4a..0fadd6f4d 100644
--- a/festim/exports/total_volume.py
+++ b/festim/exports/total_volume.py
@@ -1,9 +1,10 @@
-import festim as F
 import ufl
 from dolfinx import fem
 
+from festim.exports import VolumeQuantity
 
-class TotalVolume(F.VolumeQuantity):
+
+class TotalVolume(VolumeQuantity):
     """Computes the total value of a field in a given volume
 
     Args:
diff --git a/festim/exports/vtx.py b/festim/exports/vtx.py
index 53c7c6f13..960a1041f 100644
--- a/festim/exports/vtx.py
+++ b/festim/exports/vtx.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 
 from dolfinx.fem import Function as _Function
+
 from festim.species import Species as _Species
 from festim.subdomain.volume_subdomain import (
     VolumeSubdomain as _VolumeSubdomain,
diff --git a/festim/exports/xdmf.py b/festim/exports/xdmf.py
index d813dbe84..5c3c38cd6 100644
--- a/festim/exports/xdmf.py
+++ b/festim/exports/xdmf.py
@@ -3,6 +3,7 @@
 import mpi4py
 
 from dolfinx.io import XDMFFile
+
 from festim.species import Species as _Species
 
 from .vtx import ExportBaseClass
diff --git a/festim/heat_transfer_problem.py b/festim/heat_transfer_problem.py
index cd5b0b029..746cf278b 100644
--- a/festim/heat_transfer_problem.py
+++ b/festim/heat_transfer_problem.py
@@ -1,11 +1,13 @@
 import basix
-import festim as F
 import ufl
 from dolfinx import fem
 from dolfinx.io import VTXWriter
 
+from festim import boundary_conditions, exports, helpers, problem
+from festim import source as _source
 
-class HeatTransferProblem(F.ProblemBase):
+
+class HeatTransferProblem(problem.ProblemBase):
     def __init__(
         self,
         mesh=None,
@@ -34,7 +36,7 @@ def sources(self):
 
     @sources.setter
     def sources(self, value):
-        if not all(isinstance(source, F.HeatSource) for source in value):
+        if not all(isinstance(source, _source.HeatSource) for source in value):
             raise TypeError("sources must be a list of festim.HeatSource objects")
         self._sources = value
 
@@ -45,7 +47,14 @@ def boundary_conditions(self):
     @boundary_conditions.setter
     def boundary_conditions(self, value):
         if not all(
-            isinstance(bc, (F.FixedTemperatureBC, F.HeatFluxBC)) for bc in value
+            isinstance(
+                bc,
+                (
+                    boundary_conditions.FixedTemperatureBC,
+                    boundary_conditions.HeatFluxBC,
+                ),
+            )
+            for bc in value
         ):
             raise TypeError(
                 "boundary_conditions must be a list of festim.FixedTemperatureBC or festim.HeatFluxBC objects"
@@ -60,7 +69,7 @@ def initialise(self):
         if self.settings.transient:
             # TODO should raise error if no stepsize is provided
             # TODO Should this be an attribute of festim.Stepsize?
-            self.dt = F.as_fenics_constant(
+            self.dt = helpers.as_fenics_constant(
                 self.settings.stepsize.initial_value, self.mesh.mesh
             )
 
@@ -136,7 +145,7 @@ def create_flux_values_fenics(self):
         """For each heat flux create the value_fenics"""
         for bc in self.boundary_conditions:
             # create value_fenics for all F.HeatFluxBC objects
-            if isinstance(bc, F.HeatFluxBC):
+            if isinstance(bc, boundary_conditions.HeatFluxBC):
                 bc.create_value_fenics(
                     mesh=self.mesh.mesh,
                     temperature=self.u,
@@ -203,7 +212,7 @@ def create_formulation(self):
 
         # add fluxes
         for bc in self.boundary_conditions:
-            if isinstance(bc, F.HeatFluxBC):
+            if isinstance(bc, boundary_conditions.HeatFluxBC):
                 self.formulation -= (
                     bc.value_fenics * self.test_function * self.ds(bc.subdomain.id)
                 )
@@ -213,11 +222,11 @@ def initialise_exports(self):
         a string, find species object in self.species"""
 
         for export in self.exports:
-            if isinstance(export, F.XDMFExport):
+            if isinstance(export, exports.XDMFExport):
                 raise NotImplementedError(
                     "XDMF export is not implemented yet for heat transfer problems"
                 )
-            if isinstance(export, F.VTXTemperatureExport):
+            if isinstance(export, exports.VTXTemperatureExport):
                 self._vtxfile = VTXWriter(
                     self.u.function_space.mesh.comm,
                     export.filename,
@@ -230,7 +239,7 @@ def post_processing(self):
 
         for export in self.exports:
             # TODO if export type derived quantity
-            if isinstance(export, F.SurfaceQuantity):
+            if isinstance(export, exports.SurfaceQuantity):
                 raise NotImplementedError(
                     "SurfaceQuantity export is not implemented yet for heat transfer problems"
                 )
@@ -244,7 +253,7 @@ def post_processing(self):
                 # if filename given write export data to file
                 if export.filename is not None:
                     export.write(t=float(self.t))
-            if isinstance(export, F.XDMFExport):
+            if isinstance(export, exports.XDMFExport):
                 export.write(float(self.t))
 
         if self._vtxfile is not None:
diff --git a/festim/helpers_discontinuity.py b/festim/helpers_discontinuity.py
index 1b3aacf84..2c225a522 100644
--- a/festim/helpers_discontinuity.py
+++ b/festim/helpers_discontinuity.py
@@ -1,6 +1,5 @@
-import numpy as np
-
 import dolfinx
+import numpy as np
 
 
 def transfer_meshtags_to_submesh(
diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 6c614233d..5fd8d7d5e 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -2,19 +2,37 @@
 
 from mpi4py import MPI
 
+import basix
+import dolfinx
 import numpy as np
 import numpy.typing as npt
 import tqdm.autonotebook
-from scifem import NewtonSolver
-
-import basix
-import dolfinx
-import festim as F
 import ufl
 from dolfinx import fem
+from scifem import NewtonSolver
 
-
-class HydrogenTransportProblem(F.ProblemBase):
+import festim.boundary_conditions
+import festim.problem
+from festim import (
+    boundary_conditions,
+    exports,
+    k_B,
+    problem,
+)
+from festim import (
+    reaction as _reaction,
+)
+from festim import (
+    species as _species,
+)
+from festim import (
+    subdomain as _subdomain,
+)
+from festim.helpers import as_fenics_constant
+from festim.mesh import Mesh
+
+
+class HydrogenTransportProblem(problem.ProblemBase):
     """
     Hydrogen Transport Problem.
 
@@ -103,10 +121,12 @@ class HydrogenTransportProblem(F.ProblemBase):
 
     def __init__(
         self,
-        mesh: F.Mesh | None = None,
-        subdomains: list[F.VolumeSubdomain | F.SurfaceSubdomain] | None = None,
-        species: list[F.Species] | None = None,
-        reactions: list[F.Reaction] | None = None,
+        mesh: Mesh | None = None,
+        subdomains: (
+            list[_subdomain.VolumeSubdomain | _subdomain.SurfaceSubdomain] | None
+        ) = None,
+        species: list[_species.Species] | None = None,
+        reactions: list[_reaction.Reaction] | None = None,
         temperature: (
             float
             | int
@@ -196,14 +216,14 @@ def multispecies(self):
         return len(self.species) > 1
 
     @property
-    def species(self) -> list[F.Species]:
+    def species(self) -> list[_species.Species]:
         return self._species
 
     @species.setter
     def species(self, value):
         # check that all species are of type festim.Species
         for spe in value:
-            if not isinstance(spe, F.Species):
+            if not isinstance(spe, _species.Species):
                 raise TypeError(
                     f"elements of species must be of type festim.Species not {
                         type(spe)}"
@@ -246,7 +266,7 @@ def initialise(self):
         if self.settings.transient:
             # TODO should raise error if no stepsize is provided
             # TODO Should this be an attribute of festim.Stepsize?
-            self.dt = F.as_fenics_constant(
+            self.dt = as_fenics_constant(
                 self.settings.stepsize.initial_value, self.mesh.mesh
             )
 
@@ -281,7 +301,7 @@ def define_temperature(self):
 
         # if temperature is a float or int, create a fem.Constant
         elif isinstance(self.temperature, (float, int)):
-            self.temperature_fenics = F.as_fenics_constant(
+            self.temperature_fenics = as_fenics_constant(
                 self.temperature, self.mesh.mesh
             )
         # if temperature is a fem.Constant or function, pass it to temperature_fenics
@@ -298,7 +318,7 @@ def define_temperature(self):
                             type(self.temperature(t=float(self.t)))} "
                     )
                 # only t is an argument
-                self.temperature_fenics = F.as_fenics_constant(
+                self.temperature_fenics = as_fenics_constant(
                     mesh=self.mesh.mesh, value=self.temperature(t=float(self.t))
                 )
             else:
@@ -337,16 +357,18 @@ def initialise_exports(self):
             if isinstance(export.field, list):
                 for idx, field in enumerate(export.field):
                     if isinstance(field, str):
-                        export.field[idx] = F.find_species_from_name(
+                        export.field[idx] = _species.find_species_from_name(
                             field, self.species
                         )
             elif isinstance(export.field, str):
-                export.field = F.find_species_from_name(export.field, self.species)
+                export.field = _species.find_species_from_name(
+                    export.field, self.species
+                )
 
             # Initialize XDMFFile for writer
-            if isinstance(export, F.XDMFExport):
+            if isinstance(export, exports.XDMFExport):
                 export.define_writer(MPI.COMM_WORLD)
-            if isinstance(export, F.VTXSpeciesExport):
+            if isinstance(export, exports.VTXSpeciesExport):
                 functions = export.get_functions()
                 self._vtxfiles.append(
                     dolfinx.io.VTXWriter(
@@ -362,7 +384,7 @@ def initialise_exports(self):
         spe_to_D_global_expr = {}  # links species to D expression
 
         for export in self.exports:
-            if isinstance(export, F.SurfaceQuantity):
+            if isinstance(export, exports.SurfaceQuantity):
                 if export.field in spe_to_D_global:
                     # if already computed then use the same D
                     D = spe_to_D_global[export.field]
@@ -388,7 +410,7 @@ def define_D_global(self, species):
                 coefficient and the expression of the global diffusion coefficient
                 for a given species
         """
-        assert isinstance(species, F.Species)
+        assert isinstance(species, _species.Species)
 
         D_0 = fem.Function(self.V_DG_0)
         E_D = fem.Function(self.V_DG_0)
@@ -403,7 +425,7 @@ def define_D_global(self, species):
         D = fem.Function(self.V_DG_1)
 
         expr = D_0 * ufl.exp(
-            -E_D / F.as_fenics_constant(F.k_B, self.mesh.mesh) / self.temperature_fenics
+            -E_D / as_fenics_constant(k_B, self.mesh.mesh) / self.temperature_fenics
         )
         D_expr = fem.Expression(expr, self.V_DG_1.element.interpolation_points())
         D.interpolate(D_expr)
@@ -429,7 +451,7 @@ def define_function_spaces(self):
         else:
             elements = []
             for spe in self.species:
-                if isinstance(spe, F.Species):
+                if isinstance(spe, _species.Species):
                     elements.append(element_CG)
             element = basix.ufl.mixed_element(elements)
 
@@ -486,7 +508,7 @@ def define_boundary_conditions(self):
         for bc in self.boundary_conditions:
             if isinstance(bc.species, str):
                 # if name of species is given then replace with species object
-                bc.species = F.find_species_from_name(bc.species, self.species)
+                bc.species = _species.find_species_from_name(bc.species, self.species)
 
         super().define_boundary_conditions()
 
@@ -551,7 +573,7 @@ def create_source_values_fenics(self):
         """For each source create the value_fenics"""
         for source in self.sources:
             # create value_fenics for all F.ParticleSource objects
-            if isinstance(source, F.ParticleSource):
+            if isinstance(source, festim.source.ParticleSource):
                 source.create_value_fenics(
                     mesh=self.mesh.mesh,
                     temperature=self.temperature_fenics,
@@ -562,7 +584,7 @@ def create_flux_values_fenics(self):
         """For each particle flux create the value_fenics"""
         for bc in self.boundary_conditions:
             # create value_fenics for all F.ParticleFluxBC objects
-            if isinstance(bc, F.ParticleFluxBC):
+            if isinstance(bc, boundary_conditions.ParticleFluxBC):
                 bc.create_value_fenics(
                     mesh=self.mesh.mesh,
                     temperature=self.temperature_fenics,
@@ -628,7 +650,7 @@ def create_formulation(self):
 
         for reaction in self.reactions:
             for reactant in reaction.reactant:
-                if isinstance(reactant, F.Species):
+                if isinstance(reactant, festim.species.Species):
                     self.formulation += (
                         reaction.reaction_term(self.temperature_fenics)
                         * reactant.test_function
@@ -656,7 +678,7 @@ def create_formulation(self):
 
         # add fluxes
         for bc in self.boundary_conditions:
-            if isinstance(bc, F.ParticleFluxBC):
+            if isinstance(bc, boundary_conditions.ParticleFluxBC):
                 self.formulation -= (
                     bc.value_fenics
                     * bc.species.test_function
@@ -710,7 +732,7 @@ def post_processing(self):
             # variables time dependent
             species_not_updated = self.species.copy()  # make a copy of the species
             for export in self.exports:
-                if isinstance(export, F.SurfaceFlux):
+                if isinstance(export, exports.SurfaceFlux):
                     # if the D of the species has not been updated yet
                     if export.field in species_not_updated:
                         export.D.interpolate(export.D_expr)
@@ -718,9 +740,10 @@ def post_processing(self):
 
         for export in self.exports:
             # TODO if export type derived quantity
-            if isinstance(export, F.SurfaceQuantity):
+            if isinstance(export, exports.SurfaceQuantity):
                 if isinstance(
-                    export, (F.SurfaceFlux, F.TotalSurface, F.AverageSurface)
+                    export,
+                    (exports.SurfaceFlux, exports.TotalSurface, exports.AverageSurface),
                 ):
                     export.compute(
                         self.ds,
@@ -733,8 +756,8 @@ def post_processing(self):
                 # if filename given write export data to file
                 if export.filename is not None:
                     export.write(t=float(self.t))
-            elif isinstance(export, F.VolumeQuantity):
-                if isinstance(export, (F.TotalVolume, F.AverageVolume)):
+            elif isinstance(export, exports.VolumeQuantity):
+                if isinstance(export, (exports.TotalVolume, exports.AverageVolume)):
                     export.compute(self.dx)
                 else:
                     export.compute()
@@ -744,12 +767,12 @@ def post_processing(self):
                 # if filename given write export data to file
                 if export.filename is not None:
                     export.write(t=float(self.t))
-            if isinstance(export, F.XDMFExport):
+            if isinstance(export, exports.XDMFExport):
                 export.write(float(self.t))
 
 
 class HTransportProblemDiscontinuous(HydrogenTransportProblem):
-    interfaces: list[F.Interface]
+    interfaces: list[_subdomain.Interface]
     petsc_options: dict
     surface_to_volume: dict
 
@@ -766,9 +789,9 @@ def __init__(
         settings=None,
         exports=None,
         traps=None,
-        interfaces: list[F.Interface] = None,
-        surface_to_volume: dict = None,
-        petsc_options: dict = None,
+        interfaces: list[_subdomain.Interface] | None = None,
+        surface_to_volume: dict | None = None,
+        petsc_options: dict | None = None,
     ):
         """Class for a multi-material hydrogen transport problem
         For other arguments see ``festim.HydrogenTransportProblem``.
@@ -837,7 +860,7 @@ def initialise(self):
         if self.settings.transient:
             # TODO should raise error if no stepsize is provided
             # TODO Should this be an attribute of festim.Stepsize?
-            self.dt = F.as_fenics_constant(
+            self.dt = as_fenics_constant(
                 self.settings.stepsize.initial_value, self.mesh.mesh
             )
 
@@ -865,7 +888,7 @@ def initialise(self):
         self.create_solver()
         self.initialise_exports()
 
-    def create_dirichletbc_form(self, bc: F.FixedConcentrationBC):
+    def create_dirichletbc_form(self, bc: boundary_conditions.FixedConcentrationBC):
         """
         Creates the ``value_fenics`` attribute for a given
         ``festim.FixedConcentrationBC`` and returns the appropriate
@@ -916,7 +939,7 @@ def create_initial_conditions(self):
                 "initial conditions not yet implemented for discontinuous"
             )
 
-    def define_function_spaces(self, subdomain: F.VolumeSubdomain):
+    def define_function_spaces(self, subdomain: _subdomain.VolumeSubdomain):
         """
         Creates appropriate function space and functions for a given subdomain (submesh)
         based on the number of species existing in this subdomain. Then stores the functionspace,
@@ -971,11 +994,11 @@ def define_function_spaces(self, subdomain: F.VolumeSubdomain):
             species.subdomain_to_collapsed_function_space[subdomain] = V.sub(
                 i
             ).collapse()
-            species.subdomain_to_post_processing_solution[
-                subdomain
-            ].name = f"{species.name}_{subdomain.id}"
+            species.subdomain_to_post_processing_solution[subdomain].name = (
+                f"{species.name}_{subdomain.id}"
+            )
 
-    def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
+    def create_subdomain_formulation(self, subdomain: _subdomain.VolumeSubdomain):
         """
         Creates the variational formulation for each subdomain and stores it in ``subdomain.F``
 
@@ -1006,14 +1029,14 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
             if reaction.volume != subdomain:
                 continue
             for species in reaction.reactant + reaction.product:
-                if isinstance(species, F.Species):
+                if isinstance(species, festim.species.Species):
                     # TODO remove
                     # temporarily overide the solution to the one of the subdomain
                     species.solution = species.subdomain_to_solution[subdomain]
 
             # reactant
             for reactant in reaction.reactant:
-                if isinstance(reactant, F.Species):
+                if isinstance(reactant, festim.species.Species):
                     form += (
                         reaction.reaction_term(self.temperature_fenics)
                         * reactant.subdomain_to_test_function[subdomain]
@@ -1034,7 +1057,7 @@ def create_subdomain_formulation(self, subdomain: F.VolumeSubdomain):
 
         # add fluxes
         for bc in self.boundary_conditions:
-            if isinstance(bc, F.ParticleFluxBC):
+            if isinstance(bc, boundary_conditions.ParticleFluxBC):
                 # check that the bc is applied on a surface
                 # belonging to this subdomain
                 if subdomain == self.surface_to_volume[bc.subdomain]:
@@ -1170,7 +1193,7 @@ def create_solver(self):
     def create_flux_values_fenics(self):
         """For each particle flux create the ``value_fenics`` attribute"""
         for bc in self.boundary_conditions:
-            if isinstance(bc, F.ParticleFluxBC):
+            if isinstance(bc, boundary_conditions.ParticleFluxBC):
                 volume_subdomain = self.surface_to_volume[bc.subdomain]
                 bc.create_value_fenics(
                     mesh=volume_subdomain.submesh,
@@ -1180,7 +1203,7 @@ def create_flux_values_fenics(self):
 
     def initialise_exports(self):
         for export in self.exports:
-            if isinstance(export, F.VTXSpeciesExport):
+            if isinstance(export, exports.VTXSpeciesExport):
                 functions = export.get_functions()
                 self._vtxfiles.append(
                     dolfinx.io.VTXWriter(
@@ -1211,7 +1234,7 @@ def post_processing(self):
             vtxfile.write(float(self.t))
 
         for export in self.exports:
-            if not isinstance(export, F.VTXSpeciesExport):
+            if not isinstance(export, exports.VTXSpeciesExport):
                 raise NotImplementedError(f"Export type {type(export)} not implemented")
 
     def iterate(self):
diff --git a/festim/initial_condition.py b/festim/initial_condition.py
index 6270b59e3..b4947c521 100644
--- a/festim/initial_condition.py
+++ b/festim/initial_condition.py
@@ -1,5 +1,4 @@
 import numpy as np
-
 import ufl
 from dolfinx import fem
 
diff --git a/festim/material.py b/festim/material.py
index 8cbf3cf62..fdc176d92 100644
--- a/festim/material.py
+++ b/festim/material.py
@@ -1,6 +1,7 @@
-import festim as F
 import ufl
 
+import festim as F
+
 
 class Material:
     """
diff --git a/festim/mesh/__init__.py b/festim/mesh/__init__.py
index e69de29bb..6d931868a 100644
--- a/festim/mesh/__init__.py
+++ b/festim/mesh/__init__.py
@@ -0,0 +1,3 @@
+__all__ = ["Mesh"]
+
+from .mesh import Mesh
diff --git a/festim/mesh/mesh.py b/festim/mesh/mesh.py
index bf4d53ef6..87ce15329 100644
--- a/festim/mesh/mesh.py
+++ b/festim/mesh/mesh.py
@@ -1,6 +1,5 @@
-import numpy as np
-
 import dolfinx
+import numpy as np
 import ufl
 from dolfinx.mesh import meshtags
 
diff --git a/festim/mesh/mesh_1d.py b/festim/mesh/mesh_1d.py
index 322ead5a0..6b62684cb 100644
--- a/festim/mesh/mesh_1d.py
+++ b/festim/mesh/mesh_1d.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
-import numpy as np
-
 import basix.ufl
 import dolfinx.mesh
-import festim as F
+import numpy as np
 import ufl
 
+import festim as F
+
 
 class Mesh1D(F.Mesh):
     """
diff --git a/festim/mesh/mesh_from_xdmf.py b/festim/mesh/mesh_from_xdmf.py
index 4b744554d..6f7429ce2 100644
--- a/festim/mesh/mesh_from_xdmf.py
+++ b/festim/mesh/mesh_from_xdmf.py
@@ -1,8 +1,9 @@
 from mpi4py import MPI
 
-import festim as F
 from dolfinx.io import XDMFFile
 
+import festim as F
+
 
 class MeshFromXDMF(F.Mesh):
     """
diff --git a/festim/problem.py b/festim/problem.py
index f26084ef3..966baa8b8 100644
--- a/festim/problem.py
+++ b/festim/problem.py
@@ -4,11 +4,11 @@
 
 import numpy as np
 import tqdm.autonotebook
-
-import festim as F
 import ufl
 from dolfinx import fem
 from dolfinx.nls.petsc import NewtonSolver
+
+import festim as F
 from festim.mesh.mesh import Mesh as _Mesh
 from festim.source import SourceBase as _SourceBase
 from festim.subdomain.volume_subdomain import (
diff --git a/festim/reaction.py b/festim/reaction.py
index 23cdac11b..945a98252 100644
--- a/festim/reaction.py
+++ b/festim/reaction.py
@@ -1,10 +1,11 @@
 from typing import List, Optional, Union
 
+from ufl import exp
+
 from festim import k_B as _k_B
 from festim.species import ImplicitSpecies as _ImplicitSpecies
 from festim.species import Species as _Species
 from festim.subdomain.volume_subdomain_1d import VolumeSubdomain1D as VS1D
-from ufl import exp
 
 
 class Reaction:
diff --git a/festim/source.py b/festim/source.py
index d53ebf315..a792dda49 100644
--- a/festim/source.py
+++ b/festim/source.py
@@ -1,10 +1,10 @@
-import numpy as np
-
 import dolfinx
-import festim as F
+import numpy as np
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 class SourceBase:
     """
diff --git a/festim/subdomain/__init__.py b/festim/subdomain/__init__.py
index e69de29bb..fe4f09326 100644
--- a/festim/subdomain/__init__.py
+++ b/festim/subdomain/__init__.py
@@ -0,0 +1,6 @@
+__all__ = ["VolumeSubdomain", "SurfaceSubdomain", "Interface"]
+
+
+from .interface import Interface
+from .surface_subdomain import SurfaceSubdomain
+from .volume_subdomain import VolumeSubdomain
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index 3c3a4aac3..c2881e2e5 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -1,13 +1,13 @@
-import numpy as np
-
 import dolfinx
-import festim as F
+import numpy as np
 from dolfinx.cpp.fem import compute_integration_domains
 
+from festim.subdomain import VolumeSubdomain
+
 
 class Interface:
     id: int
-    subdomains: tuple[F.VolumeSubdomain, F.VolumeSubdomain]
+    subdomains: tuple[VolumeSubdomain, VolumeSubdomain]
     parent_mesh: dolfinx.mesh.Mesh
     mt: dolfinx.mesh.MeshTags
     restriction: list[str, str] = ("+", "-")
@@ -16,7 +16,7 @@ class Interface:
     def __init__(
         self,
         id: int,
-        subdomains: list[F.VolumeSubdomain],
+        subdomains: list[VolumeSubdomain],
         penalty_term: float = 10.0,
     ):
         """Class representing an interface between two subdomains.
diff --git a/festim/subdomain/surface_subdomain.py b/festim/subdomain/surface_subdomain.py
index b28b550c0..578545804 100644
--- a/festim/subdomain/surface_subdomain.py
+++ b/festim/subdomain/surface_subdomain.py
@@ -1,6 +1,5 @@
-import numpy as np
-
 import dolfinx.mesh
+import numpy as np
 
 
 class SurfaceSubdomain:
diff --git a/festim/subdomain/surface_subdomain_1d.py b/festim/subdomain/surface_subdomain_1d.py
index e757914ff..a14827122 100644
--- a/festim/subdomain/surface_subdomain_1d.py
+++ b/festim/subdomain/surface_subdomain_1d.py
@@ -1,6 +1,6 @@
+import dolfinx.mesh
 import numpy as np
 
-import dolfinx.mesh
 import festim as F
 
 
diff --git a/festim/subdomain/volume_subdomain.py b/festim/subdomain/volume_subdomain.py
index fb23bf520..77f6b9bb2 100644
--- a/festim/subdomain/volume_subdomain.py
+++ b/festim/subdomain/volume_subdomain.py
@@ -1,7 +1,7 @@
-import numpy as np
-
 import dolfinx
+import numpy as np
 from dolfinx.mesh import locate_entities
+
 from festim.helpers_discontinuity import transfer_meshtags_to_submesh
 
 
diff --git a/festim/subdomain/volume_subdomain_1d.py b/festim/subdomain/volume_subdomain_1d.py
index f0b4a0acd..a024ac8c1 100644
--- a/festim/subdomain/volume_subdomain_1d.py
+++ b/festim/subdomain/volume_subdomain_1d.py
@@ -1,10 +1,10 @@
 import numpy as np
-
-import festim as F
 from dolfinx.mesh import locate_entities
 
+from festim.subdomain import VolumeSubdomain
+
 
-class VolumeSubdomain1D(F.VolumeSubdomain):
+class VolumeSubdomain1D(VolumeSubdomain):
     """
     Volume subdomain class for 1D cases
 
diff --git a/pyproject.toml b/pyproject.toml
index 96f570659..25f5e3130 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,6 +39,18 @@ Issues = "https://github.com/RemDelaporteMathurin/FESTIM/issues"
 [tool.setuptools_scm]
 write_to = "festim/_version.py"
 
+[tool.mypy]
+# Suggested at https://blog.wolt.com/engineering/2021/09/30/professional-grade-mypy-configuration/
+# Goal would be to make all of the below True long-term
+disallow_untyped_defs = false
+disallow_any_unimported = false
+no_implicit_optional = false
+check_untyped_defs = false
+warn_return_any = false
+warn_unused_ignores = false
+show_error_codes = true
+ignore_missing_imports = true
+
 [tool.ruff]
 line-length = 88
 indent-width = 4
@@ -61,8 +73,19 @@ ignore = ["UP007", "RUF012"]
 allowed-confusables = ["σ"]
 
 [tool.ruff.lint.isort]
-known-first-party = ["festim", "basix", "dolfinx", "ffcx", "ufl"]
-known-third-party = ["gmsh", "numba", "numpy", "pytest", "pyvista", "pyamg"]
+known-first-party = ["festim"]
+known-third-party = [
+    "basix",
+    "dolfinx",
+    "ffcx",
+    "ufl",
+    "gmsh",
+    "numba",
+    "numpy",
+    "pytest",
+    "pyvista",
+    "pyamg",
+]
 section-order = [
     "future",
     "standard-library",
diff --git a/test/benchmark.py b/test/benchmark.py
index 0fbdec4a5..4b6b6312c 100644
--- a/test/benchmark.py
+++ b/test/benchmark.py
@@ -3,11 +3,9 @@
 from mpi4py import MPI
 from petsc4py import PETSc
 
+import basix
 import numpy as np
 import tqdm.autonotebook
-from test_permeation_problem import test_permeation_problem
-
-import basix
 from dolfinx.fem import (
     Constant,
     Function,
@@ -24,6 +22,7 @@
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import create_mesh, locate_entities, meshtags
 from dolfinx.nls.petsc import NewtonSolver
+from test_permeation_problem import test_permeation_problem
 from ufl import (
     FacetNormal,
     Measure,
diff --git a/test/system_tests/test_1_mobile_1_trap.py b/test/system_tests/test_1_mobile_1_trap.py
index 98a1e5ad5..66236cdfa 100644
--- a/test/system_tests/test_1_mobile_1_trap.py
+++ b/test/system_tests/test_1_mobile_1_trap.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 from dolfinx.mesh import create_unit_cube, create_unit_square, locate_entities
 
+import festim as F
+
 from .tools import error_L2
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
diff --git a/test/system_tests/test_1_mobile_species.py b/test/system_tests/test_1_mobile_species.py
index 64da01ac8..7e5d3a000 100644
--- a/test/system_tests/test_1_mobile_species.py
+++ b/test/system_tests/test_1_mobile_species.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 from dolfinx.mesh import create_unit_cube, create_unit_square, locate_entities
 
+import festim as F
+
 from .tools import error_L2
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
diff --git a/test/system_tests/test_fluxes.py b/test/system_tests/test_fluxes.py
index a01bfacb9..1d3237a86 100644
--- a/test/system_tests/test_fluxes.py
+++ b/test/system_tests/test_fluxes.py
@@ -1,9 +1,9 @@
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 from .tools import error_L2
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
diff --git a/test/system_tests/test_heat_transfer.py b/test/system_tests/test_heat_transfer.py
index 9486b73b4..ded5284f1 100644
--- a/test/system_tests/test_heat_transfer.py
+++ b/test/system_tests/test_heat_transfer.py
@@ -1,9 +1,9 @@
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 from .tools import error_L2
 
 test_mesh_1d = F.Mesh1D(np.linspace(0, 1, 10000))
diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index 1b2343088..fd0df647f 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
-import numpy as np
-
 import dolfinx
 import dolfinx.fem.petsc
-import festim as F
+import numpy as np
 import ufl
 
+import festim as F
+
 from .tools import error_L2
 
 
diff --git a/test/system_tests/tools.py b/test/system_tests/tools.py
index 9d13aa670..52a7df23a 100644
--- a/test/system_tests/tools.py
+++ b/test/system_tests/tools.py
@@ -1,7 +1,6 @@
 import mpi4py.MPI as MPI
 
 import numpy as np
-
 import ufl
 from dolfinx import fem
 
diff --git a/test/test_average_surface.py b/test/test_average_surface.py
index 131ff6d6a..fdf65f810 100644
--- a/test/test_average_surface.py
+++ b/test/test_average_surface.py
@@ -1,9 +1,9 @@
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 def test_average_surface_compute_1D():
     """Test that the average surface export computes the correct value"""
diff --git a/test/test_average_volume.py b/test/test_average_volume.py
index f9370c9c5..71fe6b1d6 100644
--- a/test/test_average_volume.py
+++ b/test/test_average_volume.py
@@ -1,9 +1,9 @@
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy")
 
 
diff --git a/test/test_dirichlet_bc.py b/test/test_dirichlet_bc.py
index fc49ccbad..259bf49fa 100644
--- a/test/test_dirichlet_bc.py
+++ b/test/test_dirichlet_bc.py
@@ -1,14 +1,14 @@
 from mpi4py import MPI
 
+import dolfinx.mesh
 import numpy as np
 import pytest
-
-import dolfinx.mesh
-import festim as F
 import ufl
-from dolfinx import fem
+from dolfinx import default_scalar_type, fem
 from ufl.conditional import Conditional
 
+import festim as F
+
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy_mat")
 
 mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
@@ -57,7 +57,10 @@ def test_callable_for_value():
 
     subdomain = F.SurfaceSubdomain1D(1, x=1)
     vol_subdomain = F.VolumeSubdomain1D(1, borders=[0, 1], material=dummy_mat)
-    value = lambda x, t: 1.0 + x[0] + t
+
+    def value(x, t):
+        return 1.0 + x[0] + t
+
     species = F.Species("test")
 
     bc = F.DirichletBC(subdomain, value, species)
@@ -97,7 +100,10 @@ def test_value_callable_x_t_T():
 
     subdomain = F.SurfaceSubdomain1D(1, x=1)
     vol_subdomain = F.VolumeSubdomain1D(1, borders=[0, 1], material=dummy_mat)
-    value = lambda x, t, T: 1.0 + x[0] + t + T
+
+    def value(x, t, T):
+        return 1.0 + x[0] + t + T
+
     species = F.Species("test")
 
     bc = F.DirichletBC(subdomain, value, species)
@@ -183,7 +189,10 @@ def test_callable_x_only():
     # BUILD
     subdomain = F.SurfaceSubdomain1D(1, x=1)
     vol_subdomain = F.VolumeSubdomain1D(1, borders=[0, 1], material=dummy_mat)
-    value = lambda x: 1.0 + x[0]
+
+    def value(x):
+        return 1.0 + x[0]
+
     species = F.Species("test")
 
     bc = F.DirichletBC(subdomain, value, species)
@@ -416,7 +425,7 @@ def test_integration_with_a_multispecies_HTransportProblem(value_A, value_B):
     [
         (1.0, False),
         (None, False),
-        (fem.Constant(mesh, 1.0), False),
+        (fem.Constant(mesh, default_scalar_type(1.0)), False),
         (lambda t: t, True),
         (lambda t: 1.0 + t, True),
         (lambda x: 1.0 + x[0], False),
@@ -439,7 +448,7 @@ def test_bc_time_dependent_attribute(input, expected_value):
     [
         (1.0, False),
         (None, False),
-        (fem.Constant(mesh, 1.0), False),
+        (fem.Constant(mesh, default_scalar_type(1.0)), False),
         (lambda T: T, True),
         (lambda t: 1.0 + t, False),
         (lambda x, T: 1.0 + x[0] + T, True),
diff --git a/test/test_flux_bc.py b/test/test_flux_bc.py
index ccd655161..04ff92183 100644
--- a/test/test_flux_bc.py
+++ b/test/test_flux_bc.py
@@ -1,13 +1,13 @@
 from mpi4py import MPI
 
+import dolfinx.mesh
 import numpy as np
 import pytest
-
-import dolfinx.mesh
-import festim as F
 import ufl
 import ufl.core
-from dolfinx import fem
+from dolfinx import default_scalar_type, fem
+
+import festim as F
 
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy_mat")
 mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
@@ -150,7 +150,7 @@ def my_value(t):
     "input, expected_value",
     [
         (1.0, False),
-        (fem.Constant(mesh, 1.0), False),
+        (fem.Constant(mesh, default_scalar_type(1.0)), False),
         (lambda t: t, True),
         (lambda t: 1.0 + t, True),
         (lambda x: 1.0 + x[0], False),
@@ -185,7 +185,7 @@ def test_bc_time_dependent_attribute_raises_error_when_value_none():
     [
         (1.0, False),
         (None, False),
-        (fem.Constant(mesh, 1.0), False),
+        (fem.Constant(mesh, default_scalar_type(1.0)), False),
         (lambda T: T, True),
         (lambda t: 1.0 + t, False),
         (lambda x, T: 1.0 + x[0] + T, True),
diff --git a/test/test_h_transport_problem.py b/test/test_h_transport_problem.py
index e085d266a..285c6dcd2 100644
--- a/test/test_h_transport_problem.py
+++ b/test/test_h_transport_problem.py
@@ -1,13 +1,13 @@
 import mpi4py.MPI as MPI
 
+import dolfinx.mesh
 import numpy as np
 import pytest
 import tqdm.autonotebook
+import ufl
+from dolfinx import default_scalar_type, fem, nls
 
-import dolfinx.mesh
 import festim as F
-import ufl
-from dolfinx import fem, nls
 
 test_mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
 x = ufl.SpatialCoordinate(test_mesh.mesh)
@@ -17,7 +17,13 @@
 # TODO test all the methods in the class
 @pytest.mark.parametrize(
     "value",
-    [1, fem.Constant(test_mesh.mesh, 1.0), 1.0, "coucou", lambda x: 2 * x[0]],
+    [
+        1,
+        fem.Constant(test_mesh.mesh, default_scalar_type(1.0)),
+        1.0,
+        "coucou",
+        lambda x: 2 * x[0],
+    ],
 )
 def test_temperature_setter_type(value):
     """Test that the temperature type is correctly set"""
@@ -74,7 +80,7 @@ def test_define_temperature_value_error_raised():
     [
         (1.0, fem.Constant),
         (1, fem.Constant),
-        (fem.Constant(test_mesh.mesh, 1.0), fem.Constant),
+        (fem.Constant(test_mesh.mesh, default_scalar_type(1.0)), fem.Constant),
         (lambda t: t, fem.Constant),
         (lambda t: 1.0 + t, fem.Constant),
         (lambda x: 1.0 + x[0], fem.Function),
@@ -1166,7 +1172,7 @@ def test_update_fluxes_with_time_dependent_temperature(
             assert np.isclose(computed_value, expected_values[i])
 
 
-def test_create_source_values_fenics_multispecies():
+def test_create_flux_values_fenics_multispecies():
     """Test that the create_flux_values_fenics method correctly sets the value_fenics
     attribute in a multispecies case"""
     # BUILD
diff --git a/test/test_heat_transfer_problem.py b/test/test_heat_transfer_problem.py
index b8b138e45..7a02eb0b4 100644
--- a/test/test_heat_transfer_problem.py
+++ b/test/test_heat_transfer_problem.py
@@ -2,16 +2,16 @@
 
 import mpi4py.MPI as MPI
 
+import dolfinx
 import numpy as np
 import pytest
 import tqdm.autonotebook
-
-import dolfinx
-import festim as F
 import ufl
 from dolfinx import fem
 from dolfinx.io import XDMFFile
 
+import festim as F
+
 
 def source_from_exact_solution(
     exact_solution, thermal_conductivity, density, heat_capacity
@@ -218,7 +218,9 @@ def exact_solution(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = my_problem.t.value  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = (
+        my_problem.t.value
+    )  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)
@@ -341,7 +343,9 @@ def mms_source(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = my_problem.t.value  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = (
+        my_problem.t.value
+    )  # we use the exact final time of the simulation which may differ from the one specified in the settings
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)
diff --git a/test/test_helpers.py b/test/test_helpers.py
index e9c70f138..9220b1451 100644
--- a/test/test_helpers.py
+++ b/test/test_helpers.py
@@ -1,16 +1,23 @@
 import numpy as np
 import pytest
+import ufl
+from dolfinx import default_scalar_type, fem
 
 import festim as F
-import ufl
-from dolfinx import fem
 
 test_mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
 x = ufl.SpatialCoordinate(test_mesh.mesh)
 
 
 @pytest.mark.parametrize(
-    "value", [1, fem.Constant(test_mesh.mesh, 1.0), 1.0, "coucou", 2 * x[0]]
+    "value",
+    [
+        1,
+        fem.Constant(test_mesh.mesh, default_scalar_type(1.0)),
+        1.0,
+        "coucou",
+        2 * x[0],
+    ],
 )
 def test_temperature_type_and_processing(value):
     """Test that the temperature type is correctly set"""
diff --git a/test/test_henrysbc.py b/test/test_henrysbc.py
index 0bc4ca41b..998ea78e5 100644
--- a/test/test_henrysbc.py
+++ b/test/test_henrysbc.py
@@ -1,13 +1,13 @@
 from mpi4py import MPI
 
+import dolfinx.mesh
 import numpy as np
 import pytest
-
-import dolfinx.mesh
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 def henrys_law(T, H_0, E_H, pressure):
     """Applies the Henrys law to compute the concentration at the boundary"""
diff --git a/test/test_initial_condition.py b/test/test_initial_condition.py
index f8b6eb459..b67a50c85 100644
--- a/test/test_initial_condition.py
+++ b/test/test_initial_condition.py
@@ -2,9 +2,9 @@
 
 import numpy as np
 import pytest
+from dolfinx import fem
 
 import festim as F
-from dolfinx import fem
 
 dummy_mat = F.Material(D_0=1, E_D=0.1, name="dummy_mat")
 test_mesh = F.Mesh1D(np.linspace(0, 1, 100))
diff --git a/test/test_maximum_surface.py b/test/test_maximum_surface.py
index 2e5c0c649..da1dec4f1 100644
--- a/test/test_maximum_surface.py
+++ b/test/test_maximum_surface.py
@@ -1,7 +1,7 @@
 import numpy as np
+from dolfinx import fem
 
 import festim as F
-from dolfinx import fem
 
 
 def test_maximum_surface_compute_1D():
diff --git a/test/test_maximum_volume.py b/test/test_maximum_volume.py
index 6a4fff030..3fed21f5e 100644
--- a/test/test_maximum_volume.py
+++ b/test/test_maximum_volume.py
@@ -1,7 +1,7 @@
 import numpy as np
+from dolfinx import fem
 
 import festim as F
-from dolfinx import fem
 
 
 def test_maximum_volume_compute_1D():
diff --git a/test/test_mesh.py b/test/test_mesh.py
index 74d6caeb0..5b94e6084 100644
--- a/test/test_mesh.py
+++ b/test/test_mesh.py
@@ -4,12 +4,12 @@
 
 import numpy as np
 import pytest
-
-import festim as F
 from dolfinx import mesh as fenics_mesh
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import meshtags
 
+import festim as F
+
 mesh_1D = fenics_mesh.create_unit_interval(MPI.COMM_WORLD, 10)
 mesh_2D = fenics_mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
 mesh_3D = fenics_mesh.create_unit_cube(MPI.COMM_WORLD, 10, 10, 10)
diff --git a/test/test_minimum_surface.py b/test/test_minimum_surface.py
index 9d4ec5467..eb78ee96b 100644
--- a/test/test_minimum_surface.py
+++ b/test/test_minimum_surface.py
@@ -1,7 +1,7 @@
 import numpy as np
+from dolfinx import fem
 
 import festim as F
-from dolfinx import fem
 
 
 def test_minimum_surface_export_compute_1D():
diff --git a/test/test_minimum_volume.py b/test/test_minimum_volume.py
index affa829aa..139fb8a10 100644
--- a/test/test_minimum_volume.py
+++ b/test/test_minimum_volume.py
@@ -1,7 +1,7 @@
 import numpy as np
+from dolfinx import fem
 
 import festim as F
-from dolfinx import fem
 
 
 def test_minimum_volume_compute_1D():
diff --git a/test/test_multispecies_problem.py b/test/test_multispecies_problem.py
index 552e24133..b33fc0e43 100644
--- a/test/test_multispecies_problem.py
+++ b/test/test_multispecies_problem.py
@@ -1,11 +1,11 @@
 from petsc4py import PETSc
 
 import numpy as np
-
-import festim as F
 from dolfinx.fem import Constant
 from ufl import exp
 
+import festim as F
+
 
 def relative_error_computed_to_analytical(
     D, permeability, computed_flux, L, times, P_up
diff --git a/test/test_permeation_problem.py b/test/test_permeation_problem.py
index d6c4dcc57..f03df9bd8 100644
--- a/test/test_permeation_problem.py
+++ b/test/test_permeation_problem.py
@@ -3,11 +3,11 @@
 from petsc4py import PETSc
 
 import numpy as np
-
-import festim as F
 from dolfinx.fem import Constant
 from ufl import exp
 
+import festim as F
+
 
 def relative_error_computed_to_analytical(
     D, permeability, computed_flux, L, times, P_up
diff --git a/test/test_reaction.py b/test/test_reaction.py
index 6842f29f4..92ca9c0ec 100644
--- a/test/test_reaction.py
+++ b/test/test_reaction.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
 import pytest
-
-import festim as F
 from dolfinx.fem import Function, functionspace
 from dolfinx.mesh import create_unit_cube
 from ufl import exp
 
+import festim as F
+
 my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 1], material=None)
 
 
diff --git a/test/test_sievertsbc.py b/test/test_sievertsbc.py
index c25be1404..a70d697a5 100644
--- a/test/test_sievertsbc.py
+++ b/test/test_sievertsbc.py
@@ -1,13 +1,13 @@
 from mpi4py import MPI
 
+import dolfinx.mesh
 import numpy as np
 import pytest
-
-import dolfinx.mesh
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 def sieverts_law(T, S_0, E_S, pressure):
     """Applies the Sieverts law to compute the concentration at the boundary"""
diff --git a/test/test_source.py b/test/test_source.py
index 614c5c21c..9bddca94f 100644
--- a/test/test_source.py
+++ b/test/test_source.py
@@ -1,12 +1,12 @@
 from mpi4py import MPI
 
-import pytest
-
 import dolfinx.mesh
-import festim as F
+import pytest
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy_mat")
 
 mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
@@ -92,7 +92,7 @@ def test_create_value_fenics(value, expected_type):
     [
         (1.0, False),
         (None, False),
-        (fem.Constant(mesh, 1.0), False),
+        (fem.Constant(mesh, dolfinx.default_scalar_type(1.0)), False),
         (lambda t: t, True),
         (lambda t: 1.0 + t, True),
         (lambda x: 1.0 + x[0], False),
@@ -115,7 +115,7 @@ def test_source_time_dependent_attribute(input, expected_value):
     [
         (1.0, False),
         (None, False),
-        (fem.Constant(mesh, 1.0), False),
+        (fem.Constant(mesh, dolfinx.default_scalar_type(1.0)), False),
         (lambda T: T, True),
         (lambda t: 1.0 + t, False),
         (lambda x, T: 1.0 + x[0] + T, True),
diff --git a/test/test_species.py b/test/test_species.py
index 7183587d8..827305106 100644
--- a/test/test_species.py
+++ b/test/test_species.py
@@ -1,14 +1,14 @@
 from mpi4py import MPI
 
+import dolfinx
 import numpy as np
 import pytest
-
-import dolfinx
-import festim as F
 import ufl
 from dolfinx.fem import Function, functionspace
 from dolfinx.mesh import create_unit_cube
 
+import festim as F
+
 
 def test_assign_functions_to_species():
     """Test that checks if the function assign_functions_to_species
diff --git a/test/test_subdomains.py b/test/test_subdomains.py
index 3b537900c..ca63c15c2 100644
--- a/test/test_subdomains.py
+++ b/test/test_subdomains.py
@@ -1,9 +1,9 @@
 from mpi4py import MPI
 
+import dolfinx.mesh
 import numpy as np
 import pytest
 
-import dolfinx.mesh
 import festim as F
 
 
diff --git a/test/test_surface_quantity.py b/test/test_surface_quantity.py
index 8680f356a..1dceb9c18 100644
--- a/test/test_surface_quantity.py
+++ b/test/test_surface_quantity.py
@@ -2,12 +2,12 @@
 
 import numpy as np
 import pytest
-
-import festim as F
 import ufl
 from dolfinx import fem
 from dolfinx.mesh import meshtags
 
+import festim as F
+
 
 def test_surface_flux_export_compute():
     """Test that the surface flux export computes the correct value"""
diff --git a/test/test_total_surface.py b/test/test_total_surface.py
index 18fbca27a..2f35e2c26 100644
--- a/test/test_total_surface.py
+++ b/test/test_total_surface.py
@@ -1,9 +1,9 @@
 import numpy as np
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 
 def test_total_surface_compute_1D():
     """Test that the total surface export computes the correct value"""
diff --git a/test/test_total_volume.py b/test/test_total_volume.py
index ae1908bb1..129fa2294 100644
--- a/test/test_total_volume.py
+++ b/test/test_total_volume.py
@@ -2,11 +2,11 @@
 
 import numpy as np
 import pytest
-
-import festim as F
 import ufl
 from dolfinx import fem
 
+import festim as F
+
 dummy_mat = F.Material(D_0=1, E_D=1, name="dummy")
 
 
diff --git a/test/test_vtx.py b/test/test_vtx.py
index daa597c0b..62b16af48 100644
--- a/test/test_vtx.py
+++ b/test/test_vtx.py
@@ -1,9 +1,9 @@
 from mpi4py import MPI
 
+import dolfinx
 import numpy as np
 import pytest
 
-import dolfinx
 import festim as F
 
 mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 10, 10, 10)
diff --git a/test/test_xdmf.py b/test/test_xdmf.py
index d386bd507..b4ba1e9a5 100644
--- a/test/test_xdmf.py
+++ b/test/test_xdmf.py
@@ -5,9 +5,9 @@
 
 import numpy as np
 import pytest
+from dolfinx import fem, mesh
 
 import festim as F
-from dolfinx import fem, mesh
 
 
 def test_init():
@@ -69,7 +69,10 @@ def test_field_attribute_is_always_list():
 
 @pytest.mark.parametrize("field", [["H", 2], 1, [F.Species("H"), 1]])
 def test_field_attribute_raises_error_when_invalid_type(field):
-    """Test that the field attribute raises an error if the type is not festim.Species or list"""
+    """
+    Test that the field attribute raises an error if the type
+    is not festim.Species or list.
+    """
     with pytest.raises(TypeError):
         F.XDMFExport("my_export.xdmf", field=field)
 

From d3883ec58ebe61a569137ac33a412a4953e43dab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 09:49:02 -0400
Subject: [PATCH 105/134] Import changes

---
 festim/exports/average_surface.py |  4 ++--
 festim/exports/average_volume.py  |  2 +-
 festim/exports/surface_flux.py    |  2 +-
 festim/exports/total_volume.py    |  2 +-
 festim/exports/volume_quantity.py |  4 ++--
 festim/subdomain/interface.py     |  8 +++++---
 test/benchmark.py                 | 10 +++++++---
 7 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/festim/exports/average_surface.py b/festim/exports/average_surface.py
index 5dee531aa..37b485bf2 100644
--- a/festim/exports/average_surface.py
+++ b/festim/exports/average_surface.py
@@ -1,9 +1,9 @@
 from dolfinx import fem
 
-import festim.exports.surface_quantity as sq
+from .surface_quantity import SurfaceQuantity
 
 
-class AverageSurface(sq.SurfaceQuantity):
+class AverageSurface(SurfaceQuantity):
     """Computes the average value of a field on a given surface
 
     Args:
diff --git a/festim/exports/average_volume.py b/festim/exports/average_volume.py
index e7677e341..f3e62b22f 100644
--- a/festim/exports/average_volume.py
+++ b/festim/exports/average_volume.py
@@ -1,6 +1,6 @@
 from dolfinx import fem
 
-from festim.exports import VolumeQuantity
+from .volume_quantity import VolumeQuantity
 
 
 class AverageVolume(VolumeQuantity):
diff --git a/festim/exports/surface_flux.py b/festim/exports/surface_flux.py
index 7ded96409..c08366e93 100644
--- a/festim/exports/surface_flux.py
+++ b/festim/exports/surface_flux.py
@@ -1,7 +1,7 @@
 import ufl
 from dolfinx import fem
 
-from festim.exports import SurfaceQuantity
+from .surface_quantity import SurfaceQuantity
 
 
 class SurfaceFlux(SurfaceQuantity):
diff --git a/festim/exports/total_volume.py b/festim/exports/total_volume.py
index 0fadd6f4d..7be1214c9 100644
--- a/festim/exports/total_volume.py
+++ b/festim/exports/total_volume.py
@@ -1,7 +1,7 @@
 import ufl
 from dolfinx import fem
 
-from festim.exports import VolumeQuantity
+from .volume_quantity import VolumeQuantity
 
 
 class TotalVolume(VolumeQuantity):
diff --git a/festim/exports/volume_quantity.py b/festim/exports/volume_quantity.py
index 37ccc4edb..0a5b70001 100644
--- a/festim/exports/volume_quantity.py
+++ b/festim/exports/volume_quantity.py
@@ -1,7 +1,7 @@
 import csv
 import os
 
-import festim as F
+from festim import species
 
 
 class VolumeQuantity:
@@ -50,7 +50,7 @@ def field(self):
     @field.setter
     def field(self, value):
         # check that field is festim.Species
-        if not isinstance(value, (F.Species, str)):
+        if not isinstance(value, (species.Species, str)):
             raise TypeError("field must be of type festim.Species")
 
         self._field = value
diff --git a/festim/subdomain/interface.py b/festim/subdomain/interface.py
index c2881e2e5..3d5d3b5fe 100644
--- a/festim/subdomain/interface.py
+++ b/festim/subdomain/interface.py
@@ -2,7 +2,7 @@
 import numpy as np
 from dolfinx.cpp.fem import compute_integration_domains
 
-from festim.subdomain import VolumeSubdomain
+from .volume_subdomain import VolumeSubdomain
 
 
 class Interface:
@@ -58,8 +58,10 @@ def compute_mapped_interior_facet_data(self, mesh):
         Returns
             integration_data: Integration data for interior facets
         """
-        assert (not self.subdomains[0].padded) and (not self.subdomains[1].padded)
-        mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+        assert (not self.subdomains[0].padded) and (
+            not self.subdomains[1].padded)
+        mesh.topology.create_connectivity(
+            mesh.topology.dim - 1, mesh.topology.dim)
         integration_data = compute_integration_domains(
             dolfinx.fem.IntegralType.interior_facet,
             mesh.topology._cpp_object,
diff --git a/test/benchmark.py b/test/benchmark.py
index 4b6b6312c..7fb74631b 100644
--- a/test/benchmark.py
+++ b/test/benchmark.py
@@ -71,7 +71,8 @@ def fenics_test_permeation_problem(mesh_size=1001):
     mesh_tags_volumes = meshtags(my_mesh, vdim, cells, markers)
 
     dofs_L = locate_dofs_geometrical(V, lambda x: np.isclose(x[0], 0))
-    dofs_R = locate_dofs_geometrical(V, lambda x: np.isclose(x[0], indices[-1]))
+    dofs_R = locate_dofs_geometrical(
+        V, lambda x: np.isclose(x[0], indices[-1]))
 
     dofs_facets = np.array([dofs_L[0], dofs_R[0]], dtype=np.int32)
     tags_facets = np.array([1, 2], dtype=np.int32)
@@ -89,6 +90,7 @@ def siverts_law(T, S_0, E_S, pressure):
         return S * pressure**0.5
 
     left_facets = mesh_tags_facets.find(1)
+    my_mesh.topology.create_connectivity(fdim, my_mesh.topology.dim)
     left_dofs = locate_dofs_topological(V, fdim, left_facets)
     right_facets = mesh_tags_facets.find(2)
     right_dofs = locate_dofs_topological(V, fdim, right_facets)
@@ -100,7 +102,8 @@ def siverts_law(T, S_0, E_S, pressure):
     bc_sieverts = dirichletbc(
         Constant(my_mesh, PETSc.ScalarType(surface_conc)), left_dofs, V
     )
-    bc_outgas = dirichletbc(Constant(my_mesh, PETSc.ScalarType(0)), right_dofs, V)
+    bc_outgas = dirichletbc(
+        Constant(my_mesh, PETSc.ScalarType(0)), right_dofs, V)
     bcs = [bc_sieverts, bc_outgas]
 
     D_0 = 1.9e-7
@@ -196,7 +199,8 @@ def test_festim_vs_fenics_permeation_benchmark():
     threshold = -0.1
     if diff < threshold:
         raise ValueError(
-            f"festim is {np.abs(diff):.1%} slower than fenics, current acceptable threshold of {np.abs(threshold):.1%}"
+            f"festim is {np.abs(diff):.1%} slower than fenics, current acceptable threshold of {
+                np.abs(threshold):.1%}"
         )
     else:
         print(f"avg relative diff between festim and fenics {diff:.1%}")

From e36a70cf5446942a1ad64c15b230c00d5af8f3df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 09:55:14 -0400
Subject: [PATCH 106/134] Move to source layout

---
 pyproject.toml                                             | 2 +-
 {festim => src/festim}/__init__.py                         | 0
 {festim => src/festim}/boundary_conditions/__init__.py     | 0
 {festim => src/festim}/boundary_conditions/dirichlet_bc.py | 0
 {festim => src/festim}/boundary_conditions/flux_bc.py      | 0
 {festim => src/festim}/boundary_conditions/henrys_bc.py    | 0
 {festim => src/festim}/boundary_conditions/sieverts_bc.py  | 0
 {festim => src/festim}/exports/__init__.py                 | 0
 {festim => src/festim}/exports/average_surface.py          | 0
 {festim => src/festim}/exports/average_volume.py           | 0
 {festim => src/festim}/exports/maximum_surface.py          | 0
 {festim => src/festim}/exports/maximum_volume.py           | 0
 {festim => src/festim}/exports/minimum_surface.py          | 0
 {festim => src/festim}/exports/minimum_volume.py           | 0
 {festim => src/festim}/exports/surface_flux.py             | 0
 {festim => src/festim}/exports/surface_quantity.py         | 0
 {festim => src/festim}/exports/total_surface.py            | 0
 {festim => src/festim}/exports/total_volume.py             | 0
 {festim => src/festim}/exports/volume_quantity.py          | 0
 {festim => src/festim}/exports/vtx.py                      | 0
 {festim => src/festim}/exports/xdmf.py                     | 0
 {festim => src/festim}/heat_transfer_problem.py            | 0
 {festim => src/festim}/helpers.py                          | 0
 {festim => src/festim}/helpers_discontinuity.py            | 0
 {festim => src/festim}/hydrogen_transport_problem.py       | 0
 {festim => src/festim}/initial_condition.py                | 0
 {festim => src/festim}/material.py                         | 0
 {festim => src/festim}/mesh/__init__.py                    | 0
 {festim => src/festim}/mesh/mesh.py                        | 0
 {festim => src/festim}/mesh/mesh_1d.py                     | 0
 {festim => src/festim}/mesh/mesh_from_xdmf.py              | 0
 {festim => src/festim}/problem.py                          | 0
 {festim => src/festim}/reaction.py                         | 0
 {festim => src/festim}/settings.py                         | 0
 {festim => src/festim}/source.py                           | 0
 {festim => src/festim}/species.py                          | 0
 {festim => src/festim}/stepsize.py                         | 0
 {festim => src/festim}/subdomain/__init__.py               | 0
 {festim => src/festim}/subdomain/interface.py              | 0
 {festim => src/festim}/subdomain/surface_subdomain.py      | 0
 {festim => src/festim}/subdomain/surface_subdomain_1d.py   | 0
 {festim => src/festim}/subdomain/volume_subdomain.py       | 0
 {festim => src/festim}/subdomain/volume_subdomain_1d.py    | 0
 {festim => src/festim}/trap.py                             | 0
 44 files changed, 1 insertion(+), 1 deletion(-)
 rename {festim => src/festim}/__init__.py (100%)
 rename {festim => src/festim}/boundary_conditions/__init__.py (100%)
 rename {festim => src/festim}/boundary_conditions/dirichlet_bc.py (100%)
 rename {festim => src/festim}/boundary_conditions/flux_bc.py (100%)
 rename {festim => src/festim}/boundary_conditions/henrys_bc.py (100%)
 rename {festim => src/festim}/boundary_conditions/sieverts_bc.py (100%)
 rename {festim => src/festim}/exports/__init__.py (100%)
 rename {festim => src/festim}/exports/average_surface.py (100%)
 rename {festim => src/festim}/exports/average_volume.py (100%)
 rename {festim => src/festim}/exports/maximum_surface.py (100%)
 rename {festim => src/festim}/exports/maximum_volume.py (100%)
 rename {festim => src/festim}/exports/minimum_surface.py (100%)
 rename {festim => src/festim}/exports/minimum_volume.py (100%)
 rename {festim => src/festim}/exports/surface_flux.py (100%)
 rename {festim => src/festim}/exports/surface_quantity.py (100%)
 rename {festim => src/festim}/exports/total_surface.py (100%)
 rename {festim => src/festim}/exports/total_volume.py (100%)
 rename {festim => src/festim}/exports/volume_quantity.py (100%)
 rename {festim => src/festim}/exports/vtx.py (100%)
 rename {festim => src/festim}/exports/xdmf.py (100%)
 rename {festim => src/festim}/heat_transfer_problem.py (100%)
 rename {festim => src/festim}/helpers.py (100%)
 rename {festim => src/festim}/helpers_discontinuity.py (100%)
 rename {festim => src/festim}/hydrogen_transport_problem.py (100%)
 rename {festim => src/festim}/initial_condition.py (100%)
 rename {festim => src/festim}/material.py (100%)
 rename {festim => src/festim}/mesh/__init__.py (100%)
 rename {festim => src/festim}/mesh/mesh.py (100%)
 rename {festim => src/festim}/mesh/mesh_1d.py (100%)
 rename {festim => src/festim}/mesh/mesh_from_xdmf.py (100%)
 rename {festim => src/festim}/problem.py (100%)
 rename {festim => src/festim}/reaction.py (100%)
 rename {festim => src/festim}/settings.py (100%)
 rename {festim => src/festim}/source.py (100%)
 rename {festim => src/festim}/species.py (100%)
 rename {festim => src/festim}/stepsize.py (100%)
 rename {festim => src/festim}/subdomain/__init__.py (100%)
 rename {festim => src/festim}/subdomain/interface.py (100%)
 rename {festim => src/festim}/subdomain/surface_subdomain.py (100%)
 rename {festim => src/festim}/subdomain/surface_subdomain_1d.py (100%)
 rename {festim => src/festim}/subdomain/volume_subdomain.py (100%)
 rename {festim => src/festim}/subdomain/volume_subdomain_1d.py (100%)
 rename {festim => src/festim}/trap.py (100%)

diff --git a/pyproject.toml b/pyproject.toml
index 25f5e3130..2ec87216c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,7 @@ Issues = "https://github.com/RemDelaporteMathurin/FESTIM/issues"
 
 
 [tool.setuptools_scm]
-write_to = "festim/_version.py"
+write_to = "src/festim/_version.py"
 
 [tool.mypy]
 # Suggested at https://blog.wolt.com/engineering/2021/09/30/professional-grade-mypy-configuration/
diff --git a/festim/__init__.py b/src/festim/__init__.py
similarity index 100%
rename from festim/__init__.py
rename to src/festim/__init__.py
diff --git a/festim/boundary_conditions/__init__.py b/src/festim/boundary_conditions/__init__.py
similarity index 100%
rename from festim/boundary_conditions/__init__.py
rename to src/festim/boundary_conditions/__init__.py
diff --git a/festim/boundary_conditions/dirichlet_bc.py b/src/festim/boundary_conditions/dirichlet_bc.py
similarity index 100%
rename from festim/boundary_conditions/dirichlet_bc.py
rename to src/festim/boundary_conditions/dirichlet_bc.py
diff --git a/festim/boundary_conditions/flux_bc.py b/src/festim/boundary_conditions/flux_bc.py
similarity index 100%
rename from festim/boundary_conditions/flux_bc.py
rename to src/festim/boundary_conditions/flux_bc.py
diff --git a/festim/boundary_conditions/henrys_bc.py b/src/festim/boundary_conditions/henrys_bc.py
similarity index 100%
rename from festim/boundary_conditions/henrys_bc.py
rename to src/festim/boundary_conditions/henrys_bc.py
diff --git a/festim/boundary_conditions/sieverts_bc.py b/src/festim/boundary_conditions/sieverts_bc.py
similarity index 100%
rename from festim/boundary_conditions/sieverts_bc.py
rename to src/festim/boundary_conditions/sieverts_bc.py
diff --git a/festim/exports/__init__.py b/src/festim/exports/__init__.py
similarity index 100%
rename from festim/exports/__init__.py
rename to src/festim/exports/__init__.py
diff --git a/festim/exports/average_surface.py b/src/festim/exports/average_surface.py
similarity index 100%
rename from festim/exports/average_surface.py
rename to src/festim/exports/average_surface.py
diff --git a/festim/exports/average_volume.py b/src/festim/exports/average_volume.py
similarity index 100%
rename from festim/exports/average_volume.py
rename to src/festim/exports/average_volume.py
diff --git a/festim/exports/maximum_surface.py b/src/festim/exports/maximum_surface.py
similarity index 100%
rename from festim/exports/maximum_surface.py
rename to src/festim/exports/maximum_surface.py
diff --git a/festim/exports/maximum_volume.py b/src/festim/exports/maximum_volume.py
similarity index 100%
rename from festim/exports/maximum_volume.py
rename to src/festim/exports/maximum_volume.py
diff --git a/festim/exports/minimum_surface.py b/src/festim/exports/minimum_surface.py
similarity index 100%
rename from festim/exports/minimum_surface.py
rename to src/festim/exports/minimum_surface.py
diff --git a/festim/exports/minimum_volume.py b/src/festim/exports/minimum_volume.py
similarity index 100%
rename from festim/exports/minimum_volume.py
rename to src/festim/exports/minimum_volume.py
diff --git a/festim/exports/surface_flux.py b/src/festim/exports/surface_flux.py
similarity index 100%
rename from festim/exports/surface_flux.py
rename to src/festim/exports/surface_flux.py
diff --git a/festim/exports/surface_quantity.py b/src/festim/exports/surface_quantity.py
similarity index 100%
rename from festim/exports/surface_quantity.py
rename to src/festim/exports/surface_quantity.py
diff --git a/festim/exports/total_surface.py b/src/festim/exports/total_surface.py
similarity index 100%
rename from festim/exports/total_surface.py
rename to src/festim/exports/total_surface.py
diff --git a/festim/exports/total_volume.py b/src/festim/exports/total_volume.py
similarity index 100%
rename from festim/exports/total_volume.py
rename to src/festim/exports/total_volume.py
diff --git a/festim/exports/volume_quantity.py b/src/festim/exports/volume_quantity.py
similarity index 100%
rename from festim/exports/volume_quantity.py
rename to src/festim/exports/volume_quantity.py
diff --git a/festim/exports/vtx.py b/src/festim/exports/vtx.py
similarity index 100%
rename from festim/exports/vtx.py
rename to src/festim/exports/vtx.py
diff --git a/festim/exports/xdmf.py b/src/festim/exports/xdmf.py
similarity index 100%
rename from festim/exports/xdmf.py
rename to src/festim/exports/xdmf.py
diff --git a/festim/heat_transfer_problem.py b/src/festim/heat_transfer_problem.py
similarity index 100%
rename from festim/heat_transfer_problem.py
rename to src/festim/heat_transfer_problem.py
diff --git a/festim/helpers.py b/src/festim/helpers.py
similarity index 100%
rename from festim/helpers.py
rename to src/festim/helpers.py
diff --git a/festim/helpers_discontinuity.py b/src/festim/helpers_discontinuity.py
similarity index 100%
rename from festim/helpers_discontinuity.py
rename to src/festim/helpers_discontinuity.py
diff --git a/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
similarity index 100%
rename from festim/hydrogen_transport_problem.py
rename to src/festim/hydrogen_transport_problem.py
diff --git a/festim/initial_condition.py b/src/festim/initial_condition.py
similarity index 100%
rename from festim/initial_condition.py
rename to src/festim/initial_condition.py
diff --git a/festim/material.py b/src/festim/material.py
similarity index 100%
rename from festim/material.py
rename to src/festim/material.py
diff --git a/festim/mesh/__init__.py b/src/festim/mesh/__init__.py
similarity index 100%
rename from festim/mesh/__init__.py
rename to src/festim/mesh/__init__.py
diff --git a/festim/mesh/mesh.py b/src/festim/mesh/mesh.py
similarity index 100%
rename from festim/mesh/mesh.py
rename to src/festim/mesh/mesh.py
diff --git a/festim/mesh/mesh_1d.py b/src/festim/mesh/mesh_1d.py
similarity index 100%
rename from festim/mesh/mesh_1d.py
rename to src/festim/mesh/mesh_1d.py
diff --git a/festim/mesh/mesh_from_xdmf.py b/src/festim/mesh/mesh_from_xdmf.py
similarity index 100%
rename from festim/mesh/mesh_from_xdmf.py
rename to src/festim/mesh/mesh_from_xdmf.py
diff --git a/festim/problem.py b/src/festim/problem.py
similarity index 100%
rename from festim/problem.py
rename to src/festim/problem.py
diff --git a/festim/reaction.py b/src/festim/reaction.py
similarity index 100%
rename from festim/reaction.py
rename to src/festim/reaction.py
diff --git a/festim/settings.py b/src/festim/settings.py
similarity index 100%
rename from festim/settings.py
rename to src/festim/settings.py
diff --git a/festim/source.py b/src/festim/source.py
similarity index 100%
rename from festim/source.py
rename to src/festim/source.py
diff --git a/festim/species.py b/src/festim/species.py
similarity index 100%
rename from festim/species.py
rename to src/festim/species.py
diff --git a/festim/stepsize.py b/src/festim/stepsize.py
similarity index 100%
rename from festim/stepsize.py
rename to src/festim/stepsize.py
diff --git a/festim/subdomain/__init__.py b/src/festim/subdomain/__init__.py
similarity index 100%
rename from festim/subdomain/__init__.py
rename to src/festim/subdomain/__init__.py
diff --git a/festim/subdomain/interface.py b/src/festim/subdomain/interface.py
similarity index 100%
rename from festim/subdomain/interface.py
rename to src/festim/subdomain/interface.py
diff --git a/festim/subdomain/surface_subdomain.py b/src/festim/subdomain/surface_subdomain.py
similarity index 100%
rename from festim/subdomain/surface_subdomain.py
rename to src/festim/subdomain/surface_subdomain.py
diff --git a/festim/subdomain/surface_subdomain_1d.py b/src/festim/subdomain/surface_subdomain_1d.py
similarity index 100%
rename from festim/subdomain/surface_subdomain_1d.py
rename to src/festim/subdomain/surface_subdomain_1d.py
diff --git a/festim/subdomain/volume_subdomain.py b/src/festim/subdomain/volume_subdomain.py
similarity index 100%
rename from festim/subdomain/volume_subdomain.py
rename to src/festim/subdomain/volume_subdomain.py
diff --git a/festim/subdomain/volume_subdomain_1d.py b/src/festim/subdomain/volume_subdomain_1d.py
similarity index 100%
rename from festim/subdomain/volume_subdomain_1d.py
rename to src/festim/subdomain/volume_subdomain_1d.py
diff --git a/festim/trap.py b/src/festim/trap.py
similarity index 100%
rename from festim/trap.py
rename to src/festim/trap.py

From cef405ab109381a7e954eb863a57bf8118273c5e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 10:04:33 -0400
Subject: [PATCH 107/134] Get ruff and black to agree

---
 src/festim/helpers.py                    |  2 +-
 src/festim/hydrogen_transport_problem.py |  5 ++---
 src/festim/subdomain/interface.py        |  6 ++----
 test/benchmark.py                        |  6 ++----
 test/test_heat_transfer_problem.py       | 10 ++++------
 5 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/src/festim/helpers.py b/src/festim/helpers.py
index f0e741ffc..70448dcb3 100644
--- a/src/festim/helpers.py
+++ b/src/festim/helpers.py
@@ -1,5 +1,5 @@
-from dolfinx import fem
 import dolfinx
+from dolfinx import fem
 
 
 def as_fenics_constant(
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 5e799a12e..af50b7426 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -994,9 +994,8 @@ def define_function_spaces(self, subdomain: _subdomain.VolumeSubdomain):
             species.subdomain_to_collapsed_function_space[subdomain] = V.sub(
                 i
             ).collapse()
-            species.subdomain_to_post_processing_solution[subdomain].name = (
-                f"{species.name}_{subdomain.id}"
-            )
+            name = f"{species.name}_{subdomain.id}"
+            species.subdomain_to_post_processing_solution[subdomain].name = name
 
     def create_subdomain_formulation(self, subdomain: _subdomain.VolumeSubdomain):
         """
diff --git a/src/festim/subdomain/interface.py b/src/festim/subdomain/interface.py
index 3d5d3b5fe..3792d865b 100644
--- a/src/festim/subdomain/interface.py
+++ b/src/festim/subdomain/interface.py
@@ -58,10 +58,8 @@ def compute_mapped_interior_facet_data(self, mesh):
         Returns
             integration_data: Integration data for interior facets
         """
-        assert (not self.subdomains[0].padded) and (
-            not self.subdomains[1].padded)
-        mesh.topology.create_connectivity(
-            mesh.topology.dim - 1, mesh.topology.dim)
+        assert (not self.subdomains[0].padded) and (not self.subdomains[1].padded)
+        mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
         integration_data = compute_integration_domains(
             dolfinx.fem.IntegralType.interior_facet,
             mesh.topology._cpp_object,
diff --git a/test/benchmark.py b/test/benchmark.py
index 7fb74631b..f38542d89 100644
--- a/test/benchmark.py
+++ b/test/benchmark.py
@@ -71,8 +71,7 @@ def fenics_test_permeation_problem(mesh_size=1001):
     mesh_tags_volumes = meshtags(my_mesh, vdim, cells, markers)
 
     dofs_L = locate_dofs_geometrical(V, lambda x: np.isclose(x[0], 0))
-    dofs_R = locate_dofs_geometrical(
-        V, lambda x: np.isclose(x[0], indices[-1]))
+    dofs_R = locate_dofs_geometrical(V, lambda x: np.isclose(x[0], indices[-1]))
 
     dofs_facets = np.array([dofs_L[0], dofs_R[0]], dtype=np.int32)
     tags_facets = np.array([1, 2], dtype=np.int32)
@@ -102,8 +101,7 @@ def siverts_law(T, S_0, E_S, pressure):
     bc_sieverts = dirichletbc(
         Constant(my_mesh, PETSc.ScalarType(surface_conc)), left_dofs, V
     )
-    bc_outgas = dirichletbc(
-        Constant(my_mesh, PETSc.ScalarType(0)), right_dofs, V)
+    bc_outgas = dirichletbc(Constant(my_mesh, PETSc.ScalarType(0)), right_dofs, V)
     bcs = [bc_sieverts, bc_outgas]
 
     D_0 = 1.9e-7
diff --git a/test/test_heat_transfer_problem.py b/test/test_heat_transfer_problem.py
index 7a02eb0b4..224802b7e 100644
--- a/test/test_heat_transfer_problem.py
+++ b/test/test_heat_transfer_problem.py
@@ -218,9 +218,8 @@ def exact_solution(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = (
-        my_problem.t.value
-    )  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = my_problem.t.value
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)
@@ -343,9 +342,8 @@ def mms_source(x, t):
     my_problem.run()
 
     computed_solution = my_problem.u
-    final_time_sim = (
-        my_problem.t.value
-    )  # we use the exact final time of the simulation which may differ from the one specified in the settings
+    # we use the exact final time of the simulation which may differ from the one specified in the settings
+    final_time_sim = my_problem.t.value
 
     def exact_solution_end(x):
         return exact_solution(x, final_time_sim)

From 692d1670b82503cd7222ee0af22566dc6cf85899 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 11:10:29 -0400
Subject: [PATCH 108/134] Simplify input to boundary conditions and add some
 error handling in base-class

---
 .../boundary_conditions/dirichlet_bc.py       | 189 ++++++++++++------
 src/festim/exports/vtx.py                     |   7 +-
 src/festim/heat_transfer_problem.py           |   2 -
 src/festim/hydrogen_transport_problem.py      |  15 +-
 src/festim/subdomain/__init__.py              |   2 +-
 test/test_dirichlet_bc.py                     |  12 +-
 6 files changed, 137 insertions(+), 90 deletions(-)

diff --git a/src/festim/boundary_conditions/dirichlet_bc.py b/src/festim/boundary_conditions/dirichlet_bc.py
index b7022c0c4..f11aa43ed 100644
--- a/src/festim/boundary_conditions/dirichlet_bc.py
+++ b/src/festim/boundary_conditions/dirichlet_bc.py
@@ -1,8 +1,13 @@
+from typing import Callable
+
 import numpy as np
+import numpy.typing as npt
 import ufl
 from dolfinx import fem
+from dolfinx import mesh as _mesh
 
-import festim as F
+from festim import helpers
+from festim import subdomain as _subdomain
 
 
 class DirichletBCBase:
@@ -11,23 +16,36 @@ class DirichletBCBase:
     u = value
 
     Args:
-        subdomain (festim.Subdomain): the surface subdomain where the boundary
-            condition is applied
-        value (float or fem.Constant): the value of the boundary condition
+        subdomain: The surface subdomain where the boundary condition is applied
+        value: The value of the boundary condition
 
     Attributes:
-        subdomain (festim.Subdomain): the surface subdomain where the boundary
-            condition is applied
-        value (float or fem.Constant): the value of the boundary condition
-        value_fenics (fem.Function or fem.Constant): the value of the boundary condition in
-            fenics format
-        bc_expr (fem.Expression): the expression of the boundary condition that is used to
-            update the value_fenics
-        time_dependent (bool): True if the value of the bc is time dependent
+        subdomain: The surface subdomain where the boundary condition is applied
+        value: The value of the boundary condition
+        value_fenics: The value of the boundary condition in fenics format
+        bc_expr: The expression of the boundary condition that is used to
+            update the `value_fenics`
 
     """
 
-    def __init__(self, subdomain, value) -> None:
+    subdomain: _subdomain.SurfaceSubdomain
+    value: (
+        np.ndarray
+        | fem.Constant
+        | int
+        | float
+        | Callable[[np.ndarray], np.ndarray]
+        | Callable[[np.ndarray, float], np.ndarray]
+        | Callable[[float], float]
+    )
+    value_fenics: None | fem.Function | fem.Constant | np.ndarray | float
+    bc_expr: fem.Expression
+
+    def __init__(
+        self,
+        subdomain: _subdomain.SurfaceSubdomain,
+        value: np.ndarray | fem.Constant | int | float | Callable,
+    ):
         self.subdomain = subdomain
         self.value = value
 
@@ -39,18 +57,21 @@ def value_fenics(self):
         return self._value_fenics
 
     @value_fenics.setter
-    def value_fenics(self, value):
+    def value_fenics(self, value: None | fem.Function | fem.Constant | np.ndarray):
         if value is None:
             self._value_fenics = value
             return
         if not isinstance(value, (fem.Function, fem.Constant, np.ndarray)):
+            # FIXME: Should we allow sending in a callable here?
             raise TypeError(
-                f"Value must be a dolfinx.fem.Function, dolfinx.fem.Constant, or a np.ndarray not {type(value)}"
+                f"Value must be a dolfinx.fem.Function, dolfinx.fem.Constant, or a np.ndarray not {
+                    type(value)}"
             )
         self._value_fenics = value
 
     @property
-    def time_dependent(self):
+    def time_dependent(self) -> bool:
+        """Returns true if the value of the boundary condition is time dependent"""
         if self.value is None:
             return False
         if isinstance(self.value, fem.Constant):
@@ -61,22 +82,47 @@ def time_dependent(self):
         else:
             return False
 
-    def define_surface_subdomain_dofs(self, facet_meshtags, mesh, function_space):
-        """Defines the facets and the degrees of freedom of the boundary
-        condition
+    def define_surface_subdomain_dofs(
+        self,
+        facet_meshtags: _mesh.MeshTags,
+        function_space: fem.FunctionSpace | tuple[fem.FunctionSpace, fem.FunctionSpace],
+    ) -> npt.NDArray[np.int32] | tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
+        """Defines the facets and the degrees of freedom of the boundary condition.
+
+        Given the input meshtags, find all facets matching the boundary condition subdomain ID,
+        and locate all DOFs associated with the input function space(s).
+
+        Note:
+            For sub-spaces, a tuple of sub-spaces are expected as input, and a tuple of arrays
+            associated to each of the function spaces are returned.
 
         Args:
-            facet_meshtags (ddolfinx.mesh.MeshTags): the mesh tags of the
-                surface facets
-            mesh (dolfinx.mesh.Mesh): the mesh
-            function_space (dolfinx.fem.FunctionSpaceBase): the function space
+            facet_meshtags: MeshTags describing some facets in the domain
+            mesh:
+            function_space: The function space or a tuple of function spaces: (sub, collapsed)
         """
-        bc_facets = facet_meshtags.find(self.subdomain.id)
-        bc_dofs = fem.locate_dofs_topological(function_space, mesh.fdim, bc_facets)
+        mesh = (
+            function_space[0].mesh
+            if isinstance(function_space, tuple)
+            else function_space.mesh
+        )
+        if facet_meshtags.topology != mesh.topology._cpp_object:
+            raise ValueError(
+                "Mesh of function-space is not the same as the one used for the meshtags"
+            )
+        if mesh.topology.dim - 1 != facet_meshtags.dim:
+            raise ValueError(
+                f"Meshtags of dimension {
+                    facet_meshtags.dim}, expected {mesh.topology.dim-1}"
+            )
+        bc_dofs = fem.locate_dofs_topological(
+            function_space, facet_meshtags.dim, facet_meshtags.find(
+                self.subdomain.id)
+        )
 
         return bc_dofs
 
-    def update(self, t):
+    def update(self, t: float):
         """Updates the boundary condition value
 
         Args:
@@ -95,23 +141,38 @@ class FixedConcentrationBC(DirichletBCBase):
     Args:
         subdomain (festim.Subdomain): the surface subdomain where the boundary
             condition is applied
-        value (float or fem.Constant): the value of the boundary condition
-        species (str): the name of the species
+        value: The value of the boundary condition. It can be a function of space and/or time
+        species: The name of the species
 
     Attributes:
         temperature_dependent (bool): True if the value of the bc is temperature dependent
 
-    Usage:
-        >>> from festim import FixedConcentrationBC
-        >>> FixedConcentrationBC(subdomain=my_subdomain, value=1, species="H")
-        >>> FixedConcentrationBC(subdomain=my_subdomain, value=lambda x: 1 + x[0], species="H")
-        >>> FixedConcentrationBC(subdomain=my_subdomain, value=lambda t: 1 + t, species="H")
-        >>> FixedConcentrationBC(subdomain=my_subdomain, value=lambda T: 1 + T, species="H")
-        >>> FixedConcentrationBC(subdomain=my_subdomain, value=lambda x, t: 1 + x[0] + t, species="H")
+    Examples:
+
+        .. highlight:: python
+        .. code-block:: python
+
+            from festim import FixedConcentrationBC
+            FixedConcentrationBC(subdomain=my_subdomain, value=1, species="H")
+            FixedConcentrationBC(subdomain=my_subdomain,
+                                 value=lambda x: 1 + x[0], species="H")
+            FixedConcentrationBC(subdomain=my_subdomain,
+                                 value=lambda t: 1 + t, species="H")
+            FixedConcentrationBC(subdomain=my_subdomain,
+                                 value=lambda T: 1 + T, species="H")
+            FixedConcentrationBC(subdomain=my_subdomain,
+                                 value=lambda x, t: 1 + x[0] + t, species="H")
 
     """
 
-    def __init__(self, subdomain, value, species) -> None:
+    species: str
+
+    def __init__(
+        self,
+        subdomain: _subdomain.SurfaceSubdomain,
+        value: np.ndarray | fem.Constant | int | float | Callable,
+        species: str,
+    ):
         self.species = species
         super().__init__(subdomain, value)
 
@@ -129,28 +190,28 @@ def temperature_dependent(self):
 
     def create_value(
         self,
-        mesh,
         function_space: fem.FunctionSpace,
-        temperature,
-        t: fem.Constant,
+        temperature: float | fem.Constant,
+        t: float | fem.Constant,
     ):
         """Creates the value of the boundary condition as a fenics object and sets it to
         self.value_fenics.
-        If the value is a constant, it is converted to a fenics.Constant.
-        If the value is a function of t, it is converted to a fenics.Constant.
-        Otherwise, it is converted to a fenics.Function and the
-        expression of the function is stored in self.bc_expr.
+        If the value is a constant, it is converted to a `dolfinx.fem.Constant`.
+        If the value is a function of t, it is converted to  `dolfinx.fem.Constant`.
+        Otherwise, it is converted to a `dolfinx.fem.Function`.Function and the
+        expression of the function is stored in `bc_expr`.
 
         Args:
-            mesh (dolfinx.mesh.Mesh) : the mesh
             function_space (dolfinx.fem.FunctionSpace): the function space
-            temperature (float): the temperature
+            temperature: The temperature
             t (dolfinx.fem.Constant): the time
         """
+        mesh = function_space.mesh
         x = ufl.SpatialCoordinate(mesh)
 
         if isinstance(self.value, (int, float)):
-            self.value_fenics = F.as_fenics_constant(mesh=mesh, value=self.value)
+            self.value_fenics = helpers.as_fenics_constant(
+                mesh=mesh, value=self.value)
 
         elif callable(self.value):
             arguments = self.value.__code__.co_varnames
@@ -159,9 +220,10 @@ def create_value(
                 # only t is an argument
                 if not isinstance(self.value(t=float(t)), (float, int)):
                     raise ValueError(
-                        f"self.value should return a float or an int, not {type(self.value(t=float(t)))} "
+                        f"self.value should return a float or an int, not {
+                            type(self.value(t=float(t)))} "
                     )
-                self.value_fenics = F.as_fenics_constant(
+                self.value_fenics = helpers.as_fenics_constant(
                     mesh=mesh, value=self.value(t=float(t))
                 )
             else:
@@ -184,32 +246,28 @@ def create_value(
 
 
 # alias for FixedConcentrationBC
-class DirichletBC(FixedConcentrationBC):
-    def __init__(self, subdomain, value, species) -> None:
-        super().__init__(subdomain, value, species)
+DirichletBC = FixedConcentrationBC
 
 
 class FixedTemperatureBC(DirichletBCBase):
-    def __init__(self, subdomain, value):
-        super().__init__(subdomain, value)
-
-    def create_value(self, mesh, function_space: fem.FunctionSpace, t: fem.Constant):
+    def create_value(self, function_space: fem.FunctionSpace, t: fem.Constant):
         """Creates the value of the boundary condition as a fenics object and sets it to
         self.value_fenics.
-        If the value is a constant, it is converted to a fenics.Constant.
-        If the value is a function of t, it is converted to a fenics.Constant.
-        Otherwise, it is converted to a fenics.Function and the
-        expression of the function is stored in self.bc_expr.
+        If the value is a constant, it is converted to a `dolfinx.fem.Constant`.
+        If the value is a function of t, it is converted to a `dolfinx.fem.Constant`.
+        Otherwise, it is converted to a` dolfinx.fem.Function` and the
+        expression of the function is stored in `bc_expr`.
 
         Args:
-            mesh (dolfinx.mesh.Mesh) : the mesh
-            function_space (dolfinx.fem.FunctionSpace): the function space
-            t (dolfinx.fem.Constant): the time
+            function_space: the function space
+            t: the time
         """
+        mesh = function_space.mesh
         x = ufl.SpatialCoordinate(mesh)
 
         if isinstance(self.value, (int, float)):
-            self.value_fenics = F.as_fenics_constant(mesh=mesh, value=self.value)
+            self.value_fenics = helpers.as_fenics_constant(
+                mesh=mesh, value=self.value)
 
         elif callable(self.value):
             arguments = self.value.__code__.co_varnames
@@ -218,9 +276,10 @@ def create_value(self, mesh, function_space: fem.FunctionSpace, t: fem.Constant)
                 # only t is an argument
                 if not isinstance(self.value(t=float(t)), (float, int)):
                     raise ValueError(
-                        f"self.value should return a float or an int, not {type(self.value(t=float(t)))} "
+                        f"self.value should return a float or an int, not {
+                            type(self.value(t=float(t)))} "
                     )
-                self.value_fenics = F.as_fenics_constant(
+                self.value_fenics = helpers.as_fenics_constant(
                     mesh=mesh, value=self.value(t=float(t))
                 )
             else:
diff --git a/src/festim/exports/vtx.py b/src/festim/exports/vtx.py
index 960a1041f..1522c8fb8 100644
--- a/src/festim/exports/vtx.py
+++ b/src/festim/exports/vtx.py
@@ -16,7 +16,8 @@ def __init__(self, filename: str | Path, ext: str) -> None:
         name = Path(filename)
         if name.suffix != ext:
             warnings.warn(
-                f"Filename {filename} does not have {ext} extension, adding it."
+                f"Filename {filename} does not have {
+                    ext} extension, adding it."
             )
             name = name.with_suffix(ext)
 
@@ -92,10 +93,6 @@ def field(self, value: _Species | list[_Species]):
             )
         self._field = val
 
-    @property
-    def subdomain(self) -> _VolumeSubdomain:
-        return self._subdomain
-
     def get_functions(self) -> list[_Function]:
         """
         Returns list of species for a given subdomain.
diff --git a/src/festim/heat_transfer_problem.py b/src/festim/heat_transfer_problem.py
index 746cf278b..9c49fce6c 100644
--- a/src/festim/heat_transfer_problem.py
+++ b/src/festim/heat_transfer_problem.py
@@ -112,14 +112,12 @@ def create_dirichletbc_form(self, bc):
                 the boundary condition for modifying linear systems.
         """
         bc.create_value(
-            mesh=self.mesh.mesh,
             function_space=self.function_space,
             t=self.t,
         )
 
         bc_dofs = bc.define_surface_subdomain_dofs(
             facet_meshtags=self.facet_meshtags,
-            mesh=self.mesh,
             function_space=self.function_space,
         )
 
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index af50b7426..387e0b87f 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -523,17 +523,12 @@ def create_dirichletbc_form(self, bc):
                 the boundary condition for modifying linear systems.
         """
         # create value_fenics
-        function_space_value = None
-
-        if callable(bc.value):
-            # if bc.value is a callable then need to provide a functionspace
-            if not self.multispecies:
-                function_space_value = bc.species.sub_function_space
-            else:
-                function_space_value = bc.species.collapsed_function_space
+        if not self.multispecies:
+            function_space_value = bc.species.sub_function_space
+        else:
+            function_space_value = bc.species.collapsed_function_space
 
         bc.create_value(
-            mesh=self.mesh.mesh,
             temperature=self.temperature_fenics,
             function_space=function_space_value,
             t=self.t,
@@ -550,7 +545,6 @@ def create_dirichletbc_form(self, bc):
 
         bc_dofs = bc.define_surface_subdomain_dofs(
             facet_meshtags=self.facet_meshtags,
-            mesh=self.mesh,
             function_space=function_space_dofs,
         )
 
@@ -907,7 +901,6 @@ def create_dirichletbc_form(self, bc: boundary_conditions.FixedConcentrationBC):
         collapsed_V, _ = sub_V.collapse()
 
         bc.create_value(
-            mesh=self.mesh.mesh,
             temperature=self.temperature_fenics,
             function_space=collapsed_V,
             t=self.t,
diff --git a/src/festim/subdomain/__init__.py b/src/festim/subdomain/__init__.py
index fe4f09326..0544d246b 100644
--- a/src/festim/subdomain/__init__.py
+++ b/src/festim/subdomain/__init__.py
@@ -1,4 +1,4 @@
-__all__ = ["VolumeSubdomain", "SurfaceSubdomain", "Interface"]
+__all__ = ["VolumeSubdomain", "SurfaceSubdomain", "Interface", "Subdomain"]
 
 
 from .interface import Interface
diff --git a/test/test_dirichlet_bc.py b/test/test_dirichlet_bc.py
index 259bf49fa..ff399a7ef 100644
--- a/test/test_dirichlet_bc.py
+++ b/test/test_dirichlet_bc.py
@@ -76,7 +76,7 @@ def value(x, t):
 
     T = fem.Constant(my_model.mesh.mesh, 550.0)
     t = fem.Constant(my_model.mesh.mesh, 0.0)
-    bc.create_value(my_model.mesh.mesh, my_model.function_space, T, t)
+    bc.create_value(my_model.function_space, T, t)
 
     # check that the value_fenics attribute is set correctly
     assert isinstance(bc.value_fenics, fem.Function)
@@ -119,7 +119,7 @@ def value(x, t, T):
 
     T = fem.Constant(my_model.mesh.mesh, 550.0)
     t = fem.Constant(my_model.mesh.mesh, 0.0)
-    bc.create_value(my_model.mesh.mesh, my_model.function_space, T, t)
+    bc.create_value(my_model.function_space, T, t)
 
     # check that the value_fenics attribute is set correctly
     assert isinstance(bc.value_fenics, fem.Function)
@@ -162,7 +162,7 @@ def test_callable_t_only(value):
 
     T = fem.Constant(my_model.mesh.mesh, 550.0)
     t = fem.Constant(my_model.mesh.mesh, 0.0)
-    bc.create_value(my_model.mesh.mesh, my_model.function_space, T, t)
+    bc.create_value(my_model.function_space, T, t)
 
     # check that the value_fenics attribute is set correctly
     assert isinstance(bc.value_fenics, fem.Constant)
@@ -210,7 +210,7 @@ def value(x):
     t = fem.Constant(my_model.mesh.mesh, 0.0)
 
     # TEST
-    bc.create_value(my_model.mesh.mesh, my_model.function_space, T, t)
+    bc.create_value(my_model.function_space, T, t)
 
     # check that the value_fenics attribute is set correctly
     assert isinstance(bc.value_fenics, fem.Function)
@@ -317,11 +317,11 @@ def test_define_value_error_if_ufl_conditional_t_only(value):
     bc = F.DirichletBC(subdomain, value, species)
 
     t = fem.Constant(mesh, 0.0)
-
+    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
     with pytest.raises(
         ValueError, match="self.value should return a float or an int, not "
     ):
-        bc.create_value(mesh=mesh, function_space=None, temperature=None, t=t)
+        bc.create_value(V, temperature=None, t=t)
 
 
 def test_species_predefined():

From b67aa0b1aafd5430e92cf41ea294e755504636b5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 11:23:38 -0400
Subject: [PATCH 109/134] Ruff

---
 src/festim/boundary_conditions/dirichlet_bc.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/src/festim/boundary_conditions/dirichlet_bc.py b/src/festim/boundary_conditions/dirichlet_bc.py
index f11aa43ed..095fbd9c7 100644
--- a/src/festim/boundary_conditions/dirichlet_bc.py
+++ b/src/festim/boundary_conditions/dirichlet_bc.py
@@ -116,8 +116,7 @@ def define_surface_subdomain_dofs(
                     facet_meshtags.dim}, expected {mesh.topology.dim-1}"
             )
         bc_dofs = fem.locate_dofs_topological(
-            function_space, facet_meshtags.dim, facet_meshtags.find(
-                self.subdomain.id)
+            function_space, facet_meshtags.dim, facet_meshtags.find(self.subdomain.id)
         )
 
         return bc_dofs
@@ -210,8 +209,7 @@ def create_value(
         x = ufl.SpatialCoordinate(mesh)
 
         if isinstance(self.value, (int, float)):
-            self.value_fenics = helpers.as_fenics_constant(
-                mesh=mesh, value=self.value)
+            self.value_fenics = helpers.as_fenics_constant(mesh=mesh, value=self.value)
 
         elif callable(self.value):
             arguments = self.value.__code__.co_varnames
@@ -266,8 +264,7 @@ def create_value(self, function_space: fem.FunctionSpace, t: fem.Constant):
         x = ufl.SpatialCoordinate(mesh)
 
         if isinstance(self.value, (int, float)):
-            self.value_fenics = helpers.as_fenics_constant(
-                mesh=mesh, value=self.value)
+            self.value_fenics = helpers.as_fenics_constant(mesh=mesh, value=self.value)
 
         elif callable(self.value):
             arguments = self.value.__code__.co_varnames

From edc86cf4daa10ee7c4de080b0d542fdf38415088 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 11:46:40 -0400
Subject: [PATCH 110/134] Add tests

---
 src/festim/exports/vtx.py                |  2 +-
 src/festim/hydrogen_transport_problem.py |  1 -
 test/test_vtx.py                         | 39 ++++++++++++++++++++++++
 test/test_xdmf.py                        |  6 ++++
 4 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/festim/exports/vtx.py b/src/festim/exports/vtx.py
index 1522c8fb8..30e8184bf 100644
--- a/src/festim/exports/vtx.py
+++ b/src/festim/exports/vtx.py
@@ -21,7 +21,7 @@ def __init__(self, filename: str | Path, ext: str) -> None:
             )
             name = name.with_suffix(ext)
 
-        self._filename = Path(filename)
+        self._filename = name
 
     @property
     def filename(self):
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 387e0b87f..eb6e56783 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -1108,7 +1108,6 @@ def mixed_term(u, v, n):
             if len(all_mobile_species) > 1:
                 raise NotImplementedError("Multiple mobile species not implemented")
             H = all_mobile_species[0]
-
             v_b = H.subdomain_to_test_function[subdomain_0](res[0])
             v_t = H.subdomain_to_test_function[subdomain_1](res[1])
 
diff --git a/test/test_vtx.py b/test/test_vtx.py
index 62b16af48..164905c3c 100644
--- a/test/test_vtx.py
+++ b/test/test_vtx.py
@@ -48,6 +48,45 @@ def test_vtx_export_subdomain():
     pass
 
 
+def test_vtx_suffix_converter(tmpdir):
+    filename = str(tmpdir.join("my_export.txt"))
+    my_export = F.VTXSpeciesExport(filename, field=[])
+    assert my_export.filename.suffix == ".bp"
+
+
+def test_vtx_DG(tmpdir):
+    """Test VTX export setup for DG formulation"""
+    my_model = F.HTransportProblemDiscontinuous()
+    my_model.mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
+    my_mat = F.Material(D_0=3, E_D=2, K_S_0=1, E_K_S=0, name="mat")
+
+    s0 = F.VolumeSubdomain1D(1, borders=[0.0, 2], material=my_mat)
+    s1 = F.VolumeSubdomain1D(2, borders=[2, 4], material=my_mat)
+    l0 = F.SurfaceSubdomain1D(1, x=0.0)
+    l1 = F.SurfaceSubdomain1D(2, x=4.0)
+    my_model.interfaces = [F.Interface(6, (s0, s1))]
+
+    my_model.temperature = 55
+    my_model.subdomains = [s0, s1, l0, l1]
+    my_model.surface_to_volume = {l0: s0, l1: s1}
+    # NOTE: Ask Remi why `H` has to live in both s0 and s1
+    my_model.species = [
+        F.Species("H", subdomains=[s0, s1]),
+        F.Species("T", subdomains=[s0, s1], mobile=False),
+    ]
+
+    filename = str(tmpdir.join("my_export.txt"))
+    my_export = F.VTXSpeciesExport(filename, field=my_model.species, subdomain=s0)
+    assert my_export.filename.suffix == ".bp"
+    my_model.exports = [my_export]
+    my_model.settings = F.Settings(atol=1, rtol=0.1)
+    my_model.settings.stepsize = F.Stepsize(initial_value=1)
+
+    my_model.initialise()
+    assert len(my_export.get_functions()) == 2
+    assert len(my_model._vtxfiles) == 1
+
+
 def test_vtx_integration_with_h_transport_problem(tmpdir):
     my_model = F.HydrogenTransportProblem()
     my_model.mesh = F.Mesh1D(vertices=np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
diff --git a/test/test_xdmf.py b/test/test_xdmf.py
index b4ba1e9a5..fa3387b0f 100644
--- a/test/test_xdmf.py
+++ b/test/test_xdmf.py
@@ -67,6 +67,12 @@ def test_field_attribute_is_always_list():
     assert isinstance(my_export.field, list)
 
 
+def test_vtx_suffix_converter(tmpdir):
+    filename = str(tmpdir.join("my_export.txt"))
+    my_export = F.XDMFExport(filename, field=[])
+    assert my_export.filename.suffix == ".xdmf"
+
+
 @pytest.mark.parametrize("field", [["H", 2], 1, [F.Species("H"), 1]])
 def test_field_attribute_raises_error_when_invalid_type(field):
     """

From 3e12a604bd25e712b3b6abdb3d71bcdea81fbdca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 11:48:33 -0400
Subject: [PATCH 111/134] Make mypy non-blocking

---
 .github/workflows/code_formatting.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/code_formatting.yml b/.github/workflows/code_formatting.yml
index 11576717f..f4f1163a9 100644
--- a/.github/workflows/code_formatting.yml
+++ b/.github/workflows/code_formatting.yml
@@ -25,5 +25,6 @@ jobs:
           ruff check .
 
       - name: mypy
+        continue-on-error: true
         run: |
           python -m mypy .

From 35928cfa41feab946c144758b15963ecb322cb60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 12:28:56 -0400
Subject: [PATCH 112/134] Update mesh class

---
 docs/source/api/index.rst                | 27 ++++++++---
 pyproject.toml                           |  2 +-
 src/festim/heat_transfer_problem.py      | 11 +++--
 src/festim/hydrogen_transport_problem.py | 47 ++++++++++++-------
 src/festim/mesh/mesh.py                  | 57 ++++++++++++++----------
 src/festim/mesh/mesh_1d.py               |  6 ++-
 src/festim/problem.py                    | 13 +++---
 7 files changed, 106 insertions(+), 57 deletions(-)

diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst
index 23396ca21..2d6fef3c5 100644
--- a/docs/source/api/index.rst
+++ b/docs/source/api/index.rst
@@ -1,13 +1,26 @@
-FESTIM
-======
+FESTIM API
+==========
 
-.. automodule:: festim
+
+.. automodule:: festim.boundary_conditions
     :members:
-    :undoc-members:
     :show-inheritance:
-    :inherited-members:
     :exclude-members: __weakref__
     :private-members:
-    :special-members:
     :inherited-members:
-    :exclude-members: __weakref__
\ No newline at end of file
+
+
+.. .. automodule:: festim.hydrogen_transport_problem
+..     :members:
+..     :show-inheritance:
+..     :exclude-members: __weakref__
+..     :private-members:
+..     :inherited-members:
+
+
+
+.. .. automodule:: festim.heat_transfer_problem
+..     :members:
+..     :show-inheritance:
+..     :exclude-members: __weakref__
+..     :inherited-members:
diff --git a/pyproject.toml b/pyproject.toml
index 2ec87216c..5f6870ec7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,7 +29,7 @@ classifiers = [
 [project.optional-dependencies]
 test = ["pytest >= 5.4.3", "pytest-cov", "sympy"]
 lint = ["ruff", "mypy"]
-docs = ["sphinx"]
+docs = ["sphinx", "sphinx-book-theme", "sphinx-design", "matplotlib"]
 
 [project.urls]
 Homepage = "https://github.com/RemDelaporteMathurin/FESTIM"
diff --git a/src/festim/heat_transfer_problem.py b/src/festim/heat_transfer_problem.py
index 9c49fce6c..ec322e984 100644
--- a/src/festim/heat_transfer_problem.py
+++ b/src/festim/heat_transfer_problem.py
@@ -6,6 +6,8 @@
 from festim import boundary_conditions, exports, helpers, problem
 from festim import source as _source
 
+__all__ = ["HeatTransferProblem"]
+
 
 class HeatTransferProblem(problem.ProblemBase):
     def __init__(
@@ -37,7 +39,8 @@ def sources(self):
     @sources.setter
     def sources(self, value):
         if not all(isinstance(source, _source.HeatSource) for source in value):
-            raise TypeError("sources must be a list of festim.HeatSource objects")
+            raise TypeError(
+                "sources must be a list of festim.HeatSource objects")
         self._sources = value
 
     @property
@@ -205,14 +208,16 @@ def create_formulation(self):
         # add sources
         for source in self.sources:
             self.formulation -= (
-                source.value_fenics * self.test_function * self.dx(source.volume.id)
+                source.value_fenics * self.test_function *
+                self.dx(source.volume.id)
             )
 
         # add fluxes
         for bc in self.boundary_conditions:
             if isinstance(bc, boundary_conditions.HeatFluxBC):
                 self.formulation -= (
-                    bc.value_fenics * self.test_function * self.ds(bc.subdomain.id)
+                    bc.value_fenics * self.test_function *
+                    self.ds(bc.subdomain.id)
                 )
 
     def initialise_exports(self):
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index eb6e56783..8f7362bb5 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -31,6 +31,8 @@
 from festim.helpers import as_fenics_constant
 from festim.mesh import Mesh
 
+__all__ = ["HydrogenTransportProblem", "HTransportProblemDiscontinuous"]
+
 
 class HydrogenTransportProblem(problem.ProblemBase):
     """
@@ -123,7 +125,8 @@ def __init__(
         self,
         mesh: Mesh | None = None,
         subdomains: (
-            list[_subdomain.VolumeSubdomain | _subdomain.SurfaceSubdomain] | None
+            list[_subdomain.VolumeSubdomain |
+                 _subdomain.SurfaceSubdomain] | None
         ) = None,
         species: list[_species.Species] | None = None,
         reactions: list[_reaction.Reaction] | None = None,
@@ -319,7 +322,8 @@ def define_temperature(self):
                     )
                 # only t is an argument
                 self.temperature_fenics = as_fenics_constant(
-                    mesh=self.mesh.mesh, value=self.temperature(t=float(self.t))
+                    mesh=self.mesh.mesh, value=self.temperature(
+                        t=float(self.t))
                 )
             else:
                 x = ufl.SpatialCoordinate(self.mesh.mesh)
@@ -333,7 +337,8 @@ def define_temperature(self):
                 function_space_temperature = fem.functionspace(
                     self.mesh.mesh, element_temperature
                 )
-                self.temperature_fenics = fem.Function(function_space_temperature)
+                self.temperature_fenics = fem.Function(
+                    function_space_temperature)
                 kwargs = {}
                 if "t" in arguments:
                     kwargs["t"] = self.t
@@ -415,7 +420,8 @@ def define_D_global(self, species):
         D_0 = fem.Function(self.V_DG_0)
         E_D = fem.Function(self.V_DG_0)
         for vol in self.volume_subdomains:
-            cell_indices = vol.locate_subdomain_entities(self.mesh.mesh, self.mesh.vdim)
+            cell_indices = vol.locate_subdomain_entities(
+                self.mesh.mesh, self.mesh.vdim)
 
             # replace values of D_0 and E_D by values from the material
             D_0.x.array[cell_indices] = vol.material.get_D_0(species=species)
@@ -425,9 +431,11 @@ def define_D_global(self, species):
         D = fem.Function(self.V_DG_1)
 
         expr = D_0 * ufl.exp(
-            -E_D / as_fenics_constant(k_B, self.mesh.mesh) / self.temperature_fenics
+            -E_D / as_fenics_constant(k_B, self.mesh.mesh) /
+            self.temperature_fenics
         )
-        D_expr = fem.Expression(expr, self.V_DG_1.element.interpolation_points())
+        D_expr = fem.Expression(
+            expr, self.V_DG_1.element.interpolation_points())
         D.interpolate(D_expr)
         return D, D_expr
 
@@ -508,7 +516,8 @@ def define_boundary_conditions(self):
         for bc in self.boundary_conditions:
             if isinstance(bc.species, str):
                 # if name of species is given then replace with species object
-                bc.species = _species.find_species_from_name(bc.species, self.species)
+                bc.species = _species.find_species_from_name(
+                    bc.species, self.species)
 
         super().define_boundary_conditions()
 
@@ -614,7 +623,8 @@ def create_initial_conditions(self):
 
             # assign to previous solution of species
             if not self.multispecies:
-                condition.species.prev_solution.interpolate(condition.expr_fenics)
+                condition.species.prev_solution.interpolate(
+                    condition.expr_fenics)
             else:
                 idx = self.species.index(condition.species)
                 self.u_n.sub(idx).interpolate(condition.expr_fenics)
@@ -640,7 +650,8 @@ def create_formulation(self):
                     )
 
                 if self.settings.transient:
-                    self.formulation += ((u - u_n) / self.dt) * v * self.dx(vol.id)
+                    self.formulation += ((u - u_n) / self.dt) * \
+                        v * self.dx(vol.id)
 
         for reaction in self.reactions:
             for reactant in reaction.reactant:
@@ -737,7 +748,8 @@ def post_processing(self):
             if isinstance(export, exports.SurfaceQuantity):
                 if isinstance(
                     export,
-                    (exports.SurfaceFlux, exports.TotalSurface, exports.AverageSurface),
+                    (exports.SurfaceFlux, exports.TotalSurface,
+                     exports.AverageSurface),
                 ):
                     export.compute(
                         self.ds,
@@ -1084,7 +1096,8 @@ def create_formulation(self):
             for interface in self.interfaces
         ]
         [interface.pad_parent_maps() for interface in self.interfaces]
-        dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)
+        dInterface = ufl.Measure(
+            "dS", domain=mesh, subdomain_data=integral_data)
 
         def mixed_term(u, v, n):
             return ufl.dot(ufl.grad(u), n) * v
@@ -1106,7 +1119,8 @@ def mixed_term(u, v, n):
 
             all_mobile_species = [spe for spe in self.species if spe.mobile]
             if len(all_mobile_species) > 1:
-                raise NotImplementedError("Multiple mobile species not implemented")
+                raise NotImplementedError(
+                    "Multiple mobile species not implemented")
             H = all_mobile_species[0]
             v_b = H.subdomain_to_test_function[subdomain_0](res[0])
             v_t = H.subdomain_to_test_function[subdomain_1](res[1])
@@ -1165,7 +1179,8 @@ def mixed_term(u, v, n):
                     )
                 )
             J.append(jac)
-            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
+            forms.append(dolfinx.fem.form(
+                subdomain1.F, entity_maps=entity_maps))
 
         self.forms = forms
         self.J = J
@@ -1204,7 +1219,8 @@ def initialise_exports(self):
                     )
                 )
             else:
-                raise NotImplementedError(f"Export type {type(export)} not implemented")
+                raise NotImplementedError(
+                    f"Export type {type(export)} not implemented")
 
     def post_processing(self):
         # update post-processing solutions (for each species in each subdomain)
@@ -1225,7 +1241,8 @@ def post_processing(self):
 
         for export in self.exports:
             if not isinstance(export, exports.VTXSpeciesExport):
-                raise NotImplementedError(f"Export type {type(export)} not implemented")
+                raise NotImplementedError(
+                    f"Export type {type(export)} not implemented")
 
     def iterate(self):
         """Iterates the model for a given time step"""
diff --git a/src/festim/mesh/mesh.py b/src/festim/mesh/mesh.py
index 87ce15329..67c253488 100644
--- a/src/festim/mesh/mesh.py
+++ b/src/festim/mesh/mesh.py
@@ -1,7 +1,7 @@
 import dolfinx
 import numpy as np
 import ufl
-from dolfinx.mesh import meshtags
+from dolfinx.mesh import meshtags, Mesh as dolfinx_Mesh
 
 
 class Mesh:
@@ -9,27 +9,29 @@ class Mesh:
     Mesh class
 
     Args:
-        mesh (dolfinx.mesh.Mesh, optional): the mesh. Defaults to None.
+        mesh The mesh. Defaults to None.
 
     Attributes:
-        mesh (dolfinx.mesh.Mesh): the mesh
-        vdim (int): the dimension of the mesh cells
-        fdim (int): the dimension of the mesh facets
-        n (ufl.FacetNormal): the normal vector to the facets
+        mesh The mesh
+        vdim: the dimension of the mesh cells
+        fdim: the dimension of the mesh facets
+        n: Symbolic representation of the vector normal to the facets
+            of the mesh.
     """
+    _mesh: dolfinx.mesh.Mesh
 
-    def __init__(self, mesh=None):
-        self.mesh = mesh
+    def __init__(self, mesh: dolfinx_Mesh | None = None):
 
-        if self.mesh is not None:
+        self.mesh = mesh
+        if self._mesh is not None:
             # create cell to facet connectivity
-            self.mesh.topology.create_connectivity(
-                self.mesh.topology.dim, self.mesh.topology.dim - 1
+            self._mesh.topology.create_connectivity(
+                self._mesh.topology.dim, self._mesh.topology.dim - 1
             )
 
             # create facet to cell connectivity
-            self.mesh.topology.create_connectivity(
-                self.mesh.topology.dim - 1, self.mesh.topology.dim
+            self._mesh.topology.create_connectivity(
+                self._mesh.topology.dim - 1, self._mesh.topology.dim
             )
 
     @property
@@ -45,15 +47,21 @@ def mesh(self, value):
 
     @property
     def vdim(self):
-        return self.mesh.topology.dim
+        if self._mesh is None:
+            raise RuntimeError("Mesh is not defined")
+        return self._mesh.topology.dim
 
     @property
     def fdim(self):
-        return self.mesh.topology.dim - 1
+        if self._mesh is None:
+            raise RuntimeError("Mesh is not defined")
+        return self._mesh.topology.dim - 1
 
     @property
     def n(self):
-        return ufl.FacetNormal(self.mesh)
+        if self._mesh is None:
+            raise RuntimeError("Mesh is not defined")
+        return ufl.FacetNormal(self._mesh)
 
     def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None):
         """Defines the facet and volume meshtags of the mesh
@@ -69,27 +77,27 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
             dolfinx.mesh.MeshTags: the volume meshtags
         """
         # find all cells in domain and mark them as 0
-        num_cells = self.mesh.topology.index_map(self.vdim).size_local
+        num_cells = self._mesh.topology.index_map(self.vdim).size_local
         mesh_cell_indices = np.arange(num_cells, dtype=np.int32)
         tags_volumes = np.full(num_cells, 0, dtype=np.int32)
 
         # find all facets in domain and mark them as 0
-        num_facets = self.mesh.topology.index_map(self.fdim).size_local
+        num_facets = self._mesh.topology.index_map(self.fdim).size_local
         mesh_facet_indices = np.arange(num_facets, dtype=np.int32)
         tags_facets = np.full(num_facets, 0, dtype=np.int32)
 
         for surf in surface_subdomains:
             # find all facets in subdomain and mark them as surf.id
-            entities = surf.locate_boundary_facet_indices(self.mesh)
+            entities = surf.locate_boundary_facet_indices(self._mesh)
             tags_facets[entities] = surf.id
 
         for vol in volume_subdomains:
             # find all cells in subdomain and mark them as vol.id
-            entities = vol.locate_subdomain_entities(self.mesh, self.vdim)
+            entities = vol.locate_subdomain_entities(self._mesh, self.vdim)
             tags_volumes[entities] = vol.id
 
         volume_meshtags = meshtags(
-            self.mesh, self.vdim, mesh_cell_indices, tags_volumes
+            self._mesh, self.vdim, mesh_cell_indices, tags_volumes
         )
 
         # tag interfaces
@@ -97,13 +105,13 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
         for interface in interfaces:
             (domain_0, domain_1) = interface.subdomains
             all_0_facets = dolfinx.mesh.compute_incident_entities(
-                self.mesh.topology,
+                self._mesh.topology,
                 volume_meshtags.find(domain_0.id),
                 self.vdim,
                 self.fdim,
             )
             all_1_facets = dolfinx.mesh.compute_incident_entities(
-                self.mesh.topology,
+                self._mesh.topology,
                 volume_meshtags.find(domain_1.id),
                 self.vdim,
                 self.fdim,
@@ -111,6 +119,7 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
             interface_entities = np.intersect1d(all_0_facets, all_1_facets)
             tags_facets[interface_entities] = interface.id
 
-        facet_meshtags = meshtags(self.mesh, self.fdim, mesh_facet_indices, tags_facets)
+        facet_meshtags = meshtags(
+            self._mesh, self.fdim, mesh_facet_indices, tags_facets)
 
         return facet_meshtags, volume_meshtags
diff --git a/src/festim/mesh/mesh_1d.py b/src/festim/mesh/mesh_1d.py
index 6b62684cb..da5de0a47 100644
--- a/src/festim/mesh/mesh_1d.py
+++ b/src/festim/mesh/mesh_1d.py
@@ -37,7 +37,8 @@ def generate_mesh(self):
         """Generates a 1D mesh"""
         degree = 1
         domain = ufl.Mesh(
-            basix.ufl.element(basix.ElementFamily.P, "interval", degree, shape=(1,))
+            basix.ufl.element(basix.ElementFamily.P,
+                              "interval", degree, shape=(1,))
         )
 
         mesh_points = np.reshape(self.vertices, (len(self.vertices), 1))
@@ -56,7 +57,8 @@ def check_borders(self, volume_subdomains):
             Value error: if borders outside the domain
         """
         # check that subdomains are connected
-        all_borders = [border for vol in volume_subdomains for border in vol.borders]
+        all_borders = [
+            border for vol in volume_subdomains for border in vol.borders]
         sorted_borders = np.sort(all_borders)
         for start, end in zip(sorted_borders[1:-2:2], sorted_borders[2:-1:2]):
             if start != end:
diff --git a/src/festim/problem.py b/src/festim/problem.py
index 966baa8b8..5a27c29f9 100644
--- a/src/festim/problem.py
+++ b/src/festim/problem.py
@@ -18,21 +18,24 @@
 
 class ProblemBase:
     """
-    Base class for HeatTransferProblem and HTransportProblem.
+    Base class for :py:class:`HeatTransferProblem <festim.heat_transfer_problem.HeatTransferProblem>` and
+    :py:class:`HydrogenTransportProblem <festim.hydrogen_transport_problem.HydrogenTransportProblem>`.
 
-    show_progress_bar (bool): True if a progress bar is displayed during the
-        simulation
-    progress_bar (tqdm.autonotebook.tqdm) the progress bar
+    Attributes:
+        show_progress_bar: If `True` a progress bar is displayed during the simulation
+        progress_bar: the progress bar
     """
 
     mesh: _Mesh
     sources: list[_SourceBase]
     exports: list[Any]
     subdomains: list[_VolumeSubdomain]
+    show_progress_bar: bool
+    progress_bar: None | tqdm.autonotebook.tqdm
 
     def __init__(
         self,
-        mesh=None,
+        mesh: _Mesh = None,
         sources=None,
         exports=None,
         subdomains=None,

From 584692b1b52abe2d440ff1a89daf27f111ebe036 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Thu, 24 Oct 2024 12:29:32 -0400
Subject: [PATCH 113/134] Ruff

---
 src/festim/__init__.py                   |  1 +
 src/festim/heat_transfer_problem.py      | 11 ++----
 src/festim/hydrogen_transport_problem.py | 45 ++++++++----------------
 src/festim/mesh/mesh.py                  |  8 +++--
 src/festim/mesh/mesh_1d.py               |  6 ++--
 5 files changed, 26 insertions(+), 45 deletions(-)

diff --git a/src/festim/__init__.py b/src/festim/__init__.py
index 991508e09..0e833926f 100644
--- a/src/festim/__init__.py
+++ b/src/festim/__init__.py
@@ -9,6 +9,7 @@
 R = 8.314462618  # Gas constant J.mol-1.K-1
 k_B = 8.6173303e-5  # Boltzmann constant eV.K-1
 
+
 from .boundary_conditions.dirichlet_bc import (
     DirichletBC,
     DirichletBCBase,
diff --git a/src/festim/heat_transfer_problem.py b/src/festim/heat_transfer_problem.py
index ec322e984..9c49fce6c 100644
--- a/src/festim/heat_transfer_problem.py
+++ b/src/festim/heat_transfer_problem.py
@@ -6,8 +6,6 @@
 from festim import boundary_conditions, exports, helpers, problem
 from festim import source as _source
 
-__all__ = ["HeatTransferProblem"]
-
 
 class HeatTransferProblem(problem.ProblemBase):
     def __init__(
@@ -39,8 +37,7 @@ def sources(self):
     @sources.setter
     def sources(self, value):
         if not all(isinstance(source, _source.HeatSource) for source in value):
-            raise TypeError(
-                "sources must be a list of festim.HeatSource objects")
+            raise TypeError("sources must be a list of festim.HeatSource objects")
         self._sources = value
 
     @property
@@ -208,16 +205,14 @@ def create_formulation(self):
         # add sources
         for source in self.sources:
             self.formulation -= (
-                source.value_fenics * self.test_function *
-                self.dx(source.volume.id)
+                source.value_fenics * self.test_function * self.dx(source.volume.id)
             )
 
         # add fluxes
         for bc in self.boundary_conditions:
             if isinstance(bc, boundary_conditions.HeatFluxBC):
                 self.formulation -= (
-                    bc.value_fenics * self.test_function *
-                    self.ds(bc.subdomain.id)
+                    bc.value_fenics * self.test_function * self.ds(bc.subdomain.id)
                 )
 
     def initialise_exports(self):
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 8f7362bb5..0c0607dfd 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -125,8 +125,7 @@ def __init__(
         self,
         mesh: Mesh | None = None,
         subdomains: (
-            list[_subdomain.VolumeSubdomain |
-                 _subdomain.SurfaceSubdomain] | None
+            list[_subdomain.VolumeSubdomain | _subdomain.SurfaceSubdomain] | None
         ) = None,
         species: list[_species.Species] | None = None,
         reactions: list[_reaction.Reaction] | None = None,
@@ -322,8 +321,7 @@ def define_temperature(self):
                     )
                 # only t is an argument
                 self.temperature_fenics = as_fenics_constant(
-                    mesh=self.mesh.mesh, value=self.temperature(
-                        t=float(self.t))
+                    mesh=self.mesh.mesh, value=self.temperature(t=float(self.t))
                 )
             else:
                 x = ufl.SpatialCoordinate(self.mesh.mesh)
@@ -337,8 +335,7 @@ def define_temperature(self):
                 function_space_temperature = fem.functionspace(
                     self.mesh.mesh, element_temperature
                 )
-                self.temperature_fenics = fem.Function(
-                    function_space_temperature)
+                self.temperature_fenics = fem.Function(function_space_temperature)
                 kwargs = {}
                 if "t" in arguments:
                     kwargs["t"] = self.t
@@ -420,8 +417,7 @@ def define_D_global(self, species):
         D_0 = fem.Function(self.V_DG_0)
         E_D = fem.Function(self.V_DG_0)
         for vol in self.volume_subdomains:
-            cell_indices = vol.locate_subdomain_entities(
-                self.mesh.mesh, self.mesh.vdim)
+            cell_indices = vol.locate_subdomain_entities(self.mesh.mesh, self.mesh.vdim)
 
             # replace values of D_0 and E_D by values from the material
             D_0.x.array[cell_indices] = vol.material.get_D_0(species=species)
@@ -431,11 +427,9 @@ def define_D_global(self, species):
         D = fem.Function(self.V_DG_1)
 
         expr = D_0 * ufl.exp(
-            -E_D / as_fenics_constant(k_B, self.mesh.mesh) /
-            self.temperature_fenics
+            -E_D / as_fenics_constant(k_B, self.mesh.mesh) / self.temperature_fenics
         )
-        D_expr = fem.Expression(
-            expr, self.V_DG_1.element.interpolation_points())
+        D_expr = fem.Expression(expr, self.V_DG_1.element.interpolation_points())
         D.interpolate(D_expr)
         return D, D_expr
 
@@ -516,8 +510,7 @@ def define_boundary_conditions(self):
         for bc in self.boundary_conditions:
             if isinstance(bc.species, str):
                 # if name of species is given then replace with species object
-                bc.species = _species.find_species_from_name(
-                    bc.species, self.species)
+                bc.species = _species.find_species_from_name(bc.species, self.species)
 
         super().define_boundary_conditions()
 
@@ -623,8 +616,7 @@ def create_initial_conditions(self):
 
             # assign to previous solution of species
             if not self.multispecies:
-                condition.species.prev_solution.interpolate(
-                    condition.expr_fenics)
+                condition.species.prev_solution.interpolate(condition.expr_fenics)
             else:
                 idx = self.species.index(condition.species)
                 self.u_n.sub(idx).interpolate(condition.expr_fenics)
@@ -650,8 +642,7 @@ def create_formulation(self):
                     )
 
                 if self.settings.transient:
-                    self.formulation += ((u - u_n) / self.dt) * \
-                        v * self.dx(vol.id)
+                    self.formulation += ((u - u_n) / self.dt) * v * self.dx(vol.id)
 
         for reaction in self.reactions:
             for reactant in reaction.reactant:
@@ -748,8 +739,7 @@ def post_processing(self):
             if isinstance(export, exports.SurfaceQuantity):
                 if isinstance(
                     export,
-                    (exports.SurfaceFlux, exports.TotalSurface,
-                     exports.AverageSurface),
+                    (exports.SurfaceFlux, exports.TotalSurface, exports.AverageSurface),
                 ):
                     export.compute(
                         self.ds,
@@ -1096,8 +1086,7 @@ def create_formulation(self):
             for interface in self.interfaces
         ]
         [interface.pad_parent_maps() for interface in self.interfaces]
-        dInterface = ufl.Measure(
-            "dS", domain=mesh, subdomain_data=integral_data)
+        dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=integral_data)
 
         def mixed_term(u, v, n):
             return ufl.dot(ufl.grad(u), n) * v
@@ -1119,8 +1108,7 @@ def mixed_term(u, v, n):
 
             all_mobile_species = [spe for spe in self.species if spe.mobile]
             if len(all_mobile_species) > 1:
-                raise NotImplementedError(
-                    "Multiple mobile species not implemented")
+                raise NotImplementedError("Multiple mobile species not implemented")
             H = all_mobile_species[0]
             v_b = H.subdomain_to_test_function[subdomain_0](res[0])
             v_t = H.subdomain_to_test_function[subdomain_1](res[1])
@@ -1179,8 +1167,7 @@ def mixed_term(u, v, n):
                     )
                 )
             J.append(jac)
-            forms.append(dolfinx.fem.form(
-                subdomain1.F, entity_maps=entity_maps))
+            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
 
         self.forms = forms
         self.J = J
@@ -1219,8 +1206,7 @@ def initialise_exports(self):
                     )
                 )
             else:
-                raise NotImplementedError(
-                    f"Export type {type(export)} not implemented")
+                raise NotImplementedError(f"Export type {type(export)} not implemented")
 
     def post_processing(self):
         # update post-processing solutions (for each species in each subdomain)
@@ -1241,8 +1227,7 @@ def post_processing(self):
 
         for export in self.exports:
             if not isinstance(export, exports.VTXSpeciesExport):
-                raise NotImplementedError(
-                    f"Export type {type(export)} not implemented")
+                raise NotImplementedError(f"Export type {type(export)} not implemented")
 
     def iterate(self):
         """Iterates the model for a given time step"""
diff --git a/src/festim/mesh/mesh.py b/src/festim/mesh/mesh.py
index 67c253488..560205aff 100644
--- a/src/festim/mesh/mesh.py
+++ b/src/festim/mesh/mesh.py
@@ -1,7 +1,8 @@
 import dolfinx
 import numpy as np
 import ufl
-from dolfinx.mesh import meshtags, Mesh as dolfinx_Mesh
+from dolfinx.mesh import Mesh as dolfinx_Mesh
+from dolfinx.mesh import meshtags
 
 
 class Mesh:
@@ -18,10 +19,10 @@ class Mesh:
         n: Symbolic representation of the vector normal to the facets
             of the mesh.
     """
+
     _mesh: dolfinx.mesh.Mesh
 
     def __init__(self, mesh: dolfinx_Mesh | None = None):
-
         self.mesh = mesh
         if self._mesh is not None:
             # create cell to facet connectivity
@@ -120,6 +121,7 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
             tags_facets[interface_entities] = interface.id
 
         facet_meshtags = meshtags(
-            self._mesh, self.fdim, mesh_facet_indices, tags_facets)
+            self._mesh, self.fdim, mesh_facet_indices, tags_facets
+        )
 
         return facet_meshtags, volume_meshtags
diff --git a/src/festim/mesh/mesh_1d.py b/src/festim/mesh/mesh_1d.py
index da5de0a47..6b62684cb 100644
--- a/src/festim/mesh/mesh_1d.py
+++ b/src/festim/mesh/mesh_1d.py
@@ -37,8 +37,7 @@ def generate_mesh(self):
         """Generates a 1D mesh"""
         degree = 1
         domain = ufl.Mesh(
-            basix.ufl.element(basix.ElementFamily.P,
-                              "interval", degree, shape=(1,))
+            basix.ufl.element(basix.ElementFamily.P, "interval", degree, shape=(1,))
         )
 
         mesh_points = np.reshape(self.vertices, (len(self.vertices), 1))
@@ -57,8 +56,7 @@ def check_borders(self, volume_subdomains):
             Value error: if borders outside the domain
         """
         # check that subdomains are connected
-        all_borders = [
-            border for vol in volume_subdomains for border in vol.borders]
+        all_borders = [border for vol in volume_subdomains for border in vol.borders]
         sorted_borders = np.sort(all_borders)
         for start, end in zip(sorted_borders[1:-2:2], sorted_borders[2:-1:2]):
             if start != end:

From 682a4f3477a6b364fb3fa4f1cece565e64b0b42e Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 25 Oct 2024 12:18:01 -0400
Subject: [PATCH 114/134] added tolerances in the solver

---
 src/festim/hydrogen_transport_problem.py | 2 +-
 test/system_tests/test_multi_material.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 0c0607dfd..660ee2fbf 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -1240,7 +1240,7 @@ def iterate(self):
         self.update_time_dependent_values()
 
         # solve main problem
-        self.solver.solve(self.settings.rtol)
+        self.solver.solve(self.settings.atol, self.settings.rtol)
 
         # post processing
         self.post_processing()
diff --git a/test/system_tests/test_multi_material.py b/test/system_tests/test_multi_material.py
index fd0df647f..7d2720d09 100644
--- a/test/system_tests/test_multi_material.py
+++ b/test/system_tests/test_multi_material.py
@@ -312,7 +312,9 @@ def test_3_materials_transient():
 
     my_model.temperature = lambda x: 300 + 100 * x[0]
 
-    my_model.settings = F.Settings(atol=None, rtol=1e-5, transient=True, final_time=100)
+    my_model.settings = F.Settings(
+        atol=1e-10, rtol=1e-10, transient=True, final_time=100
+    )
     my_model.settings.stepsize = 1
 
     my_model.exports = [

From 99f57afdaed9f9b947d2a1c24ae5734c0eb2a044 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 16:30:31 +0000
Subject: [PATCH 115/134] Fix switch

---
 src/festim/subdomain/interface.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/festim/subdomain/interface.py b/src/festim/subdomain/interface.py
index 3792d865b..85e66212c 100644
--- a/src/festim/subdomain/interface.py
+++ b/src/festim/subdomain/interface.py
@@ -75,8 +75,8 @@ def compute_mapped_interior_facet_data(self, mesh):
         switch = mapped_cell_1 > mapped_cell_0
         # Order restriction on one side
         if True in switch:
-            ordered_integration_data[switch, [0, 1, 2, 3]] = ordered_integration_data[
-                switch, [2, 3, 0, 1]
+            ordered_integration_data[switch, :] = ordered_integration_data[switch][
+                :, [2, 3, 0, 1]
             ]
 
         # Check that other restriction lies in other interface

From e713ef5681b87cd1c9e3d05adb89653da7d1a9cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 13:02:59 -0400
Subject: [PATCH 116/134] Use 3.10 syntax

---
 pyproject.toml                                 |  1 +
 src/festim/boundary_conditions/dirichlet_bc.py | 17 ++++++++---------
 src/festim/hydrogen_transport_problem.py       |  2 +-
 src/festim/reaction.py                         |  4 ++--
 4 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 5f6870ec7..64f5d7d16 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,6 +54,7 @@ ignore_missing_imports = true
 [tool.ruff]
 line-length = 88
 indent-width = 4
+target-version = "py310"
 
 
 [tool.ruff.lint]
diff --git a/src/festim/boundary_conditions/dirichlet_bc.py b/src/festim/boundary_conditions/dirichlet_bc.py
index 095fbd9c7..492024d8a 100644
--- a/src/festim/boundary_conditions/dirichlet_bc.py
+++ b/src/festim/boundary_conditions/dirichlet_bc.py
@@ -1,4 +1,4 @@
-from typing import Callable
+from collections.abc import Callable
 
 import numpy as np
 import numpy.typing as npt
@@ -64,8 +64,8 @@ def value_fenics(self, value: None | fem.Function | fem.Constant | np.ndarray):
         if not isinstance(value, (fem.Function, fem.Constant, np.ndarray)):
             # FIXME: Should we allow sending in a callable here?
             raise TypeError(
-                f"Value must be a dolfinx.fem.Function, dolfinx.fem.Constant, or a np.ndarray not {
-                    type(value)}"
+                "Value must be a dolfinx.fem.Function, dolfinx.fem.Constant, or a np.ndarray not"
+                + f"{type(value)}"
             )
         self._value_fenics = value
 
@@ -112,8 +112,7 @@ def define_surface_subdomain_dofs(
             )
         if mesh.topology.dim - 1 != facet_meshtags.dim:
             raise ValueError(
-                f"Meshtags of dimension {
-                    facet_meshtags.dim}, expected {mesh.topology.dim-1}"
+                f"Meshtags of dimension {facet_meshtags.dim}, expected {mesh.topology.dim-1}"
             )
         bc_dofs = fem.locate_dofs_topological(
             function_space, facet_meshtags.dim, facet_meshtags.find(self.subdomain.id)
@@ -218,8 +217,8 @@ def create_value(
                 # only t is an argument
                 if not isinstance(self.value(t=float(t)), (float, int)):
                     raise ValueError(
-                        f"self.value should return a float or an int, not {
-                            type(self.value(t=float(t)))} "
+                        "self.value should return a float or an int, not"
+                        + f"{type(self.value(t=float(t)))} "
                     )
                 self.value_fenics = helpers.as_fenics_constant(
                     mesh=mesh, value=self.value(t=float(t))
@@ -273,8 +272,8 @@ def create_value(self, function_space: fem.FunctionSpace, t: fem.Constant):
                 # only t is an argument
                 if not isinstance(self.value(t=float(t)), (float, int)):
                     raise ValueError(
-                        f"self.value should return a float or an int, not {
-                            type(self.value(t=float(t)))} "
+                        "self.value should return a float or an int, not"
+                        + f"{type(self.value(t=float(t)))} "
                     )
                 self.value_fenics = helpers.as_fenics_constant(
                     mesh=mesh, value=self.value(t=float(t))
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 660ee2fbf..41d19421b 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -1,4 +1,4 @@
-from typing import Callable
+from collections.abc import Callable
 
 from mpi4py import MPI
 
diff --git a/src/festim/reaction.py b/src/festim/reaction.py
index 945a98252..5577cb96e 100644
--- a/src/festim/reaction.py
+++ b/src/festim/reaction.py
@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import Optional, Union
 
 from ufl import exp
 
@@ -53,7 +53,7 @@ def __init__(
         k_0: float,
         E_k: float,
         volume: VS1D,
-        product: Optional[Union[_Species, List[_Species]]] = [],
+        product: Optional[Union[_Species, list[_Species]]] = [],
         p_0: float = None,
         E_p: float = None,
     ) -> None:

From ce344e8fd35334b9f45873de25184e38b5b3fe73 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Fri, 25 Oct 2024 13:12:37 -0400
Subject: [PATCH 117/134] fixed a few more multilines

---
 src/festim/exports/vtx.py                | 3 +--
 src/festim/exports/xdmf.py               | 4 ++--
 src/festim/hydrogen_transport_problem.py | 8 ++++----
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/festim/exports/vtx.py b/src/festim/exports/vtx.py
index 30e8184bf..1ef15a60b 100644
--- a/src/festim/exports/vtx.py
+++ b/src/festim/exports/vtx.py
@@ -16,8 +16,7 @@ def __init__(self, filename: str | Path, ext: str) -> None:
         name = Path(filename)
         if name.suffix != ext:
             warnings.warn(
-                f"Filename {filename} does not have {
-                    ext} extension, adding it."
+                f"Filename {filename} does not have {ext} extension, adding it."
             )
             name = name.with_suffix(ext)
 
diff --git a/src/festim/exports/xdmf.py b/src/festim/exports/xdmf.py
index 5c3c38cd6..68871cc5d 100644
--- a/src/festim/exports/xdmf.py
+++ b/src/festim/exports/xdmf.py
@@ -50,8 +50,8 @@ def field(self, value: _Species | list[_Species]):
             val = [value]
         else:
             raise TypeError(
-                f"field must be of type festim.Species or a list of festim.Species, got {
-                    type(value)}."
+                f"field must be of type festim.Species or a list of festim.Species, got "
+                f"{type(value)}."
             )
         self._field = val
 
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 41d19421b..511a202a9 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -227,8 +227,8 @@ def species(self, value):
         for spe in value:
             if not isinstance(spe, _species.Species):
                 raise TypeError(
-                    f"elements of species must be of type festim.Species not {
-                        type(spe)}"
+                    f"elements of species must be of type festim.Species not "
+                    f"{type(spe)}"
                 )
         self._species = value
 
@@ -316,8 +316,8 @@ def define_temperature(self):
             if "t" in arguments and "x" not in arguments:
                 if not isinstance(self.temperature(t=float(self.t)), (float, int)):
                     raise ValueError(
-                        f"self.temperature should return a float or an int, not {
-                            type(self.temperature(t=float(self.t)))} "
+                        f"self.temperature should return a float or an int, not "
+                        f"{type(self.temperature(t=float(self.t)))} "
                     )
                 # only t is an argument
                 self.temperature_fenics = as_fenics_constant(

From 085bd30bdee2ac814416aa4736f2a79e47b17003 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 18:01:28 +0000
Subject: [PATCH 118/134] Last 3.10 backports of fstring

---
 src/festim/reaction.py | 12 ++++++------
 src/festim/species.py  |  4 ++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/festim/reaction.py b/src/festim/reaction.py
index 5577cb96e..ab61956b3 100644
--- a/src/festim/reaction.py
+++ b/src/festim/reaction.py
@@ -80,8 +80,8 @@ def reactant(self, value):
         for i in value:
             if not isinstance(i, (_Species, _ImplicitSpecies)):
                 raise TypeError(
-                    f"reactant must be an F.Species or F.ImplicitSpecies, not {
-                        type(i)}"
+                    "reactant must be an F.Species or F.ImplicitSpecies, not"
+                    + f"{type(i)}"
                 )
         self._reactant = value
 
@@ -112,13 +112,13 @@ def reaction_term(self, temperature):
         if self.product == []:
             if self.p_0 is not None:
                 raise ValueError(
-                    f"p_0 must be None, not {
-                        self.p_0} when no products are present."
+                    f"p_0 must be None, not {self.p_0}"
+                    + " when no products are present."
                 )
             if self.E_p is not None:
                 raise ValueError(
-                    f"E_p must be None, not {
-                        self.E_p} when no products are present."
+                    f"E_p must be None, not {self.E_p}"
+                    + " when no products are present."
                 )
         else:
             if self.p_0 == None:
diff --git a/src/festim/species.py b/src/festim/species.py
index ca9d28921..477fff660 100644
--- a/src/festim/species.py
+++ b/src/festim/species.py
@@ -141,8 +141,8 @@ def concentration(self):
             for other in self.others:
                 if other.solution is None:
                     raise ValueError(
-                        f"Cannot compute concentration of {
-                            self.name} because {other.name} has no solution"
+                        f"Cannot compute concentration of {self.name} " +
+                        f"because {other.name} has no solution."
                     )
         return self.n - sum([other.solution for other in self.others])
 

From f989dbe9e1d5a15cb5b8010ba1478d4ff9ef89cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 18:01:39 +0000
Subject: [PATCH 119/134] One more fstring backoport

---
 src/festim/species.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/festim/species.py b/src/festim/species.py
index 477fff660..104dea82c 100644
--- a/src/festim/species.py
+++ b/src/festim/species.py
@@ -141,8 +141,8 @@ def concentration(self):
             for other in self.others:
                 if other.solution is None:
                     raise ValueError(
-                        f"Cannot compute concentration of {self.name} " +
-                        f"because {other.name} has no solution."
+                        f"Cannot compute concentration of {self.name} "
+                        + f"because {other.name} has no solution."
                     )
         return self.n - sum([other.solution for other in self.others])
 

From 25a97903fcdf7d45f4fac384ca86002653e37cac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 18:10:42 +0000
Subject: [PATCH 120/134] Lasts tests

---
 src/festim/boundary_conditions/dirichlet_bc.py | 4 ++--
 src/festim/reaction.py                         | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/festim/boundary_conditions/dirichlet_bc.py b/src/festim/boundary_conditions/dirichlet_bc.py
index 492024d8a..0a5af01d6 100644
--- a/src/festim/boundary_conditions/dirichlet_bc.py
+++ b/src/festim/boundary_conditions/dirichlet_bc.py
@@ -217,7 +217,7 @@ def create_value(
                 # only t is an argument
                 if not isinstance(self.value(t=float(t)), (float, int)):
                     raise ValueError(
-                        "self.value should return a float or an int, not"
+                        "self.value should return a float or an int, not "
                         + f"{type(self.value(t=float(t)))} "
                     )
                 self.value_fenics = helpers.as_fenics_constant(
@@ -272,7 +272,7 @@ def create_value(self, function_space: fem.FunctionSpace, t: fem.Constant):
                 # only t is an argument
                 if not isinstance(self.value(t=float(t)), (float, int)):
                     raise ValueError(
-                        "self.value should return a float or an int, not"
+                        "self.value should return a float or an int, not "
                         + f"{type(self.value(t=float(t)))} "
                     )
                 self.value_fenics = helpers.as_fenics_constant(
diff --git a/src/festim/reaction.py b/src/festim/reaction.py
index ab61956b3..64a4effea 100644
--- a/src/festim/reaction.py
+++ b/src/festim/reaction.py
@@ -80,7 +80,7 @@ def reactant(self, value):
         for i in value:
             if not isinstance(i, (_Species, _ImplicitSpecies)):
                 raise TypeError(
-                    "reactant must be an F.Species or F.ImplicitSpecies, not"
+                    "reactant must be an F.Species or F.ImplicitSpecies, not "
                     + f"{type(i)}"
                 )
         self._reactant = value

From 6fd66db6fd250dd49076bfffe6bdd8ea97ac756d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 18:46:34 +0000
Subject: [PATCH 121/134] Simplify form compilation

---
 src/festim/hydrogen_transport_problem.py | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 511a202a9..37ccddd70 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -1155,22 +1155,18 @@ def mixed_term(u, v, n):
             subdomain_1.F += F_1
 
         J = []
-        forms = []
         for subdomain1 in self.volume_subdomains:
             jac = []
-            form = subdomain1.F
             for subdomain2 in self.volume_subdomains:
                 jac.append(
-                    dolfinx.fem.form(
-                        ufl.derivative(form, subdomain2.u),
-                        entity_maps=entity_maps,
-                    )
+                    ufl.derivative(subdomain1.F, subdomain2.u),
                 )
             J.append(jac)
-            forms.append(dolfinx.fem.form(subdomain1.F, entity_maps=entity_maps))
-
-        self.forms = forms
-        self.J = J
+        self.forms = dolfinx.fem.form(
+            [subdomain.F for subdomain in self.volume_subdomains],
+            entity_maps=entity_maps,
+        )
+        self.J = dolfinx.fem.form(J, entity_maps=entity_maps)
 
     def create_solver(self):
         self.solver = NewtonSolver(

From dbde7e3bb6f9bffc49af0e0e64cd341a4dd6c347 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 14:49:45 -0400
Subject: [PATCH 122/134] Minor fixes to transfer meshtags

---
 src/festim/helpers_discontinuity.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/festim/helpers_discontinuity.py b/src/festim/helpers_discontinuity.py
index 2c225a522..5395dc806 100644
--- a/src/festim/helpers_discontinuity.py
+++ b/src/festim/helpers_discontinuity.py
@@ -9,6 +9,7 @@ def transfer_meshtags_to_submesh(
     Transfer a meshtag from a parent mesh to a sub-mesh.
     """
 
+    # Invert map from sub-mesh to parent-mesh to parent-mesh to sub-mesh
     tdim = mesh.topology.dim
     cell_imap = mesh.topology.index_map(tdim)
     num_cells = cell_imap.size_local + cell_imap.num_ghosts
@@ -16,8 +17,9 @@ def transfer_meshtags_to_submesh(
     mesh_to_submesh[sub_cell_to_parent] = np.arange(
         len(sub_cell_to_parent), dtype=np.int32
     )
-    sub_vertex_to_parent = np.asarray(sub_vertex_to_parent)
 
+    # Compute and access various entity to vertex and vertex to
+    # entity connectivity for parent and sub-mesh.
     submesh.topology.create_connectivity(entity_tag.dim, 0)
 
     num_child_entities = (

From f05d785ae73039e07940e2eeb46b3a373abf385e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 14:56:49 -0400
Subject: [PATCH 123/134] Unify formatting for integration measures with Rem(i)

---
 src/festim/hydrogen_transport_problem.py | 22 +++++-----------------
 1 file changed, 5 insertions(+), 17 deletions(-)

diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 37ccddd70..b8ec50b2c 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -4,7 +4,6 @@
 
 import basix
 import dolfinx
-import numpy as np
 import numpy.typing as npt
 import tqdm.autonotebook
 import ufl
@@ -864,18 +863,8 @@ def initialise(self):
         self.create_source_values_fenics()
         self.create_flux_values_fenics()
         self.create_initial_conditions()
-
         for subdomain in self.volume_subdomains:
             self.define_function_spaces(subdomain)
-            ct_r = dolfinx.mesh.meshtags(
-                self.mesh.mesh,
-                self.mesh.mesh.topology.dim,
-                subdomain.submesh_to_mesh,
-                np.full_like(subdomain.submesh_to_mesh, 1, dtype=np.int32),
-            )
-            subdomain.dx = ufl.Measure(
-                "dx", domain=self.mesh.mesh, subdomain_data=ct_r, subdomain_id=1
-            )
             self.create_subdomain_formulation(subdomain)
             subdomain.u.name = f"u_{subdomain.id}"
 
@@ -1007,16 +996,15 @@ def create_subdomain_formulation(self, subdomain: _subdomain.VolumeSubdomain):
             u = spe.subdomain_to_solution[subdomain]
             u_n = spe.subdomain_to_prev_solution[subdomain]
             v = spe.subdomain_to_test_function[subdomain]
-            dx = subdomain.dx
 
             D = subdomain.material.get_diffusion_coefficient(
                 self.mesh.mesh, self.temperature_fenics, spe
             )
             if self.settings.transient:
-                form += ((u - u_n) / self.dt) * v * dx
+                form += ((u - u_n) / self.dt) * v * self.dx(subdomain.id)
 
             if spe.mobile:
-                form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * dx
+                form += ufl.inner(D * ufl.grad(u), ufl.grad(v)) * self.dx(subdomain.id)
 
         # add reaction terms
         for reaction in self.reactions:
@@ -1034,7 +1022,7 @@ def create_subdomain_formulation(self, subdomain: _subdomain.VolumeSubdomain):
                     form += (
                         reaction.reaction_term(self.temperature_fenics)
                         * reactant.subdomain_to_test_function[subdomain]
-                        * dx
+                        * self.dx(subdomain.id)
                     )
 
             # product
@@ -1046,7 +1034,7 @@ def create_subdomain_formulation(self, subdomain: _subdomain.VolumeSubdomain):
                 form += (
                     -reaction.reaction_term(self.temperature_fenics)
                     * product.subdomain_to_test_function[subdomain]
-                    * dx
+                    * self.dx(subdomain.id)
                 )
 
         # add fluxes
@@ -1062,7 +1050,7 @@ def create_subdomain_formulation(self, subdomain: _subdomain.VolumeSubdomain):
         for source in self.sources:
             v = source.species.subdomain_to_test_function[subdomain]
             if source.volume == subdomain:
-                form -= source.value_fenics * v * dx
+                form -= source.value_fenics * v * self.dx(subdomain.id)
 
         # store the form in the subdomain object
         subdomain.F = form

From 18e4cad2edd028ce48270853cefbf7520900e5e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 15:21:57 -0400
Subject: [PATCH 124/134] Various fixes to post-processing

---
 src/festim/exports/maximum_surface.py        |  6 +++++-
 src/festim/exports/maximum_volume.py         |  7 ++++++-
 src/festim/exports/minimum_surface.py        |  4 +++-
 src/festim/exports/minimum_volume.py         |  6 +++++-
 src/festim/hydrogen_transport_problem.py     |  2 +-
 src/festim/mesh/mesh.py                      |  2 +-
 src/festim/subdomain/surface_subdomain.py    | 14 ++++++++------
 src/festim/subdomain/surface_subdomain_1d.py | 19 +++++++++++--------
 src/festim/subdomain/volume_subdomain.py     | 12 ++++--------
 src/festim/subdomain/volume_subdomain_1d.py  | 13 ++++++-------
 10 files changed, 50 insertions(+), 35 deletions(-)

diff --git a/src/festim/exports/maximum_surface.py b/src/festim/exports/maximum_surface.py
index b398f91c0..0d13eb400 100644
--- a/src/festim/exports/maximum_surface.py
+++ b/src/festim/exports/maximum_surface.py
@@ -24,5 +24,9 @@ def compute(self):
         Computes the maximum value of the field on the defined surface
         subdomain, and appends it to the data list
         """
-        self.value = np.max(self.field.solution.x.array[self.surface.indices])
+        solution = self.field.solution
+
+        indices = self.surf.locate_boundary_facet_indices(solution.function_space.mesh)
+
+        self.value = np.max(self.field.solution.x.array[indices])
         self.data.append(self.value)
diff --git a/src/festim/exports/maximum_volume.py b/src/festim/exports/maximum_volume.py
index 80534e777..73bb51e36 100644
--- a/src/festim/exports/maximum_volume.py
+++ b/src/festim/exports/maximum_volume.py
@@ -25,5 +25,10 @@ def compute(self):
         Computes the maximum value of solution function within the defined volume
         subdomain, and appends it to the data list
         """
-        self.value = np.max(self.field.solution.x.array[self.volume.entities])
+
+        solution = self.field.solution
+        indices = self.volume.locate_subdomain_entities(
+            solution.function_space.mesh, solution.function_space.mesh.topology.dim
+        )
+        self.value = np.max(self.field.solution.x.array[indices])
         self.data.append(self.value)
diff --git a/src/festim/exports/minimum_surface.py b/src/festim/exports/minimum_surface.py
index c5b53e3a7..5f1fbf720 100644
--- a/src/festim/exports/minimum_surface.py
+++ b/src/festim/exports/minimum_surface.py
@@ -24,5 +24,7 @@ def compute(self):
         Computes the minimum value of the field on the defined surface
         subdomain, and appends it to the data list
         """
-        self.value = np.min(self.field.solution.x.array[self.surface.indices])
+        solution = self.field.solution
+        indices = self.surf.locate_boundary_facet_indices(solution.function_space.mesh)
+        self.value = np.min(self.field.solution.x.array[indices])
         self.data.append(self.value)
diff --git a/src/festim/exports/minimum_volume.py b/src/festim/exports/minimum_volume.py
index 0a29e17f8..4c3b5f84e 100644
--- a/src/festim/exports/minimum_volume.py
+++ b/src/festim/exports/minimum_volume.py
@@ -24,5 +24,9 @@ def compute(self):
         Computes the minimum value of solution function within the defined volume
         subdomain, and appends it to the data list
         """
-        self.value = np.min(self.field.solution.x.array[self.volume.entities])
+        solution = self.field.solution
+        indices = self.volume.locate_subdomain_entities(solution.function_space.mesh)
+        # FIXME: np.min/np.max is not parallel safe (unique value per process)
+        # Needs to use a reduction operation (MPI.comm.allreduce(..., op=...))
+        self.value = np.min(self.field.solution.x.array[indices])
         self.data.append(self.value)
diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index b8ec50b2c..3366fb3fe 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -416,7 +416,7 @@ def define_D_global(self, species):
         D_0 = fem.Function(self.V_DG_0)
         E_D = fem.Function(self.V_DG_0)
         for vol in self.volume_subdomains:
-            cell_indices = vol.locate_subdomain_entities(self.mesh.mesh, self.mesh.vdim)
+            cell_indices = vol.locate_subdomain_entities(self.mesh.mesh)
 
             # replace values of D_0 and E_D by values from the material
             D_0.x.array[cell_indices] = vol.material.get_D_0(species=species)
diff --git a/src/festim/mesh/mesh.py b/src/festim/mesh/mesh.py
index 560205aff..757a4ed92 100644
--- a/src/festim/mesh/mesh.py
+++ b/src/festim/mesh/mesh.py
@@ -94,7 +94,7 @@ def define_meshtags(self, surface_subdomains, volume_subdomains, interfaces=None
 
         for vol in volume_subdomains:
             # find all cells in subdomain and mark them as vol.id
-            entities = vol.locate_subdomain_entities(self._mesh, self.vdim)
+            entities = vol.locate_subdomain_entities(self._mesh)
             tags_volumes[entities] = vol.id
 
         volume_meshtags = meshtags(
diff --git a/src/festim/subdomain/surface_subdomain.py b/src/festim/subdomain/surface_subdomain.py
index 578545804..7064e890e 100644
--- a/src/festim/subdomain/surface_subdomain.py
+++ b/src/festim/subdomain/surface_subdomain.py
@@ -1,5 +1,6 @@
 import dolfinx.mesh
 import numpy as np
+import numpy.typing as npt
 
 
 class SurfaceSubdomain:
@@ -10,10 +11,14 @@ class SurfaceSubdomain:
         id (int): the id of the surface subdomain
     """
 
+    id: int
+
     def __init__(self, id):
         self.id = id
 
-    def locate_boundary_facet_indices(self, mesh):
+    def locate_boundary_facet_indices(
+        self, mesh: dolfinx.mesh.Mesh
+    ) -> npt.NDArray[np.int32]:
         """Locates the dof of the surface subdomain within the function space
         and return the index of the dof
 
@@ -25,11 +30,8 @@ def locate_boundary_facet_indices(self, mesh):
                 indices of the subdomain
         """
         fdim = mesh.topology.dim - 1
-        # By default, all entities are included
-        indices = dolfinx.mesh.locate_entities_boundary(
-            mesh, fdim, lambda x: np.full(x.shape[1], True, dtype=bool)
-        )
-        return indices
+        mesh.topology.create_connectivity(fdim, fdim + 1)
+        return dolfinx.mesh.exterior_facet_indices(mesh.topology)
 
 
 def find_surface_from_id(id: int, surfaces: list):
diff --git a/src/festim/subdomain/surface_subdomain_1d.py b/src/festim/subdomain/surface_subdomain_1d.py
index a14827122..03e5c421c 100644
--- a/src/festim/subdomain/surface_subdomain_1d.py
+++ b/src/festim/subdomain/surface_subdomain_1d.py
@@ -1,10 +1,10 @@
 import dolfinx.mesh
 import numpy as np
 
-import festim as F
+from .surface_subdomain import SurfaceSubdomain
 
 
-class SurfaceSubdomain1D(F.SurfaceSubdomain):
+class SurfaceSubdomain1D(SurfaceSubdomain):
     """
     Surface subdomain class for 1D cases
 
@@ -20,11 +20,15 @@ class SurfaceSubdomain1D(F.SurfaceSubdomain):
         >>> surf_subdomain = F.SurfaceSubdomain1D(id=1, x=1)
     """
 
-    def __init__(self, id, x) -> None:
+    # FIXME: Rename this to _id and use getter/setter
+    id: int
+    x: float
+
+    def __init__(self, id: int, x: float) -> None:
         super().__init__(id)
         self.x = x
 
-    def locate_boundary_facet_indices(self, mesh):
+    def locate_boundary_facet_indices(self, mesh: dolfinx.mesh.Mesh):
         """Locates the dof of the surface subdomain within the function space
         and return the index of the dof
 
@@ -36,8 +40,7 @@ def locate_boundary_facet_indices(self, mesh):
                 indices of the subdomain
         """
         assert mesh.geometry.dim == 1, "This method is only for 1D meshes"
-        fdim = 0
-        self.indices = dolfinx.mesh.locate_entities_boundary(
-            mesh, fdim, lambda x: np.isclose(x[0], self.x)
+        indices = dolfinx.mesh.locate_entities_boundary(
+            mesh, 0, lambda x: np.isclose(x[0], self.x)
         )
-        return self.indices
+        return indices
diff --git a/src/festim/subdomain/volume_subdomain.py b/src/festim/subdomain/volume_subdomain.py
index 77f6b9bb2..d30ae65f1 100644
--- a/src/festim/subdomain/volume_subdomain.py
+++ b/src/festim/subdomain/volume_subdomain.py
@@ -1,6 +1,5 @@
 import dolfinx
 import numpy as np
-from dolfinx.mesh import locate_entities
 
 from festim.helpers_discontinuity import transfer_meshtags_to_submesh
 
@@ -30,21 +29,18 @@ def __init__(self, id, material):
         self.id = id
         self.material = material
 
-    def locate_subdomain_entities(self, mesh, vdim):
+    def locate_subdomain_entities(self, mesh: dolfinx.mesh.Mesh):
         """Locates all cells in subdomain borders within domain
 
         Args:
             mesh (dolfinx.mesh.Mesh): the mesh of the model
-            vdim (int): the dimension of the volumes of the mesh,
-                for 1D this is always 1
 
         Returns:
             entities (np.array): the entities of the subdomain
         """
-        # By default, all entities are included
-        # return array like x full of True
-        entities = locate_entities(mesh, vdim, lambda x: np.full(x.shape[1], True))
-        return entities
+        cell_map = mesh.topology.index_map(mesh.topology.dim)
+        num_cells_local = cell_map.size_local + cell_map.num_ghosts
+        return np.arange(num_cells_local, dtype=np.int32)
 
     def create_subdomain(self, mesh: dolfinx.mesh.Mesh, marker: dolfinx.mesh.MeshTags):
         """
diff --git a/src/festim/subdomain/volume_subdomain_1d.py b/src/festim/subdomain/volume_subdomain_1d.py
index a024ac8c1..bdf4a9ecd 100644
--- a/src/festim/subdomain/volume_subdomain_1d.py
+++ b/src/festim/subdomain/volume_subdomain_1d.py
@@ -1,5 +1,6 @@
 import numpy as np
-from dolfinx.mesh import locate_entities
+import numpy.typing as npt
+from dolfinx.mesh import Mesh, locate_entities
 
 from festim.subdomain import VolumeSubdomain
 
@@ -27,20 +28,18 @@ def __init__(self, id, borders, material) -> None:
         super().__init__(id, material)
         self.borders = borders
 
-    def locate_subdomain_entities(self, mesh, vdim):
+    def locate_subdomain_entities(self, mesh: Mesh) -> npt.NDArray[np.int32]:
         """Locates all cells in subdomain borders within domain
 
         Args:
             mesh (dolfinx.mesh.Mesh): the mesh of the model
-            vdim (int): the dimension of the volumes of the mesh,
-                for 1D this is always 1
 
         Returns:
             entities (np.array): the entities of the subdomain
         """
-        self.entities = locate_entities(
+        entities = locate_entities(
             mesh,
-            vdim,
+            mesh.topology.dim,
             lambda x: np.logical_and(x[0] >= self.borders[0], x[0] <= self.borders[1]),
         )
-        return self.entities
+        return entities

From 9e149e3f1173144584c9b645a8152fc0acd518af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 15:22:27 -0400
Subject: [PATCH 125/134] Fix variable name

---
 src/festim/exports/maximum_surface.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/festim/exports/maximum_surface.py b/src/festim/exports/maximum_surface.py
index 0d13eb400..f94cc8105 100644
--- a/src/festim/exports/maximum_surface.py
+++ b/src/festim/exports/maximum_surface.py
@@ -25,8 +25,6 @@ def compute(self):
         subdomain, and appends it to the data list
         """
         solution = self.field.solution
-
-        indices = self.surf.locate_boundary_facet_indices(solution.function_space.mesh)
-
+        indices = self.surface.locate_boundary_facet_indices(solution.function_space.mesh)
         self.value = np.max(self.field.solution.x.array[indices])
         self.data.append(self.value)

From 7271ec9ac66fb7f514f3a2a1706138faacde074c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 15:23:24 -0400
Subject: [PATCH 126/134] Fix variable name

---
 src/festim/exports/maximum_volume.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/festim/exports/maximum_volume.py b/src/festim/exports/maximum_volume.py
index 73bb51e36..60ac1c24c 100644
--- a/src/festim/exports/maximum_volume.py
+++ b/src/festim/exports/maximum_volume.py
@@ -25,10 +25,7 @@ def compute(self):
         Computes the maximum value of solution function within the defined volume
         subdomain, and appends it to the data list
         """
-
         solution = self.field.solution
-        indices = self.volume.locate_subdomain_entities(
-            solution.function_space.mesh, solution.function_space.mesh.topology.dim
-        )
+        indices = self.volume.locate_subdomain_entities(solution.function_space.mesh)
         self.value = np.max(self.field.solution.x.array[indices])
         self.data.append(self.value)

From 0ea4875ae9756a4b5d515e61129ef44ce35d4745 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 15:24:09 -0400
Subject: [PATCH 127/134] More fixes

---
 src/festim/exports/maximum_surface.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/festim/exports/maximum_surface.py b/src/festim/exports/maximum_surface.py
index f94cc8105..5275c6f49 100644
--- a/src/festim/exports/maximum_surface.py
+++ b/src/festim/exports/maximum_surface.py
@@ -25,6 +25,8 @@ def compute(self):
         subdomain, and appends it to the data list
         """
         solution = self.field.solution
-        indices = self.surface.locate_boundary_facet_indices(solution.function_space.mesh)
+        indices = self.surface.locate_boundary_facet_indices(
+            solution.function_space.mesh
+        )
         self.value = np.max(self.field.solution.x.array[indices])
         self.data.append(self.value)

From a4cdb826980bc888505557f95dbed02087314b1e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Fri, 25 Oct 2024 15:28:07 -0400
Subject: [PATCH 128/134] Fix test

---
 src/festim/exports/minimum_surface.py | 2 +-
 test/test_maximum_volume.py           | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/festim/exports/minimum_surface.py b/src/festim/exports/minimum_surface.py
index 5f1fbf720..4f2918789 100644
--- a/src/festim/exports/minimum_surface.py
+++ b/src/festim/exports/minimum_surface.py
@@ -25,6 +25,6 @@ def compute(self):
         subdomain, and appends it to the data list
         """
         solution = self.field.solution
-        indices = self.surf.locate_boundary_facet_indices(solution.function_space.mesh)
+        indices = self.surface.locate_boundary_facet_indices(solution.function_space.mesh)
         self.value = np.min(self.field.solution.x.array[indices])
         self.data.append(self.value)
diff --git a/test/test_maximum_volume.py b/test/test_maximum_volume.py
index 3fed21f5e..6ff25ea3f 100644
--- a/test/test_maximum_volume.py
+++ b/test/test_maximum_volume.py
@@ -12,7 +12,7 @@ def test_maximum_volume_compute_1D():
     dummy_material = F.Material(D_0=1.5, E_D=1, name="dummy")
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
     dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_material)
-    dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh, vdim=0)
+    dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh)
 
     # give function to species
     V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))

From f3df3b0839c92d447665c38c31ebafb00495ec73 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Mon, 28 Oct 2024 10:18:07 -0400
Subject: [PATCH 129/134] fixed test

---
 test/test_minimum_volume.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_minimum_volume.py b/test/test_minimum_volume.py
index 139fb8a10..4318da330 100644
--- a/test/test_minimum_volume.py
+++ b/test/test_minimum_volume.py
@@ -12,7 +12,7 @@ def test_minimum_volume_compute_1D():
     dummy_material = F.Material(D_0=1.5, E_D=1, name="dummy")
     my_mesh = F.Mesh1D(np.linspace(0, L, 10000))
     dummy_volume = F.VolumeSubdomain1D(id=1, borders=[0, L], material=dummy_material)
-    dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh, vdim=0)
+    dummy_volume.locate_subdomain_entities(mesh=my_mesh.mesh)
 
     # give function to species
     V = fem.functionspace(my_mesh.mesh, ("Lagrange", 1))

From 403192e0a612adea3241486dbd9f4917a211a1cc Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Mon, 28 Oct 2024 10:25:01 -0400
Subject: [PATCH 130/134] format

---
 src/festim/exports/minimum_surface.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/festim/exports/minimum_surface.py b/src/festim/exports/minimum_surface.py
index 4f2918789..948af5816 100644
--- a/src/festim/exports/minimum_surface.py
+++ b/src/festim/exports/minimum_surface.py
@@ -25,6 +25,8 @@ def compute(self):
         subdomain, and appends it to the data list
         """
         solution = self.field.solution
-        indices = self.surface.locate_boundary_facet_indices(solution.function_space.mesh)
+        indices = self.surface.locate_boundary_facet_indices(
+            solution.function_space.mesh
+        )
         self.value = np.min(self.field.solution.x.array[indices])
         self.data.append(self.value)

From ae176983bd1171be105388ec882032aeabe101c3 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Mon, 28 Oct 2024 21:44:12 -0400
Subject: [PATCH 131/134] vtxfile write in HydrogenTransportProblem

---
 src/festim/hydrogen_transport_problem.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 3366fb3fe..7884a78b4 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -765,6 +765,10 @@ def post_processing(self):
             if isinstance(export, exports.XDMFExport):
                 export.write(float(self.t))
 
+        # should we move this to problem.ProblemBase?
+        for vtxfile in self._vtxfiles:
+            vtxfile.write(float(self.t))
+
 
 class HTransportProblemDiscontinuous(HydrogenTransportProblem):
     interfaces: list[_subdomain.Interface]

From d2f77877073d1dd1501c3a11b94951a23d57668a Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Wed, 30 Oct 2024 09:55:06 -0400
Subject: [PATCH 132/134] moved convergence rates

---
 convergence_rates.py => examples/convergence_rates.py | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename convergence_rates.py => examples/convergence_rates.py (100%)

diff --git a/convergence_rates.py b/examples/convergence_rates.py
similarity index 100%
rename from convergence_rates.py
rename to examples/convergence_rates.py

From 6371bb617a3444bf3e4114844da60beb09dba608 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Delaporte-Mathurin?=
 <40028739+RemDelaporteMathurin@users.noreply.github.com>
Date: Wed, 30 Oct 2024 09:55:42 -0400
Subject: [PATCH 133/134] comment

---
 src/festim/hydrogen_transport_problem.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 7884a78b4..6d62aa9c4 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -1147,6 +1147,7 @@ def mixed_term(u, v, n):
             subdomain_1.F += F_1
 
         J = []
+        # this is the symbolic differentiation of the Jacobian
         for subdomain1 in self.volume_subdomains:
             jac = []
             for subdomain2 in self.volume_subdomains:

From 7d0a6d459c4371f3d07758f2a5e0dd35d7e7d0b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Delaporte-Mathurin?=
 <40028739+RemDelaporteMathurin@users.noreply.github.com>
Date: Wed, 30 Oct 2024 09:55:51 -0400
Subject: [PATCH 134/134] comment

---
 src/festim/hydrogen_transport_problem.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py
index 6d62aa9c4..ddefd98f1 100644
--- a/src/festim/hydrogen_transport_problem.py
+++ b/src/festim/hydrogen_transport_problem.py
@@ -1155,6 +1155,7 @@ def mixed_term(u, v, n):
                     ufl.derivative(subdomain1.F, subdomain2.u),
                 )
             J.append(jac)
+        # compile jacobian (J) and residual (F)
         self.forms = dolfinx.fem.form(
             [subdomain.F for subdomain in self.volume_subdomains],
             entity_maps=entity_maps,