From 8e649e3cb0ca45909acb91cd164f5a18e4dbf4e1 Mon Sep 17 00:00:00 2001
From: J Dark <james.dark.95@gmail.com>
Date: Tue, 10 Oct 2023 15:27:56 -0400
Subject: [PATCH] create meshtags and meaures from subdomains

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

diff --git a/festim/hydrogen_transport_problem.py b/festim/hydrogen_transport_problem.py
index 20a037fec..2ecf91e33 100644
--- a/festim/hydrogen_transport_problem.py
+++ b/festim/hydrogen_transport_problem.py
@@ -3,12 +3,9 @@
 import ufl
 from mpi4py import MPI
 from dolfinx.fem import Function
-from ufl import (
-    TestFunction,
-    dot,
-    grad,
-    exp,
-)
+from dolfinx.mesh import meshtags, locate_entities
+from ufl import TestFunction, dot, grad, exp, Measure
+import numpy as np
 
 import festim as F
 
@@ -98,12 +95,7 @@ def initialise(self):
         """
 
         self.define_function_space()
-        (
-            self.facet_tags,
-            self.volume_tags,
-            self.dx,
-            self.ds,
-        ) = self.mesh.create_measures_and_tags(self.function_space)
+        self.define_markers_and_measures()
         self.assign_functions_to_species()
         self.create_formulation()
 
@@ -111,6 +103,47 @@ def define_function_space(self):
         elements = ufl.FiniteElement("CG", self.mesh.mesh.ufl_cell(), 1)
         self.function_space = fem.FunctionSpace(self.mesh.mesh, elements)
 
+    def define_markers_and_measures(self):
+        dofs_facets = []
+        tags_facets = []
+
+        # find all cells in domain and mark them as 1
+        num_cells = self.mesh.mesh.topology.index_map(vdim).size_local
+        mesh_cell_indicies = np.arange(num_cells, dtype=np.int32)
+        tags_volumes = np.full(num_cells, 0, dtype=np.int32)
+
+        # TODO this should be a property of mesh
+        fdim = self.mesh.mesh.topology.dim - 1
+        vdim = self.mesh.mesh.topology.dim
+
+        for sub_dom in self.subdomains:
+            if isinstance(sub_dom, F.SurfaceSubdomain1D):
+                dof = sub_dom.locate_dof(self.function_space)
+                dofs_facets.append(dof)
+                tags_facets.append(sub_dom.id)
+            if isinstance(sub_dom, F.VolumeSubdomain):
+                # find all cells in subdomain and mark them as sub_dom.id
+                entities = sub_dom.locate_subdomain_entities(self.mesh.mesh, vdim)
+                tags_volumes[entities] = sub_dom.id
+
+        # dofs and tags need to be in np.in32 format for meshtags
+        dofs_facets = np.array(dofs_facets, dtype=np.int32)
+        tags_facets = np.array(tags_facets, dtype=np.int32)
+
+        # define mesh tags
+        self.mesh_tags_facets = meshtags(self.mesh.mesh, fdim, dofs_facets, tags_facets)
+        self.mesh_tags_volumes = meshtags(
+            self.mesh.mesh, vdim, mesh_cell_indicies, tags_volumes
+        )
+
+        # define measures
+        self.ds = Measure(
+            "ds", domain=self.mesh.mesh, subdomain_data=self.mesh_tags_facets
+        )
+        self.dx = Measure(
+            "dx", domain=self.mesh.mesh, subdomain_data=self.mesh_tags_volumes
+        )
+
     def assign_functions_to_species(self):
         """Creates for each species the solution, prev solution and test function"""
         if len(self.species) > 1: