diff --git a/notebooks/03-gravity/fwd_gravity_gradiometry_3d.ipynb b/notebooks/03-gravity/fwd_gravity_gradiometry_3d.ipynb
index c31c3b9..eb587e4 100644
--- a/notebooks/03-gravity/fwd_gravity_gradiometry_3d.ipynb
+++ b/notebooks/03-gravity/fwd_gravity_gradiometry_3d.ipynb
@@ -354,7 +354,7 @@
     "hx = [(dx, nbcx)]\n",
     "hy = [(dy, nbcy)]\n",
     "hz = [(dz, nbcz)]\n",
-    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift vertically to top same as maximum topography\n",
     "mesh.origin += np.r_[0.0, 0.0, z_topo.max()]\n",
diff --git a/notebooks/03-gravity/inv_gravity_anomaly_3d.ipynb b/notebooks/03-gravity/inv_gravity_anomaly_3d.ipynb
index 1d6ae27..0b25ad0 100644
--- a/notebooks/03-gravity/inv_gravity_anomaly_3d.ipynb
+++ b/notebooks/03-gravity/inv_gravity_anomaly_3d.ipynb
@@ -1193,7 +1193,7 @@
     "hx = [(dx, nbcx)]\n",
     "hy = [(dy, nbcy)]\n",
     "hz = [(dz, nbcz)]\n",
-    "tree_mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "tree_mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift vertically to top same as maximum topography\n",
     "tree_mesh.origin += np.r_[0.0, 0.0, topo_xyz[:, -1].max()]\n",
diff --git a/notebooks/03-gravity/weighting_strategies.ipynb b/notebooks/03-gravity/weighting_strategies.ipynb
index 157acf0..a428e2f 100644
--- a/notebooks/03-gravity/weighting_strategies.ipynb
+++ b/notebooks/03-gravity/weighting_strategies.ipynb
@@ -1537,11 +1537,11 @@
     "\n",
     "# Defines the directives for the IRLS regularization. This includes setting\n",
     "# the cooling schedule for the trade-off parameter.\n",
-    "update_IRLS = directives.Update_IRLS(\n",
+    "update_IRLS = directives.UpdateIRLS(\n",
     "    f_min_change=1e-4,\n",
     "    max_irls_iterations=30,\n",
-    "    coolEpsFact=1.5,\n",
-    "    beta_tol=1e-2,\n",
+    "    irls_cooling_factor=1.5,\n",
+    "    misfit_tolerance=1e-2,\n",
     ")\n",
     "\n",
     "# Options for outputting recovered models and predicted data for each beta.\n",
diff --git a/notebooks/04-magnetics/fwd_magnetics_mvi_3d.ipynb b/notebooks/04-magnetics/fwd_magnetics_mvi_3d.ipynb
index 42542a5..c139105 100644
--- a/notebooks/04-magnetics/fwd_magnetics_mvi_3d.ipynb
+++ b/notebooks/04-magnetics/fwd_magnetics_mvi_3d.ipynb
@@ -358,7 +358,7 @@
     "hx = [(dx, nbcx)]\n",
     "hy = [(dy, nbcy)]\n",
     "hz = [(dz, nbcz)]\n",
-    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift vertically to top same as maximum topography\n",
     "mesh.origin += np.r_[0.0, 0.0, z_topo.max()]\n",
diff --git a/notebooks/04-magnetics/inv_magnetics_induced_3d.ipynb b/notebooks/04-magnetics/inv_magnetics_induced_3d.ipynb
index 90514fd..deb7250 100644
--- a/notebooks/04-magnetics/inv_magnetics_induced_3d.ipynb
+++ b/notebooks/04-magnetics/inv_magnetics_induced_3d.ipynb
@@ -1171,7 +1171,7 @@
     "hx = [(dx, nbcx)]\n",
     "hy = [(dy, nbcy)]\n",
     "hz = [(dz, nbcz)]\n",
-    "tree_mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "tree_mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift vertically to top same as maximum topography\n",
     "tree_mesh.origin += np.r_[0.0, 0.0, topo_xyz[:, -1].max()]\n",
diff --git a/notebooks/05-dcr/fwd_dcr_2d.ipynb b/notebooks/05-dcr/fwd_dcr_2d.ipynb
index 127526b..c42568a 100644
--- a/notebooks/05-dcr/fwd_dcr_2d.ipynb
+++ b/notebooks/05-dcr/fwd_dcr_2d.ipynb
@@ -483,7 +483,7 @@
     "# Define the base mesh with top at z = 0 m.\n",
     "hx = [(dh, nbcx)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hz], x0=\"CN\")\n",
+    "mesh = TreeMesh([hx, hz], x0=\"CN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, z_topo.max()]\n",
diff --git a/notebooks/05-dcr/fwd_dcr_3d.ipynb b/notebooks/05-dcr/fwd_dcr_3d.ipynb
index ae4ccf1..3268d74 100644
--- a/notebooks/05-dcr/fwd_dcr_3d.ipynb
+++ b/notebooks/05-dcr/fwd_dcr_3d.ipynb
@@ -383,7 +383,7 @@
     "hx = [(dh, nbcx)]\n",
     "hy = [(dh, nbcy)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, 0.0, z_topo.max()]\n",
diff --git a/notebooks/05-dcr/inv_dcr_2d.ipynb b/notebooks/05-dcr/inv_dcr_2d.ipynb
index 4ae011b..2fdf9c2 100644
--- a/notebooks/05-dcr/inv_dcr_2d.ipynb
+++ b/notebooks/05-dcr/inv_dcr_2d.ipynb
@@ -462,7 +462,7 @@
     "# Define the base mesh with top at z = 0 m\n",
     "hx = [(dh, nbcx)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hz], x0=\"CN\")\n",
+    "mesh = TreeMesh([hx, hz], x0=\"CN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, topo_2d[:, -1].max()]\n",
diff --git a/notebooks/05-dcr/inv_dcr_3d.ipynb b/notebooks/05-dcr/inv_dcr_3d.ipynb
index 026cef2..b4c5ed2 100644
--- a/notebooks/05-dcr/inv_dcr_3d.ipynb
+++ b/notebooks/05-dcr/inv_dcr_3d.ipynb
@@ -538,7 +538,7 @@
     "hx = [(dh, nbcx)]\n",
     "hy = [(dh, nbcy)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, 0.0, topo_xyz[:, -1].max()]\n",
diff --git a/notebooks/06-ip/fwd_ip_2d.ipynb b/notebooks/06-ip/fwd_ip_2d.ipynb
index 4a77a1e..c799e79 100644
--- a/notebooks/06-ip/fwd_ip_2d.ipynb
+++ b/notebooks/06-ip/fwd_ip_2d.ipynb
@@ -305,7 +305,7 @@
     "# Define the base mesh with top at z = 0 m.\n",
     "hx = [(dh, nbcx)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hz], x0=\"CN\")\n",
+    "mesh = TreeMesh([hx, hz], x0=\"CN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, z_topo.max()]\n",
diff --git a/notebooks/06-ip/fwd_ip_3d.ipynb b/notebooks/06-ip/fwd_ip_3d.ipynb
index a8b8539..8859967 100644
--- a/notebooks/06-ip/fwd_ip_3d.ipynb
+++ b/notebooks/06-ip/fwd_ip_3d.ipynb
@@ -356,7 +356,7 @@
     "hx = [(dh, nbcx)]\n",
     "hy = [(dh, nbcy)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, 0.0, z_topo.max()]\n",
diff --git a/notebooks/06-ip/inv_ip_2d.ipynb b/notebooks/06-ip/inv_ip_2d.ipynb
index 6857ace..ff16531 100644
--- a/notebooks/06-ip/inv_ip_2d.ipynb
+++ b/notebooks/06-ip/inv_ip_2d.ipynb
@@ -429,7 +429,7 @@
     "# Define the base mesh with top at z = 0 m\n",
     "hx = [(dh, nbcx)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hz], x0=\"CN\")\n",
+    "mesh = TreeMesh([hx, hz], x0=\"CN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, topo_2d[:, -1].max()]\n",
diff --git a/notebooks/06-ip/inv_ip_3d.ipynb b/notebooks/06-ip/inv_ip_3d.ipynb
index 4082db2..c99fd43 100644
--- a/notebooks/06-ip/inv_ip_3d.ipynb
+++ b/notebooks/06-ip/inv_ip_3d.ipynb
@@ -532,7 +532,7 @@
     "hx = [(dh, nbcx)]\n",
     "hy = [(dh, nbcy)]\n",
     "hz = [(dh, nbcz)]\n",
-    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\")\n",
+    "mesh = TreeMesh([hx, hy, hz], x0=\"CCN\", diagonal_balance=True)\n",
     "\n",
     "# Shift top to maximum topography\n",
     "mesh.origin = mesh.origin + np.r_[0.0, 0.0, topo_xyz[:, -1].max()]\n",
diff --git a/notebooks/07-fdem/fwd_fdem_3d.ipynb b/notebooks/07-fdem/fwd_fdem_3d.ipynb
index 06d32c5..a264660 100644
--- a/notebooks/07-fdem/fwd_fdem_3d.ipynb
+++ b/notebooks/07-fdem/fwd_fdem_3d.ipynb
@@ -370,7 +370,7 @@
     "nbc = 2 ** int(np.round(np.log(dom_width / dh) / np.log(2.0)))  # num. base cells\n",
     "\n",
     "h = [(dh, nbc)]\n",
-    "mesh = TreeMesh([h, h, h], x0=\"CCC\")\n",
+    "mesh = TreeMesh([h, h, h], x0=\"CCC\", diagonal_balance=True)\n",
     "\n",
     "# Mesh refinement based on surface topography (starting with 20 cells)\n",
     "k = np.sqrt(np.sum(topo_xyz[:, 0:2] ** 2, axis=1)) < 400\n",
diff --git a/notebooks/08-tdem/fwd_tdem_3d.ipynb b/notebooks/08-tdem/fwd_tdem_3d.ipynb
index ee38063..9f1dc4a 100644
--- a/notebooks/08-tdem/fwd_tdem_3d.ipynb
+++ b/notebooks/08-tdem/fwd_tdem_3d.ipynb
@@ -407,7 +407,7 @@
     "nbc = 2 ** int(np.round(np.log(dom_width / dh) / np.log(2.0)))  # num. base cells\n",
     "\n",
     "h = [(dh, nbc)]\n",
-    "mesh = TreeMesh([h, h, h], x0=\"CCC\")\n",
+    "mesh = TreeMesh([h, h, h], x0=\"CCC\", diagonal_balance=True)\n",
     "\n",
     "# Mesh refinement based on surface topography (starting with 20 cells)\n",
     "k = np.sqrt(np.sum(topo_xyz[:, 0:2] ** 2, axis=1)) < 400\n",
diff --git a/notebooks/08-tdem/fwd_utem_3d.ipynb b/notebooks/08-tdem/fwd_utem_3d.ipynb
index 8deaf82..8c395b1 100644
--- a/notebooks/08-tdem/fwd_utem_3d.ipynb
+++ b/notebooks/08-tdem/fwd_utem_3d.ipynb
@@ -598,7 +598,7 @@
     "L = 80000.0\n",
     "nbc = 2 ** int(np.round(np.log(L / dh) / np.log(2.0)))\n",
     "h = [(dh, nbc)]\n",
-    "mesh = TreeMesh([h, h, h], origin=\"CCC\")\n",
+    "mesh = TreeMesh([h, h, h], origin=\"CCC\", diagonal_balance=True)\n",
     "\n",
     "# Shift vertically to top same as maximum topography\n",
     "mesh.origin += np.r_[0.0, 0.0, z_topo.min()]"