Some refactoring

.gitignore

@@ -393,3 +393,5 @@ venv.bak/
 .mypy_cache/
 
 results/
+
+.vscode

build_grid.py

@@ -12,50 +12,64 @@
 THRESHOLD = 0.85
 
 
-def build_grid(size, location_probability = THRESHOLD, output_format=(int, int), seed=None):
-    """ Build a square grid of elements, where each element may or may not contain oil
-
-     size: The number of elements along one edge of the grid
-     seed: Random seed to be used to generate the grid
-     """
-    random.seed(seed)
-
-    if location_probability is None:
-        location_probability = THRESHOLD
-
-    initial_grid = set()
-    for location in itertools.product(range(0, size), repeat=2):
-        if random.random() > location_probability:
-            initial_grid.add(location)
-
-    grid = set()
-    # Cluster the grid. If an active element is not isolated,
-    # or if an inactive element has at least 4 active neighbors
-    for location in itertools.product(range(0, size), repeat=2):
-        state = location in initial_grid
-        sites_nearby = get_neighbors(location, initial_grid)
-        neighbor_count = len(sites_nearby)
-        if (state and neighbor_count != 0) or neighbor_count >= 4:
-            grid.add(location)
-
-    if output_format == (int, int):
-        return grid
-
-    if output_format == ET._Element: # pylint: disable=W0143
-        root = ET.Element('grid')
-        grid = list(grid)
-        grid.sort()
-        for location in grid:
-            y_element = ET.Element('y')
-            x_element = ET.Element('x')
-            x_element.text, y_element.text = (map(str, location))
-            location_element = ET.Element('el')
+class Grid:
+    def __init__(self, size, location_probability=THRESHOLD, seed=None):
+        """Build a square grid of elements, where each element may or may not contain oil
+
+        size: The number of elements along one edge of the grid
+        seed: Random seed to be used to generate the grid
+        """
+        self.size = size
+        self.location_probability = location_probability
+        self.seed = seed
+        random.seed(seed)
+
+        if location_probability is None:
+            location_probability = THRESHOLD
+
+        initial_grid = set()
+        for location in itertools.product(range(0, size), repeat=2):
+            if random.random() > location_probability:
+                initial_grid.add(location)
+
+        self._grid = set()
+        # Cluster the grid. If an active element is not isolated,
+        # or if an inactive element has at least 4 active neighbors
+        for location in itertools.product(range(0, size), repeat=2):
+            state = location in initial_grid
+            sites_nearby = get_neighbors(location, initial_grid)
+            neighbor_count = len(sites_nearby)
+            if (state and neighbor_count != 0) or neighbor_count >= 4:
+                self._grid.add(location)
+
+    @property
+    def number_of_sites(self) -> int:
+        return len(self._grid)
+
+    @property
+    def density(self) -> int:
+        return self.number_of_sites / (self.size * self.size)
+
+    @property
+    def to_set_of_tuples(self) -> set[tuple[int, int]]:
+        return self._grid
+
+    @property
+    def to_xml(self) -> ET.Element:
+        root = ET.Element("grid")
+        sortable_grid = list(self._grid)
+        sortable_grid.sort()
+        for location in sortable_grid:
+            y_element = ET.Element("y")
+            x_element = ET.Element("x")
+            x_element.text, y_element.text = map(str, location)
+            location_element = ET.Element("el")
             location_element.extend([x_element, y_element])
             root.append(location_element)
         return root
 
-    raise NotImplementedError
 
 if __name__ == "__main__":
+    grid = Grid(25).to_xml
     with open("../surveying_xslt/input.xml", "w") as f:
-        f.write((ET.tostring(build_grid(25, ET._Element), pretty_print=True).decode()))
+        f.write((ET.tostring(grid, pretty_print=True).decode()))

display.py

@@ -2,40 +2,49 @@
 
 import colorama
 
-TOP_LEFT_CHAR = u'\u250c'
-TOP_RIGHT_CHAR = u'\u2510'
-TOP_CHAR = u'\u252c'
-BOTTOM_LEFT_CHAR = u'\u2514'
-BOTTOM_RIGHT_CHAR = u'\u2518'
-BOTTOM_CHAR = u'\u2534'
-LEFT_CHAR = u'\u251c'
-RIGHT_CHAR = u'\u2524'
-VERTICAL_CHAR = u'\u2502'
-HORIZ_CHAR = u'\u2500'
-VERTEX_CHAR = u'\u253c'
+TOP_LEFT_CHAR = u"\u250c"
+TOP_RIGHT_CHAR = u"\u2510"
+TOP_CHAR = u"\u252c"
+BOTTOM_LEFT_CHAR = u"\u2514"
+BOTTOM_RIGHT_CHAR = u"\u2518"
+BOTTOM_CHAR = u"\u2534"
+LEFT_CHAR = u"\u251c"
+RIGHT_CHAR = u"\u2524"
+VERTICAL_CHAR = u"\u2502"
+HORIZ_CHAR = u"\u2500"
+VERTEX_CHAR = u"\u253c"
 
-def display_set_as_grid(size, locations):
-    """ Render the grid. Relies on a fixed-width font. """
 
-    if size > 100:
+def display_set_as_grid(grid):
+    """Render the grid. Relies on a fixed-width font."""
+
+    if grid.size > 100:
         print("Too big to render!")
         return
 
     # Use the size of the grid to figure out how much to pad row and column indices by
-    num_chars_in_label = len(str(size))
+    num_chars_in_label = len(str(grid.size))
 
     # Generate column headers
-    print(TOP_LEFT_CHAR +
-          (HORIZ_CHAR* num_chars_in_label + TOP_CHAR)*size +
-          HORIZ_CHAR* num_chars_in_label +
-          TOP_RIGHT_CHAR)
-    column_heading_numbers = [str(y).zfill(num_chars_in_label)
-                              for y in list(range(0, size))]
+    print(
+        TOP_LEFT_CHAR
+        + (HORIZ_CHAR * num_chars_in_label + TOP_CHAR) * grid.size
+        + HORIZ_CHAR * num_chars_in_label
+        + TOP_RIGHT_CHAR
+    )
+    column_heading_numbers = [
+        str(y).zfill(num_chars_in_label) for y in list(range(0, grid.size))
+    ]
     header = VERTICAL_CHAR.join(column_heading_numbers) + VERTICAL_CHAR
     print(VERTICAL_CHAR + " " * num_chars_in_label + VERTICAL_CHAR + header)
 
     cell_separator = VERTEX_CHAR + HORIZ_CHAR * (num_chars_in_label)
-    row_divider = LEFT_CHAR + HORIZ_CHAR * (num_chars_in_label) + cell_separator * size + RIGHT_CHAR
+    row_divider = (
+        LEFT_CHAR
+        + HORIZ_CHAR * (num_chars_in_label)
+        + cell_separator * grid.size
+        + RIGHT_CHAR
+    )
 
     print(row_divider)
 
@@ -44,21 +53,31 @@ def display_set_as_grid(size, locations):
     current_row = 0
 
     # TODO: Make this not rely on a particular creation order
-    for y_coord in list(range(0, size)):
-        for x_coord in list(range(0, size)):
+    for y_coord in list(range(0, grid.size)):
+        for x_coord in list(range(0, grid.size)):
             # TODO: Somehow highlight different wells with different colors
-            marker = colorama.Fore.RED + "x" + colorama.Style.RESET_ALL \
-                if (x_coord, y_coord) in locations else " "
+            marker = (
+                colorama.Fore.RED + "x" + colorama.Style.RESET_ALL
+                if (x_coord, y_coord) in grid.to_set_of_tuples
+                else " "
+            )
             if y_coord == current_row:
                 row = row + marker * num_chars_in_label + VERTICAL_CHAR
             else:
                 current_row = y_coord
                 print(row)
                 print(row_divider)
-                row = VERTICAL_CHAR + str(y_coord).zfill(num_chars_in_label) \
-                      + VERTICAL_CHAR + marker * num_chars_in_label + VERTICAL_CHAR
+                row = (
+                    VERTICAL_CHAR
+                    + str(y_coord).zfill(num_chars_in_label)
+                    + VERTICAL_CHAR
+                    + marker * num_chars_in_label
+                    + VERTICAL_CHAR
+                )
     print(row)
-    print(BOTTOM_LEFT_CHAR +
-          (HORIZ_CHAR* num_chars_in_label + BOTTOM_CHAR)*size +
-          HORIZ_CHAR* num_chars_in_label +
-          BOTTOM_RIGHT_CHAR)
+    print(
+        BOTTOM_LEFT_CHAR
+        + (HORIZ_CHAR * num_chars_in_label + BOTTOM_CHAR) * grid.size
+        + HORIZ_CHAR * num_chars_in_label
+        + BOTTOM_RIGHT_CHAR
+    )

methods/graphtool_method.py

@@ -10,8 +10,8 @@
 METHOD_NAME = "Graph-Tool Method"
 
 
-def find_reservoirs(locations):
-    """ Uses a graph approach to find how many wells are needed, making the assumption that
+def find_reservoirs(locations) -> list[set[tuple[int, int]]]:
+    """Uses a graph approach to find how many wells are needed, making the assumption that
     only one well is needed per contiguous field
 
     locations: Set containing all locations with oil
@@ -28,8 +28,9 @@ def find_reservoirs(locations):
     edge_list = []
     for location in locations:
         neighbor_coords = get_neighbors(location, locations)
-        edge_list.extend([(locations[location], locations[neighbor])
-                          for neighbor in neighbor_coords])
+        edge_list.extend(
+            [(locations[location], locations[neighbor]) for neighbor in neighbor_coords]
+        )
 
     locations_graph.add_edge_list(edge_list)
 

methods/igraph_method.py

@@ -9,27 +9,38 @@
 METHOD_NAME = "IGraph Method"
 
 
-def find_reservoirs(locations):
-    """ Uses a graph approach to find how many wells are needed, making the assumption that
+def find_reservoirs(locations) -> list[set[tuple[int, int]]]:
+    """Uses a graph approach to find how many wells are needed, making the assumption that
     only one well is needed per contiguous field
 
     locations: Set containing all locations with oil
     """
 
     locations_igraph = igraph.Graph()
-    locations_str = [','.join(map(str, location)) for location in locations]
+    locations_str = [",".join(map(str, location)) for location in locations]
     locations_igraph.add_vertices(locations_str)
 
     edge_list = set()
     for coords in locations:
         neighbor_coords = get_neighbors(coords, locations)
-        edge_list.update({(','.join(map(str, coords)), ','.join(map(str, neighbor_coord)))
-                          for neighbor_coord in neighbor_coords})
+        edge_list.update(
+            {
+                (",".join(map(str, coords)), ",".join(map(str, neighbor_coord)))
+                for neighbor_coord in neighbor_coords
+            }
+        )
 
     locations_igraph.add_edges(edge_list)
 
-    clusters = locations_igraph.clusters(mode='STRONG')
-    wells = [{locations_igraph.vs[vertex]['name'] for vertex in cluster} # pylint: disable=E1136
-             for cluster in clusters]
-
-    return [{tuple(int(value) for value in site.split(',')) for site in igraph_result} for igraph_result in wells]
+    clusters = locations_igraph.clusters(mode="STRONG")
+    wells = [
+        {
+            locations_igraph.vs[vertex]["name"] for vertex in cluster
+        }  # pylint: disable=E1136
+        for cluster in clusters
+    ]
+
+    return [
+        {tuple(int(value) for value in site.split(",")) for site in igraph_result}
+        for igraph_result in wells
+    ]

methods/networkx_method.py

@@ -10,8 +10,8 @@
 METHOD_NAME = "NetworkX Method"
 
 
-def find_reservoirs(grid):
-    """ Uses a graph approach to find how many wells are needed, making the assumption that
+def find_reservoirs(grid) -> list[set[tuple[int, int]]]:
+    """Uses a graph approach to find how many wells are needed, making the assumption that
     only one well is needed per contiguous field
 
     locations: Set containing all locations with oil

methods/recursive_method.py

@@ -6,20 +6,17 @@
 from tools import get_neighbors
 
 METHOD_NAME = "Recursive Method"
-REQUIRED_GRID_TYPE = (int, int)
 
 
-def find_reservoirs(this_grid, reservoir=None, original_grid=None):
-    """ Recursively determines how many wells are needed, making the assumption that
+def find_reservoirs(this_grid, reservoir=None, original_grid=None) -> list[set[tuple[int, int]]]:
+    """Recursively determines how many wells are needed, making the assumption that
     only one well is needed per contiguous field
 
     this_grid: This is the list of locations to be checked for the current reservoir
     reservoir: If being called recursively, this contains the current reservoir that is being built
     original_grid: If being called recursively, this is the full grid to find neighbor elements
     """
 
-
-
     # well is None iff this is the 'outer' call of this function
     if reservoir is None:
         # Initialize variables

methods/stack_method.py

@@ -7,8 +7,8 @@
 METHOD_NAME = "Stack Method (List)"
 
 
-def find_reservoirs(this_grid):
-    """ Recursively determines how many wells are needed, making the assumption that
+def find_reservoirs(this_grid) -> list[set[tuple[int, int]]]:
+    """Recursively determines how many wells are needed, making the assumption that
     only one well is needed per contiguous field
 
     this_grid: This is the list of locations to be checked for the current reservoir

methods/stack_method_deque.py

@@ -9,8 +9,8 @@
 METHOD_NAME = "Stack Method (deque)"
 
 
-def find_reservoirs(this_grid):
-    """ Recursively determines how many wells are needed, making the assumption that
+def find_reservoirs(this_grid) -> list[set[tuple[int, int]]]:
+    """Recursively determines how many wells are needed, making the assumption that
     only one well is needed per contiguous field
 
     this_grid: This is the list of locations to be checked for the current reservoir