From fee12b58ebf404d72bb34d67c4d02992115bd57f Mon Sep 17 00:00:00 2001
From: Santiago Soler <santisoler@fastmail.com>
Date: Wed, 20 Nov 2024 11:42:11 -0800
Subject: [PATCH] Set no diagonal balance when reading UBC tree meshes

Set `diagonal_balance=False` when creating a TreeMesh after reading
a UBC file in `read_UBC`. This prevents raising the future warning
related to the default value of `diagonal_balance`. Add test that checks
that no warning is being raised by that method.
---
 discretize/mixins/mesh_io.py |  2 +-
 tests/tree/test_tree_io.py   | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/discretize/mixins/mesh_io.py b/discretize/mixins/mesh_io.py
index 947eec959..a9218db0b 100644
--- a/discretize/mixins/mesh_io.py
+++ b/discretize/mixins/mesh_io.py
@@ -492,7 +492,7 @@ def read_UBC(TreeMesh, file_name, directory=None):
         else:
             max_level = min(ls) + 1
 
-        mesh = TreeMesh(hs, origin=origin)
+        mesh = TreeMesh(hs, origin=origin, diagonal_balance=False)
         levels = indArr[:, -1]
         indArr = indArr[:, :-1]
 
diff --git a/tests/tree/test_tree_io.py b/tests/tree/test_tree_io.py
index e149edcb6..d19f420c7 100644
--- a/tests/tree/test_tree_io.py
+++ b/tests/tree/test_tree_io.py
@@ -1,3 +1,4 @@
+import warnings
 import numpy as np
 import discretize
 import pickle
@@ -49,6 +50,19 @@ def test_UBCfiles(mesh, tmp_path):
     np.testing.assert_array_equal(vec, vecUBC2)
 
 
+def test_ubc_files_no_warning_diagonal_balance(mesh, tmp_path):
+    """
+    Test that reading UBC files don't trigger the diagonal balance warning.
+    """
+    # Save the sample mesh into a UBC file
+    fname = tmp_path / "temp.msh"
+    mesh.write_UBC(fname)
+    # Make sure that no warning is raised when reading the mesh
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        discretize.TreeMesh.read_UBC(fname)
+
+
 if has_vtk:
 
     def test_write_VTU_files(mesh, tmp_path):