Start fixing entity_maps with backwards compatibility

Author: jorgensd

examples/advection_diffusion_multi_material.py

@@ -80,9 +80,9 @@ def velocity_func(x):
         return values
 
     mesh2, mt2, ct2 = generate_mesh(n=10)
-    submesh, submesh_to_mesh, v_map = dolfinx.mesh.create_submesh(
-        mesh2, ct2.dim, ct2.find(4)
-    )[0:3]
+    submesh, cell_map, v_map = dolfinx.mesh.create_submesh(mesh2, ct2.dim, ct2.find(4))[
+        0:3
+    ]
     v_cg = basix.ufl.element(
         "Lagrange", submesh.topology.cell_name(), 2, shape=(submesh.geometry.dim,)
     )

examples/multi_material_2d.py

@@ -107,7 +107,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=1e-5, transient=False)
+my_model.settings = F.Settings(atol=1e-5, rtol=1e-5, transient=False)
 
 my_model.exports = [
     F.VTXSpeciesExport(f"u_{subdomain.id}.bp", field=H, subdomain=subdomain)
@@ -120,7 +120,20 @@ def half(x):
 # -------------------- post processing --------------------
 
 # derived quantities
-entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
+# Check for DOLFINx version compatibility
+try:
+    from dolfinx.mesh import EntityMap  # noqa: F401
+
+    legacy_entity_map = False
+    entity_map_wrapper = lambda e_map: list(e_map.values())
+    entity_maps = [sd.cell_map for sd in my_model.volume_subdomains]
+except ImportError:
+    legacy_entity_map = True
+    entity_map_wrapper = lambda e_map: e_map
+    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)
@@ -143,18 +156,26 @@ def half(x):
     * ds(1),
     entity_maps=entity_maps,
 )
-print(dolfinx.fem.assemble_scalar(form))
 
-entity_maps[mesh] = bottom_domain.submesh_to_mesh
+
+def assemble_and_print(form: dolfinx.fem.Form):
+    loc_val = dolfinx.fem.assemble_scalar(form)
+    glob_val = mesh.comm.allreduce(loc_val, op=MPI.SUM)
+    if MPI.COMM_WORLD.rank == 0:
+        print(glob_val, flush=True)
+
+
+assemble_and_print(form)
+
 
 form = dolfinx.fem.form(bottom_domain.u.sub(0) * dx_b)
-print(dolfinx.fem.assemble_scalar(form))
+assemble_and_print(form)
 
 form = dolfinx.fem.form(
     my_model.temperature_fenics * dx_b,
-    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+    entity_maps=entity_map_wrapper({mesh: bottom_domain.cell_map}),
 )
-print(dolfinx.fem.assemble_scalar(form))
+assemble_and_print(form)
 
 D = bottom_domain.material.get_diffusion_coefficient(
     my_model.mesh.mesh, my_model.temperature_fenics, H
@@ -166,15 +187,15 @@ def half(x):
         n_b,
     )
     * ds_b(id_interface),
-    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+    entity_maps=entity_map_wrapper({mesh: bottom_domain.cell_map}),
 )
-print(dolfinx.fem.assemble_scalar(form))
+assemble_and_print(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},
+    entity_maps=entity_map_wrapper({mesh: top_domain.cell_map}),
 )
-print(dolfinx.fem.assemble_scalar(form))
+assemble_and_print(form)

examples/multi_material_2d_bis.py

@@ -146,8 +146,7 @@ def half(x):
 
 
 # derived quantities
-entity_maps = {sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains}
-entity_maps[mesh] = bottom_domain.submesh_to_mesh
+
 
 ds = ufl.Measure("ds", domain=mesh, subdomain_data=my_model.facet_meshtags)
 ds_b = ufl.Measure("ds", domain=bottom_domain.submesh, subdomain_data=bottom_domain.ft)
@@ -164,23 +163,35 @@ def half(x):
 form = dolfinx.fem.form(bottom_domain.u.sub(1) * dx_b)
 print(dolfinx.fem.assemble_scalar(form))
 
+# Check for DOLFINx version compatibility
+try:
+    from dolfinx.mesh import EntityMap  # noqa: F401
+
+    legacy_entity_map = False
+    entity_map_wrapper = lambda e_map: list(e_map.values())
+except ImportError:
+    legacy_entity_map = True
+    entity_map_wrapper = lambda e_map: e_map
+
+
 form = dolfinx.fem.form(
     my_model.temperature_fenics * dx_b,
-    entity_maps={mesh: bottom_domain.submesh_to_mesh},
+    entity_maps=entity_map_wrapper({mesh: bottom_domain.cell_map}),
 )
 print(dolfinx.fem.assemble_scalar(form))
 
 D = bottom_domain.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},
+    entity_maps=entity_map_wrapper({mesh: bottom_domain.cell_map}),
 )
 print(dolfinx.fem.assemble_scalar(form))
 form = dolfinx.fem.form(
@@ -189,19 +200,26 @@ def half(x):
         n_t,
     )
     * ds_t(id_interface),
-    entity_maps={mesh: top_domain.submesh_to_mesh},
+    entity_maps=entity_map_wrapper({mesh: top_domain.cell_map}),
 )
 print(dolfinx.fem.assemble_scalar(form))
 
 
 # using submesh function
 u = H.subdomain_to_post_processing_solution[bottom_domain]
+if legacy_entity_map:
+    entity_maps = {
+        sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains
+    }
+else:
+    entity_maps = [sd.cell_map for sd in my_model.volume_subdomains]
+
 form = dolfinx.fem.form(
     ufl.dot(
         D * ufl.grad(u),
         n_b,
     )
     * ds(bottom_surface.id),
-    entity_maps={sd.submesh: sd.parent_to_submesh for sd in my_model.volume_subdomains},
+    entity_maps=entity_maps,
 )
 print(dolfinx.fem.assemble_scalar(form))

src/festim/helpers_discontinuity.py

@@ -12,11 +12,19 @@ def transfer_meshtags_to_submesh(
     # 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
-    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
-    )
+    try:
+        from dolfinx.mesh import EntityMap  # noqa: F401
+
+        legacy_entity_map = False
+    except ImportError:
+        legacy_entity_map = True
+
+    if legacy_entity_map:
+        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
+        )
 
     # Compute and access various entity to vertex and vertex to
     # entity connectivity for parent and sub-mesh.
@@ -38,22 +46,51 @@ def transfer_meshtags_to_submesh(
     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
+
+    if legacy_entity_map:
+        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
+    else:
+        for facet, value in zip(entity_tag.indices, entity_tag.values):
+            facet_found = False
+            parent_cells = p_f_to_c.links(facet)
+            sub_cells = sub_cell_to_parent.sub_topology_to_topology(
+                parent_cells, inverse=True
+            )
+
+            for child_cell in sub_cells:
+                if facet_found:
+                    break
+                if child_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.sub_topology_to_topology(
+                                child_vertices, inverse=False
+                            )
+                        )
+                        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,

src/festim/hydrogen_transport_problem.py

@@ -1392,10 +1392,14 @@ def mixed_term(u, v, n):
 
         n = ufl.FacetNormal(mesh)
         cr = ufl.Circumradius(mesh)
-
-        entity_maps = {
-            sd.submesh: sd.parent_to_submesh for sd in self.volume_subdomains
-        }
+        try:
+            from dolfinx.mesh import EntityMap
+
+            entity_maps = [sd.cell_map for sd in self.volume_subdomains]
+        except ImportError:
+            entity_maps = {
+                sd.submesh: sd.parent_to_submesh for sd in self.volume_subdomains
+            }
         for interface in self.interfaces:
             gamma = interface.penalty_term
 
@@ -1532,8 +1536,16 @@ def create_solver(self):
         )
         self.solver.max_iterations = self.settings.max_iterations
         self.solver.convergence_criterion = self.settings.convergence_criterion
-        self.solver.atol = self.settings.atol
-        self.solver.rtol = self.settings.rtol
+        self.solver.atol = (
+            self.settings.atol
+            if not callable(self.settings.rtol)
+            else self.settings.rtol(float(self.t))
+        )
+        self.solver.rtol = (
+            self.settings.rtol
+            if not callable(self.settings.rtol)
+            else self.settings.rtol(float(self.t))
+        )
 
     def create_flux_values_fenics(self):
         """For each particle flux create the ``value_fenics`` attribute"""

src/festim/subdomain/interface.py

@@ -34,6 +34,13 @@ def __init__(
 
     def pad_parent_maps(self):
         """Workaround to make sparsity-pattern work without skips"""
+        try:
+            # No padding needed for latest version of DOLFINx
+            from dolfinx.mesh import EntityMap  # noqa: F401
+
+            return
+        except ImportError:
+            pass
 
         if Version(dolfinx.__version__) == Version("0.9.0"):
             args = (
@@ -68,7 +75,7 @@ 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)
 
         if Version(dolfinx.__version__) == Version("0.9.0"):
@@ -89,8 +96,23 @@ def compute_mapped_interior_facet_data(self, mesh):
 
         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]]
+        try:
+            # No padding needed for latest version of DOLFINx
+            from dolfinx.mesh import EntityMap  # noqa: F401
+
+            mapped_cell_0 = self.subdomains[0].cell_map.sub_topology_to_topology(
+                integration_data[0::4], inverse=True
+            )
+            mapped_cell_1 = self.subdomains[0].cell_map.sub_topology_to_topology(
+                integration_data[2::4], inverse=True
+            )
+            legacy_entity_map = False
+
+        except ImportError:
+            assert (not self.subdomains[0].padded) and (not self.subdomains[1].padded)
+            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]]
+            legacy_entity_map = True
 
         switch = mapped_cell_1 > mapped_cell_0
         # Order restriction on one side
@@ -100,9 +122,14 @@ def compute_mapped_interior_facet_data(self, mesh):
             ]
 
         # Check that other restriction lies in other interface
-        domain1_cell = self.subdomains[1].parent_to_submesh[
-            ordered_integration_data[:, 2]
-        ]
+        if legacy_entity_map:
+            domain1_cell = self.subdomains[1].parent_to_submesh[
+                ordered_integration_data[:, 2]
+            ]
+        else:
+            domain1_cell = self.subdomains[1].cell_map.sub_topology_to_topology(
+                ordered_integration_data[:, 2], inverse=True
+            )
         assert (domain1_cell >= 0).all()
 
         return (self.id, ordered_integration_data.reshape(-1))