diff --git a/discretize/tree_mesh.py b/discretize/tree_mesh.py
index cfa817a90..a9518862b 100644
--- a/discretize/tree_mesh.py
+++ b/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 = [
+ "
", # need to close this tag
+ "",
+ f"{mesh_name} | ",
+ f"(non finalized) | ",
+ "
",
+ "", # need to close this tag
+ "",
+ " | ",
+ f'Mesh extent | ',
+ "
",
+ "",
+ " | ",
+ f"min | ",
+ f"max | ",
+ "
",
+ ]
+ for dim in range(self.dim):
+ n_vector = getattr(self, "nodes_" + dim_label[dim])
+ output += [
+ "",
+ f"{dim_label[dim]} | ",
+ f"{n_vector[0]} | ",
+ f"{n_vector[-1]} | ",
+ "
",
+ ]
+ output += ["
", "
"]
+ 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 = "\n"
cel_tbl += "\n"
diff --git a/tests/tree/test_tree.py b/tests/tree/test_tree.py
index 6fd92936f..2a27fb19d 100644
--- a/tests/tree/test_tree.py
+++ b/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()