Merge pull request #393 from simpeg/repr-non-finalized-mesh

Allow `TreeMesh.__repr__` to run when non finalized

Author: jcapriot

discretize/tree_mesh.py

@@ -291,6 +291,11 @@ def __repr__(self):
                 "{}: {:^13},{:^13}".format(dim_label[dim], n_vector[0], n_vector[-1])
             )
 
+        # Return partial information if mesh is not finalized
+        if not self.finalized:
+            top = f"\n {mesh_name} (non finalized)\n\n"
+            return top + "\n".join(extent_display)
+
         for i, line in enumerate(extent_display):
             if i == len(cell_display):
                 cell_display.append(" " * (len(cell_display[0]) - 3 - len(line)))
@@ -315,14 +320,47 @@ def __repr__(self):
     def _repr_html_(self):
         """HTML representation."""
         mesh_name = "{0!s}TreeMesh".format(("Oc" if self.dim == 3 else "Quad"))
+        style = " style='padding: 5px 20px 5px 20px;'"
+        dim_label = {0: "x", 1: "y", 2: "z"}
+
+        if not self.finalized:
+            style_bold = '"font-weight: bold; font-size: 1.2em; text-align: center;"'
+            style_regular = '"font-size: 1.2em; text-align: center;"'
+            output = [
+                "<table>",  # need to close this tag
+                "<tr>",
+                f"<td style={style_bold}>{mesh_name}</td>",
+                f"<td style={style_regular} colspan='2'>(non finalized)</td>",
+                "</tr>",
+                "<table>",  # need to close this tag
+                "<tr>",
+                "<th></th>",
+                f'<th {style} colspan="2">Mesh extent</th>',
+                "</tr>",
+                "<tr>",
+                "<th></th>",
+                f"<th {style}>min</th>",
+                f"<th {style}>max</th>",
+                "</tr>",
+            ]
+            for dim in range(self.dim):
+                n_vector = getattr(self, "nodes_" + dim_label[dim])
+                output += [
+                    "<tr>",
+                    f"<td {style}>{dim_label[dim]}</td>",
+                    f"<td {style}>{n_vector[0]}</td>",
+                    f"<td {style}>{n_vector[-1]}</td>",
+                    "</tr>",
+                ]
+            output += ["</table>", "</table>"]
+            return "\n".join(output)
+
         level_count = self._count_cells_per_index()
         non_zero_levels = np.nonzero(level_count)[0]
-        dim_label = {0: "x", 1: "y", 2: "z"}
         h_gridded = self.h_gridded
         mins = np.min(h_gridded, axis=0)
         maxs = np.max(h_gridded, axis=0)
 
-        style = " style='padding: 5px 20px 5px 20px;'"
         # Cell level table:
         cel_tbl = "<table>\n"
         cel_tbl += "<tr>\n"

tests/tree/test_tree.py

@@ -590,5 +590,64 @@ def test_cell_locator(dim):
     assert found
 
 
+class TestRepr:
+    """
+    Test repr methods on TreeMesh.
+
+    Check if no error is raised when calling repr methods on a finalized and
+    non finalized meshes.
+    """
+
+    @pytest.fixture(params=["2D", "3D"])
+    def mesh(self, request):
+        """Return a sample TreeMesh"""
+        nc = 16
+        if request.param == "2D":
+            h = [nc, nc]
+            origin = (-32.4, 245.4)
+        else:
+            h = [nc, nc, nc]
+            origin = (-32.4, 245.4, 192.3)
+        mesh = discretize.TreeMesh(h, origin, diagonal_balance=True)
+        return mesh
+
+    def finalize(self, mesh):
+        """
+        Finalize the sample tree mesh.
+        """
+        origin = mesh.origin
+        if mesh.dim == 2:
+            p1 = (origin[0] + 0.4, origin[1] + 0.4)
+            p2 = (origin[0] + 0.6, origin[1] + 0.6)
+            mesh.refine_box(p1, p2, levels=5)
+        else:
+            p1 = (origin[0] + 0.4, origin[1] + 0.4, origin[2] + 0.7)
+            p2 = (origin[0] + 0.6, origin[1] + 0.6, origin[2] + 0.9)
+            mesh.refine_box(p1, p2, levels=5)
+        mesh.finalize()
+
+    @pytest.mark.parametrize("finalize", [True, False])
+    def test_repr(self, mesh, finalize):
+        """
+        Test if __repr__ doesn't raise errors on any TreeMesh.
+        """
+        if finalize:
+            self.finalize(mesh)
+        output = mesh.__repr__()
+        assert type(output) is str
+        assert len(output) != 0
+
+    @pytest.mark.parametrize("finalize", [True, False])
+    def test_repr_html(self, mesh, finalize):
+        """
+        Test if _repr_html_ doesn't raise errors on any TreeMesh.
+        """
+        if finalize:
+            self.finalize(mesh)
+        output = mesh._repr_html_()
+        assert type(output) is str
+        assert len(output) != 0
+
+
 if __name__ == "__main__":
     unittest.main()