simplify main loop to build connectivity etc..

gprMax/grid.py

@@ -29,6 +29,9 @@ class Grid(object):
         self.nx = grid.shape[0]
         self.ny = grid.shape[1]
         self.nz = grid.shape[2]
+        self.i_max = self.nx - 1
+        self.j_max = self.ny - 1
+        self.k_max = self.nz - 1
         self.grid = grid
 
     def n_edges(self):

gprMax/xdmf.py

@@ -16,7 +16,7 @@ class ListCounter():
         self.array[self.count] = item
         self.count += 1
 
-class Edges:
+class EdgeLabels:
 
     def __init__(self, grid):
 
@@ -37,6 +37,20 @@ class Edges:
         edge = np.array([in_label, out_label])
         self.edge_counter.add(edge)
 
+class EdgeMaterials:
+
+    def __init__(self, fdtd_grid):
+        self.fdtd_grid = fdtd_grid
+        self.n_edges = fdtd_grid.n_edges()
+        self.materials = np.zeros((self.n_edges), np.int8)
+        self.materialCounter = ListCounter(self.materials)
+
+    # direction x->0 y->1 z->2
+    def add_material(self, i, j, k, direction):
+
+        material = self.fdtd_grid.ID[direction, i, j, k]
+        self.materialCounter.add(material)
+
 
 class Coordinates:
 
@@ -110,66 +124,48 @@ class SolidLabels():
         solid_labels = self.hexCellPicker(self.label_grid.grid, i, j, k)
         self.label_counter(solid_labels)
 
+class SolidManager(Grid):
 
-class Materials:
+    def __init__(self, label_grid, fdtd_grid):
 
-    def __init__(self, fdtd_grid):
-        self.fdtd_grid = fdtd_grid
-        self.n_edges = fdtd_grid.n_edges()
-        self.materials = np.zeros((self.n_edges), np.int8)
-        self.materialCounter = ListCounter(self.materials)
+        super().__init__(label_grid)
+        self.solids = Solids(fdtd_grid)
+        self.solid_labels = SolidLabels(label_grid)
 
-    # direction x->0 y->1 z->2
-    def add_material(self, i, j, k, direction):
+    def createSolid(self, i, j, k):
+        if i < self.i_max and j < self.j_max and k < self.k_max:
+            self.solids.add_solid(i, j, k)
+            self.solid_labels.add(i, j, k)
 
-        material = self.fdtd_grid.ID[direction, i, j, k]
-        self.materialCounter.add(material)
 
-def process_grid(fdtd_grid, res):
+class EdgeManager(Grid):
+    """
+    Class to manage the creation of edges and matching edge materials.
+    """
 
-    # Dimensions of the problem domain.
-    nx = fdtd_grid.nx
-    ny = fdtd_grid.ny
-    nz = fdtd_grid.nz
+    def __init__(self, label_grid, fdtd_grid):
+        super().__init__(label_grid)
+        self.edges = EdgeLabels(label_grid)
+        self.edge_materials = EdgeMaterials(fdtd_grid)
 
-    # useful indices
-    i_max = nx - 1
-    j_max = ny - 1
-    k_max = nz - 1
+    def createEdges(self, i, j, k):
+        """
+            Create the relevant edges and corresponding edge materials.
+            Args:
+                i (int): i index of label in labels_grid
+                j (int): j index of label in labels_grid
+                k (int): k index of label in labels_grid
 
-    # label each node in the space
-    labels = np.arange(nx * ny * nz).reshape(nx, ny, nz)
+        """
+        edges = self.edges
+        edge_materials = self.edge_materials
+        i_max = self.i_max
+        j_max = self.j_max
+        k_max = self.k_max
+        label = self.edges.grid.get(i, j, k)
 
-    label_grid = Grid(labels)
-
-    # Define coordinates for each node
-    coordinates = Coordinates(fdtd_grid)
-
-    # Material for each solid
-    solids = Solids(fdtd_grid)
-
-    # Connectivity for hexhahedron grid
-    solid_labels = SolidLabels(label_grid)
-
-    if res == 'f':
-        # Edges define the connectivity of the grid.
-        edges = Edges(label_grid)
-
-        # Material for each edge
-        edge_materials = Materials(fdtd_grid)
-
-    for i, ix in enumerate(labels):
-        for j, jx in enumerate(ix):
-            for k, kx in enumerate(jx):
-
-                label = labels[i][j][k]
-
-                if i < i_max and j < j_max and k < k_max:
-                    solids.add_solid(i, j, k)
-                    solid_labels.add(i, j, k)
-
-                if res == 'f':
         # Each vertex can have varying numbers of edges
+
         # Type 1 vertex
         if i < i_max and j < j_max and k < k_max:
             edges.add_edge(label, i + 1, j, k)
@@ -225,18 +221,44 @@ def process_grid(fdtd_grid, res):
         else:
             print('oh no')
 
+
+def process_grid(fdtd_grid, res):
+
+    # Create a grid of labels with equal dimension to fdtd grid
+    labels = np.arange(fdtd_grid.n_nodes()).reshape(fdtd_grid.nx, fdtd_grid.ny, fdtd_grid.nz)
+
+    label_grid = Grid(labels)
+
+    # Define coordinates for each node
+    coordinates = Coordinates(fdtd_grid)
+
+    solid_manager = SolidManager(label_grid, fdtd_grid)
+
+    if res == 'f':
+        edge_manager = EdgeManager()
+
+    # Iterate through the label and create relevant edges and solids.
+    for i, ix in enumerate(labels):
+        for j, jx in enumerate(ix):
+            for k, kx in enumerate(jx):
+
+                if res == 'f':
+                    edge_manager.createEdges(i, j, k)
+
+                solid_manager.createSolid(i, j, k)
+
                 # Add the coordinates
                 coordinates.add_coordinate(i, j, k)
 
     data = {
         'coordinates': coordinates,
-        'solids': solids,
-        'solid_labels': solid_labels,
+        'solids': solid_manager.solids,
+        'solid_labels': solid_manager.solid_labels,
     }
 
     if res == 'f':
-        data['edges'] = edges
-        data['edge_materials'] = edge_materials
+        data['edges'] = edge_manager.edges
+        data['edge_materials'] = edge_manager.edge_materials
 
     return data