diff --git a/.github/workflows/pytest-unix-os.yaml b/.github/workflows/pytest-unix-os.yaml
index b4a45dbbf..ad800e30a 100644
--- a/.github/workflows/pytest-unix-os.yaml
+++ b/.github/workflows/pytest-unix-os.yaml
@@ -40,11 +40,11 @@ jobs:
       CONDA_LOCK_ENV_FILE: environments/conda-py-${{ matrix.python_ver }}-${{ startsWith(matrix.os, 'macos') && 'osx' || 'linux' }}-64-dev.lock.yml
       PIP_EXTRA_INDEX_URL: https://test.pypi.org/simple/
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           lfs: true
       - name: Setup conda env
-        uses: mamba-org/provision-with-micromamba@main
+        uses: mamba-org/setup-micromamba@v1
         with:
           environment-file: ${{ env.CONDA_LOCK_ENV_FILE }}
           environment-name: test_env
diff --git a/.github/workflows/pytest-windows.yaml b/.github/workflows/pytest-windows.yaml
index 86853dc17..9d8dff5b5 100644
--- a/.github/workflows/pytest-windows.yaml
+++ b/.github/workflows/pytest-windows.yaml
@@ -39,14 +39,15 @@ jobs:
       CONDA_LOCK_ENV_FILE: environments/conda-py-${{ matrix.python_ver }}-win-64-dev.lock.yml
       PIP_EXTRA_INDEX_URL: https://test.pypi.org/simple/
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           lfs: true
       - name: Setup conda env
-        uses: mamba-org/provision-with-micromamba@main
+        uses: mamba-org/setup-micromamba@v1
         with:
           environment-file: ${{ env.CONDA_LOCK_ENV_FILE }}
           environment-name: test_env
+          init-shell: powershell
           cache-downloads: true
       - name: pytest
         run: |
diff --git a/.github/workflows/static-analysis.yml b/.github/workflows/static-analysis.yml
index 8ef8a2124..99edd0f5d 100644
--- a/.github/workflows/static-analysis.yml
+++ b/.github/workflows/static-analysis.yml
@@ -26,18 +26,18 @@ env:
 jobs:
   pylint:
     name: pylint
-    runs-on: windows-latest
+    runs-on: ubuntu-latest
     defaults:
       run:
         shell: bash -l {0}
     env:
       PYTHONUTF8: 1
-      CONDA_LOCK_ENV_FILE: environments/conda-py-3.9-win-64-dev.lock.yml
+      CONDA_LOCK_ENV_FILE: environments/conda-py-3.9-linux-64-dev.lock.yml
       PIP_EXTRA_INDEX_URL: https://test.pypi.org/simple/
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Setup conda env
-        uses: mamba-org/provision-with-micromamba@main
+        uses: mamba-org/setup-micromamba@v1
         with:
           environment-file: ${{ env.CONDA_LOCK_ENV_FILE }}
           environment-name: linter_env
diff --git a/geoapps/inversion/components/models.py b/geoapps/inversion/components/models.py
index 64a2c6b38..ef6f4e74a 100644
--- a/geoapps/inversion/components/models.py
+++ b/geoapps/inversion/components/models.py
@@ -114,42 +114,39 @@ def active_cells(self, active_cells):
         self._active_cells = active_cells
 
     @property
-    def starting(self):
-        mstart = self._starting.model
+    def starting(self) -> np.ndarray | None:
+        if self._starting.model is None:
+            return None
+
+        mstart = self._starting.model.copy()
 
         if mstart is not None and self.is_sigma:
-            mstart = mstart.copy()
             mstart = np.log(mstart)
 
         return mstart
 
     @property
-    def reference(self):
+    def reference(self) -> np.ndarray | None:
         mref = self._reference.model
 
         if self.driver.params.forward_only:
             return mref
 
-        if mref is None:
+        if mref is None or (self.is_sigma and all(mref == 0)):
             mref = self.starting
             self.driver.params.alpha_s = 0.0
-        else:
-            mref = mref.copy()
-            if self.is_sigma & (all(mref == 0)):
-                mref = self.starting
-                self.driver.params.alpha_s = 0.0
-            else:
-                mref = np.log(mref) if self.is_sigma else mref
-        return mref
 
-    @property
-    def lower_bound(self):
-        lbound = self._lower_bound.model
+        ref_model = mref.copy()
+        ref_model = np.log(ref_model) if self.is_sigma else ref_model
+
+        return ref_model
 
-        if lbound is None:
+    @property
+    def lower_bound(self) -> np.ndarray | None:
+        if self._lower_bound.model is None:
             return -np.inf
 
-        lbound = lbound.copy()
+        lbound = self._lower_bound.model.copy()
 
         if self.is_sigma:
             is_finite = np.isfinite(lbound)
@@ -157,13 +154,12 @@ def lower_bound(self):
         return lbound
 
     @property
-    def upper_bound(self):
-        ubound = self._upper_bound.model
-
-        if ubound is None:
+    def upper_bound(self) -> np.ndarray | None:
+        if self._upper_bound.model is None:
             return np.inf
 
-        ubound = ubound.copy()
+        ubound = self._upper_bound.model.copy()
+
         if self.is_sigma:
             is_finite = np.isfinite(ubound)
             ubound[is_finite] = np.log(ubound[is_finite])
@@ -171,14 +167,16 @@ def upper_bound(self):
         return ubound
 
     @property
-    def conductivity(self):
-        mstart = self._conductivity.model
+    def conductivity(self) -> np.ndarray | None:
+        if self._conductivity.model is None:
+            return None
 
-        if mstart is not None and self.is_sigma:
-            mstart = mstart.copy()
-            mstart = np.log(mstart)
+        cond_model = self._conductivity.model.copy()
 
-        return mstart
+        if cond_model is not None and self.is_sigma:
+            cond_model = np.log(cond_model)
+
+        return cond_model
 
     def _initialize(self, driver):
         self.driver = driver
diff --git a/geoapps/inversion/joint/joint_surveys/driver.py b/geoapps/inversion/joint/joint_surveys/driver.py
index ff7acbb42..46aab255b 100644
--- a/geoapps/inversion/joint/joint_surveys/driver.py
+++ b/geoapps/inversion/joint/joint_surveys/driver.py
@@ -47,6 +47,9 @@ def validate_create_models(self):
                 norm = np.array(np.sum(projection, axis=1)).flatten()
                 model = (projection * model_local_values) / (norm + 1e-8)
 
+                if self.drivers[0].models.is_sigma:
+                    model = np.exp(model)
+
                 setattr(
                     getattr(self.models, f"_{model_type}"),
                     "model",
diff --git a/geoapps/octree_creation/application.py b/geoapps/octree_creation/application.py
index 29dc93dab..6c0ee4ff7 100644
--- a/geoapps/octree_creation/application.py
+++ b/geoapps/octree_creation/application.py
@@ -23,13 +23,23 @@
 
 from geoapps.base.application import BaseApplication
 from geoapps.base.selection import ObjectDataSelection
+from geoapps.octree_creation.constants import app_initializer
+from geoapps.octree_creation.driver import OctreeDriver
+from geoapps.octree_creation.params import OctreeParams
 from geoapps.utils import warn_module_not_found
 
-from . import OctreeParams, app_initializer
-from .driver import OctreeDriver
-
 with warn_module_not_found():
-    from ipywidgets import Dropdown, FloatText, Label, Layout, Text, VBox, Widget
+    from ipywidgets import (
+        Checkbox,
+        Dropdown,
+        FloatText,
+        IntText,
+        Label,
+        Layout,
+        Text,
+        VBox,
+        Widget,
+    )
     from ipywidgets.widgets.widget_selection import TraitError
 
 
@@ -46,6 +56,8 @@ class OctreeMesh(ObjectDataSelection):
     _depth_core = None
     _horizontal_padding = None
     _vertical_padding = None
+    _diagonal_balance = None
+    _minimum_level = None
 
     def __init__(self, ui_json=None, **kwargs):
         app_initializer.update(kwargs)
@@ -79,6 +91,9 @@ def __init__(self, ui_json=None, **kwargs):
                         Label("Padding distance"),
                         self.horizontal_padding,
                         self.vertical_padding,
+                        Label("Basic"),
+                        self.diagonal_balance,
+                        self.minimum_level,
                     ],
                     layout=Layout(border="solid"),
                 ),
@@ -180,6 +195,28 @@ def vertical_padding(self) -> FloatText:
             )
         return self._vertical_padding
 
+    @property
+    def diagonal_balance(self) -> Checkbox:
+        """
+        Widget controlling the diagonal balance.
+        """
+        if getattr(self, "_diagonal_balance", None) is None:
+            self._diagonal_balance = Checkbox(
+                description="UBC compatible",
+            )
+        return self._diagonal_balance
+
+    @property
+    def minimum_level(self) -> IntText:
+        """
+        Widget controlling the minimum refinement level.
+        """
+        if getattr(self, "_minimum_level", None) is None:
+            self._minimum_level = IntText(
+                description="Minimum refinement level",
+            )
+        return self._minimum_level
+
     @property
     def workspace(self):
         """
diff --git a/geoapps/octree_creation/constants.py b/geoapps/octree_creation/constants.py
index d784dae75..027c9e7b9 100644
--- a/geoapps/octree_creation/constants.py
+++ b/geoapps/octree_creation/constants.py
@@ -99,6 +99,12 @@
             "main": True,
             "value": 500.0,
         },
+        "diagonal_balance": {
+            "group": "Basic",
+            "label": "UBC Compatible",
+            "main": True,
+            "value": True,
+        },
         "minimum_level": {
             "enabled": True,
             "group": "Basic",
diff --git a/geoapps/octree_creation/driver.py b/geoapps/octree_creation/driver.py
index a613b3eac..eb8aa62ad 100644
--- a/geoapps/octree_creation/driver.py
+++ b/geoapps/octree_creation/driver.py
@@ -65,7 +65,9 @@ def octree_from_params(params: OctreeParams):
             depth_core=params.depth_core,
         )
         minimum_level = OctreeDriver.minimum_level(mesh, params.minimum_level)
-        mesh.refine(minimum_level, finalize=False)
+        mesh.refine(
+            minimum_level, finalize=False, diagonal_balance=params.diagonal_balance
+        )
 
         for label, value in params.free_parameter_dict.items():
             refinement_object = getattr(params, value["object"])
@@ -77,12 +79,12 @@ def octree_from_params(params: OctreeParams):
 
             if isinstance(refinement_object, Curve):
                 mesh = OctreeDriver.refine_tree_from_curve(
-                    mesh, refinement_object, levels
+                    mesh, refinement_object, levels, params.diagonal_balance
                 )
 
             elif isinstance(refinement_object, Surface):
                 mesh = OctreeDriver.refine_tree_from_triangulation(
-                    mesh, refinement_object, levels
+                    mesh, refinement_object, levels, params.diagonal_balance
                 )
 
             elif getattr(params, value["type"]) == "surface":
@@ -90,6 +92,7 @@ def octree_from_params(params: OctreeParams):
                     mesh,
                     refinement_object,
                     levels,
+                    params.diagonal_balance,
                     max_distance=getattr(params, value["distance"]),
                 )
 
@@ -98,6 +101,7 @@ def octree_from_params(params: OctreeParams):
                     mesh,
                     refinement_object,
                     levels,
+                    diagonal_balance=params.diagonal_balance,
                 )
 
             else:
@@ -116,6 +120,7 @@ def refine_tree_from_curve(
         mesh: TreeMesh,
         curve: Curve,
         levels: list[int] | np.ndarray,
+        diagonal_balance: bool = True,
         finalize: bool = False,
     ) -> TreeMesh:
         """
@@ -126,6 +131,7 @@ def refine_tree_from_curve(
         :param curve: Curve object to use for refinement.
         :param levels: Number of cells requested at each refinement level.
             Defined in reversed order from the highest octree to lowest.
+        :param diagonal_balance: Whether to balance cells along the diagonal of the tree during construction.
         :param finalize: Finalize the tree mesh after refinement.
 
         """
@@ -137,7 +143,7 @@ def refine_tree_from_curve(
 
         locations = densify_curve(curve, mesh.h[0][0])
         mesh = OctreeDriver.refine_tree_from_points(
-            mesh, locations, levels, finalize=False
+            mesh, locations, levels, diagonal_balance=diagonal_balance, finalize=False
         )
 
         if finalize:
@@ -150,6 +156,7 @@ def refine_tree_from_points(
         mesh: TreeMesh,
         points: ObjectBase | np.ndarray,
         levels: list[int] | np.ndarray,
+        diagonal_balance: bool = True,
         finalize: bool = False,
     ) -> TreeMesh:
         """
@@ -159,6 +166,7 @@ def refine_tree_from_points(
         :param points: Object to use for refinement.
         :param levels: Number of cells requested at each refinement level.
             Defined in reversed order from the highest octree to lowest.
+        :param diagonal_balance: Whether to balance cells along the diagonal of the tree during construction.
         :param finalize: Finalize the tree mesh after refinement.
 
         :return: Refined tree mesh.
@@ -177,7 +185,13 @@ def refine_tree_from_points(
         distance = 0
         for ii, n_cells in enumerate(levels):
             distance += n_cells * OctreeDriver.cell_size_from_level(mesh, ii)
-            mesh.refine_ball(locations, distance, mesh.max_level - ii, finalize=False)
+            mesh.refine_ball(
+                locations,
+                distance,
+                mesh.max_level - ii,
+                diagonal_balance=diagonal_balance,
+                finalize=False,
+            )
 
         if finalize:
             mesh.finalize()
@@ -189,6 +203,7 @@ def refine_tree_from_surface(
         mesh: TreeMesh,
         surface: ObjectBase,
         levels: list[int] | np.ndarray,
+        diagonal_balance: bool = True,
         max_distance: float = np.inf,
         finalize: bool = False,
     ) -> TreeMesh:
@@ -200,6 +215,7 @@ def refine_tree_from_surface(
         :param levels: Number of cells requested at each refinement level.
             Defined in reversed order from the highest octree to lowest.
         :param max_distance: Maximum distance from the surface to refine.
+        :param diagonal_balance: Whether to balance cells along the diagonal of the tree during construction.
         :param finalize: Finalize the tree mesh after refinement.
 
         :return: Refined tree mesh.
@@ -247,6 +263,7 @@ def refine_tree_from_surface(
                 mesh.insert_cells(
                     np.c_[xy[keeper], elevation - depth],
                     np.ones(nnz) * mesh.max_level - ind,
+                    diagonal_balance=diagonal_balance,
                     finalize=False,
                 )
 
@@ -257,7 +274,11 @@ def refine_tree_from_surface(
 
     @staticmethod
     def refine_tree_from_triangulation(
-        mesh: TreeMesh, surface, levels: list[int] | np.ndarray, finalize=False
+        mesh: TreeMesh,
+        surface,
+        levels: list[int] | np.ndarray,
+        diagonal_balance: bool = True,
+        finalize=False,
     ) -> TreeMesh:
         """
         Refine a tree mesh along the simplicies of a surface.
@@ -266,6 +287,7 @@ def refine_tree_from_triangulation(
         :param surface: Surface object to use for refinement.
         :param levels: Number of cells requested at each refinement level.
             Defined in reversed order from highest octree to lowest.
+        :param diagonal_balance: Whether to balance cells along the diagonal of the tree during construction.
         :param finalize: Finalize the tree mesh after refinement.
 
         :return: Refined tree mesh.
@@ -290,6 +312,7 @@ def refine_tree_from_triangulation(
                 (surface.vertices, surface.cells),
                 -ind[0] - 1,
                 paddings,
+                diagonal_balance=diagonal_balance,
                 finalize=finalize,
             )
         return mesh
diff --git a/geoapps/octree_creation/params.py b/geoapps/octree_creation/params.py
index 4a327683f..53d844f7e 100644
--- a/geoapps/octree_creation/params.py
+++ b/geoapps/octree_creation/params.py
@@ -31,6 +31,7 @@ def __init__(self, input_file=None, **kwargs):
         self._u_cell_size = None
         self._v_cell_size = None
         self._w_cell_size = None
+        self._diagonal_balance = None
         self._minimum_level = None
         self._horizontal_padding = None
         self._vertical_padding = None
@@ -143,6 +144,14 @@ def depth_core(self):
     def depth_core(self, val):
         self.setter_validator("depth_core", val)
 
+    @property
+    def diagonal_balance(self):
+        return self._diagonal_balance
+
+    @diagonal_balance.setter
+    def diagonal_balance(self, val):
+        self.setter_validator("diagonal_balance", val)
+
     @property
     def minimum_level(self):
         return self._minimum_level
diff --git a/geoapps/utils/testing.py b/geoapps/utils/testing.py
index 44f77b1fa..638b86263 100644
--- a/geoapps/utils/testing.py
+++ b/geoapps/utils/testing.py
@@ -413,6 +413,7 @@ def topo_drape(x, y):
             mesh,
             topography,
             levels=refinement,
+            diagonal_balance=False,
             finalize=False,
         )
 
@@ -421,6 +422,7 @@ def topo_drape(x, y):
                 mesh,
                 vertices,
                 levels=[2],
+                diagonal_balance=False,
                 finalize=False,
             )
 
diff --git a/tests/data_test.py b/tests/data_test.py
index fa3b2d67f..2c8b4be62 100644
--- a/tests/data_test.py
+++ b/tests/data_test.py
@@ -81,6 +81,7 @@ def test_survey_data(tmp_path: Path):
             mesh,
             test_topo_object,
             levels=[2],
+            diagonal_balance=False,
             finalize=True,
         )
 
diff --git a/tests/run_tests/octree_creation_test.py b/tests/run_tests/octree_creation_test.py
index 331aac5e4..60885478f 100644
--- a/tests/run_tests/octree_creation_test.py
+++ b/tests/run_tests/octree_creation_test.py
@@ -19,7 +19,9 @@
 from scipy.spatial import Delaunay
 
 from geoapps.driver_base.utils import treemesh_2_octree
+from geoapps.octree_creation import OctreeParams
 from geoapps.octree_creation.application import OctreeDriver, OctreeMesh
+from geoapps.shared_utils.utils import octree_2_treemesh
 from geoapps.utils.testing import get_output_workspace
 
 # pylint: disable=redefined-outer-name
@@ -88,6 +90,7 @@ def test_create_octree_radial(tmp_path: Path, setup_test_octree):
             treemesh,
             points,
             str2list(refinement),
+            diagonal_balance=False,
             finalize=True,
         )
         octree = treemesh_2_octree(workspace, treemesh, name="Octree_Mesh")
@@ -113,6 +116,7 @@ def test_create_octree_radial(tmp_path: Path, setup_test_octree):
             w_cell_size=cell_sizes[2],
             horizontal_padding=horizontal_padding,
             vertical_padding=vertical_padding,
+            diagonal_balance=False,
             depth_core=depth_core,
             **refinements,
         )
@@ -139,11 +143,16 @@ def test_create_octree_curve(tmp_path: Path, setup_test_octree):
     with Workspace.create(tmp_path / "testOctree.geoh5") as workspace:
         curve = Curve.create(workspace, vertices=locations)
         curve.remove_cells([-1])
-        treemesh.refine(treemesh.max_level - minimum_level + 1, finalize=False)
+        treemesh.refine(
+            treemesh.max_level - minimum_level + 1,
+            diagonal_balance=False,
+            finalize=False,
+        )
         treemesh = OctreeDriver.refine_tree_from_curve(
             treemesh,
             curve,
             str2list(refinement),
+            diagonal_balance=False,
             finalize=True,
         )
         octree = treemesh_2_octree(workspace, treemesh, name="Octree_Mesh")
@@ -167,6 +176,7 @@ def test_create_octree_curve(tmp_path: Path, setup_test_octree):
             horizontal_padding=horizontal_padding,
             vertical_padding=vertical_padding,
             depth_core=depth_core,
+            diagonal_balance=False,
             **refinements,
         )
         app.trigger_click(None)
@@ -191,11 +201,16 @@ def test_create_octree_surface(tmp_path: Path, setup_test_octree):
 
     with Workspace.create(tmp_path / "testOctree.geoh5") as workspace:
         points = Points.create(workspace, vertices=locations)
-        treemesh.refine(treemesh.max_level - minimum_level + 1, finalize=False)
+        treemesh.refine(
+            treemesh.max_level - minimum_level + 1,
+            diagonal_balance=False,
+            finalize=False,
+        )
         treemesh = OctreeDriver.refine_tree_from_surface(
             treemesh,
             points,
             str2list(refinement),
+            diagonal_balance=False,
             finalize=True,
         )
         octree = treemesh_2_octree(workspace, treemesh, name="Octree_Mesh")
@@ -222,6 +237,7 @@ def test_create_octree_surface(tmp_path: Path, setup_test_octree):
             horizontal_padding=horizontal_padding,
             vertical_padding=vertical_padding,
             depth_core=depth_core,
+            diagonal_balance=False,
             **refinements,
         )
         app.trigger_click(None)
@@ -270,11 +286,16 @@ def test_create_octree_triangulation(tmp_path: Path, setup_test_octree):
             vertices=np.c_[x.flatten(), y.flatten(), z.flatten()],
             cells=surf.simplices,
         )
-        treemesh.refine(treemesh.max_level - minimum_level + 1, finalize=False)
+        treemesh.refine(
+            treemesh.max_level - minimum_level + 1,
+            diagonal_balance=False,
+            finalize=False,
+        )
         treemesh = OctreeDriver.refine_tree_from_triangulation(
             treemesh,
             sphere,
             str2list(refinement),
+            diagonal_balance=False,
             finalize=True,
         )
         octree = treemesh_2_octree(workspace, treemesh, name="Octree_Mesh")
@@ -306,3 +327,74 @@ def test_create_octree_triangulation(tmp_path: Path, setup_test_octree):
         with Workspace(get_output_workspace(tmp_path)) as workspace:
             rec_octree = workspace.get_entity("Octree_Mesh")[0]
             compare_entities(octree, rec_octree, ignore=["_uid"])
+
+
+@pytest.mark.parametrize(
+    "diagonal_balance, exp_values, exp_counts",
+    [(True, [0, 1], [22, 10]), (False, [0, 1, 2], [22, 8, 2])],
+)
+def test_octree_diagonal_balance(
+    tmp_path: Path, diagonal_balance, exp_values, exp_counts
+):
+    workspace = Workspace.create(tmp_path / "testDiagonalBalance.geoh5")
+    with workspace.open(mode="r+"):
+        point = [0, 0, 0]
+        points = Points.create(workspace, vertices=np.array([[150, 0, 150], point]))
+
+        # Repeat the creation using the app
+        params_dict = {
+            "geoh5": workspace,
+            "objects": str(points.uid),
+            "u_cell_size": 10.0,
+            "v_cell_size": 10.0,
+            "w_cell_size": 10.0,
+            "horizontal_padding": 500.0,
+            "vertical_padding": 200.0,
+            "depth_core": 400.0,
+            "Refinement A object": points.uid,
+            "Refinement A levels": "1",
+            "Refinement A type": "radial",
+            "Refinement A distance": 200,
+        }
+
+        params = OctreeParams(
+            **params_dict, diagonal_balance=diagonal_balance, ga_group_name="mesh"
+        )
+        filename = "diag_balance.ui.json"
+
+        params.write_input_file(name=filename, path=tmp_path, validate=False)
+
+        OctreeDriver.start(tmp_path / filename)
+
+    with workspace.open(mode="r"):
+        results = []
+        treemesh = octree_2_treemesh(workspace.get_entity("mesh")[0])
+
+        ind = treemesh._get_containing_cell_indexes(  # pylint: disable=protected-access
+            point
+        )
+        starting_cell = treemesh[ind]
+
+        level = starting_cell._level  # pylint: disable=protected-access
+        for first_neighbor in starting_cell.neighbors:
+            neighbors = []
+            for neighbor in treemesh[first_neighbor].neighbors:
+                if isinstance(neighbor, list):
+                    neighbors += neighbor
+                else:
+                    neighbors.append(neighbor)
+
+            for second_neighbor in neighbors:
+                compare_cell = treemesh[second_neighbor]
+                if set(starting_cell.nodes) & set(compare_cell.nodes):
+                    results.append(
+                        np.abs(
+                            level
+                            - compare_cell._level  # pylint: disable=protected-access
+                        )
+                    )
+
+        values, counts = np.unique(results, return_counts=True)
+
+        assert (values == np.array(exp_values)).all()
+        assert (counts == np.array(exp_counts)).all()
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 7f4d31c8b..cc49c16de 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -106,13 +106,14 @@ def test_drape_to_octree(tmp_path: Path):
         mesh_type="TREE",
     )
     tree = OctreeDriver.refine_tree_from_points(
-        tree, locs, levels=[4, 2], finalize=False
+        tree, locs, levels=[4, 2], diagonal_balance=False, finalize=False
     )
     topography = Points.create(ws, vertices=topo)
     tree = OctreeDriver.refine_tree_from_surface(
         tree,
         topography,
         levels=[2, 2],
+        diagonal_balance=False,
         finalize=True,
     )
     # interp and save common models into the octree