diff --git a/.coveragerc b/.coveragerc
index 1eed499..6b5acd9 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,2 +1,2 @@
 [run]
-omit = pipedream_solver/nsuperlink.py,pipedream_solver/ninfiltration.py,pipedream_solver/nquality.py,pipedream_solver/ngeometry.py,pipedream_solver/nutils.py
\ No newline at end of file
+omit = pipedream_solver/geometry.py, pipedream_solver/infiltration.py, pipedream_solver/_nsuperlink.py, pipedream_solver/_ninfiltration.py
\ No newline at end of file
diff --git a/pipedream_solver/_ninfiltration.py b/pipedream_solver/_ninfiltration.py
new file mode 100644
index 0000000..26e7619
--- /dev/null
+++ b/pipedream_solver/_ninfiltration.py
@@ -0,0 +1,121 @@
+import numpy as np
+from numba import njit
+from pipedream_solver.nutils import newton_raphson, bounded_newton_raphson, numba_any
+
+@njit
+def run_green_ampt_newton(F_2, x0, F_1, dt, Ks, theta_d, psi_f, ia, max_iter=50,
+                          atol=1.48e-8, rtol=0.0, bounded=True):
+    """
+    Use Newton-Raphson iteration to find cumulative infiltration at next time step (F_2).
+
+    Inputs:
+    -------
+    F_2 : np.ndarray (float)
+        Cumulative infiltration at next time step (meters).
+    x0 : np.ndarray (float)
+        Initial guess for cumulative infiltration (meters)
+    F_1 : np.ndarray (float)
+        Cumulative infiltration at previous time step (meters)
+    dt : np.ndarray (float)
+        Time step (seconds)
+    Ks : np.ndarray (float)
+        Saturated hydraulic conductivity (m/s)
+    theta_d : np.ndarray (float)
+        Soil moisture deficit (-)
+    psi_f : np.ndarray (float)
+        Matric potential of the wetting front (m)
+    ia : np.ndarray (float)
+        Available rainfall depth (meters)
+    max_iter : int
+        Maximum number of Newton-Raphson iterations
+    atol : float
+        Allowable (absolute) error of the zero value
+    rtol : float
+        Allowable (relative) error of the zero value
+    bounded : bool
+        If True, use bounded Newton-Raphson iteration
+    """
+    n = F_2.size
+    for i in range(n):
+        x_0_i = x0[i]
+        F_1_i = F_1[i]
+        dt_i = dt[i]
+        nargs = np.zeros(5)
+        nargs[0] = F_1_i
+        nargs[1] = dt_i
+        nargs[2] = Ks[i]
+        nargs[3] = theta_d[i]
+        nargs[4] = psi_f[i]
+        if bounded:
+            min_F = 0
+            max_F = F_1_i + ia[i] * dt_i
+            F_est = bounded_newton_raphson(numba_integrated_green_ampt,
+                                        numba_derivative_green_ampt,
+                                        x_0_i, min_F, max_F,
+                                        nargs, max_iter=max_iter,
+                                        atol=atol, rtol=rtol)
+        else:
+            F_est = newton_raphson(numba_integrated_green_ampt,
+                                   numba_derivative_green_ampt,
+                                   x_0_i, nargs, max_iter=max_iter,
+                                   atol=atol, rtol=rtol)
+        F_2[i] = F_est
+    return F_2
+
+@njit
+def numba_integrated_green_ampt(F_2, args):
+    """
+    Solve integrated form of Green Ampt equation for cumulative infiltration.
+
+    Inputs:
+    -------
+    F_2: np.ndarray (float)
+        Cumulative infiltration at current timestep (m)
+    F_1: np.ndarray (float)
+        Cumulative infiltration at next timestep (m)
+    dt: float
+        Time step (seconds)
+    Ks: np.ndarray (float)
+        Saturated hydraulic conductivity (m/s)
+    theta_d: np.ndarray (float)
+        Soil moisture deficit
+    psi_s: np.ndarray (float)
+        Soil suction head (m)
+    """
+    F_1 = args[0]
+    dt = args[1]
+    Ks = args[2]
+    theta_d = args[3]
+    psi_s = args[4]
+    C = Ks * dt + F_1 - psi_s * theta_d * np.log(F_1 + np.abs(psi_s) * theta_d)
+    zero = C + psi_s * theta_d * np.log(F_2 + np.abs(psi_s) * theta_d) - F_2
+    return zero
+
+@njit
+def numba_derivative_green_ampt(F_2, args):
+    """
+    Derivative of Green Ampt equation for cumulative infiltration.
+
+    Inputs:
+    -------
+    F_2: np.ndarray (float)
+        Cumulative infiltration at current timestep (m)
+    F_1: np.ndarray (float)
+        Cumulative infiltration at next timestep (m)
+    dt: float
+        Time step (seconds)
+    Ks: np.ndarray (float)
+        Saturated hydraulic conductivity (m/s)
+    theta_d: np.ndarray (float)
+        Soil moisture deficit
+    psi_s: np.ndarray (float)
+        Soil suction head (m)
+    """
+    F_1 = args[0]
+    dt = args[1]
+    Ks = args[2]
+    theta_d = args[3]
+    psi_s = args[4]
+    zero = (psi_s * theta_d / (psi_s * theta_d + F_2)) - 1
+    return zero
+
diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py
new file mode 100644
index 0000000..d9d8175
--- /dev/null
+++ b/pipedream_solver/_nsuperlink.py
@@ -0,0 +1,1678 @@
+import numpy as np
+from numba import njit, prange
+from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, DictType, void
+import pipedream_solver.ngeometry
+
+SMALLEST_NORMAL = np.finfo(np.float64).smallest_normal
+
+MIN_SJ_AREA = 1e-8
+
+CIRCULAR = 1
+RECT_CLOSED = 2
+RECT_OPEN = 3
+TRIANGULAR = 4
+TRAPEZOIDAL = 5
+PARABOLIC = 6
+ELLIPTICAL = 7
+WIDE = 8
+FORCE_MAIN = 9
+FLOODPLAIN = 10
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_hydraulic_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik,
+                             _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
+                             _geom_codes, _Ik, _ik):
+    n = len(_ik)
+    for i in range(n):
+        I = _Ik[i]
+        Ip1 = I + 1
+        geom_code = _geom_codes[i]
+        h_I = _h_Ik[I]
+        h_Ip1 = _h_Ik[Ip1]
+        h_i = (h_I + h_Ip1) / 2
+        g1_i = _g1_ik[i]
+        g2_i = _g2_ik[i]
+        g3_i = _g3_ik[i]
+        g4_i = _g4_ik[i]
+        g5_i = _g5_ik[i]
+        g6_i = _g6_ik[i]
+        g7_i = _g7_ik[i]
+        if geom_code:
+            if geom_code == CIRCULAR:
+                _A_ik[i] = pipedream_solver.ngeometry.Circular_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Circular_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Circular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == RECT_CLOSED:
+                _A_ik[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Rect_Closed_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Rect_Closed_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == RECT_OPEN:
+                _A_ik[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Rect_Open_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Rect_Open_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRIANGULAR:
+                _A_ik[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Triangular_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Triangular_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Triangular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRAPEZOIDAL:
+                _A_ik[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_i, g1_i, g2_i, g3_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Trapezoidal_Pe_ik(h_i, g1_i, g2_i, g3_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Trapezoidal_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == PARABOLIC:
+                _A_ik[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Parabolic_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Parabolic_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Parabolic_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == ELLIPTICAL:
+                raise NotImplementedError
+            elif geom_code == WIDE:
+                _A_ik[i] = pipedream_solver.ngeometry.Wide_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Wide_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Wide_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Wide_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FORCE_MAIN:
+                _A_ik[i] = pipedream_solver.ngeometry.Force_Main_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Force_Main_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Force_Main_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Force_Main_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FLOODPLAIN:
+                _A_ik[i] = pipedream_solver.ngeometry.Floodplain_A_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Floodplain_Pe_ik(h_i, g1_i, g2_i,
+                                                                        g3_i, g4_i, g5_i, g6_i, g7_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Floodplain_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Floodplain_B_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, _theta_bk,
+                            _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
+                            _geom_codes, _i_bk, _I_bk, _J_bk):
+    n = len(_i_bk)
+    for k in range(n):
+        i = _i_bk[k]
+        I = _I_bk[k]
+        j = _J_bk[k]
+        # TODO: If theta is zero, should return critical depth
+        h_I = _h_Ik[I]
+        h_Ip1 = _theta_bk[k] * (_H_j[j] - _z_inv_bk[k])
+        h_i = (h_I + h_Ip1) / 2
+        geom_code = _geom_codes[i]
+        g1_i = _g1_ik[i]
+        g2_i = _g2_ik[i]
+        g3_i = _g3_ik[i]
+        g4_i = _g4_ik[i]
+        g5_i = _g5_ik[i]
+        g6_i = _g6_ik[i]
+        g7_i = _g7_ik[i]
+        if geom_code:
+            if geom_code == CIRCULAR:
+                _A_bk[k] = pipedream_solver.ngeometry.Circular_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Circular_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Circular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == RECT_CLOSED:
+                _A_bk[k] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Rect_Closed_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Rect_Closed_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == RECT_OPEN:
+                _A_bk[k] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Rect_Open_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Rect_Open_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRIANGULAR:
+                _A_bk[k] = pipedream_solver.ngeometry.Triangular_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Triangular_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Triangular_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Triangular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRAPEZOIDAL:
+                _A_bk[k] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_i, g1_i, g2_i, g3_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Trapezoidal_Pe_ik(h_i, g1_i, g2_i, g3_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Trapezoidal_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == PARABOLIC:
+                _A_bk[k] = pipedream_solver.ngeometry.Parabolic_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Parabolic_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Parabolic_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Parabolic_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == ELLIPTICAL:
+                raise NotImplementedError
+            elif geom_code == WIDE:
+                _A_bk[k] = pipedream_solver.ngeometry.Wide_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Wide_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Wide_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Wide_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FORCE_MAIN:
+                _A_bk[k] = pipedream_solver.ngeometry.Force_Main_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Force_Main_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Force_Main_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Force_Main_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FLOODPLAIN:
+                _A_bk[k] = pipedream_solver.ngeometry.Floodplain_A_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Floodplain_Pe_ik(h_i, g1_i, g2_i,
+                                                                        g3_i, g4_i, g5_i, g6_i, g7_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Floodplain_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Floodplain_B_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64),
+      cache=True)
+def numba_orifice_geometry(_Ao, h_eo, u_o, _g1_o, _g2_o, _g3_o, _geom_codes_o, n_o):
+    for i in range(n_o):
+        geom_code = _geom_codes_o[i]
+        g1 = _g1_o[i]
+        g2 = _g2_o[i]
+        g3 = _g3_o[i]
+        u = u_o[i]
+        h_e = h_eo[i]
+        if geom_code:
+            if geom_code == CIRCULAR:
+                _Ao[i] = pipedream_solver.ngeometry.Circular_A_ik(h_e, g1 * u, g2)
+            elif geom_code == RECT_CLOSED:
+                _Ao[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_e, g1 * u, g2)
+            elif geom_code == RECT_OPEN:
+                _Ao[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_e, g1 * u, g2)
+            elif geom_code == TRIANGULAR:
+                _Ao[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_e, g1 * u, g2)
+            elif geom_code == TRAPEZOIDAL:
+                _Ao[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_e, g1 * u, g2, g3)
+            elif geom_code == PARABOLIC:
+                _Ao[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_e, g1 * u, g2)
+            elif geom_code == ELLIPTICAL:
+                raise NotImplementedError
+            elif geom_code == WIDE:
+                _Ao[i] = pipedream_solver.ngeometry.Wide_A_ik(h_e, g1 * u, g2)
+            elif geom_code == FORCE_MAIN:
+                _Ao[i] = pipedream_solver.ngeometry.Force_Main_A_ik(h_e, g1 * u, g2)
+            elif geom_code == FLOODPLAIN:
+                # TODO: This can be implemented
+                raise NotImplementedError
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_transect_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik, _is_irregular, _transect_zs,
+                            _transect_As, _transect_Bs, _transect_Pes, _transect_Rs,
+                            _transect_codes, _transect_inds, _transect_lens, _Ik, _ik):
+    n = len(_ik)
+    for i in range(n):
+        is_irregular = _is_irregular[i]
+        if is_irregular:
+            I = _Ik[i]
+            Ip1 = I + 1
+            transect_code = _transect_codes[i]
+            h_I = _h_Ik[I]
+            h_Ip1 = _h_Ik[Ip1]
+            h_i = (h_I + h_Ip1) / 2
+            start = _transect_inds[transect_code]
+            size = _transect_lens[transect_code]
+            end = start + size
+            _z_range = _transect_zs[start:end]
+            _A_range = _transect_As[start:end]
+            _B_range = _transect_Bs[start:end]
+            _Pe_range = _transect_Pes[start:end]
+            _R_range = _transect_Rs[start:end]
+            _A_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _A_range, 1)
+            _B_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _B_range, 1)
+            _Pe_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _Pe_range, 1)
+            _R_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _R_range, 1)
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], boolean[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], int64[:], int64[:],
+            int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_boundary_transect(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk,
+                            _theta_bk, _is_irregular, _transect_zs, _transect_As, _transect_Bs,
+                            _transect_Pes, _transect_Rs, _transect_codes, _transect_inds,
+                            _transect_lens, _I_bk, _i_bk, _J_bk):
+    n = len(_i_bk)
+    for k in range(n):
+        i = _i_bk[k]
+        I = _I_bk[k]
+        j = _J_bk[k]
+        is_irregular_bk = _is_irregular[i]
+        if is_irregular_bk:
+            h_I = _h_Ik[I]
+            # TODO: This should incorporate theta
+            h_Ip1 = _theta_bk[k] * (_H_j[j] - _z_inv_bk[k])
+            transect_code = _transect_codes[i]
+            h_i = (h_I + h_Ip1) / 2
+            start = _transect_inds[transect_code]
+            size = _transect_lens[transect_code]
+            end = start + size
+            _z_range = _transect_zs[start:end]
+            _A_range = _transect_As[start:end]
+            _B_range = _transect_Bs[start:end]
+            _Pe_range = _transect_Pes[start:end]
+            _R_range = _transect_Rs[start:end]
+            _A_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _A_range, 1)
+            _B_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _B_range, 1)
+            _Pe_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _Pe_range, 1)
+            _R_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _R_range, 1)
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]),
+      cache=True)
+def numba_compute_functional_storage_areas(h, A, a, b, c, _functional):
+    M = h.size
+    for j in range(M):
+        if _functional[j]:
+            if h[j] < 0:
+                A[j] = 0
+            else:
+                A[j] = max(a[j] * (h[j]**b[j]) + c[j], MIN_SJ_AREA)
+    return A
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]),
+      cache=True)
+def numba_compute_functional_storage_volumes(h, V, a, b, c, _functional):
+    M = h.size
+    for j in range(M):
+        if _functional[j]:
+            if h[j] < 0:
+                V[j] = 0
+            else:
+                # TODO: Ensure this works when functional storage area is negative
+                V[j] = (a[j] / (b[j] + 1)) * h[j] ** (b[j] + 1) + c[j] * h[j]
+    return V
+
+@njit
+def numba_compute_tabular_storage_areas(h_j, A_sj, hs, As, sjs, sts, inds, lens):
+    n = sjs.size
+    for i in range(n):
+        sj = sjs[i]
+        st = sts[i]
+        ind = inds[st]
+        size = lens[st]
+        h_range = hs[ind:ind+size]
+        A_range = As[ind:ind+size]
+        Amin = A_range.min()
+        Amax = A_range.max()
+        h_search = h_j[sj]
+        ix = np.searchsorted(h_range, h_search)
+        # NOTE: np.interp not supported in this version of numba
+        # A_result = np.interp(h_search, h_range, A_range)
+        # A_out[i] = A_result
+        if (ix == 0):
+            A_sj[sj] = Amin
+        elif (ix >= size):
+            A_sj[sj] = Amax
+        else:
+            dx_0 = h_search - h_range[ix - 1]
+            dx_1 = h_range[ix] - h_search
+            frac = dx_0 / (dx_0 + dx_1)
+            A_sj[sj] = (1 - frac) * A_range[ix - 1] + (frac) * A_range[ix]
+    return A_sj
+
+@njit
+def numba_compute_tabular_storage_volumes(h_j, V_sj, hs, As, Vs, sjs, sts, inds, lens):
+    n = sjs.size
+    for i in range(n):
+        sj = sjs[i]
+        st = sts[i]
+        ind = inds[st]
+        size = lens[st]
+        h_range = hs[ind:ind+size]
+        A_range = As[ind:ind+size]
+        V_range = Vs[ind:ind+size]
+        hmax = h_range.max()
+        Vmin = V_range.min()
+        Vmax = V_range.max()
+        Amax = A_range.max()
+        h_search = h_j[sj]
+        ix = np.searchsorted(h_range, h_search)
+        # NOTE: np.interp not supported in this version of numba
+        # A_result = np.interp(h_search, h_range, A_range)
+        # A_out[i] = A_result
+        if (ix == 0):
+            V_sj[sj] = Vmin
+        elif (ix >= size):
+            V_sj[sj] = Vmax + Amax * (h_search - hmax)
+        else:
+            dx_0 = h_search - h_range[ix - 1]
+            dx_1 = h_range[ix] - h_search
+            frac = dx_0 / (dx_0 + dx_1)
+            V_sj[sj] = (1 - frac) * V_range[ix - 1] + (frac) * V_range[ix]
+    return V_sj
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64))
+def friction_slope(Q_ik_t, dx_ik, A_ik, R_ik, n_ik, Sf_method_ik, g=9.81):
+    if A_ik > 0:
+        # Chezy-Manning eq.
+        if Sf_method_ik == 0:
+            t_1 = (g * n_ik**2 * np.abs(Q_ik_t) * dx_ik
+                   / A_ik / R_ik**(4/3))
+        # Hazen-Williams eq.
+        elif Sf_method_ik == 1:
+            t_1 = (1.354 * g * np.abs(Q_ik_t)**0.85 * dx_ik
+                   / A_ik**0.85 / n_ik**1.85 / R_ik**1.1655)
+        # Darcy-Weisbach eq.
+        elif Sf_method_ik == 2:
+            # kinematic viscosity(meter^2/sec), we can consider this is constant.
+            nu = 0.0000010034
+            Re = (np.abs(Q_ik_t) / A_ik) * 4 * R_ik / nu
+            if Re > 2000:
+                f = 0.25 / (np.log10(n_ik / (3.7 * 4 * R_ik) + 5.74 / (Re**0.9)))**2
+            elif (Re > 0) and (Re <= 2000):
+                f = 64 / Re
+            elif (Re == 0):
+                f = 0.1
+            else:
+                raise ValueError('Reynolds number outside allowable range')
+            t_1 = (0.01274 * g * f * np.abs(Q_ik_t) * dx_ik
+                   / (A_ik * R_ik))
+        else:
+            raise ValueError('Invalid friction method.')
+        return t_1
+    else:
+        return 0.
+
+@njit(float64[:](float64[:], float64[:]),
+      cache=True)
+def numba_a_ik(u_Ik, sigma_ik):
+    """
+    Compute link coefficient 'a' for link i, superlink k.
+    """
+    return -np.maximum(u_Ik, 0) * sigma_ik
+
+@njit(float64[:](float64[:], float64[:]),
+      cache=True)
+def numba_c_ik(u_Ip1k, sigma_ik):
+    """
+    Compute link coefficient 'c' for link i, superlink k.
+    """
+    return -np.maximum(-u_Ip1k, 0) * sigma_ik
+
+@njit(float64[:](float64[:], float64, float64[:], float64[:], float64[:], float64[:],
+                 float64[:], float64[:], float64[:], float64[:], boolean[:], float64[:], int64[:], float64),
+      cache=True)
+def numba_b_ik(dx_ik, dt, n_ik, Q_ik_t, A_ik, R_ik,
+               A_c_ik, C_ik, a_ik, c_ik, ctrl, sigma_ik, Sf_method_ik, g=9.81):
+    """
+    Compute link coefficient 'b' for link i, superlink k.
+    """
+    # TODO: Clean up
+    t_0 = (dx_ik / dt)
+    t_1 = np.zeros(Q_ik_t.size)
+    k = len(Sf_method_ik)
+    for n in range(k):
+        t_1[n] = friction_slope(Q_ik_t[n], dx_ik[n], A_ik[n], R_ik[n],
+                                n_ik[n], Sf_method_ik[n], g)
+    t_2 = np.zeros(ctrl.size)
+    cond = ctrl
+    t_2[cond] = C_ik[cond] * np.abs(Q_ik_t[cond]) / 2 / g / A_c_ik[cond]**2
+    t_3 = a_ik * sigma_ik
+    t_4 = c_ik * sigma_ik
+    return t_0 + t_1 + t_2 - t_3 - t_4
+
+#@njit(float64[:](float64[:], float64, float64[:], float64[:], float64[:], float64[:],
+#                 float64[:], float64[:], float64[:], boolean[:], float64[:], int64[:], float64),
+#      cache=True)
+#def numba_b_ik(dx_ik, dt, n_ik, Q_ik_t, A_ik, R_ik,
+#               A_c_ik, C_ik, u_ik, ctrl, sigma_ik, Sf_method_ik, g=9.81):
+#    """
+#    Compute link coefficient 'b' for link i, superlink k.
+#    """
+#    # TODO: Clean up
+#    t_0 = (dx_ik / dt)
+#    t_1 = np.zeros(Q_ik_t.size)
+#    k = len(Sf_method_ik)
+#    for n in range(k):
+#        t_1[n] = friction_slope(Q_ik_t[n], dx_ik[n], A_ik[n], R_ik[n],
+#                                n_ik[n], Sf_method_ik[n], g)
+#    t_2 = np.zeros(ctrl.size)
+#    cond = ctrl
+#    t_2[cond] = C_ik[cond] * np.abs(Q_ik_t[cond]) / 2 / g / A_c_ik[cond]**2
+#    t_3 = u_ik * sigma_ik
+#    return t_0 + t_1 + t_2 + t_3
+
+@njit(float64[:](float64[:], float64[:], float64, float64[:], float64[:], float64[:], float64),
+      cache=True)
+def numba_P_ik(Q_ik_t, dx_ik, dt, A_ik, S_o_ik, sigma_ik, g=9.81):
+    """
+    Compute link coefficient 'P' for link i, superlink k.
+    """
+    t_0 = (Q_ik_t * dx_ik / dt) * sigma_ik
+    t_1 = g * A_ik * S_o_ik * dx_ik
+    return t_0 + t_1
+
+@njit(float64(float64, float64, float64, float64, float64, float64),
+      cache=True)
+def E_Ik(B_ik, dx_ik, B_im1k, dx_im1k, A_SIk, dt):
+    """
+    Compute node coefficient 'E' for node I, superlink k.
+    """
+    t_0 = B_ik * dx_ik / 2
+    t_1 = B_im1k * dx_im1k / 2
+    t_2 = A_SIk
+    t_3 = dt
+    return (t_0 + t_1 + t_2) / t_3
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def D_Ik(Q_0IK, B_ik, dx_ik, B_im1k, dx_im1k, A_SIk, h_Ik_t, dt):
+    """
+    Compute node coefficient 'D' for node I, superlink k.
+    """
+    t_0 = Q_0IK
+    t_1 = B_ik * dx_ik / 2
+    t_2 = B_im1k * dx_im1k / 2
+    t_3 = A_SIk
+    t_4 = h_Ik_t / dt
+    return t_0 + ((t_1 + t_2 + t_3) * t_4)
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], int64[:], float64,
+            int64[:], int64[:], boolean[:], boolean[:]),
+      cache=True)
+def numba_node_coeffs(_D_Ik, _E_Ik, _Q_0Ik, _B_ik, _h_Ik, _dx_ik, _A_SIk,
+                      _B_uk, _B_dk, _dx_uk, _dx_dk, _kI, _dt,
+                      _forward_I_i, _backward_I_i, _is_start, _is_end):
+    N = _h_Ik.size
+    for I in range(N):
+        k = _kI[I]
+        if _is_start[I]:
+            i = _forward_I_i[I]
+            _E_Ik[I] = E_Ik(_B_ik[i], _dx_ik[i], _B_uk[k], _dx_uk[k], _A_SIk[I], _dt)
+            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_ik[i], _dx_ik[i], _B_uk[k], _dx_uk[k], _A_SIk[I],
+                            _h_Ik[I], _dt)
+        elif _is_end[I]:
+            im1 = _backward_I_i[I]
+            _E_Ik[I] = E_Ik(_B_dk[k], _dx_dk[k], _B_ik[im1], _dx_ik[im1],
+                            _A_SIk[I], _dt)
+            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_dk[k], _dx_dk[k], _B_ik[im1],
+                            _dx_ik[im1], _A_SIk[I], _h_Ik[I], _dt)
+        else:
+            i = _forward_I_i[I]
+            im1 = i - 1
+            _E_Ik[I] = E_Ik(_B_ik[i], _dx_ik[i], _B_ik[im1], _dx_ik[im1],
+                            _A_SIk[I], _dt)
+            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_ik[i], _dx_ik[i], _B_ik[im1],
+                            _dx_ik[im1], _A_SIk[I], _h_Ik[I], _dt)
+    return 1
+
+@njit(float64(float64, float64),
+      cache=True)
+def safe_divide(num, den):
+    #if (den == 0):
+    #    return 0
+    #else:
+    #    return num / den
+    return (num + SMALLEST_NORMAL) / (den + SMALLEST_NORMAL)
+
+# TODO: This is really just meant for orifice, weir, and pump equations
+# Where division by zero implies entire expression is zero
+@njit(float64[:](float64[:], float64[:]),
+      cache=True)
+def safe_divide_vec(num, den):
+    result = np.zeros_like(num)
+    is_safe = (np.abs(den) > SMALLEST_NORMAL)
+    result[is_safe] = num[is_safe] / den[is_safe]
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def Q_i_f(h_Ip1k, h_1k, U_Ik, V_Ik, W_Ik):
+    t_0 = U_Ik * h_Ip1k
+    t_1 = V_Ik
+    t_2 = W_Ik * h_1k
+    return t_0 + t_1 + t_2
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def Q_i_b(h_Ik, h_Np1k, X_Ik, Y_Ik, Z_Ik):
+    t_0 = X_Ik * h_Ik
+    t_1 = Y_Ik
+    t_2 = Z_Ik * h_Np1k
+    return t_0 + t_1 + t_2
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def h_i_b(Q_ik, h_Np1k, X_Ik, Y_Ik, Z_Ik):
+    num = Q_ik - Y_Ik - Z_Ik * h_Np1k
+    den = X_Ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64[:], int64[:], int64[:], int64,
+            float64, float64[:], boolean),
+      cache=True)
+def numba_solve_internals(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik,
+                          _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, nk, NK,
+                          min_depth, max_depth_k, first_link_backwards=True):
+    for k in range(NK):
+        n = nk[k]
+        i_1 = _i_1k[k]
+        I_1 = _I_1k[k]
+        i_n = i_1 + n - 1
+        I_Np1 = I_1 + n
+        I_N = I_Np1 - 1
+        # Set boundary depths
+        _h_1k = _h_uk[k]
+        _h_Np1k = _h_dk[k]
+        _h_Ik[I_1] = _h_1k
+        _h_Ik[I_Np1] = _h_Np1k
+        # Set max depth
+        max_depth = max_depth_k[k]
+        # Compute internal depths and flows (except first link flow)
+        for j in range(n - 1):
+            I = I_N - j
+            Ip1 = I + 1
+            i = i_n - j
+            _Q_ik[i] = Q_i_f(_h_Ik[Ip1], _h_1k, _U_Ik[I], _V_Ik[I], _W_Ik[I])
+            _h_Ik[I] = h_i_b(_Q_ik[i], _h_Np1k, _X_Ik[I], _Y_Ik[I], _Z_Ik[I])
+            #if _h_Ik[I] < min_depth:
+            #    _h_Ik[I] = min_depth
+            #if _h_Ik[I] > max_depth:
+            #    _h_Ik[I] = max_depth
+        if first_link_backwards:
+            _Q_ik[i_1] = Q_i_b(_h_Ik[I_1], _h_Np1k, _X_Ik[I_1], _Y_Ik[I_1],
+                            _Z_Ik[I_1])
+        else:
+            # Not theoretically correct, but seems to be more stable sometimes
+            _Q_ik[i_1] = Q_i_f(_h_Ik[I_1 + 1], _h_1k, _U_Ik[I_1], _V_Ik[I_1],
+                            _W_Ik[I_1])
+    return 1
+
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64[:], int64[:], int64[:], int64,
+            float64, float64[:], boolean),
+      cache=True)
+def numba_solve_internals_matrix(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik,
+                          _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, nk, NK,
+                          min_depth, max_depth_k, first_link_backwards=True):
+    for k in range(NK):
+        n = nk[k]
+        i_1 = _i_1k[k]
+        I_1 = _I_1k[k]
+        i_n = i_1 + n - 1
+        I_2 = I_1 + 1
+        I_Np1 = I_1 + n
+        I_N = I_Np1 - 1
+        # Set boundary depths
+        _h_1k = _h_uk[k]
+        _h_Np1k = _h_dk[k]
+        _h_Ik[I_1] = _h_1k
+        _h_Ik[I_Np1] = _h_Np1k
+        # Set up matrix
+        dim = 2 * n - 1
+        A = np.zeros((dim, dim))
+        b = np.zeros(dim)
+        # Fill linear system
+        #### h indices
+        for j, i in enumerate(range(1, dim, 2)):
+            A[i, i] = -_X_Ik[I_2+j]
+            A[i, i+1] = 1.
+            b[i] = _Y_Ik[I_2+j] + _h_Np1k * _Z_Ik[I_2+j]
+        #### Q indices
+        for j, i in enumerate(range(2, dim, 2)):
+            A[i, i] = 1.
+            if i < dim - 1:
+                A[i, i+1] = -_U_Ik[I_2+j]
+            b[i] = _V_Ik[I_2+j] + _h_1k * _W_Ik[I_2+j]
+        #### Other indices
+        A[0 ,0] = 1.
+        A[0, 1] = -_X_Ik[I_1]
+        b[0] = _Y_Ik[I_1] + _h_Np1k * _Z_Ik[I_1]
+        b[-1] += _h_1k * _U_Ik[I_N]
+        # Solve linear system
+        x = np.linalg.solve(A, b)
+        # Write to output arrays
+        _Q_ik[i_1:i_1+n] = x[::2]
+        _h_Ik[I_2:I_2+n-1] = x[1::2]
+    return 1
+
+
+@njit(float64[:](float64[:], int64, int64[:], int64[:], int64[:], int64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_solve_internals_ls(_h_Ik, NK, nk, _k_1k, _i_1k, _I_1k, _U, _X, _b):
+    for k in range(NK):
+        nlinks = nk[k]
+        lstart = _k_1k[k]
+        rstart = _i_1k[k]
+        jstart = _I_1k[k]
+        _Ak = np.zeros((nlinks, nlinks - 1))
+        for i in range(nlinks - 1):
+            _Ak[i, i] = _U[lstart + i]
+            _Ak[i + 1, i] = -_X[lstart + i]
+        _AkT = _Ak.T.copy()
+        _bk = _b[rstart:rstart+nlinks].copy()
+        _AA = _AkT @ _Ak
+        _Ab = _AkT @ _bk
+        # If want to prevent singular matrix, set ( diag == 0 ) = 1
+        for i in range(nlinks - 1):
+            if (_AA[i, i] == 0.0):
+                _AA[i, i] = 1.0
+        _h_inner = np.linalg.solve(_AA, _Ab)
+        _h_Ik[jstart+1:jstart+nlinks] = _h_inner
+    return _h_Ik
+
+# NOTE: This writes to u_ik in place
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_u_ik(_Q_ik, _A_ik, _A_ik_min, _u_ik):
+    n = _u_ik.size
+    for i in range(n):
+        _Q_i = _Q_ik[i]
+        _A_i = _A_ik[i]
+        _A_i_min = _A_ik_min[i]
+        _u_ik[i] = _Q_i / (_A_i + _A_i_min**2 / _A_i)
+    return _u_ik
+
+#@njit(float64[:](float64[:], float64[:], float64[:]),
+#      cache=True)
+#def numba_u_ik(_Q_ik, _A_ik, _u_ik):
+#    n = _u_ik.size
+#    for i in range(n):
+#        _Q_i = _Q_ik[i]
+#        _A_i = _A_ik[i]
+#        if _A_i > SMALLEST_NORMAL:
+#            _u_ik[i] = _Q_i / _A_i
+#        else:
+#            _u_ik[i] = 0.
+#    return _u_ik
+
+# NOTE: This writes to u_Ik in place
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]),
+      cache=True)
+def numba_u_Ik(_dx_ik, _u_ik, _dx_uk, _u_uk, _link_start, _ki, _u_Ik):
+    n = _u_Ik.size
+    for i in range(n):
+        k = _ki[i]
+        if _link_start[i]:
+            num = _dx_ik[i] * _u_uk[k] + _dx_uk[k] * _u_ik[i]
+            den = _dx_ik[i] + _dx_uk[k]
+            if den > SMALLEST_NORMAL:
+                _u_Ik[i] = num / den
+            else:
+                _u_Ik[i] = 0.
+        else:
+            im1 = i - 1
+            num = _dx_ik[i] * _u_ik[im1] + _dx_ik[im1] * _u_ik[i]
+            den = _dx_ik[i] + _dx_ik[im1]
+            if den > SMALLEST_NORMAL:
+                _u_Ik[i] = num / den
+            else:
+                _u_Ik[i] = 0.
+    return _u_Ik
+
+# NOTE: This writes to u_Ip1k in place
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]),
+      cache=True)
+def numba_u_Ip1k(_dx_ik, _u_ik, _dx_dk, _u_dk, _link_end, _ki, _u_Ip1k):
+    n = _u_Ip1k.size
+    for i in range(n):
+        k = _ki[i]
+        if _link_end[i]:
+            num = _dx_ik[i] * _u_dk[k] + _dx_dk[k] * _u_ik[i]
+            den = _dx_ik[i] + _dx_dk[k]
+            if den > SMALLEST_NORMAL:
+                _u_Ip1k[i] = num / den
+            else:
+                _u_Ip1k[i] = 0.
+        else:
+            ip1 = i + 1
+            num = _dx_ik[i] * _u_ik[ip1] + _dx_ik[ip1] * _u_ik[i]
+            den = _dx_ik[i] + _dx_ik[ip1]
+            if den > SMALLEST_NORMAL:
+                _u_Ip1k[i] = num / den
+            else:
+                _u_Ip1k[i] = 0.
+    return _u_Ip1k
+
+#@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]),
+#      cache=True)
+#def numba_u_Ik(_dx_ik, _u_ik, _dx_uk, _u_uk, _link_start, _ki, _u_Ik):
+#    n = _u_Ik.size
+#    for i in range(n):
+#        k = _ki[i]
+#        if _link_start[i]:
+#            _u_Ik[i] = _u_uk[k]
+#        else:
+#            im1 = i - 1
+#            _u_Ik[i] = _u_ik[im1]
+#    return _u_Ik
+
+#@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]),
+#      cache=True)
+#def numba_u_Ip1k(_dx_ik, _u_ik, _dx_dk, _u_dk, _link_end, _ki, _u_Ip1k):
+#    n = _u_Ip1k.size
+#    for i in range(n):
+#        k = _ki[i]
+#        if _link_end[i]:
+#            _u_Ip1k[i] = _u_dk[k]
+#        else:
+#            ip1 = i + 1
+#            _u_Ip1k[i] = _u_ik[ip1]
+#    return _u_Ip1k
+
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def U_1k(E_2k, c_1k, A_1k, T_1k, g=9.81):
+    """
+    Compute forward recurrence coefficient 'U' for node 1, superlink k.
+    """
+    num = E_2k * c_1k - g * A_1k
+    den = T_1k
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64),
+      cache=True)
+def V_1k(P_1k, D_2k, c_1k, T_1k, a_1k=0.0, D_1k=0.0):
+    """
+    Compute forward recurrence coefficient 'V' for node 1, superlink k.
+    """
+    num = P_1k - D_2k * c_1k + D_1k * a_1k
+    den = T_1k
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def W_1k(A_1k, T_1k, a_1k=0.0, E_1k=0.0, g=9.81):
+    """
+    Compute forward recurrence coefficient 'W' for node 1, superlink k.
+    """
+    num = g * A_1k - E_1k * a_1k
+    den = T_1k
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64),
+      cache=True)
+def T_1k(a_1k, b_1k, c_1k):
+    """
+    Compute forward recurrence coefficient 'T' for link 1, superlink k.
+    """
+    return a_1k + b_1k + c_1k
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def U_Ik(E_Ip1k, c_ik, A_ik, T_ik, g=9.81):
+    """
+    Compute forward recurrence coefficient 'U' for node I, superlink k.
+    """
+    num = E_Ip1k * c_ik - g * A_ik
+    den = T_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def V_Ik(P_ik, a_ik, D_Ik, D_Ip1k, c_ik, A_ik, E_Ik, V_Im1k, U_Im1k, T_ik, g=9.81):
+    """
+    Compute forward recurrence coefficient 'V' for node I, superlink k.
+    """
+    t_0 = P_ik
+    t_1 = a_ik * D_Ik
+    t_2 = D_Ip1k * c_ik
+    t_3 = (g * A_ik - E_Ik * a_ik)
+    t_4 = V_Im1k + D_Ik
+    t_5 = U_Im1k - E_Ik
+    t_6 = T_ik
+    # TODO: There is still a divide by zero here
+    num = (t_0 + t_1 - t_2 - (t_3 * t_4 / t_5))
+    den = t_6
+    result = safe_divide(num, den)
+    return  result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def W_Ik(A_ik, E_Ik, a_ik, W_Im1k, U_Im1k, T_ik, g=9.81):
+    """
+    Compute forward recurrence coefficient 'W' for node I, superlink k.
+    """
+    num = -(g * A_ik - E_Ik * a_ik) * W_Im1k
+    den = (U_Im1k - E_Ik) * T_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def T_ik(a_ik, b_ik, c_ik, A_ik, E_Ik, U_Im1k, g=9.81):
+    """
+    Compute forward recurrence coefficient 'T' for link i, superlink k.
+    """
+    t_0 = a_ik + b_ik + c_ik
+    t_1 = g * A_ik - E_Ik * a_ik
+    t_2 = U_Im1k - E_Ik
+    result = t_0 - safe_divide(t_1, t_2)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def X_Nk(A_nk, E_Nk, a_nk, O_nk, g=9.81):
+    """
+    Compute backward recurrence coefficient 'X' for node N, superlink k.
+    """
+    num = g * A_nk - E_Nk * a_nk
+    den = O_nk
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64),
+      cache=True)
+def Y_Nk(P_nk, D_Nk, a_nk, O_nk, c_nk=0.0, D_Np1k=0.0):
+    """
+    Compute backward recurrence coefficient 'Y' for node N, superlink k.
+    """
+    num = P_nk + D_Nk * a_nk - D_Np1k * c_nk
+    den = O_nk
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def Z_Nk(A_nk, O_nk, c_nk=0.0, E_Np1k=0.0, g=9.81):
+    """
+    Compute backward recurrence coefficient 'Z' for node N, superlink k.
+    """
+    num = E_Np1k * c_nk - g * A_nk
+    den = O_nk
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64),
+      cache=True)
+def O_nk(a_nk, b_nk, c_nk):
+    """
+    Compute backward recurrence coefficient 'O' for link n, superlink k.
+    """
+    return a_nk + b_nk + c_nk
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def X_Ik(A_ik, E_Ik, a_ik, O_ik, g=9.81):
+    """
+    Compute backward recurrence coefficient 'X' for node I, superlink k.
+    """
+    num = g * A_ik - E_Ik * a_ik
+    den = O_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def Y_Ik(P_ik, a_ik, D_Ik, D_Ip1k, c_ik, A_ik, E_Ip1k, Y_Ip1k, X_Ip1k, O_ik, g=9.81):
+    """
+    Compute backward recurrence coefficient 'Y' for node I, superlink k.
+    """
+    t_0 = P_ik
+    t_1 = a_ik * D_Ik
+    t_2 = D_Ip1k * c_ik
+    t_3 = (g * A_ik - E_Ip1k * c_ik)
+    t_4 = D_Ip1k - Y_Ip1k
+    t_5 = X_Ip1k + E_Ip1k
+    t_6 = O_ik
+    # TODO: There is still a divide by zero here
+    num = (t_0 + t_1 - t_2 - (t_3 * t_4 / t_5))
+    den = t_6
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def Z_Ik(A_ik, E_Ip1k, c_ik, Z_Ip1k, X_Ip1k, O_ik, g=9.81):
+    """
+    Compute backward recurrence coefficient 'Z' for node I, superlink k.
+    """
+    num = (g * A_ik - E_Ip1k * c_ik) * Z_Ip1k
+    den = (X_Ip1k + E_Ip1k) * O_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def O_ik(a_ik, b_ik, c_ik, A_ik, E_Ip1k, X_Ip1k, g=9.81):
+    """
+    Compute backward recurrence coefficient 'O' for link i, superlink k.
+    """
+    t_0 = a_ik + b_ik + c_ik
+    t_1 = g * A_ik - E_Ip1k * c_ik
+    t_2 = X_Ip1k + E_Ip1k
+    result = t_0 + safe_divide(t_1, t_2)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_D_k_star_old(X_1k, kappa_uk, U_Nk, kappa_dk, Z_1k, W_Nk):
+    """
+    Compute superlink boundary condition coefficient 'D_k_star'.
+    """
+    t_0 = (X_1k * kappa_uk - 1) * (U_Nk * kappa_dk - 1)
+    t_1 = (Z_1k * kappa_dk) * (W_Nk * kappa_uk)
+    result = t_0 - t_1
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_alpha_uk_old(U_Nk, kappa_dk, X_1k, Z_1k, W_Nk, D_k_star, lambda_uk):
+    """
+    Compute superlink boundary condition coefficient 'alpha' for upstream end
+    of superlink k.
+    """
+    num = lambda_uk * ((1 - U_Nk * kappa_dk) * X_1k + (Z_1k * kappa_dk * W_Nk))
+    den = D_k_star
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_beta_uk_old(U_Nk, kappa_dk, Z_1k, W_Nk, D_k_star, lambda_dk):
+    """
+    Compute superlink boundary condition coefficient 'beta' for upstream end
+    of superlink k.
+    """
+    num = lambda_dk * ((1 - U_Nk * kappa_dk) * Z_1k + (Z_1k * kappa_dk * U_Nk))
+    den = D_k_star
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+                 float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_chi_uk_old(U_Nk, kappa_dk, Y_1k, X_1k, mu_uk, Z_1k,
+                 mu_dk, V_Nk, W_Nk, D_k_star):
+    """
+    Compute superlink boundary condition coefficient 'chi' for upstream end
+    of superlink k.
+    """
+    t_0 = (1 - U_Nk * kappa_dk) * (Y_1k + X_1k * mu_uk + Z_1k * mu_dk)
+    t_1 = (Z_1k * kappa_dk) * (V_Nk + W_Nk * mu_uk + U_Nk * mu_dk)
+    num = t_0 + t_1
+    den = D_k_star
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_alpha_dk_old(X_1k, kappa_uk, W_Nk, D_k_star, lambda_uk):
+    """
+    Compute superlink boundary condition coefficient 'alpha' for downstream end
+    of superlink k.
+    """
+    num = lambda_uk * ((1 - X_1k * kappa_uk) * W_Nk + (W_Nk * kappa_uk * X_1k))
+    den = D_k_star
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_beta_dk_old(X_1k, kappa_uk, U_Nk, W_Nk, Z_1k, D_k_star, lambda_dk):
+    """
+    Compute superlink boundary condition coefficient 'beta' for downstream end
+    of superlink k.
+    """
+    num = lambda_dk * ((1 - X_1k * kappa_uk) * U_Nk + (W_Nk * kappa_uk * Z_1k))
+    den = D_k_star
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+                 float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_chi_dk_old(X_1k, kappa_uk, V_Nk, W_Nk, mu_uk, U_Nk,
+                    mu_dk, Y_1k, Z_1k, D_k_star):
+    """
+    Compute superlink boundary condition coefficient 'chi' for downstream end
+    of superlink k.
+    """
+    t_0 = (1 - X_1k * kappa_uk) * (V_Nk + W_Nk * mu_uk + U_Nk * mu_dk)
+    t_1 = (W_Nk * kappa_uk) * (Y_1k + X_1k * mu_uk + Z_1k * mu_dk)
+    num = t_0 + t_1
+    den = D_k_star
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64),
+      cache=True)
+def gamma_o(Q_o_t, Ao, Co, g=9.81):
+    """
+    Compute flow coefficient 'gamma' for orifice o.
+    """
+    num = 2 * g * Co**2 * Ao**2
+    den = np.abs(Q_o_t)
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def gamma_w(Q_w_t, H_w_t, L_w, s_w, Cwr, Cwt):
+    """
+    Compute flow coefficient 'gamma' for weir w.
+    """
+    num = (Cwr * L_w * H_w_t + Cwt * s_w * H_w_t**2)**2
+    den = np.abs(Q_w_t)
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def gamma_p(Q_p_t, b_p, c_p, u):
+    """
+    Compute flow coefficient 'gamma' for pump p.
+    """
+    num = u
+    den = b_p * np.abs(Q_p_t)**(c_p - 1)
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:]),
+      cache=True)
+def xi_uk(dx_uk, B_uk, theta_uk):
+    result = (dx_uk * B_uk * theta_uk) / 2
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:]),
+      cache=True)
+def xi_dk(dx_dk, B_dk, theta_dk):
+    result = (dx_dk * B_dk * theta_dk) / 2
+    return result
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:],
+             int64[:], int64[:]),
+      cache=True)
+def numba_orifice_flow_coefficients(_alpha_o, _beta_o, _chi_o, H_j, _Qo, u, _z_inv_j,
+                                    _z_o, _tau_o, _Co, _Ao, _y_max_o, _unidir_o,
+                                     _J_uo, _J_do):
+    g = 9.81
+    _H_uo = H_j[_J_uo]
+    _H_do = H_j[_J_do]
+    _z_inv_uo = _z_inv_j[_J_uo]
+    # Create indicator functions
+    _omega_o = np.zeros_like(_H_uo)
+    _omega_o[_H_uo >= _H_do] = 1.0
+    # Compute universal coefficients
+    _gamma_o = gamma_o(_Qo, _Ao, _Co, g)
+    # Create conditionals
+    cond_0 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u))
+    cond_1 = ((1 - _omega_o) * _H_uo + _omega_o * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u / 2))
+    cond_2 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo)
+    cond_3 = (_H_do >= _H_uo) & _unidir_o
+    # Fill coefficient arrays
+    # Submerged on both sides
+    a = (cond_0 & cond_1)
+    _alpha_o[a] = _gamma_o[a]
+    _beta_o[a] = -_gamma_o[a]
+    _chi_o[a] = 0.0
+    # Submerged on one side
+    b = (cond_0 & ~cond_1)
+    _alpha_o[b] = _gamma_o[b] * _omega_o[b] * (-1)**(1 - _omega_o[b])
+    _beta_o[b] = _gamma_o[b] * (1 - _omega_o[b]) * (-1)**(1 - _omega_o[b])
+    _chi_o[b] = (_gamma_o[b] * (-1)**(1 - _omega_o[b])
+                                    * (- _z_inv_uo[b] - _z_o[b] -
+                                        _tau_o[b] * _y_max_o[b] * u[b] / 2))
+    # Weir flow
+    c = (~cond_0 & cond_2)
+    _alpha_o[c] = _gamma_o[c] * _omega_o[c] * (-1)**(1 - _omega_o[c])
+    _beta_o[c] = _gamma_o[c] * (1 - _omega_o[c]) * (-1)**(1 - _omega_o[c])
+    _chi_o[c] = (_gamma_o[c] * (-1)**(1 - _omega_o[c])
+                                    * (- _z_inv_uo[c] - _z_o[c]))
+    # No flow
+    d = (~cond_0 & ~cond_2) | cond_3
+    _alpha_o[d] = 0.0
+    _beta_o[d] = 0.0
+    _chi_o[d] = 0.0
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], boolean[:],
+            int64[:], int64[:], float64),
+      cache=True)
+def numba_solve_orifice_flows(H_j, u, _z_inv_j, _z_o,
+                              _tau_o, _y_max_o, _Co, _Ao, _unidir_o, _J_uo, _J_do, g=9.81):
+    # Specify orifice heads at previous timestep
+    _H_uo = H_j[_J_uo]
+    _H_do = H_j[_J_do]
+    _z_inv_uo = _z_inv_j[_J_uo]
+    # Create indicator functions
+    _omega_o = np.zeros_like(_H_uo)
+    _omega_o[_H_uo >= _H_do] = 1.0
+    # Create arrays to store flow coefficients for current time step
+    _alpha_o = np.zeros_like(_H_uo)
+    _beta_o = np.zeros_like(_H_uo)
+    _chi_o = np.zeros_like(_H_uo)
+    # Compute universal coefficients
+    _gamma_o = 2 * g * _Co**2 * _Ao**2
+    # Create conditionals
+    cond_0 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u))
+    cond_1 = ((1 - _omega_o) * _H_uo + _omega_o * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u / 2))
+    cond_2 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo)
+    cond_3 = (_H_do >= _H_uo) & _unidir_o
+    # Fill coefficient arrays
+    # Submerged on both sides
+    a = (cond_0 & cond_1)
+    _alpha_o[a] = _gamma_o[a]
+    _beta_o[a] = -_gamma_o[a]
+    _chi_o[a] = 0.0
+    # Submerged on one side
+    b = (cond_0 & ~cond_1)
+    _alpha_o[b] = _gamma_o[b] * _omega_o[b] * (-1)**(1 - _omega_o[b])
+    _beta_o[b] = _gamma_o[b] * (1 - _omega_o[b]) * (-1)**(1 - _omega_o[b])
+    _chi_o[b] = (_gamma_o[b] * (-1)**(1 - _omega_o[b])
+                                    * (- _z_inv_uo[b] - _z_o[b]
+                                        - _tau_o[b] * _y_max_o[b] * u[b] / 2))
+    # Weir flow on one side
+    c = (~cond_0 & cond_2)
+    _alpha_o[c] = _gamma_o[c] * _omega_o[c] * (-1)**(1 - _omega_o[c])
+    _beta_o[c] = _gamma_o[c] * (1 - _omega_o[c]) * (-1)**(1 - _omega_o[c])
+    _chi_o[c] = (_gamma_o[c] * (-1)**(1 - _omega_o[c])
+                                    * (- _z_inv_uo[c] - _z_o[c]))
+    # No flow
+    d = (~cond_0 & ~cond_2) | cond_3
+    _alpha_o[d] = 0.0
+    _beta_o[d] = 0.0
+    _chi_o[d] = 0.0
+    # Compute flow
+    _Qo_next = (-1)**(1 - _omega_o) * np.sqrt(np.abs(
+               _alpha_o * _H_uo + _beta_o * _H_do + _chi_o))
+    # Export instance variables
+    return _Qo_next
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:]),
+      cache=True)
+def numba_weir_flow_coefficients(_Hw, _Qw, _alpha_w, _beta_w, _chi_w, H_j, _z_inv_j, _z_w,
+                                 _y_max_w, u, _L_w, _s_w, _Cwr, _Cwt, _J_uw, _J_dw):
+    # Specify weir heads at previous timestep
+    _H_uw = H_j[_J_uw]
+    _H_dw = H_j[_J_dw]
+    _z_inv_uw = _z_inv_j[_J_uw]
+    # Create indicator functions
+    _omega_w = np.zeros(_H_uw.size)
+    _omega_w[_H_uw >= _H_dw] = 1.0
+    # Create conditionals
+    cond_0 = (_omega_w * _H_uw + (1 - _omega_w) * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    cond_1 = ((1 - _omega_w) * _H_uw + _omega_w * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    # Effective heads
+    a = (cond_0 & cond_1)
+    b = (cond_0 & ~cond_1)
+    c = (~cond_0)
+    _Hw[a] = _H_uw[a] - _H_dw[a]
+    _Hw[b] = (_omega_w[b] * _H_uw[b] + (1 - _omega_w[b]) * _H_dw[b]
+                    + (-_z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
+    _Hw[c] = 0.0
+    _Hw = np.abs(_Hw)
+    # Compute universal coefficients
+    _gamma_w = gamma_w(_Qw, _Hw, _L_w, _s_w, _Cwr, _Cwt)
+    # Fill coefficient arrays
+    # Submerged on both sides
+    a = (cond_0 & cond_1)
+    _alpha_w[a] = _gamma_w[a]
+    _beta_w[a] = -_gamma_w[a]
+    _chi_w[a] = 0.0
+    # Submerged on one side
+    b = (cond_0 & ~cond_1)
+    _alpha_w[b] = _gamma_w[b] * _omega_w[b] * (-1)**(1 - _omega_w[b])
+    _beta_w[b] = _gamma_w[b] * (1 - _omega_w[b]) * (-1)**(1 - _omega_w[b])
+    _chi_w[b] = (_gamma_w[b] * (-1)**(1 - _omega_w[b]) *
+                                (- _z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
+    # No flow
+    c = (~cond_0)
+    _alpha_w[c] = 0.0
+    _beta_w[c] = 0.0
+    _chi_w[c] = 0.0
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64[:], int64[:]),
+      cache=True)
+def numba_solve_weir_flows(_Hw, _Qw, H_j, _z_inv_j, _z_w, _y_max_w, u, _L_w,
+                           _s_w, _Cwr, _Cwt, _J_uw, _J_dw):
+    _H_uw = H_j[_J_uw]
+    _H_dw = H_j[_J_dw]
+    _z_inv_uw = _z_inv_j[_J_uw]
+    # Create indicator functions
+    _omega_w = np.zeros(_H_uw.size)
+    _omega_w[_H_uw >= _H_dw] = 1.0
+    # Create conditionals
+    cond_0 = (_omega_w * _H_uw + (1 - _omega_w) * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    cond_1 = ((1 - _omega_w) * _H_uw + _omega_w * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    # TODO: Is this being recalculated for a reason?
+    # Effective heads
+    a = (cond_0 & cond_1)
+    b = (cond_0 & ~cond_1)
+    c = (~cond_0)
+    _Hw[a] = _H_uw[a] - _H_dw[a]
+    _Hw[b] = (_omega_w[b] * _H_uw[b] + (1 - _omega_w[b]) * _H_dw[b]
+                    + (-_z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
+    _Hw[c] = 0.0
+    _Hw = np.abs(_Hw)
+    # Compute universal coefficient
+    _gamma_ww = (_Cwr * _L_w * _Hw + _Cwt * _s_w * _Hw**2)**2
+    # Compute flow
+    _Qw_next = (-1)**(1 - _omega_w) * np.sqrt(_gamma_ww * _Hw)
+    return _Qw_next
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:]),
+      cache=True)
+def numba_pump_flow_coefficients(_alpha_p, _beta_p, _chi_p, H_j, _z_inv_j, _Qp, u,
+                                 _z_p, _dHp_max, _dHp_min, _a_p, _b_p, _c_p,
+                                 _J_up, _J_dp):
+    # Get upstream and downstream heads and invert elevation
+    _H_up = H_j[_J_up]
+    _H_dp = H_j[_J_dp]
+    _z_inv_up = _z_inv_j[_J_up]
+    # Compute effective head
+    _dHp = _H_dp - _H_up
+    # Condition 0: Upstream head is above inlet height
+    cond_0 = _H_up > _z_inv_up + _z_p
+    # Condition 1: Head difference is within range of pump curve
+    cond_1 = (_dHp > _dHp_min) & (_dHp < _dHp_max)
+    _dHp[_dHp > _dHp_max] = _dHp_max[_dHp > _dHp_max]
+    _dHp[_dHp < _dHp_min] = _dHp_min[_dHp < _dHp_min]
+    # Compute universal coefficients
+    _gamma_p = gamma_p(_Qp, _b_p, _c_p, u)
+    # Fill coefficient arrays
+    # Head in pump curve range
+    a = (cond_0 & cond_1)
+    _alpha_p[a] = _gamma_p[a]
+    _beta_p[a] = -_gamma_p[a]
+    _chi_p[a] = _gamma_p[a] * _a_p[a]
+    # Head outside of pump curve range
+    b = (cond_0 & ~cond_1)
+    _alpha_p[b] = 0.0
+    _beta_p[b] = 0.0
+    _chi_p[b] = _gamma_p[b] * (_a_p[b] - _dHp[b])
+    # Depth below inlet
+    c = (~cond_0)
+    _alpha_p[c] = 0.0
+    _beta_p[c] = 0.0
+    _chi_p[c] = 0.0
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+                 float64[:], float64[:], float64[:], int64[:], int64[:]),
+      cache=True)
+def numba_solve_pump_flows(H_j, u, _z_inv_j, _z_p, _dHp_max, _dHp_min, _a_p, _b_p, _c_p,
+                           _J_up, _J_dp):
+    _H_up = H_j[_J_up]
+    _H_dp = H_j[_J_dp]
+    _z_inv_up = _z_inv_j[_J_up]
+    # Create conditionals
+    _dHp = _H_dp - _H_up
+    _dHp[_dHp > _dHp_max] = _dHp_max[_dHp > _dHp_max]
+    _dHp[_dHp < _dHp_min] = _dHp_min[_dHp < _dHp_min]
+    cond_0 = _H_up > _z_inv_up + _z_p
+    # Compute universal coefficients
+    _Qp_next = (u / _b_p * (_a_p - _dHp))**(1 / _c_p)
+    _Qp_next[~cond_0] = 0.0
+    return _Qp_next
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], int64, int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_forward_recurrence(_T_ik, _U_Ik, _V_Ik, _W_Ik, _a_ik, _b_ik, _c_ik,
+                             _P_ik, _A_ik, _E_Ik, _D_Ik, NK, nk, _I_1k, _i_1k):
+    g = 9.81
+    for k in range(NK):
+        # Start at junction 1
+        _I_1 = _I_1k[k]
+        _i_1 = _i_1k[k]
+        _I_2 = _I_1 + 1
+        _i_2 = _i_1 + 1
+        nlinks = nk[k]
+        _T_ik[_i_1] = T_1k(_a_ik[_i_1], _b_ik[_i_1], _c_ik[_i_1])
+        _U_Ik[_I_1] = U_1k(_E_Ik[_I_2], _c_ik[_i_1], _A_ik[_i_1], _T_ik[_i_1], g)
+        _V_Ik[_I_1] = V_1k(_P_ik[_i_1], _D_Ik[_I_2], _c_ik[_i_1], _T_ik[_i_1],
+                            _a_ik[_i_1], _D_Ik[_I_1])
+        _W_Ik[_I_1] = W_1k(_A_ik[_i_1], _T_ik[_i_1], _a_ik[_i_1], _E_Ik[_I_1], g)
+        # Loop from junction 2 -> Nk
+        for i in range(nlinks - 1):
+            _i_next = _i_2 + i
+            _I_next = _I_2 + i
+            _Im1_next = _I_next - 1
+            _Ip1_next = _I_next + 1
+            _T_ik[_i_next] = T_ik(_a_ik[_i_next], _b_ik[_i_next], _c_ik[_i_next],
+                                  _A_ik[_i_next], _E_Ik[_I_next], _U_Ik[_Im1_next], g)
+            _U_Ik[_I_next] = U_Ik(_E_Ik[_Ip1_next], _c_ik[_i_next],
+                                  _A_ik[_i_next], _T_ik[_i_next], g)
+            _V_Ik[_I_next] = V_Ik(_P_ik[_i_next], _a_ik[_i_next], _D_Ik[_I_next],
+                                  _D_Ik[_Ip1_next], _c_ik[_i_next], _A_ik[_i_next],
+                                  _E_Ik[_I_next], _V_Ik[_Im1_next], _U_Ik[_Im1_next],
+                                  _T_ik[_i_next], g)
+            _W_Ik[_I_next] = W_Ik(_A_ik[_i_next], _E_Ik[_I_next], _a_ik[_i_next],
+                                  _W_Ik[_Im1_next], _U_Ik[_Im1_next], _T_ik[_i_next], g)
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], int64, int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_backward_recurrence(_O_ik, _X_Ik, _Y_Ik, _Z_Ik, _a_ik, _b_ik, _c_ik,
+                              _P_ik, _A_ik, _E_Ik, _D_Ik, NK, nk, _I_Nk, _i_nk):
+    g = 9.81
+    for k in range(NK):
+        _I_N = _I_Nk[k]
+        _i_n = _i_nk[k]
+        _I_Nm1 = _I_N - 1
+        _i_nm1 = _i_n - 1
+        _I_Np1 = _I_N + 1
+        nlinks = nk[k]
+        _O_ik[_i_n] = O_nk(_a_ik[_i_n], _b_ik[_i_n], _c_ik[_i_n])
+        _X_Ik[_I_N] = X_Nk(_A_ik[_i_n], _E_Ik[_I_N], _a_ik[_i_n], _O_ik[_i_n], g)
+        _Y_Ik[_I_N] = Y_Nk(_P_ik[_i_n], _D_Ik[_I_N], _a_ik[_i_n], _O_ik[_i_n],
+                            _c_ik[_i_n], _D_Ik[_I_Np1])
+        _Z_Ik[_I_N] = Z_Nk(_A_ik[_i_n], _O_ik[_i_n], _c_ik[_i_n], _E_Ik[_I_Np1], g)
+        for i in range(nlinks - 1):
+            _i_next = _i_nm1 - i
+            _I_next = _I_Nm1 - i
+            _Ip1_next = _I_next + 1
+            _O_ik[_i_next] = O_ik(_a_ik[_i_next], _b_ik[_i_next], _c_ik[_i_next],
+                                  _A_ik[_i_next], _E_Ik[_Ip1_next], _X_Ik[_Ip1_next], g)
+            _X_Ik[_I_next] = X_Ik(_A_ik[_i_next], _E_Ik[_I_next], _a_ik[_i_next],
+                                  _O_ik[_i_next], g)
+            _Y_Ik[_I_next] = Y_Ik(_P_ik[_i_next], _a_ik[_i_next], _D_Ik[_I_next],
+                                  _D_Ik[_Ip1_next], _c_ik[_i_next], _A_ik[_i_next],
+                                  _E_Ik[_Ip1_next], _Y_Ik[_Ip1_next], _X_Ik[_Ip1_next],
+                                  _O_ik[_i_next], g)
+            _Z_Ik[_I_next] = Z_Ik(_A_ik[_i_next], _E_Ik[_Ip1_next], _c_ik[_i_next],
+                                  _Z_Ik[_Ip1_next], _X_Ik[_Ip1_next], _O_ik[_i_next], g)
+    return 1
+
+@njit(float64[:,:](float64[:,:], int64, int64),
+      cache=True)
+def numba_create_banded(l, bandwidth, M):
+    AB = np.zeros((2*bandwidth + 1, M))
+    for i in range(M):
+        AB[bandwidth, i] = l[i, i]
+    for n in range(bandwidth):
+        for j in range(M - n - 1):
+            AB[bandwidth - n - 1, -j - 1] = l[-j - 2 - n, -j - 1]
+            AB[bandwidth + n + 1, j] = l[j + n + 1, j]
+    return AB
+
+@njit(void(float64[:], int64[:], float64[:]),
+      cache=True)
+def numba_add_at(a, indices, b):
+    n = len(indices)
+    for k in range(n):
+        i = indices[k]
+        a[i] += b[k]
+
+@njit(void(float64[:, :], boolean[:], int64[:], int64[:], int64),
+      cache=True)
+def numba_clear_off_diagonals(A, bc, _J_uk, _J_dk, NK):
+    for k in range(NK):
+        _J_u = _J_uk[k]
+        _J_d = _J_dk[k]
+        _bc_u = bc[_J_u]
+        _bc_d = bc[_J_d]
+        if not _bc_u:
+            A[_J_u, _J_d] = 0.0
+        if not _bc_d:
+            A[_J_d, _J_u] = 0.0
+
+def numba_B_j(_B_j, _J_uk, _J_dk, _xi_uk, _xi_dk, _A_sj):
+    numba_add_at(_B_j, _J_uk, _xi_uk)
+    numba_add_at(_B_j, _J_dk, _xi_dk)
+    _B_j += _A_sj
+    return _B_j
+
+@njit(void(float64[:, :], boolean[:], float64[:], float64, int64),
+      cache=True)
+def numba_create_J_matrix(J, bc, _A_sj, _dt, M):
+    for i in range(M):
+        if bc[i]:
+            J[i,i] = 1.0
+        else:
+            J[i,i] = (_A_sj[i] / _dt)
+
+@njit(void(float64[:, :], float64[:], boolean[:], int64[:], int64[:], float64[:],
+           float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+           float64, int64, int64),
+      cache=True)
+def numba_create_K_matrix(K, _F_jj, bc, _J_uk, _J_dk, _alpha_uk,
+                          _alpha_dk, _beta_uk, _beta_dk, _xi_uk, _xi_dk,
+                          _A_sj, _dt, M, NK):
+    numba_add_at(_F_jj, _J_uk, _alpha_uk)
+    numba_add_at(_F_jj, _J_dk, -_beta_dk)
+    numba_add_at(_F_jj, _J_uk, _xi_uk / _dt)
+    numba_add_at(_F_jj, _J_dk, _xi_dk / _dt)
+    # Set diagonal of K matrix
+    for i in range(M):
+        if bc[i]:
+            K[i,i] = 0.0
+        else:
+            K[i,i] = _F_jj[i]
+    for k in range(NK):
+        _J_u = _J_uk[k]
+        _J_d = _J_dk[k]
+        _bc_u = bc[_J_u]
+        _bc_d = bc[_J_d]
+        if not _bc_u:
+            K[_J_u, _J_d] += _beta_uk[k]
+        if not _bc_d:
+            K[_J_d, _J_u] -= _alpha_dk[k]
+
+@njit(void(float64[:, :], float64[:], boolean[:], int64[:], int64[:], float64[:],
+           float64[:], float64[:], float64[:], int64, int64),
+      cache=True)
+def numba_create_OWP_matrix(X, diag, bc, _J_uc, _J_dc, _alpha_uc,
+                            _alpha_dc, _beta_uc, _beta_dc, M, NC):
+    # Set diagonal
+    numba_add_at(diag, _J_uc, _alpha_uc)
+    numba_add_at(diag, _J_dc, -_beta_dc)
+    for i in range(M):
+        if bc[i]:
+            X[i,i] = 0.0
+        else:
+            X[i,i] = diag[i]
+    # Set off-diagonal
+    for c in range(NC):
+        _J_u = _J_uc[c]
+        _J_d = _J_dc[c]
+        _bc_u = bc[_J_u]
+        _bc_d = bc[_J_d]
+        if not _bc_u:
+            X[_J_u, _J_d] += _beta_uc[c]
+        if not _bc_d:
+            X[_J_d, _J_u] -= _alpha_dc[c]
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64),
+      cache=True)
+def numba_Q_i_next_b(X_Ik, h_Ik, Y_Ik, Z_Ik, h_Np1k, _Ik, _ki, n):
+    _Q_i = np.zeros(n)
+    for i in range(n):
+        I = _Ik[i]
+        k = _ki[i]
+        t_0 = X_Ik[I] * h_Ik[I]
+        t_1 = Y_Ik[I]
+        t_2 = Z_Ik[I] * h_Np1k[k]
+        _Q_i[i] = t_0 + t_1 + t_2
+    return _Q_i
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64),
+      cache=True)
+def numba_Q_im1k_next_f(U_Ik, h_Ik, V_Ik, W_Ik, h_1k, _Ik, _ki, n):
+    _Q_i = np.zeros(n)
+    for i in range(n):
+        I = _Ik[i]
+        Ip1 = I + 1
+        k = _ki[i]
+        t_0 = U_Ik[I] * h_Ik[Ip1]
+        t_1 = V_Ik[I]
+        t_2 = W_Ik[I] * h_1k[k]
+        _Q_i[i] = t_0 + t_1 + t_2
+    return _Q_i
+
+@njit(void(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+           float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64[:],
+           int64[:], int64[:], int64, boolean[:]),
+      cache=True)
+def numba_reposition_junctions(_x_Ik, _z_inv_Ik, _h_Ik, _dx_ik, _Q_ik, _H_dk,
+                               _b0, _zc, _xc, _m, _elem_pos, _i_1k, _I_1k,
+                               _I_Np1k, nk, NK, reposition):
+    for k in range(NK):
+        if reposition[k]:
+            _i_1 = _i_1k[k]
+            _I_1 = _I_1k[k]
+            _I_Np1 = _I_Np1k[k]
+            nlinks = nk[k]
+            njunctions = nlinks + 1
+            _i_end = _i_1 + nlinks
+            _I_end = _I_1 + njunctions
+            _H_d = _H_dk[k]
+            _z_inv_1 = _z_inv_Ik[_I_1]
+            _z_inv_Np1 = _z_inv_Ik[_I_Np1]
+            pos_prev = _elem_pos[k]
+            # Junction arrays for superlink k
+            _x_I = _x_Ik[_I_1:_I_end]
+            _z_inv_I = _z_inv_Ik[_I_1:_I_end]
+            _h_I = _h_Ik[_I_1:_I_end]
+            _dx_i = _dx_ik[_i_1:_i_end]
+            # Move junction if downstream head is within range
+            move_junction = (_H_d > _z_inv_Np1) & (_H_d < _z_inv_1)
+            if move_junction:
+                z_m = _H_d
+                _x0 = _x_I[_I_1]
+                x_m = (_H_d - _b0[k]) / _m[k] + _x0
+            else:
+                z_m = _zc[k]
+                x_m = _xc[k]
+            # Determine new x-position of junction
+            c = np.searchsorted(_x_I, x_m)
+            cm1 = c - 1
+            # Compute fractional x-position along superlink k
+            frac = (x_m - _x_I[cm1]) / (_x_I[c] - _x_I[cm1])
+            # Interpolate depth at new position
+            h_m = (1 - frac) * _h_I[cm1] + (frac) * _h_I[c]
+            # Link length ratio
+            r = _dx_i[pos_prev - 1] / (_dx_i[pos_prev - 1]
+                                    + _dx_i[pos_prev])
+            # Set new values
+            _x_I[pos_prev] = x_m
+            _z_inv_I[pos_prev] = z_m
+            _h_I[pos_prev] = h_m
+            Ix = np.argsort(_x_I)
+            _dx_i = np.diff(_x_I[Ix])
+            _x_Ik[_I_1:_I_end] = _x_I[Ix]
+            _z_inv_Ik[_I_1:_I_end] = _z_inv_I[Ix]
+            _h_Ik[_I_1:_I_end] = _h_I[Ix]
+            _dx_ik[_i_1:_i_end] = _dx_i
+            # Set position to new position
+            pos_change = np.argsort(Ix)
+            pos_next = pos_change[pos_prev]
+            _elem_pos[k] = pos_next
+            shifted = (pos_prev != pos_next)
+            # If position has shifted interpolate flow
+            if shifted:
+                ix = np.arange(nlinks)
+                ix[pos_prev] = pos_next
+                ix.sort()
+                _Q_i = _Q_ik[_i_1:_i_end]
+                _Q_i[pos_prev - 1] = (1 - r) * _Q_i[pos_prev - 1] + r * _Q_i[pos_prev]
+                _Q_ik[_i_1:_i_end] = _Q_i[ix]
+
+@njit(float64[:](float64[:], float64[:], float64[:], int64[:], 
+                 int64[:], int64[:], boolean[:], boolean[:]),
+      cache=True)
+def junction_numerator(_dQ_ik, _dQ_uk, _dQ_dk, _kI, 
+                       _forward_I_i, _backward_I_i, _is_start, _is_end):
+    N = _kI.size
+    num = np.zeros(N, dtype=np.float64)
+    for I in range(N):
+        k = _kI[I]
+        if _is_start[I]:
+            i = _forward_I_i[I]
+            num[I] = _dQ_ik[i] - _dQ_uk[k]
+        elif _is_end[I]:
+            im1 = _backward_I_i[I]
+            num[I] = _dQ_dk[k] - _dQ_ik[im1]
+        else:
+            i = _forward_I_i[I]
+            im1 = i - 1
+            num[I] = _dQ_ik[i] - _dQ_ik[im1]
+    return num
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], int64[:],
+                 int64[:], int64[:], boolean[:], boolean[:]),
+      cache=True)
+def junction_denominator(_D_Ik, _Q_ik, _Q_uk, _Q_dk, _kI, 
+                         _forward_I_i, _backward_I_i, _is_start, _is_end):
+    N = _kI.size
+    denom = np.zeros(N, dtype=np.float64)
+    for I in range(N):
+        k = _kI[I]
+        if _is_start[I]:
+            i = _forward_I_i[I]
+            denom[I] = _D_Ik[I] + _Q_uk[k] - _Q_ik[i]
+        elif _is_end[I]:
+            im1 = _backward_I_i[I]
+            denom[I] = _D_Ik[I] + _Q_ik[im1] - _Q_dk[k]
+        else:
+            i = _forward_I_i[I]
+            im1 = i - 1
+            denom[I] = _D_Ik[I] + _Q_ik[im1] - _Q_ik[i]
+    return denom
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+                 int64[:], int64[:], int64[:], int64[:], int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def superjunction_numerator(D_j, _dQ_uk, _dQ_dk, _dQ_o, _dQ_w, _dQ_p, 
+                            _J_uk, _J_dk, _J_uo, _J_do, _J_uw, _J_dw, _J_up, _J_dp):
+    M = D_j.size
+    num = np.zeros(M, dtype=np.float64)
+    numba_add_at(num, _J_uk, _dQ_uk)
+    numba_add_at(num, _J_dk, -_dQ_dk)
+    numba_add_at(num, _J_uo, _dQ_o)
+    numba_add_at(num, _J_do, -_dQ_o)
+    numba_add_at(num, _J_uw, _dQ_w)
+    numba_add_at(num, _J_dw, -_dQ_w)
+    numba_add_at(num, _J_up, _dQ_p)
+    numba_add_at(num, _J_dp, -_dQ_p)
+    return num
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+                 int64[:], int64[:], int64[:], int64[:], int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def superjunction_denominator(D_j, _Q_uk, _Q_dk, _Q_o, _Q_w, _Q_p, 
+                            _J_uk, _J_dk, _J_uo, _J_do, _J_uw, _J_dw, _J_up, _J_dp):
+    M = D_j.size
+    denom = np.zeros(M, dtype=np.float64)
+    denom[:] += D_j
+    numba_add_at(denom, _J_dk, _Q_dk)
+    numba_add_at(denom, _J_uk, -_Q_uk)
+    numba_add_at(denom, _J_do, _Q_o)
+    numba_add_at(denom, _J_uo, -_Q_o)
+    numba_add_at(denom, _J_dw, _Q_w)
+    numba_add_at(denom, _J_uw, -_Q_w)
+    numba_add_at(denom, _J_dp, _Q_p)
+    numba_add_at(denom, _J_up, -_Q_p)
+    return denom
diff --git a/pipedream_solver/_tsuperlink.py b/pipedream_solver/_tsuperlink.py
new file mode 100644
index 0000000..eb6d5ae
--- /dev/null
+++ b/pipedream_solver/_tsuperlink.py
@@ -0,0 +1,1442 @@
+import numpy as np
+from numba import njit, prange
+from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, DictType, void
+import pipedream_solver.ngeometry
+
+SMALLEST_NORMAL = np.finfo(np.float64).smallest_normal
+
+MIN_SJ_AREA = 1e-8
+
+CIRCULAR = 1
+RECT_CLOSED = 2
+RECT_OPEN = 3
+TRIANGULAR = 4
+TRAPEZOIDAL = 5
+PARABOLIC = 6
+ELLIPTICAL = 7
+WIDE = 8
+FORCE_MAIN = 9
+FLOODPLAIN = 10
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_hydraulic_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik,
+                             _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
+                             _geom_codes, _Ik, _ik):
+    n = len(_ik)
+    for i in range(n):
+        I = _Ik[i]
+        Ip1 = I + 1
+        geom_code = _geom_codes[i]
+        h_I = _h_Ik[I]
+        h_Ip1 = _h_Ik[Ip1]
+        h_i = (h_I + h_Ip1) / 2
+        g1_i = _g1_ik[i]
+        g2_i = _g2_ik[i]
+        g3_i = _g3_ik[i]
+        g4_i = _g4_ik[i]
+        g5_i = _g5_ik[i]
+        g6_i = _g6_ik[i]
+        g7_i = _g7_ik[i]
+        if geom_code:
+            if geom_code == CIRCULAR:
+                _A_ik[i] = pipedream_solver.ngeometry.Circular_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Circular_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Circular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == RECT_CLOSED:
+                _A_ik[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Rect_Closed_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Rect_Closed_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == RECT_OPEN:
+                _A_ik[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Rect_Open_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Rect_Open_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRIANGULAR:
+                _A_ik[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Triangular_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Triangular_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Triangular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRAPEZOIDAL:
+                _A_ik[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_i, g1_i, g2_i, g3_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Trapezoidal_Pe_ik(h_i, g1_i, g2_i, g3_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Trapezoidal_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == PARABOLIC:
+                _A_ik[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Parabolic_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Parabolic_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Parabolic_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == ELLIPTICAL:
+                raise NotImplementedError
+            elif geom_code == WIDE:
+                _A_ik[i] = pipedream_solver.ngeometry.Wide_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Wide_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Wide_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Wide_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FORCE_MAIN:
+                _A_ik[i] = pipedream_solver.ngeometry.Force_Main_A_ik(h_i, g1_i, g2_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Force_Main_Pe_ik(h_i, g1_i, g2_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Force_Main_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Force_Main_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FLOODPLAIN:
+                _A_ik[i] = pipedream_solver.ngeometry.Floodplain_A_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+                _Pe_ik[i] = pipedream_solver.ngeometry.Floodplain_Pe_ik(h_i, g1_i, g2_i,
+                                                                        g3_i, g4_i, g5_i, g6_i, g7_i)
+                _R_ik[i] = pipedream_solver.ngeometry.Floodplain_R_ik(_A_ik[i], _Pe_ik[i])
+                _B_ik[i] = pipedream_solver.ngeometry.Floodplain_B_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, _theta_bk,
+                            _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
+                            _geom_codes, _i_bk, _I_bk, _J_bk):
+    n = len(_i_bk)
+    for k in range(n):
+        i = _i_bk[k]
+        I = _I_bk[k]
+        j = _J_bk[k]
+        # TODO: If theta is zero, should return critical depth
+        h_I = _h_Ik[I]
+        h_Ip1 = _theta_bk[k] * (_H_j[j] - _z_inv_bk[k])
+        h_i = (h_I + h_Ip1) / 2
+        geom_code = _geom_codes[i]
+        g1_i = _g1_ik[i]
+        g2_i = _g2_ik[i]
+        g3_i = _g3_ik[i]
+        g4_i = _g4_ik[i]
+        g5_i = _g5_ik[i]
+        g6_i = _g6_ik[i]
+        g7_i = _g7_ik[i]
+        if geom_code:
+            if geom_code == CIRCULAR:
+                _A_bk[k] = pipedream_solver.ngeometry.Circular_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Circular_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Circular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == RECT_CLOSED:
+                _A_bk[k] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Rect_Closed_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Rect_Closed_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == RECT_OPEN:
+                _A_bk[k] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Rect_Open_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Rect_Open_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRIANGULAR:
+                _A_bk[k] = pipedream_solver.ngeometry.Triangular_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Triangular_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Triangular_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Triangular_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == TRAPEZOIDAL:
+                _A_bk[k] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_i, g1_i, g2_i, g3_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Trapezoidal_Pe_ik(h_i, g1_i, g2_i, g3_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Trapezoidal_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_i, g1_i, g2_i, g3_i)
+            elif geom_code == PARABOLIC:
+                _A_bk[k] = pipedream_solver.ngeometry.Parabolic_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Parabolic_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Parabolic_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Parabolic_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == ELLIPTICAL:
+                raise NotImplementedError
+            elif geom_code == WIDE:
+                _A_bk[k] = pipedream_solver.ngeometry.Wide_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Wide_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Wide_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Wide_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FORCE_MAIN:
+                _A_bk[k] = pipedream_solver.ngeometry.Force_Main_A_ik(h_i, g1_i, g2_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Force_Main_Pe_ik(h_i, g1_i, g2_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Force_Main_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Force_Main_B_ik(h_i, g1_i, g2_i)
+            elif geom_code == FLOODPLAIN:
+                _A_bk[k] = pipedream_solver.ngeometry.Floodplain_A_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+                _Pe_bk[k] = pipedream_solver.ngeometry.Floodplain_Pe_ik(h_i, g1_i, g2_i,
+                                                                        g3_i, g4_i, g5_i, g6_i, g7_i)
+                _R_bk[k] = pipedream_solver.ngeometry.Floodplain_R_ik(_A_bk[k], _Pe_bk[k])
+                _B_bk[k] = pipedream_solver.ngeometry.Floodplain_B_ik(h_i, g1_i, g2_i,
+                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64),
+      cache=True)
+def numba_orifice_geometry(_Ao, h_eo, u_o, _g1_o, _g2_o, _g3_o, _geom_codes_o, n_o):
+    for i in range(n_o):
+        geom_code = _geom_codes_o[i]
+        g1 = _g1_o[i]
+        g2 = _g2_o[i]
+        g3 = _g3_o[i]
+        u = u_o[i]
+        h_e = h_eo[i]
+        if geom_code:
+            if geom_code == CIRCULAR:
+                _Ao[i] = pipedream_solver.ngeometry.Circular_A_ik(h_e, g1 * u, g2)
+            elif geom_code == RECT_CLOSED:
+                _Ao[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_e, g1 * u, g2)
+            elif geom_code == RECT_OPEN:
+                _Ao[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_e, g1 * u, g2)
+            elif geom_code == TRIANGULAR:
+                _Ao[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_e, g1 * u, g2)
+            elif geom_code == TRAPEZOIDAL:
+                _Ao[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_e, g1 * u, g2, g3)
+            elif geom_code == PARABOLIC:
+                _Ao[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_e, g1 * u, g2)
+            elif geom_code == ELLIPTICAL:
+                raise NotImplementedError
+            elif geom_code == WIDE:
+                _Ao[i] = pipedream_solver.ngeometry.Wide_A_ik(h_e, g1 * u, g2)
+            elif geom_code == FORCE_MAIN:
+                _Ao[i] = pipedream_solver.ngeometry.Force_Main_A_ik(h_e, g1 * u, g2)
+            elif geom_code == FLOODPLAIN:
+                # TODO: This can be implemented
+                raise NotImplementedError
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_transect_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik, _is_irregular, _transect_zs,
+                            _transect_As, _transect_Bs, _transect_Pes, _transect_Rs,
+                            _transect_codes, _transect_inds, _transect_lens, _Ik, _ik):
+    n = len(_ik)
+    for i in range(n):
+        is_irregular = _is_irregular[i]
+        if is_irregular:
+            I = _Ik[i]
+            Ip1 = I + 1
+            transect_code = _transect_codes[i]
+            h_I = _h_Ik[I]
+            h_Ip1 = _h_Ik[Ip1]
+            h_i = (h_I + h_Ip1) / 2
+            start = _transect_inds[transect_code]
+            size = _transect_lens[transect_code]
+            end = start + size
+            _z_range = _transect_zs[start:end]
+            _A_range = _transect_As[start:end]
+            _B_range = _transect_Bs[start:end]
+            _Pe_range = _transect_Pes[start:end]
+            _R_range = _transect_Rs[start:end]
+            _A_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _A_range, 1)
+            _B_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _B_range, 1)
+            _Pe_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _Pe_range, 1)
+            _R_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _R_range, 1)
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], boolean[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], int64[:], int64[:],
+            int64[:], int64[:], int64[:], int64[:]),
+      cache=True)
+def numba_boundary_transect(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk,
+                            _theta_bk, _is_irregular, _transect_zs, _transect_As, _transect_Bs,
+                            _transect_Pes, _transect_Rs, _transect_codes, _transect_inds,
+                            _transect_lens, _I_bk, _i_bk, _J_bk):
+    n = len(_i_bk)
+    for k in range(n):
+        i = _i_bk[k]
+        I = _I_bk[k]
+        j = _J_bk[k]
+        is_irregular_bk = _is_irregular[i]
+        if is_irregular_bk:
+            h_I = _h_Ik[I]
+            # TODO: This should incorporate theta
+            h_Ip1 = _theta_bk[k] * (_H_j[j] - _z_inv_bk[k])
+            transect_code = _transect_codes[i]
+            h_i = (h_I + h_Ip1) / 2
+            start = _transect_inds[transect_code]
+            size = _transect_lens[transect_code]
+            end = start + size
+            _z_range = _transect_zs[start:end]
+            _A_range = _transect_As[start:end]
+            _B_range = _transect_Bs[start:end]
+            _Pe_range = _transect_Pes[start:end]
+            _R_range = _transect_Rs[start:end]
+            _A_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _A_range, 1)
+            _B_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _B_range, 1)
+            _Pe_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _Pe_range, 1)
+            _R_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _R_range, 1)
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]),
+      cache=True)
+def numba_compute_functional_storage_areas(h, A, a, b, c, _functional):
+    M = h.size
+    for j in range(M):
+        if _functional[j]:
+            if h[j] < 0:
+                A[j] = 0
+            else:
+                A[j] = max(a[j] * (h[j]**b[j]) + c[j], MIN_SJ_AREA)
+    return A
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]),
+      cache=True)
+def numba_compute_functional_storage_volumes(h, V, a, b, c, _functional):
+    M = h.size
+    for j in range(M):
+        if _functional[j]:
+            if h[j] < 0:
+                V[j] = 0
+            else:
+                # TODO: Ensure this works when functional storage area is negative
+                V[j] = (a[j] / (b[j] + 1)) * h[j] ** (b[j] + 1) + c[j] * h[j]
+    return V
+
+@njit
+def numba_compute_tabular_storage_areas(h_j, A_sj, hs, As, sjs, sts, inds, lens):
+    n = sjs.size
+    for i in range(n):
+        sj = sjs[i]
+        st = sts[i]
+        ind = inds[st]
+        size = lens[st]
+        h_range = hs[ind:ind+size]
+        A_range = As[ind:ind+size]
+        Amin = A_range.min()
+        Amax = A_range.max()
+        h_search = h_j[sj]
+        ix = np.searchsorted(h_range, h_search)
+        # NOTE: np.interp not supported in this version of numba
+        # A_result = np.interp(h_search, h_range, A_range)
+        # A_out[i] = A_result
+        if (ix == 0):
+            A_sj[sj] = Amin
+        elif (ix >= size):
+            A_sj[sj] = Amax
+        else:
+            dx_0 = h_search - h_range[ix - 1]
+            dx_1 = h_range[ix] - h_search
+            frac = dx_0 / (dx_0 + dx_1)
+            A_sj[sj] = (1 - frac) * A_range[ix - 1] + (frac) * A_range[ix]
+    return A_sj
+
+@njit
+def numba_compute_tabular_storage_volumes(h_j, V_sj, hs, As, Vs, sjs, sts, inds, lens):
+    n = sjs.size
+    for i in range(n):
+        sj = sjs[i]
+        st = sts[i]
+        ind = inds[st]
+        size = lens[st]
+        h_range = hs[ind:ind+size]
+        A_range = As[ind:ind+size]
+        V_range = Vs[ind:ind+size]
+        hmax = h_range.max()
+        Vmin = V_range.min()
+        Vmax = V_range.max()
+        Amax = A_range.max()
+        h_search = h_j[sj]
+        ix = np.searchsorted(h_range, h_search)
+        # NOTE: np.interp not supported in this version of numba
+        # A_result = np.interp(h_search, h_range, A_range)
+        # A_out[i] = A_result
+        if (ix == 0):
+            V_sj[sj] = Vmin
+        elif (ix >= size):
+            V_sj[sj] = Vmax + Amax * (h_search - hmax)
+        else:
+            dx_0 = h_search - h_range[ix - 1]
+            dx_1 = h_range[ix] - h_search
+            frac = dx_0 / (dx_0 + dx_1)
+            V_sj[sj] = (1 - frac) * V_range[ix - 1] + (frac) * V_range[ix]
+    return V_sj
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64))
+def friction_slope(Q_ik_t, dx_ik, A_ik, R_ik, n_ik, Sf_method_ik, g=9.81):
+    if A_ik > 0:
+        # Chezy-Manning eq.
+        if Sf_method_ik == 0:
+            t_1 = (g * n_ik**2 * np.abs(Q_ik_t) * dx_ik
+                   / A_ik / R_ik**(4/3))
+        # Hazen-Williams eq.
+        elif Sf_method_ik == 1:
+            t_1 = (1.354 * g * np.abs(Q_ik_t)**0.85 * dx_ik
+                   / A_ik**0.85 / n_ik**1.85 / R_ik**1.1655)
+        # Darcy-Weisbach eq.
+        elif Sf_method_ik == 2:
+            # kinematic viscosity(meter^2/sec), we can consider this is constant.
+            nu = 0.0000010034
+            Re = (np.abs(Q_ik_t) / A_ik) * 4 * R_ik / nu
+            if Re > 2000:
+                f = 0.25 / (np.log10(n_ik / (3.7 * 4 * R_ik) + 5.74 / (Re**0.9)))**2
+            elif (Re > 0) and (Re <= 2000):
+                f = 64 / Re
+            elif (Re == 0):
+                f = 0.1
+            else:
+                raise ValueError('Reynolds number outside allowable range')
+            t_1 = (0.01274 * g * f * np.abs(Q_ik_t) * dx_ik
+                   / (A_ik * R_ik))
+        else:
+            raise ValueError('Invalid friction method.')
+        return t_1
+    else:
+        return 0.
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64, float64),
+      cache=True)
+def tnumba_a_ik(u_ik, B_ik, A_ik, dx_ik, dt, g=9.81):
+    t_0 = -u_ik * B_ik * dx_ik / dt
+    t_1 = u_ik**2 * B_ik
+    t_2 = -g * A_ik
+    result = t_0 + t_1 + t_2
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64, float64),
+      cache=True)
+def tnumba_c_ik(u_ik, B_ik, A_ik, dx_ik, dt, g=9.81):
+    t_0 = -u_ik * B_ik * dx_ik / dt
+    t_1 = -u_ik**2 * B_ik
+    t_2 = g * A_ik
+    result = t_0 + t_1 + t_2
+    return result
+
+@njit(float64[:](float64[:], float64, float64[:], float64[:],
+                 float64[:], float64[:], float64[:], int64[:], float64),
+      cache=True)
+def tnumba_b_ik(dx_ik, dt, n_ik, Q_ik_t, 
+               A_ik, R_ik, C_ik, Sf_method_ik, g=9.81):
+    t_0 = dx_ik / dt
+    t_1 = np.zeros(Q_ik_t.size)
+    k = len(Sf_method_ik)
+    for n in range(k):
+        t_1[n] = friction_slope(Q_ik_t[n], dx_ik[n], A_ik[n], R_ik[n],
+                                n_ik[n], Sf_method_ik[n], g)
+    t_2 = C_ik * np.abs(Q_ik_t) / 2 / g
+    result = t_0 + t_1 + t_2
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64, float64[:], float64[:], float64[:],
+                 float64[:], float64[:], float64[:], float64),
+      cache=True)
+def tnumba_P_ik(Q_ik_t, dx_ik, dt, u_ik, B_ik, A_ik, 
+         h_Ik_t, h_Ip1k_t, S_o_ik, g=9.81):
+    t0 = Q_ik_t * dx_ik / dt
+    t1 = - u_ik * B_ik * dx_ik * (h_Ik_t + h_Ip1k_t) / dt
+    t2 = g * A_ik * S_o_ik * dx_ik
+    result = t0 + t1 + t2
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64),
+      cache=True)
+def E_Ik(B_ik, dx_ik, B_im1k, dx_im1k, A_SIk, dt):
+    """
+    Compute node coefficient 'E' for node I, superlink k.
+    """
+    t_0 = B_ik * dx_ik / 2
+    t_1 = B_im1k * dx_im1k / 2
+    t_2 = A_SIk
+    t_3 = dt
+    return (t_0 + t_1 + t_2) / t_3
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def D_Ik(Q_0IK, B_ik, dx_ik, B_im1k, dx_im1k, A_SIk, h_Ik_t, dt):
+    """
+    Compute node coefficient 'D' for node I, superlink k.
+    """
+    t_0 = Q_0IK
+    t_1 = B_ik * dx_ik / 2
+    t_2 = B_im1k * dx_im1k / 2
+    t_3 = A_SIk
+    t_4 = h_Ik_t / dt
+    return t_0 + ((t_1 + t_2 + t_3) * t_4)
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], int64[:], float64,
+            int64[:], int64[:], boolean[:], boolean[:]),
+      cache=True)
+def numba_node_coeffs(_D_Ik, _E_Ik, _Q_0Ik, _B_ik, _h_Ik, _dx_ik, _A_SIk,
+                      _B_uk, _B_dk, _dx_uk, _dx_dk, _kI, _dt,
+                      _forward_I_i, _backward_I_i, _is_start, _is_end):
+    N = _h_Ik.size
+    for I in range(N):
+        k = _kI[I]
+        if _is_start[I]:
+            i = _forward_I_i[I]
+            _E_Ik[I] = E_Ik(_B_ik[i], _dx_ik[i], _B_uk[k], _dx_uk[k], _A_SIk[I], _dt)
+            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_ik[i], _dx_ik[i], _B_uk[k], _dx_uk[k], _A_SIk[I],
+                            _h_Ik[I], _dt)
+        elif _is_end[I]:
+            im1 = _backward_I_i[I]
+            _E_Ik[I] = E_Ik(_B_dk[k], _dx_dk[k], _B_ik[im1], _dx_ik[im1],
+                            _A_SIk[I], _dt)
+            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_dk[k], _dx_dk[k], _B_ik[im1],
+                            _dx_ik[im1], _A_SIk[I], _h_Ik[I], _dt)
+        else:
+            i = _forward_I_i[I]
+            im1 = i - 1
+            _E_Ik[I] = E_Ik(_B_ik[i], _dx_ik[i], _B_ik[im1], _dx_ik[im1],
+                            _A_SIk[I], _dt)
+            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_ik[i], _dx_ik[i], _B_ik[im1],
+                            _dx_ik[im1], _A_SIk[I], _h_Ik[I], _dt)
+    return 1
+
+@njit(float64(float64, float64),
+      cache=True)
+def safe_divide(num, den):
+    return (num + SMALLEST_NORMAL) / (den + SMALLEST_NORMAL)
+
+# TODO: This is really just meant for orifice, weir, and pump equations
+# Where division by zero implies entire expression is zero
+@njit(float64[:](float64[:], float64[:]),
+      cache=True)
+def safe_divide_vec(num, den):
+    result = np.zeros_like(num)
+    is_safe = (np.abs(den) > SMALLEST_NORMAL)
+    result[is_safe] = num[is_safe] / den[is_safe]
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def Q_i_f(h_Ip1k, h_1k, U_Ik, V_Ik, W_Ik):
+    t_0 = U_Ik * h_Ip1k
+    t_1 = V_Ik
+    t_2 = W_Ik * h_1k
+    return t_0 + t_1 + t_2
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def Q_i_b(h_Ik, h_Np1k, X_Ik, Y_Ik, Z_Ik):
+    t_0 = X_Ik * h_Ik
+    t_1 = Y_Ik
+    t_2 = Z_Ik * h_Np1k
+    return t_0 + t_1 + t_2
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def h_i_f(Q_ik, h_1k, U_Ik, V_Ik, W_Ik):
+    num = Q_ik - V_Ik - W_Ik * h_1k
+    den = U_Ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64),
+      cache=True)
+def h_i_b(Q_ik, h_Np1k, X_Ik, Y_Ik, Z_Ik):
+    num = Q_ik - Y_Ik - Z_Ik * h_Np1k
+    den = X_Ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64[:], int64[:], int64[:], int64[:], int64),
+      cache=True)
+def tnumba_solve_internals(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik,
+                          _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, _I_Nk, nk, NK):
+    for k in range(NK):
+        n = nk[k]
+        i_1 = _i_1k[k]
+        I_1 = _I_1k[k]
+        I_N = _I_Nk[k]
+        i_n = i_1 + n - 1
+        I_Np1 = I_1 + n
+        # Set boundary depths
+        _h_1k = _h_uk[k]
+        _h_Np1k = _h_dk[k]
+        _h_Ik[I_1] = _h_1k
+        _h_Ik[I_Np1] = _h_Np1k
+        # Determine direction (proxy for Froude number)
+        ratio = 0.
+        # Todo: check this range
+        for j in range(n):
+            I = I_1 + j
+            ratio += (np.log10(_U_Ik[I]) - np.log10(_X_Ik[I]))
+        go_forwards = ratio > 0.
+        go_backwards = (ratio <= 0.) or (not np.isfinite(ratio))
+        # If ratio > 0, go forwards
+        if go_forwards:
+            _Q_ik[i_1] = _X_Ik[I_1] * _h_1k + _Z_Ik[I_1] * _h_Np1k + _Y_Ik[I_1]
+            for j in range(I_N - I_1):
+                I = I_1 + j
+                Ip1 = I + 1
+                i = i_1 + j
+                ip1 = i + 1
+                _h_Ik[Ip1] = h_i_f(_Q_ik[i], _h_1k, _U_Ik[I], _V_Ik[I], _W_Ik[I])
+                _Q_ik[ip1] = Q_i_b(_h_Ik[Ip1], _h_Np1k, _X_Ik[Ip1], _Y_Ik[Ip1], _Z_Ik[Ip1])
+                #_h_Ik[Ip1] = (_Q_ik[i] - _W_Ik[I] * _h_1k - _V_Ik[I]) / _U_Ik[I]
+                #_Q_ik[ip1] = _X_Ik[Ip1] * _h_Ik[Ip1] + _Z_Ik[Ip1] * _h_Np1k + _Y_Ik[Ip1]
+            assert ip1 == i_n
+            assert Ip1 == I_N
+        # If ratio <= 0, go backwards
+        elif go_backwards:
+            _Q_ik[i_n] = _U_Ik[I_N] * _h_Np1k + _W_Ik[I_N] * _h_1k + _V_Ik[I_N]
+            for j in range(I_N - I_1):
+                I = I_N - j
+                Im1 = I - 1
+                i = i_n - j
+                im1 = i - 1
+                _h_Ik[I] = h_i_b(_Q_ik[i], _h_Np1k, _X_Ik[I], _Y_Ik[I], _Z_Ik[I])
+                _Q_ik[im1] = Q_i_f(_h_Ik[I], _h_1k, _U_Ik[Im1], _V_Ik[Im1], _W_Ik[Im1])
+                #_h_Ik[I] = (_Q_ik[i] - _Z_Ik[I] * _h_Np1k - _Y_Ik[I]) / _X_Ik[I]
+                #_Q_ik[im1] = _U_Ik[Im1] * _h_Ik[I] + _W_Ik[Im1] * _h_1k + _V_Ik[Im1]
+            assert im1 == i_1
+            assert I == I_1 + 1
+        else:
+            raise ValueError('Invalid Froude Number.')
+    return 1
+
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64[:], int64[:], int64[:], int64,
+            float64, float64[:], boolean),
+      cache=True)
+def numba_solve_internals_matrix(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik,
+                          _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, nk, NK,
+                          min_depth, max_depth_k, first_link_backwards=True):
+    for k in range(NK):
+        n = nk[k]
+        i_1 = _i_1k[k]
+        I_1 = _I_1k[k]
+        i_n = i_1 + n - 1
+        I_2 = I_1 + 1
+        I_Np1 = I_1 + n
+        I_N = I_Np1 - 1
+        # Set boundary depths
+        _h_1k = _h_uk[k]
+        _h_Np1k = _h_dk[k]
+        _h_Ik[I_1] = _h_1k
+        _h_Ik[I_Np1] = _h_Np1k
+        # Set up matrix
+        dim = 2 * n - 1
+        A = np.zeros((dim, dim))
+        b = np.zeros(dim)
+        # Fill linear system
+        #### h indices
+        for j, i in enumerate(range(1, dim, 2)):
+            A[i, i] = -_X_Ik[I_2+j]
+            A[i, i+1] = 1.
+            b[i] = _Y_Ik[I_2+j] + _h_Np1k * _Z_Ik[I_2+j]
+        #### Q indices
+        for j, i in enumerate(range(2, dim, 2)):
+            A[i, i] = 1.
+            if i < dim - 1:
+                A[i, i+1] = -_U_Ik[I_2+j]
+            b[i] = _V_Ik[I_2+j] + _h_1k * _W_Ik[I_2+j]
+        #### Other indices
+        A[0 ,0] = 1.
+        A[0, 1] = -_X_Ik[I_1]
+        b[0] = _Y_Ik[I_1] + _h_Np1k * _Z_Ik[I_1]
+        b[-1] += _h_1k * _U_Ik[I_N]
+        # Solve linear system
+        x = np.linalg.solve(A, b)
+        # Write to output arrays
+        _Q_ik[i_1:i_1+n] = x[::2]
+        _h_Ik[I_2:I_2+n-1] = x[1::2]
+    return 1
+
+
+@njit(float64[:](float64[:], int64, int64[:], int64[:], int64[:], int64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_solve_internals_ls(_h_Ik, NK, nk, _k_1k, _i_1k, _I_1k, _U, _X, _b):
+    for k in range(NK):
+        nlinks = nk[k]
+        lstart = _k_1k[k]
+        rstart = _i_1k[k]
+        jstart = _I_1k[k]
+        _Ak = np.zeros((nlinks, nlinks - 1))
+        for i in range(nlinks - 1):
+            _Ak[i, i] = _U[lstart + i]
+            _Ak[i + 1, i] = -_X[lstart + i]
+        _AkT = _Ak.T.copy()
+        _bk = _b[rstart:rstart+nlinks].copy()
+        _AA = _AkT @ _Ak
+        _Ab = _AkT @ _bk
+        # If want to prevent singular matrix, set ( diag == 0 ) = 1
+        for i in range(nlinks - 1):
+            if (_AA[i, i] == 0.0):
+                _AA[i, i] = 1.0
+        _h_inner = np.linalg.solve(_AA, _Ab)
+        _h_Ik[jstart+1:jstart+nlinks] = _h_inner
+    return _h_Ik
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def numba_u_ik(_Q_ik, _A_ik, _A_ik_min, _u_ik):
+    n = _u_ik.size
+    for i in range(n):
+        _Q_i = _Q_ik[i]
+        _A_i = _A_ik[i]
+        _A_i_min = _A_ik_min[i]
+        _u_ik[i] = _Q_i / (_A_i + _A_i_min**2 / _A_i)
+    return _u_ik
+
+#@njit(float64[:](float64[:], float64[:], float64[:]),
+#      cache=True)
+#def numba_u_ik(_Q_ik, _A_ik, _u_ik):
+#    n = _u_ik.size
+#    for i in range(n):
+#        _Q_i = _Q_ik[i]
+#        _A_i = _A_ik[i]
+#        if _A_i > SMALLEST_NORMAL:
+#            _u_ik[i] = _Q_i / _A_i
+#        else:
+#            _u_ik[i] = 0.
+#    return _u_ik
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]),
+      cache=True)
+def numba_u_Ik(_dx_ik, _u_ik, _dx_uk, _u_uk, _link_start, _ki, _u_Ik):
+    n = _u_Ik.size
+    for i in range(n):
+        k = _ki[i]
+        if _link_start[i]:
+            num = _dx_ik[i] * _u_uk[k] + _dx_uk[k] * _u_ik[i]
+            den = _dx_ik[i] + _dx_uk[k]
+            if den > SMALLEST_NORMAL:
+                _u_Ik[i] = num / den
+            else:
+                _u_Ik[i] = 0.
+        else:
+            im1 = i - 1
+            num = _dx_ik[i] * _u_ik[im1] + _dx_ik[im1] * _u_ik[i]
+            den = _dx_ik[i] + _dx_ik[im1]
+            if den > SMALLEST_NORMAL:
+                _u_Ik[i] = num / den
+            else:
+                _u_Ik[i] = 0.
+    return _u_Ik
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]),
+      cache=True)
+def numba_u_Ip1k(_dx_ik, _u_ik, _dx_dk, _u_dk, _link_end, _ki, _u_Ip1k):
+    n = _u_Ip1k.size
+    for i in range(n):
+        k = _ki[i]
+        if _link_end[i]:
+            num = _dx_ik[i] * _u_dk[k] + _dx_dk[k] * _u_ik[i]
+            den = _dx_ik[i] + _dx_dk[k]
+            if den > SMALLEST_NORMAL:
+                _u_Ip1k[i] = num / den
+            else:
+                _u_Ip1k[i] = 0.
+        else:
+            ip1 = i + 1
+            num = _dx_ik[i] * _u_ik[ip1] + _dx_ik[ip1] * _u_ik[i]
+            den = _dx_ik[i] + _dx_ik[ip1]
+            if den > SMALLEST_NORMAL:
+                _u_Ip1k[i] = num / den
+            else:
+                _u_Ip1k[i] = 0.
+    return _u_Ip1k
+
+@njit(float64(float64, float64),
+      cache=True)
+def U_1k(c_1k, T_1k):
+    """
+    Compute forward recurrence coefficient 'U' for node 1, superlink k.
+    """
+    num = -c_1k
+    den = T_1k
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64),
+      cache=True)
+def V_1k(P_1k, T_1k):
+    """
+    Compute forward recurrence coefficient 'V' for node 1, superlink k.
+    """
+    num = P_1k
+    den = T_1k
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64),
+      cache=True)
+def W_1k(a_1k, T_1k):
+    """
+    Compute forward recurrence coefficient 'W' for node 1, superlink k.
+    """
+    num = -a_1k
+    den = T_1k
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64),
+      cache=True)
+def T_1k(b_1k):
+    """
+    Compute forward recurrence coefficient 'T' for link 1, superlink k.
+    """
+    return b_1k
+
+@njit(float64(float64, float64, float64, float64),
+      cache=True)
+def U_Ik(E_Ik, c_ik, U_Im1k, T_ik):
+    """
+    Compute forward recurrence coefficient 'U' for node I, superlink k.
+    """
+    num = c_ik * (E_Ik - U_Im1k)
+    den = T_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def V_Ik(P_ik, a_ik, D_Ik, E_Ik, V_Im1k, U_Im1k, T_ik):
+    """
+    Compute forward recurrence coefficient 'V' for node I, superlink k.
+    """
+    num = a_ik * (D_Ik + V_Im1k) - P_ik * (E_Ik - U_Im1k)
+    den = T_ik
+    result = safe_divide(num, den)
+    return  result
+
+@njit(float64(float64, float64, float64),
+      cache=True)
+def W_Ik(a_ik, W_Im1k, T_ik):
+    """
+    Compute forward recurrence coefficient 'W' for node I, superlink k.
+    """
+    num = a_ik * W_Im1k
+    den = T_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64),
+      cache=True)
+def T_ik(a_ik, b_ik, E_Ik, U_Im1k):
+    """
+    Compute forward recurrence coefficient 'T' for link i, superlink k.
+    """
+    result = a_ik - b_ik * (E_Ik - U_Im1k)
+    return result
+
+@njit(float64(float64, float64),
+      cache=True)
+def X_Nk(a_nk, O_nk):
+    """
+    Compute backward recurrence coefficient 'X' for node N, superlink k.
+    """
+    num = -a_nk
+    den = O_nk
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64),
+      cache=True)
+def Y_Nk(P_nk, O_nk):
+    """
+    Compute backward recurrence coefficient 'Y' for node N, superlink k.
+    """
+    num = P_nk
+    den = O_nk
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64),
+      cache=True)
+def Z_Nk(c_nk, O_nk):
+    """
+    Compute backward recurrence coefficient 'Z' for node N, superlink k.
+    """
+    num = -c_nk
+    den = O_nk
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64),
+      cache=True)
+def O_nk(b_nk):
+    """
+    Compute backward recurrence coefficient 'O' for link n, superlink k.
+    """
+    return b_nk
+
+@njit(float64(float64, float64, float64, float64),
+      cache=True)
+def X_Ik(a_ik, E_Ip1k, X_Ip1k, O_ik):
+    """
+    Compute backward recurrence coefficient 'X' for node I, superlink k.
+    """
+    num = -a_ik * (E_Ip1k + X_Ip1k)
+    den = O_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64, float64, float64, float64),
+      cache=True)
+def Y_Ik(P_ik, c_ik, D_Ip1k, E_Ip1k, Y_Ip1k, X_Ip1k, O_ik):
+    """
+    Compute backward recurrence coefficient 'Y' for node I, superlink k.
+    """
+    num = -c_ik * (D_Ip1k - Y_Ip1k) + P_ik * (E_Ip1k + X_Ip1k)
+    den = O_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64),
+      cache=True)
+def Z_Ik(c_ik, Z_Ip1k, O_ik):
+    """
+    Compute backward recurrence coefficient 'Z' for node I, superlink k.
+    """
+    num = c_ik * Z_Ip1k
+    den = O_ik
+    result = safe_divide(num, den)
+    return result
+
+@njit(float64(float64, float64, float64, float64),
+      cache=True)
+def O_ik(b_ik, c_ik, E_Ip1k, X_Ip1k):
+    """
+    Compute backward recurrence coefficient 'O' for link i, superlink k.
+    """
+    result = c_ik + b_ik * (E_Ip1k + X_Ip1k)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64),
+      cache=True)
+def gamma_o(Q_o_t, Ao, Co, g=9.81):
+    """
+    Compute flow coefficient 'gamma' for orifice o.
+    """
+    num = 2 * g * Co**2 * Ao**2
+    den = np.abs(Q_o_t)
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def gamma_w(Q_w_t, H_w_t, L_w, s_w, Cwr, Cwt):
+    """
+    Compute flow coefficient 'gamma' for weir w.
+    """
+    num = (Cwr * L_w * H_w_t + Cwt * s_w * H_w_t**2)**2
+    den = np.abs(Q_w_t)
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:]),
+      cache=True)
+def gamma_p(Q_p_t, b_p, c_p, u):
+    """
+    Compute flow coefficient 'gamma' for pump p.
+    """
+    num = u
+    den = b_p * np.abs(Q_p_t)**(c_p - 1)
+    result = safe_divide_vec(num, den)
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:]),
+      cache=True)
+def xi_uk(dx_uk, B_uk, theta_uk):
+    result = (dx_uk * B_uk * theta_uk) / 2
+    return result
+
+@njit(float64[:](float64[:], float64[:], float64[:]),
+      cache=True)
+def xi_dk(dx_dk, B_dk, theta_dk):
+    result = (dx_dk * B_dk * theta_dk) / 2
+    return result
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:],
+             int64[:], int64[:]),
+      cache=True)
+def numba_orifice_flow_coefficients(_alpha_o, _beta_o, _chi_o, H_j, _Qo, u, _z_inv_j,
+                                    _z_o, _tau_o, _Co, _Ao, _y_max_o, _unidir_o,
+                                     _J_uo, _J_do):
+    g = 9.81
+    _H_uo = H_j[_J_uo]
+    _H_do = H_j[_J_do]
+    _z_inv_uo = _z_inv_j[_J_uo]
+    # Create indicator functions
+    _omega_o = np.zeros_like(_H_uo)
+    _omega_o[_H_uo >= _H_do] = 1.0
+    # Compute universal coefficients
+    _gamma_o = gamma_o(_Qo, _Ao, _Co, g)
+    # Create conditionals
+    cond_0 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u))
+    cond_1 = ((1 - _omega_o) * _H_uo + _omega_o * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u / 2))
+    cond_2 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo)
+    cond_3 = (_H_do >= _H_uo) & _unidir_o
+    # Fill coefficient arrays
+    # Submerged on both sides
+    a = (cond_0 & cond_1)
+    _alpha_o[a] = _gamma_o[a]
+    _beta_o[a] = -_gamma_o[a]
+    _chi_o[a] = 0.0
+    # Submerged on one side
+    b = (cond_0 & ~cond_1)
+    _alpha_o[b] = _gamma_o[b] * _omega_o[b] * (-1)**(1 - _omega_o[b])
+    _beta_o[b] = _gamma_o[b] * (1 - _omega_o[b]) * (-1)**(1 - _omega_o[b])
+    _chi_o[b] = (_gamma_o[b] * (-1)**(1 - _omega_o[b])
+                                    * (- _z_inv_uo[b] - _z_o[b] -
+                                        _tau_o[b] * _y_max_o[b] * u[b] / 2))
+    # Weir flow
+    c = (~cond_0 & cond_2)
+    _alpha_o[c] = _gamma_o[c] * _omega_o[c] * (-1)**(1 - _omega_o[c])
+    _beta_o[c] = _gamma_o[c] * (1 - _omega_o[c]) * (-1)**(1 - _omega_o[c])
+    _chi_o[c] = (_gamma_o[c] * (-1)**(1 - _omega_o[c])
+                                    * (- _z_inv_uo[c] - _z_o[c]))
+    # No flow
+    d = (~cond_0 & ~cond_2) | cond_3
+    _alpha_o[d] = 0.0
+    _beta_o[d] = 0.0
+    _chi_o[d] = 0.0
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], boolean[:],
+            int64[:], int64[:], float64),
+      cache=True)
+def numba_solve_orifice_flows(H_j, u, _z_inv_j, _z_o,
+                              _tau_o, _y_max_o, _Co, _Ao, _unidir_o, _J_uo, _J_do, g=9.81):
+    # Specify orifice heads at previous timestep
+    _H_uo = H_j[_J_uo]
+    _H_do = H_j[_J_do]
+    _z_inv_uo = _z_inv_j[_J_uo]
+    # Create indicator functions
+    _omega_o = np.zeros_like(_H_uo)
+    _omega_o[_H_uo >= _H_do] = 1.0
+    # Create arrays to store flow coefficients for current time step
+    _alpha_o = np.zeros_like(_H_uo)
+    _beta_o = np.zeros_like(_H_uo)
+    _chi_o = np.zeros_like(_H_uo)
+    # Compute universal coefficients
+    _gamma_o = 2 * g * _Co**2 * _Ao**2
+    # Create conditionals
+    cond_0 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u))
+    cond_1 = ((1 - _omega_o) * _H_uo + _omega_o * _H_do >
+                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u / 2))
+    cond_2 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
+                _z_o + _z_inv_uo)
+    cond_3 = (_H_do >= _H_uo) & _unidir_o
+    # Fill coefficient arrays
+    # Submerged on both sides
+    a = (cond_0 & cond_1)
+    _alpha_o[a] = _gamma_o[a]
+    _beta_o[a] = -_gamma_o[a]
+    _chi_o[a] = 0.0
+    # Submerged on one side
+    b = (cond_0 & ~cond_1)
+    _alpha_o[b] = _gamma_o[b] * _omega_o[b] * (-1)**(1 - _omega_o[b])
+    _beta_o[b] = _gamma_o[b] * (1 - _omega_o[b]) * (-1)**(1 - _omega_o[b])
+    _chi_o[b] = (_gamma_o[b] * (-1)**(1 - _omega_o[b])
+                                    * (- _z_inv_uo[b] - _z_o[b]
+                                        - _tau_o[b] * _y_max_o[b] * u[b] / 2))
+    # Weir flow on one side
+    c = (~cond_0 & cond_2)
+    _alpha_o[c] = _gamma_o[c] * _omega_o[c] * (-1)**(1 - _omega_o[c])
+    _beta_o[c] = _gamma_o[c] * (1 - _omega_o[c]) * (-1)**(1 - _omega_o[c])
+    _chi_o[c] = (_gamma_o[c] * (-1)**(1 - _omega_o[c])
+                                    * (- _z_inv_uo[c] - _z_o[c]))
+    # No flow
+    d = (~cond_0 & ~cond_2) | cond_3
+    _alpha_o[d] = 0.0
+    _beta_o[d] = 0.0
+    _chi_o[d] = 0.0
+    # Compute flow
+    _Qo_next = (-1)**(1 - _omega_o) * np.sqrt(np.abs(
+               _alpha_o * _H_uo + _beta_o * _H_do + _chi_o))
+    # Export instance variables
+    return _Qo_next
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:]),
+      cache=True)
+def numba_weir_flow_coefficients(_Hw, _Qw, _alpha_w, _beta_w, _chi_w, H_j, _z_inv_j, _z_w,
+                                 _y_max_w, u, _L_w, _s_w, _Cwr, _Cwt, _J_uw, _J_dw):
+    # Specify weir heads at previous timestep
+    _H_uw = H_j[_J_uw]
+    _H_dw = H_j[_J_dw]
+    _z_inv_uw = _z_inv_j[_J_uw]
+    # Create indicator functions
+    _omega_w = np.zeros(_H_uw.size)
+    _omega_w[_H_uw >= _H_dw] = 1.0
+    # Create conditionals
+    cond_0 = (_omega_w * _H_uw + (1 - _omega_w) * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    cond_1 = ((1 - _omega_w) * _H_uw + _omega_w * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    # Effective heads
+    a = (cond_0 & cond_1)
+    b = (cond_0 & ~cond_1)
+    c = (~cond_0)
+    _Hw[a] = _H_uw[a] - _H_dw[a]
+    _Hw[b] = (_omega_w[b] * _H_uw[b] + (1 - _omega_w[b]) * _H_dw[b]
+                    + (-_z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
+    _Hw[c] = 0.0
+    _Hw = np.abs(_Hw)
+    # Compute universal coefficients
+    _gamma_w = gamma_w(_Qw, _Hw, _L_w, _s_w, _Cwr, _Cwt)
+    # Fill coefficient arrays
+    # Submerged on both sides
+    a = (cond_0 & cond_1)
+    _alpha_w[a] = _gamma_w[a]
+    _beta_w[a] = -_gamma_w[a]
+    _chi_w[a] = 0.0
+    # Submerged on one side
+    b = (cond_0 & ~cond_1)
+    _alpha_w[b] = _gamma_w[b] * _omega_w[b] * (-1)**(1 - _omega_w[b])
+    _beta_w[b] = _gamma_w[b] * (1 - _omega_w[b]) * (-1)**(1 - _omega_w[b])
+    _chi_w[b] = (_gamma_w[b] * (-1)**(1 - _omega_w[b]) *
+                                (- _z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
+    # No flow
+    c = (~cond_0)
+    _alpha_w[c] = 0.0
+    _beta_w[c] = 0.0
+    _chi_w[c] = 0.0
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64[:], int64[:]),
+      cache=True)
+def numba_solve_weir_flows(_Hw, _Qw, H_j, _z_inv_j, _z_w, _y_max_w, u, _L_w,
+                           _s_w, _Cwr, _Cwt, _J_uw, _J_dw):
+    _H_uw = H_j[_J_uw]
+    _H_dw = H_j[_J_dw]
+    _z_inv_uw = _z_inv_j[_J_uw]
+    # Create indicator functions
+    _omega_w = np.zeros(_H_uw.size)
+    _omega_w[_H_uw >= _H_dw] = 1.0
+    # Create conditionals
+    cond_0 = (_omega_w * _H_uw + (1 - _omega_w) * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    cond_1 = ((1 - _omega_w) * _H_uw + _omega_w * _H_dw >
+                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
+    # TODO: Is this being recalculated for a reason?
+    # Effective heads
+    a = (cond_0 & cond_1)
+    b = (cond_0 & ~cond_1)
+    c = (~cond_0)
+    _Hw[a] = _H_uw[a] - _H_dw[a]
+    _Hw[b] = (_omega_w[b] * _H_uw[b] + (1 - _omega_w[b]) * _H_dw[b]
+                    + (-_z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
+    _Hw[c] = 0.0
+    _Hw = np.abs(_Hw)
+    # Compute universal coefficient
+    _gamma_ww = (_Cwr * _L_w * _Hw + _Cwt * _s_w * _Hw**2)**2
+    # Compute flow
+    _Qw_next = (-1)**(1 - _omega_w) * np.sqrt(_gamma_ww * _Hw)
+    return _Qw_next
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            int64[:], int64[:]),
+      cache=True)
+def numba_pump_flow_coefficients(_alpha_p, _beta_p, _chi_p, H_j, _z_inv_j, _Qp, u,
+                                 _z_p, _dHp_max, _dHp_min, _a_p, _b_p, _c_p,
+                                 _J_up, _J_dp):
+    # Get upstream and downstream heads and invert elevation
+    _H_up = H_j[_J_up]
+    _H_dp = H_j[_J_dp]
+    _z_inv_up = _z_inv_j[_J_up]
+    # Compute effective head
+    _dHp = _H_dp - _H_up
+    # Condition 0: Upstream head is above inlet height
+    cond_0 = _H_up > _z_inv_up + _z_p
+    # Condition 1: Head difference is within range of pump curve
+    cond_1 = (_dHp > _dHp_min) & (_dHp < _dHp_max)
+    _dHp[_dHp > _dHp_max] = _dHp_max[_dHp > _dHp_max]
+    _dHp[_dHp < _dHp_min] = _dHp_min[_dHp < _dHp_min]
+    # Compute universal coefficients
+    _gamma_p = gamma_p(_Qp, _b_p, _c_p, u)
+    # Fill coefficient arrays
+    # Head in pump curve range
+    a = (cond_0 & cond_1)
+    _alpha_p[a] = _gamma_p[a]
+    _beta_p[a] = -_gamma_p[a]
+    _chi_p[a] = _gamma_p[a] * _a_p[a]
+    # Head outside of pump curve range
+    b = (cond_0 & ~cond_1)
+    _alpha_p[b] = 0.0
+    _beta_p[b] = 0.0
+    _chi_p[b] = _gamma_p[b] * (_a_p[b] - _dHp[b])
+    # Depth below inlet
+    c = (~cond_0)
+    _alpha_p[c] = 0.0
+    _beta_p[c] = 0.0
+    _chi_p[c] = 0.0
+    return 1
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+                 float64[:], float64[:], float64[:], int64[:], int64[:]),
+      cache=True)
+def numba_solve_pump_flows(H_j, u, _z_inv_j, _z_p, _dHp_max, _dHp_min, _a_p, _b_p, _c_p,
+                           _J_up, _J_dp):
+    _H_up = H_j[_J_up]
+    _H_dp = H_j[_J_dp]
+    _z_inv_up = _z_inv_j[_J_up]
+    # Create conditionals
+    _dHp = _H_dp - _H_up
+    _dHp[_dHp > _dHp_max] = _dHp_max[_dHp > _dHp_max]
+    _dHp[_dHp < _dHp_min] = _dHp_min[_dHp < _dHp_min]
+    cond_0 = _H_up > _z_inv_up + _z_p
+    # Compute universal coefficients
+    _Qp_next = (u / _b_p * (_a_p - _dHp))**(1 / _c_p)
+    _Qp_next[~cond_0] = 0.0
+    return _Qp_next
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64, int64[:], int64[:], int64[:]),
+      cache=True)
+def tnumba_forward_recurrence(_T_ik, _U_Ik, _V_Ik, _W_Ik, _a_ik, _b_ik, _c_ik,
+                             _P_ik, _E_Ik, _D_Ik, NK, nk, _I_1k, _i_1k):
+    for k in range(NK):
+        # Start at junction 1
+        _I_1 = _I_1k[k]
+        _i_1 = _i_1k[k]
+        _I_2 = _I_1 + 1
+        _i_2 = _i_1 + 1
+        nlinks = nk[k]
+        _T_ik[_i_1] = T_1k(_b_ik[_i_1])
+        _U_Ik[_I_1] = U_1k(_c_ik[_i_1], _T_ik[_i_1])
+        _V_Ik[_I_1] = V_1k(_P_ik[_i_1], _T_ik[_i_1])
+        _W_Ik[_I_1] = W_1k(_a_ik[_i_1], _T_ik[_i_1])
+        # Loop from junction 2 -> Nk
+        for i in range(nlinks - 1):
+            _i_next = _i_2 + i
+            _I_next = _I_2 + i
+            _Im1_next = _I_next - 1
+            _T_ik[_i_next] = T_ik(_a_ik[_i_next], _b_ik[_i_next], _E_Ik[_I_next], _U_Ik[_Im1_next])
+            _U_Ik[_I_next] = U_Ik(_E_Ik[_I_next], _c_ik[_i_next],
+                                  _U_Ik[_Im1_next], _T_ik[_i_next])
+            _V_Ik[_I_next] = V_Ik(_P_ik[_i_next], _a_ik[_i_next], _D_Ik[_I_next],
+                                  _E_Ik[_I_next], _V_Ik[_Im1_next], _U_Ik[_Im1_next],
+                                  _T_ik[_i_next])
+            _W_Ik[_I_next] = W_Ik(_a_ik[_i_next], _W_Ik[_Im1_next], _T_ik[_i_next])
+    return 1
+
+@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+            float64[:], float64[:], float64[:], int64, int64[:], int64[:], int64[:]),
+      cache=True)
+def tnumba_backward_recurrence(_O_ik, _X_Ik, _Y_Ik, _Z_Ik, _a_ik, _b_ik, _c_ik,
+                              _P_ik, _E_Ik, _D_Ik, NK, nk, _I_Nk, _i_nk):
+    for k in range(NK):
+        _I_N = _I_Nk[k]
+        _i_n = _i_nk[k]
+        _I_Nm1 = _I_N - 1
+        _i_nm1 = _i_n - 1
+        _I_Np1 = _I_N + 1
+        nlinks = nk[k]
+        _O_ik[_i_n] = O_nk(_b_ik[_i_n])
+        _X_Ik[_I_N] = X_Nk(_a_ik[_i_n], _O_ik[_i_n])
+        _Y_Ik[_I_N] = Y_Nk(_P_ik[_i_n], _O_ik[_i_n])
+        _Z_Ik[_I_N] = Z_Nk(_c_ik[_i_n], _O_ik[_i_n])
+        for i in range(nlinks - 1):
+            _i_next = _i_nm1 - i
+            _I_next = _I_Nm1 - i
+            _Ip1_next = _I_next + 1
+            _O_ik[_i_next] = O_ik(_b_ik[_i_next], _c_ik[_i_next], _E_Ik[_Ip1_next], _X_Ik[_Ip1_next])
+            _X_Ik[_I_next] = X_Ik(_a_ik[_i_next], _E_Ik[_Ip1_next], _X_Ik[_Ip1_next], _O_ik[_i_next])
+            _Y_Ik[_I_next] = Y_Ik(_P_ik[_i_next], _c_ik[_i_next], _D_Ik[_Ip1_next], 
+                                  _E_Ik[_Ip1_next], _Y_Ik[_Ip1_next], _X_Ik[_Ip1_next],
+                                  _O_ik[_i_next])
+            _Z_Ik[_I_next] = Z_Ik(_c_ik[_i_next], _Z_Ik[_Ip1_next], _O_ik[_i_next])
+    return 1
+
+@njit(float64[:,:](float64[:,:], int64, int64),
+      cache=True)
+def numba_create_banded(l, bandwidth, M):
+    AB = np.zeros((2*bandwidth + 1, M))
+    for i in range(M):
+        AB[bandwidth, i] = l[i, i]
+    for n in range(bandwidth):
+        for j in range(M - n - 1):
+            AB[bandwidth - n - 1, -j - 1] = l[-j - 2 - n, -j - 1]
+            AB[bandwidth + n + 1, j] = l[j + n + 1, j]
+    return AB
+
+@njit(void(float64[:], int64[:], float64[:]),
+      cache=True)
+def numba_add_at(a, indices, b):
+    n = len(indices)
+    for k in range(n):
+        i = indices[k]
+        a[i] += b[k]
+
+@njit(void(float64[:, :], boolean[:], int64[:], int64[:], int64),
+      cache=True)
+def numba_clear_off_diagonals(A, bc, _J_uk, _J_dk, NK):
+    for k in range(NK):
+        _J_u = _J_uk[k]
+        _J_d = _J_dk[k]
+        _bc_u = bc[_J_u]
+        _bc_d = bc[_J_d]
+        if not _bc_u:
+            A[_J_u, _J_d] = 0.0
+        if not _bc_d:
+            A[_J_d, _J_u] = 0.0
+
+def numba_B_j(_B_j, _J_uk, _J_dk, _xi_uk, _xi_dk, _A_sj):
+    numba_add_at(_B_j, _J_uk, _xi_uk)
+    numba_add_at(_B_j, _J_dk, _xi_dk)
+    _B_j += _A_sj
+    return _B_j
+
+@njit(void(float64[:, :], boolean[:], float64[:], float64, int64),
+      cache=True)
+def numba_create_J_matrix(J, bc, _A_sj, _dt, M):
+    for i in range(M):
+        if bc[i]:
+            J[i,i] = 1.0
+        else:
+            J[i,i] = (_A_sj[i] / _dt)
+
+@njit(void(float64[:, :], float64[:], boolean[:], int64[:], int64[:], float64[:],
+           float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+           float64, int64, int64),
+      cache=True)
+def numba_create_K_matrix(K, _F_jj, bc, _J_uk, _J_dk, _alpha_uk,
+                          _alpha_dk, _beta_uk, _beta_dk, _xi_uk, _xi_dk,
+                          _A_sj, _dt, M, NK):
+    numba_add_at(_F_jj, _J_uk, _alpha_uk)
+    numba_add_at(_F_jj, _J_dk, -_beta_dk)
+    numba_add_at(_F_jj, _J_uk, _xi_uk / _dt)
+    numba_add_at(_F_jj, _J_dk, _xi_dk / _dt)
+    # Set diagonal of K matrix
+    for i in range(M):
+        if bc[i]:
+            K[i,i] = 0.0
+        else:
+            K[i,i] = _F_jj[i]
+    for k in range(NK):
+        _J_u = _J_uk[k]
+        _J_d = _J_dk[k]
+        _bc_u = bc[_J_u]
+        _bc_d = bc[_J_d]
+        if not _bc_u:
+            K[_J_u, _J_d] += _beta_uk[k]
+        if not _bc_d:
+            K[_J_d, _J_u] -= _alpha_dk[k]
+
+@njit(void(float64[:, :], float64[:], boolean[:], int64[:], int64[:], float64[:],
+           float64[:], float64[:], float64[:], int64, int64),
+      cache=True)
+def numba_create_OWP_matrix(X, diag, bc, _J_uc, _J_dc, _alpha_uc,
+                            _alpha_dc, _beta_uc, _beta_dc, M, NC):
+    # Set diagonal
+    numba_add_at(diag, _J_uc, _alpha_uc)
+    numba_add_at(diag, _J_dc, -_beta_dc)
+    for i in range(M):
+        if bc[i]:
+            X[i,i] = 0.0
+        else:
+            X[i,i] = diag[i]
+    # Set off-diagonal
+    for c in range(NC):
+        _J_u = _J_uc[c]
+        _J_d = _J_dc[c]
+        _bc_u = bc[_J_u]
+        _bc_d = bc[_J_d]
+        if not _bc_u:
+            X[_J_u, _J_d] += _beta_uc[c]
+        if not _bc_d:
+            X[_J_d, _J_u] -= _alpha_dc[c]
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64),
+      cache=True)
+def numba_Q_i_next_b(X_Ik, h_Ik, Y_Ik, Z_Ik, h_Np1k, _Ik, _ki, n):
+    _Q_i = np.zeros(n)
+    for i in range(n):
+        I = _Ik[i]
+        k = _ki[i]
+        t_0 = X_Ik[I] * h_Ik[I]
+        t_1 = Y_Ik[I]
+        t_2 = Z_Ik[I] * h_Np1k[k]
+        _Q_i[i] = t_0 + t_1 + t_2
+    return _Q_i
+
+@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64),
+      cache=True)
+def numba_Q_im1k_next_f(U_Ik, h_Ik, V_Ik, W_Ik, h_1k, _Ik, _ki, n):
+    _Q_i = np.zeros(n)
+    for i in range(n):
+        I = _Ik[i]
+        Ip1 = I + 1
+        k = _ki[i]
+        t_0 = U_Ik[I] * h_Ik[Ip1]
+        t_1 = V_Ik[I]
+        t_2 = W_Ik[I] * h_1k[k]
+        _Q_i[i] = t_0 + t_1 + t_2
+    return _Q_i
+
+@njit(void(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
+           float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64[:],
+           int64[:], int64[:], int64, boolean[:]),
+      cache=True)
+def numba_reposition_junctions(_x_Ik, _z_inv_Ik, _h_Ik, _dx_ik, _Q_ik, _H_dk,
+                               _b0, _zc, _xc, _m, _elem_pos, _i_1k, _I_1k,
+                               _I_Np1k, nk, NK, reposition):
+    for k in range(NK):
+        if reposition[k]:
+            _i_1 = _i_1k[k]
+            _I_1 = _I_1k[k]
+            _I_Np1 = _I_Np1k[k]
+            nlinks = nk[k]
+            njunctions = nlinks + 1
+            _i_end = _i_1 + nlinks
+            _I_end = _I_1 + njunctions
+            _H_d = _H_dk[k]
+            _z_inv_1 = _z_inv_Ik[_I_1]
+            _z_inv_Np1 = _z_inv_Ik[_I_Np1]
+            pos_prev = _elem_pos[k]
+            # Junction arrays for superlink k
+            _x_I = _x_Ik[_I_1:_I_end]
+            _z_inv_I = _z_inv_Ik[_I_1:_I_end]
+            _h_I = _h_Ik[_I_1:_I_end]
+            _dx_i = _dx_ik[_i_1:_i_end]
+            # Move junction if downstream head is within range
+            move_junction = (_H_d > _z_inv_Np1) & (_H_d < _z_inv_1)
+            if move_junction:
+                z_m = _H_d
+                _x0 = _x_I[_I_1]
+                x_m = (_H_d - _b0[k]) / _m[k] + _x0
+            else:
+                z_m = _zc[k]
+                x_m = _xc[k]
+            # Determine new x-position of junction
+            c = np.searchsorted(_x_I, x_m)
+            cm1 = c - 1
+            # Compute fractional x-position along superlink k
+            frac = (x_m - _x_I[cm1]) / (_x_I[c] - _x_I[cm1])
+            # Interpolate depth at new position
+            h_m = (1 - frac) * _h_I[cm1] + (frac) * _h_I[c]
+            # Link length ratio
+            r = _dx_i[pos_prev - 1] / (_dx_i[pos_prev - 1]
+                                    + _dx_i[pos_prev])
+            # Set new values
+            _x_I[pos_prev] = x_m
+            _z_inv_I[pos_prev] = z_m
+            _h_I[pos_prev] = h_m
+            Ix = np.argsort(_x_I)
+            _dx_i = np.diff(_x_I[Ix])
+            _x_Ik[_I_1:_I_end] = _x_I[Ix]
+            _z_inv_Ik[_I_1:_I_end] = _z_inv_I[Ix]
+            _h_Ik[_I_1:_I_end] = _h_I[Ix]
+            _dx_ik[_i_1:_i_end] = _dx_i
+            # Set position to new position
+            pos_change = np.argsort(Ix)
+            pos_next = pos_change[pos_prev]
+            _elem_pos[k] = pos_next
+            shifted = (pos_prev != pos_next)
+            # If position has shifted interpolate flow
+            if shifted:
+                ix = np.arange(nlinks)
+                ix[pos_prev] = pos_next
+                ix.sort()
+                _Q_i = _Q_ik[_i_1:_i_end]
+                _Q_i[pos_prev - 1] = (1 - r) * _Q_i[pos_prev - 1] + r * _Q_i[pos_prev]
+                _Q_ik[_i_1:_i_end] = _Q_i[ix]
+
diff --git a/pipedream_solver/callbacks.py b/pipedream_solver/callbacks.py
new file mode 100644
index 0000000..1232d39
--- /dev/null
+++ b/pipedream_solver/callbacks.py
@@ -0,0 +1,22 @@
+class BaseCallback():
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __on_step_start__(self, *args, **kwargs):
+        return None
+    
+    def __on_step_end__(self, *args, **kwargs):
+        return None
+
+    def __on_save_state__(self, *args, **kwargs):
+        return None
+
+    def __on_load_state__(self, *args, **kwargs):
+        return None
+
+    def __on_simulation_start__(self, *args, **kwargs):
+        return None
+
+    def __on_simulation_end__(self, *args, **kwargs):
+        return None
+ 
\ No newline at end of file
diff --git a/pipedream_solver/convergence.py b/pipedream_solver/convergence.py
new file mode 100644
index 0000000..8394ce4
--- /dev/null
+++ b/pipedream_solver/convergence.py
@@ -0,0 +1,70 @@
+import numpy as np
+from pipedream_solver.callbacks import BaseCallback
+
+class BaseConvergenceManager(BaseCallback):
+    def __init__(self, model, num_iter):
+        self.model = model
+        self.num_iter = num_iter
+        self.iter_elapsed = 0
+
+    def __on_step_start__(self, *args, **kwargs):
+        if (self.iter_elapsed == 0):
+            self.model.save_state()
+
+    def __on_step_end__(self, *args, **kwargs):
+        dt = kwargs['dt']
+        self.iter_elapsed += 1
+        while (self.iter_elapsed < self.num_iter):
+            condition_met = self.convergence_condition()
+            if not condition_met:
+                self.H_j_prev = self.model.H_j.copy()
+                self.model.iter_count -= 1
+                self.model.t -= dt
+                self.model._setup_step(*args, **kwargs)
+                self.model._solve_step(*args, **kwargs)
+                self.H_j_next = self.model.H_j.copy()
+                self.iter_elapsed += 1
+            else:
+                break
+        self.iter_elapsed = 0
+
+    def convergence_condition(self):
+        return False
+
+class DefaultConvergenceManager(BaseConvergenceManager):
+    def __init__(self, model, head_tol, num_iter):
+        self.model = model
+        self.head_tol = head_tol
+        self.num_iter = num_iter
+        self.iter_elapsed = 0
+
+    def __on_step_start__(self, *args, **kwargs):
+        self.H_j_prev = self.model.H_j.copy()
+        if (self.iter_elapsed == 0):
+            self.model.save_state()
+
+    def __on_step_end__(self, *args, **kwargs):
+        dt = kwargs['dt']
+        self.H_j_next = self.model.H_j.copy()
+        self.iter_elapsed += 1
+        while (self.iter_elapsed < self.num_iter):
+            condition_met = self.convergence_condition()
+            if not condition_met:
+                self.H_j_prev = self.model.H_j.copy()
+                self.model.iter_count -= 1
+                self.model.t -= dt
+                self.model._setup_step(*args, **kwargs)
+                self.model._solve_step(*args, **kwargs)
+                self.H_j_next = self.model.H_j.copy()
+                self.iter_elapsed += 1
+            else:
+                break
+        self.iter_elapsed = 0
+
+    def convergence_condition(self):
+        head_tol = self.head_tol
+        H_j_prev = self.H_j_prev
+        H_j_next = self.H_j_next
+        residual = np.abs(H_j_next - H_j_prev)
+        condition_met = (residual < head_tol).all()
+        return condition_met
diff --git a/pipedream_solver/coupling.py b/pipedream_solver/coupling.py
new file mode 100644
index 0000000..989ac95
--- /dev/null
+++ b/pipedream_solver/coupling.py
@@ -0,0 +1,32 @@
+class SpeciesCoupling():
+    def __init__(self, species={}, functions=[], parameters={}):
+        self.species = {}
+        self.functions = []
+        self.parameters = {}
+        for key, value in species.items():
+            if hasattr(self, key):
+                raise NameError
+            setattr(self, key, value)
+            self.species[key] = value
+        for function in functions:
+            if hasattr(self, function.__name__):
+                raise NameError
+            setattr(self, function.__name__, function)
+            self.functions.append(function)
+        for key, value in parameters.items():
+            if hasattr(self, key):
+                raise NameError
+            setattr(self, key, value)
+            self.parameters[key] = value
+
+    def step(self, **kwargs):
+        '''
+        Steps through all functions in `self.functions` in order.
+        '''
+        for function in self.functions:
+            arg_keys = function.__code__.co_varnames
+            # NOTE: Has potential to fail silently
+            arg_values = (getattr(self, key, None) for key in arg_keys)
+            function_inputs = dict(zip(arg_keys, arg_values))
+            function_inputs.update(kwargs)
+            result = function(**function_inputs)
diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py
new file mode 100644
index 0000000..c4dee67
--- /dev/null
+++ b/pipedream_solver/diagnostics.py
@@ -0,0 +1,784 @@
+import numpy as np
+from pipedream_solver._nsuperlink import numba_compute_functional_storage_volumes, numba_compute_tabular_storage_volumes
+from pipedream_solver._nsuperlink import junction_numerator, junction_denominator, superjunction_numerator, superjunction_denominator
+from pipedream_solver.callbacks import BaseCallback
+
+from numba import njit, prange
+from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, DictType, void
+
+class ErrorTracker(BaseCallback):
+    """
+    Tracks the error between the left-hand and right-hand side of each equation
+    """
+    def __init__(self, model, rtol=1e-1, atol=1e-10):
+        self.model = model
+        self.errors = {}
+        self.continuity_error_j = np.zeros(model.M)
+        self.continuity_error_Ik = np.zeros(model._I.size)
+        self.momentum_error_ik = np.zeros(model._i.size)
+        self.momentum_error_uk = np.zeros(model.NK)
+        self.momentum_error_dk = np.zeros(model.NK)
+        self.momentum_error_o = np.zeros(model.n_o)
+        self.momentum_error_w = np.zeros(model.n_w)
+        self.momentum_error_p = np.zeros(model.n_p)
+        self.continuity_magnitude_j = np.zeros(model.M)
+        self.continuity_magnitude_Ik = np.zeros(model._I.size)
+        self.momentum_magnitude_ik = np.zeros(model._i.size)
+        self.momentum_magnitude_uk = np.zeros(model.NK)
+        self.momentum_magnitude_dk = np.zeros(model.NK)
+        self.momentum_magnitude_o = np.zeros(model.n_o)
+        self.momentum_magnitude_w = np.zeros(model.n_w)
+        self.momentum_magnitude_p = np.zeros(model.n_p)
+        self.rtol = rtol
+        self.atol = atol
+
+    @property
+    def continuity_error(self):
+        return np.concatenate([self.continuity_error_j, self.continuity_error_Ik])
+
+    @property
+    def momentum_error(self):
+        return np.concatenate([self.momentum_error_uk, self.momentum_error_dk,
+                               self.momentum_error_ik])
+
+    @property
+    def error(self):
+        return np.concatenate([self.continuity_error_j, self.momentum_error_uk,
+                               self.momentum_error_dk, self.continuity_error_Ik,
+                               self.momentum_error_ik])
+    @property
+    def magnitude(self):
+        return np.concatenate([self.continuity_magnitude_j, self.momentum_magnitude_uk,
+                               self.momentum_magnitude_dk, self.continuity_magnitude_Ik,
+                               self.momentum_magnitude_ik])
+
+    @property
+    def error_metric(self):
+        rtol = self.rtol
+        atol = self.atol
+        dt = self.model._dt
+        e = np.abs(self.error) * dt
+        # TODO: Is maximum correct here?
+        ewt = np.maximum(rtol * np.abs(self.magnitude), atol)
+        metric = (e / ewt).max()
+        return metric
+
+    @property
+    def success(self):
+        error_metric = self.error_metric
+        condition = error_metric <= 1.
+        return condition
+    
+    def _record_error(self, *args, **kwargs):
+        t = self.model.t
+        self.errors[t] = self.error
+        
+    def _compute_error(self, *args, **kwargs):
+        model = self.model
+        # TODO: This could cause problems, need to make sure this stays updated at each step
+        dt = model._dt
+        self.continuity_error_j = continuity_error_j(model, dt)
+        self.continuity_error_Ik = continuity_error_Ik(model, dt)
+        self.momentum_error_ik = momentum_error_ik(model, dt)
+        self.momentum_error_uk = momentum_error_uk_2(model, dt)
+        self.momentum_error_dk = momentum_error_dk_2(model, dt)
+        # Use these for tsuperlink
+        #self.momentum_error_ik = momentum_error_ik_2(model, dt)
+        #self.momentum_error_uk = momentum_error_uk_3(model, dt)
+        #self.momentum_error_dk = momentum_error_dk_3(model, dt)
+        #self.momentum_error_o = 
+        #self.momentum_error_w = 
+        #self.momentum_error_p = 
+
+    def _compute_magnitudes(self, *args, **kwargs):
+        model = self.model
+        # TODO: This dt could cause problems, need to make sure this stays updated at each step
+        dt = model._dt
+        self.continuity_magnitude_j = continuity_magnitude_j(model, dt)
+        self.continuity_magnitude_Ik = continuity_magnitude_Ik(model, dt)
+        self.momentum_magnitude_uk = momentum_magnitude_uk(model, dt)
+        self.momentum_magnitude_dk = momentum_magnitude_dk(model, dt)
+        self.momentum_magnitude_ik = momentum_magnitude_ik(model, dt)
+
+    def __on_step_end__(self, *args, **kwargs):
+        self._compute_error(*args, **kwargs)
+        self._compute_magnitudes(*args, **kwargs)
+
+
+class ConditionTracker(BaseCallback):
+    def __init__(self, model):
+        self.model = model
+
+    def _superlink_condition_number(self):
+        model = self.model
+        singular_values = (model._X_Ik[model._Ik])**2
+        max_singular_value = max(1., singular_values.max())
+        min_singular_value = min(1., singular_values.min())
+        condition_number = max_singular_value / min_singular_value
+
+class VolumeTracker(BaseCallback):
+    def __init__(self, model):
+        self.model = model
+        self.volume_j = np.zeros(model.M)
+        self.volume_Ik = np.zeros(model._I.size)
+        self.volume_ik = np.zeros(model._i.size)
+        self.volume_uk = np.zeros(model.NK)
+        self.volume_dk = np.zeros(model.NK)
+        self.volume_flux_j = np.zeros(model.M)
+        self.volume_flux_Ik = np.zeros(model._I.size)
+        self.cumulative_vol_flux_j = np.zeros(model.M)
+        self.cumulative_vol_flux_Ik = np.zeros(model._I.size)
+        self.set_init_volume()
+
+    @property
+    def init_volume(self):
+        result = (self.init_volume_j.sum() + self.init_volume_Ik.sum() + self.init_volume_ik.sum() 
+                  + self.init_volume_uk.sum() + self.init_volume_dk.sum())
+        return result
+
+    @property
+    def total_volume(self):
+        result = (self.volume_j.sum() + self.volume_Ik.sum() + self.volume_ik.sum() 
+                  + self.volume_uk.sum() + self.volume_dk.sum())
+        return result
+
+    @property
+    def total_volume_flux(self):
+        result = (self.volume_flux_j.sum() + self.volume_flux_Ik.sum())
+        return result
+
+    @property
+    def cumulative_volume_flux(self):
+        result = (self.cumulative_vol_flux_j.sum() + self.cumulative_vol_flux_Ik.sum())
+        return result
+
+    def set_init_volume(self):
+        model = self.model
+        self.init_volume_j = volume_j(model)
+        self.init_volume_Ik = volume_Ik(model)
+        self.init_volume_ik = volume_ik(model)
+        self.init_volume_uk = volume_uk(model)
+        self.init_volume_dk = volume_dk(model)
+
+    def __on_step_end__(self, *args, **kwargs):
+        model = self.model
+        dt = model._dt
+        Q_in = model._Q_in
+        Q_0Ik = model._Q_0Ik
+        Q_bc = model._Q_bc
+        self.volume_j = volume_j(model)
+        self.volume_Ik = volume_Ik(model)
+        self.volume_ik = volume_ik(model)
+        self.volume_uk = volume_uk(model)
+        self.volume_dk = volume_dk(model)
+        self.volume_flux_j = volume_flux_j(Q_in, Q_bc, dt)
+        self.volume_flux_Ik = volume_flux_Ik(Q_0Ik, dt)
+        self.cumulative_vol_flux_j += self.volume_flux_j
+        self.cumulative_vol_flux_Ik += self.volume_flux_Ik
+
+
+class ConvergenceTracker(BaseCallback):
+    def __init__(self, model, xtol=1e-6, max_learning_rate=0.5, min_learning_rate=0.01, beta=2.):
+        self.model = model
+        self.x_old = self._compute_prior_guess()
+        self.x_new = self._compute_next_guess()
+        self.dx = self._compute_guess_difference(self.x_old, self.x_new)
+        self.xtol = xtol
+        self.max_learning_rate = max_learning_rate
+        self.min_learning_rate = min_learning_rate
+        self.learning_rate = max_learning_rate
+        self.beta = beta
+        self.success = None
+        self.convergence_queue = []
+        self.learning_queue = []
+        self.convergence_metric = np.inf
+
+    def _compute_step_ratio(self, dx):
+        # TODO: Need to add weirs, orifices, pumps
+        # TODO: Need to account for elements with multiple outflows
+        model = self.model
+        # Compute bounds on allowable discharges
+        Q_ik_ub = model._E_Ik[model._Ik] * model._h_Ik[model._Ik]
+        Q_ik_lb = -model._E_Ik[model._Ip1k] * model._h_Ik[model._Ip1k]
+        Q_uk_ub = (model._A_sj[model._J_uk] * model.H_j[model._J_uk] 
+                + (model._B_uk * model._dx_uk / 2) 
+                * ((model._theta_uk * (model.H_j[model._J_uk] - model._z_inv_j[model._J_uk]) 
+                    + model._h_Ik[model._I_1k]) / 2)) / model._dt
+        Q_uk_lb = -model._E_Ik[model._I_1k] * model._h_Ik[model._I_1k]
+        Q_dk_ub = model._E_Ik[model._I_Np1k] * model._h_Ik[model._I_Np1k]
+        Q_dk_lb = -(model._A_sj[model._J_dk] * model.H_j[model._J_dk] 
+                + (model._B_dk * model._dx_dk / 2) 
+                * ((model._theta_dk * (model.H_j[model._J_dk] - model._z_inv_j[model._J_dk]) 
+                    + model._h_Ik[model._I_Np1k]) / 2)) / model._dt
+        Q_w_ub = (model._A_sj[model._J_uw] * model.H_j[model._J_uw]) / model._dt
+        Q_w_lb = -(model._A_sj[model._J_dw] * model.H_j[model._J_dw]) / model._dt
+        Q_o_ub = (model._A_sj[model._J_uo] * model.H_j[model._J_uo]) / model._dt
+        Q_o_lb = -(model._A_sj[model._J_do] * model.H_j[model._J_do]) / model._dt
+        Q_p_ub = (model._A_sj[model._J_up] * model.H_j[model._J_up]) / model._dt
+        Q_p_lb = -(model._A_sj[model._J_dp] * model.H_j[model._J_dp]) / model._dt
+
+
+        # Compute binding step ratio
+        ratio_ub = max(max(max_if(dx.get('Q_ik', 0.) / Q_ik_ub), 0.), 
+                       max(max_if(dx.get('Q_uk', 0.) / Q_uk_ub), 0.),
+                       max(max_if(dx.get('Q_dk', 0.) / Q_dk_ub), 0.),
+                       max(max_if(dx.get('Q_w', 0.) / Q_w_ub), 0.),
+                       max(max_if(dx.get('Q_o', 0.) / Q_o_ub), 0.),
+                       max(max_if(dx.get('Q_p', 0.) / Q_p_ub), 0.)
+                       )
+        # TODO: Check if this should be max or min for maxif
+        ####################################################
+        # Negative signs removed before max_if because lb should be negative
+        ratio_lb = max(max(max_if(dx.get('Q_ik', 0.) / Q_ik_lb), 0.), 
+                       max(max_if(dx.get('Q_uk', 0.) / Q_uk_lb), 0.), 
+                       max(max_if(dx.get('Q_dk', 0.) / Q_dk_lb), 0.),
+                       max(max_if(dx.get('Q_w', 0.) / Q_w_lb), 0.),
+                       max(max_if(dx.get('Q_o', 0.) / Q_o_lb), 0.),
+                       max(max_if(dx.get('Q_p', 0.) / Q_p_lb), 0.)
+                       )
+        step_ratio = max(ratio_ub, ratio_lb)
+        return step_ratio
+
+    def _compute_learning_rate(self, step_ratio):
+        max_learning_rate = self.max_learning_rate
+        min_learning_rate = self.min_learning_rate
+        beta = self.beta
+        if step_ratio > 0:
+            learning_rate = min(1 / step_ratio / beta, 1.)
+        else:
+            learning_rate = 1.
+        learning_rate = max(min(learning_rate, max_learning_rate), min_learning_rate)
+        return learning_rate
+
+    def _convergence_metric(self, dx, x_old, x_new):
+        assert x_old.keys() == x_new.keys() == dx.keys()
+        abs_dx = {k : np.abs(dx[k]) for k in dx}
+        abs_mag = {k : 1. + np.maximum(np.abs(x_new[k]), np.abs(x_old[k])) for k in x_new}
+        max_step = max([(abs_dx[k] / abs_mag[k]).max() for k in abs_dx])
+        return max_step
+
+    def _convergence_met(self, convergence_metric, xtol):
+        condition = (convergence_metric <= xtol)
+        return condition
+
+    def _compute_prior_guess(self):
+        state = self.model.return_state()
+        x_old = {k : v for k, v in state.items() if v.size > 0}
+        return x_old
+    
+    def _compute_next_guess(self):
+        state = self.model.return_state()
+        x_new = {k : v for k, v in state.items() if v.size > 0}
+        return x_new
+
+    def _compute_guess_difference(self, x_old, x_new):
+        assert x_old.keys() == x_new.keys()
+        dx = {k : x_new[k] - x_old[k] for k in x_new}
+        return dx
+
+    def _set_states(self, x_old, x_new, learning_rate):
+        model = self.model
+        for state_name in x_new:
+            prior_state = x_old[state_name]
+            next_state = x_new[state_name]
+            updated_state = (1 - learning_rate) * prior_state + (learning_rate) * next_state
+            setattr(model, state_name, updated_state)
+
+    def __on_step_start__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, dt=None,
+                          first_time=False, implicit=True, banded=None, first_iter=True,
+                          num_iter=0, rtol=None, atol=None, head_tol=0.0015):
+       self.x_old = self._compute_prior_guess() 
+       self.success = False
+    
+    def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, dt=None,
+                        first_time=False, implicit=True, banded=None, first_iter=True,
+                        num_iter=1, rtol=None, atol=None, head_tol=0.0015):
+        # Perform fixed-point iteration until convergence
+        iter_elapsed = 1
+        if (num_iter > 0):
+            self.convergence_queue = []
+            self.learning_queue = []
+            self.x_new = self._compute_next_guess()
+            self.dx = self._compute_guess_difference(self.x_old, self.x_new)
+            self.convergence_metric = self._convergence_metric(self.dx, self.x_old, self.x_new)
+            self.success = self._convergence_met(self.convergence_metric, self.xtol)
+            if not self.success:
+                for _ in range(num_iter):
+                    # TODO: Rename this to step count
+                    self.model.iter_count -= 1
+                    self.model.t -= dt
+                    self.x_old = self._compute_prior_guess()
+                    # Enforce minimum depth
+                    self.model.H_j = np.maximum(self.model.H_j, self.model._z_inv_j + self.model.min_depth)
+                    self.model.h_Ik = np.maximum(self.model.h_Ik, self.model.min_depth)
+                    try:
+                        self.model._setup_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                                            first_time=first_time, implicit=implicit, banded=banded,
+                                            first_iter=False)
+                        self.model._solve_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                                            first_time=first_time, implicit=implicit, banded=banded,
+                                            first_iter=False)
+                    except:
+                        self.model.iter_elapsed = iter_elapsed 
+                        #self.model.load_state()
+                        raise
+                    self.x_new = self._compute_next_guess()
+                    self.dx = self._compute_guess_difference(self.x_old, self.x_new)
+                    self.step_ratio = self._compute_step_ratio(self.dx)
+                    self.learning_rate = self._compute_learning_rate(self.step_ratio)
+                    self._set_states(self.x_old, self.x_new, self.learning_rate)
+                    self.convergence_metric = self._convergence_metric(self.dx, self.x_old, self.x_new)
+                    self.success = self._convergence_met(self.convergence_metric, self.xtol)
+                    self.convergence_queue.append(self.convergence_metric)
+                    self.learning_queue.append(self.learning_rate)
+                    iter_elapsed += 1
+                    if self.success:
+                        break
+        self.model.iter_elapsed = iter_elapsed 
+        # Enforce minimum depth
+        self.model.H_j = np.maximum(self.model.H_j, self.model._z_inv_j + self.model.min_depth)
+        self.model.h_Ik = np.maximum(self.model.h_Ik, self.model.min_depth)
+
+
+class ExperimentalConvergenceTracker(ConvergenceTracker):
+
+    def _compute_step_ratio(self, dx):
+        # TODO: Need to add weirs, orifices, pumps
+        model = self.model
+        _D_Ik = model._D_Ik
+        _Q_ik = model.Q_ik
+        _Q_uk = model.Q_uk
+        _Q_dk = model.Q_dk
+        _Q_o = model.Q_o
+        _Q_w = model.Q_w
+        _Q_p = model.Q_p
+        _kI = model._kI
+        _forward_I_i = model.forward_I_i
+        _backward_I_i = model.backward_I_i
+        _is_start = model._is_start
+        _is_end = model._is_end
+        _D_j = model.b
+        _J_uk = model._J_uk
+        _J_dk = model._J_dk
+        _J_uo = model._J_uo
+        _J_do = model._J_do
+        _J_uw = model._J_uw
+        _J_dw = model._J_dw
+        _J_up = model._J_up
+        _J_dp = model._J_dp
+        default = np.array([], dtype=np.float64)
+        _dQ_ik = self.dx.get('Q_ik', default)
+        _dQ_uk = self.dx.get('Q_uk', default)
+        _dQ_dk = self.dx.get('Q_dk', default)
+        _dQ_o = self.dx.get('Q_o', default)
+        _dQ_w = self.dx.get('Q_w', default)
+        _dQ_p = self.dx.get('Q_p', default)
+        num_I = junction_numerator(_dQ_ik, _dQ_uk, _dQ_dk, _kI, 
+                                   _forward_I_i, _backward_I_i, _is_start, _is_end)
+        denom_I = junction_denominator(_D_Ik, _Q_ik, _Q_uk, _Q_dk, _kI, 
+                                       _forward_I_i, _backward_I_i, _is_start, _is_end)
+        num_j = superjunction_numerator(_D_j, _dQ_uk, _dQ_dk, _dQ_o, _dQ_w, _dQ_p, 
+                                        _J_uk, _J_dk, _J_uo, _J_do, _J_uw, _J_dw, _J_up, _J_dp)
+        denom_j = superjunction_denominator(_D_j, _Q_uk, _Q_dk, _Q_o, _Q_w, _Q_p, 
+                                            _J_uk, _J_dk, _J_uo, _J_do, _J_uw, _J_dw, _J_up, _J_dp)
+        beta_I = (num_I / denom_I).max()
+        beta_j = (num_j / denom_j).max()
+        step_ratio = max(beta_I, beta_j)
+        return step_ratio
+
+    def _compute_learning_rate(self, step_ratio):
+        max_learning_rate = self.max_learning_rate
+        min_learning_rate = self.min_learning_rate
+        beta = self.beta
+        if step_ratio > 0:
+            learning_rate = min(1 / step_ratio / beta, 1.)
+        else:
+            learning_rate = 1.
+        learning_rate = max(min(learning_rate, max_learning_rate), min_learning_rate)
+        return learning_rate
+
+
+class LegacyConvergenceTracker(BaseCallback):
+    def __init__(self, model):
+        self.model = model
+        self.prior_guess = 0.
+        self.next_guess = 0.
+        self.learning_rate = 0.5
+        self.min_depth = 1e-5
+
+    def _convergence_met(self, prior_guess, next_guess, head_tol=0.0015):
+        e = np.abs(next_guess - prior_guess)
+        condition = e.max() <= head_tol
+        return condition
+
+    def _compute_prior_guess(self):
+        return self.model.H_j
+    
+    def _compute_next_guess(self):
+        return self.model.H_j
+
+    def _set_states(self, prior_states, next_states):
+        model = self.model
+        alpha = self.learning_rate
+        for state_name in next_states:
+            prior_state = prior_states[state_name]
+            next_state = next_states[state_name]
+            updated_state = (1 - alpha) * prior_state + (alpha) * next_state
+            setattr(model, state_name, updated_state)
+
+    def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, dt=None,
+                        first_time=False, implicit=True, banded=None, first_iter=True,
+                        num_iter=1, rtol=None, atol=None, head_tol=0.0015):
+        self.prior_guess = self._compute_prior_guess() 
+
+    def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, dt=None,
+                        first_time=False, implicit=True, banded=None, first_iter=True,
+                        num_iter=1, rtol=None, atol=None, head_tol=0.0015):
+        # Perform fixed-point iteration until convergence
+        iter_elapsed = 1
+        if (num_iter > 0):
+            self.next_guess = self._compute_next_guess()
+            convergence_met = self._convergence_met(self.prior_guess, self.next_guess, head_tol=head_tol)
+            if not convergence_met:
+                for _ in range(num_iter):
+                    # TODO: Rename this to step count
+                    self.model.iter_count -= 1
+                    self.model.t -= dt
+                    self.prior_guess = self._compute_prior_guess()
+                    self.prior_states = self.model.return_state()
+                    self.model.H_j = np.maximum(self.model.H_j, self.model._z_inv_j + self.min_depth)
+                    self.model.h_Ik = np.maximum(self.model.h_Ik, self.min_depth)
+                    try:
+                        self.model._setup_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                                            first_time=first_time, implicit=implicit, banded=banded,
+                                            first_iter=False)
+                        self.model._solve_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                                            first_time=first_time, implicit=implicit, banded=banded,
+                                            first_iter=False)
+                    except:
+                        self.model.load_state()
+                        raise
+                    self.next_guess = self._compute_next_guess()
+                    self.next_states = self.model.return_state()
+                    iter_elapsed += 1
+                    convergence_met = self._convergence_met(self.prior_guess, self.next_guess, head_tol=head_tol)
+                    self._set_states(self.prior_states, self.next_states)
+                    if convergence_met:
+                        break
+        self.model.iter_elapsed = iter_elapsed 
+
+def continuity_error_j(model, dt):
+    """
+    err = [A_sj (dH_j / dt) + Q_dk - Q_uk] - [Q_in] 
+    """
+    error = np.zeros(model.M)
+    H_j_next = model.H_j
+    H_j_prev = model.states['H_j']
+    bc = model.bc
+    Q_in = model._Q_in
+    Q_bc = model._Q_bc
+    error += model.A_sj / dt
+    np.add.at(error, model._J_uk, model._B_uk * model._dx_uk * model._theta_uk / 2 / dt)
+    np.add.at(error, model._J_dk, model._B_dk * model._dx_dk * model._theta_dk / 2 / dt)
+    error *= (H_j_next - H_j_prev)
+    np.add.at(error, model._J_uk, model._Q_uk)
+    np.subtract.at(error, model._J_dk, model._Q_dk)
+    np.add.at(error, model._J_uo, model._Qo)
+    np.subtract.at(error, model._J_do, model._Qo)
+    np.add.at(error, model._J_uw, model._Qw)
+    np.subtract.at(error, model._J_dw, model._Qw)
+    np.add.at(error, model._J_up, model._Qp)
+    np.subtract.at(error, model._J_dp, model._Qp)
+    error -= Q_in
+    #error -= Q_bc
+    error[bc] = 0.
+    return error
+
+def continuity_magnitude_j(model, dt):
+    mag = np.zeros(model.M)
+    H_j_next = model.H_j
+    H_j_prev = model.states['H_j']
+    mag += model.A_sj / dt
+    np.add.at(mag, model._J_uk, model._B_uk * model._dx_uk * model._theta_uk / 2 / dt)
+    np.add.at(mag, model._J_dk, model._B_dk * model._dx_dk * model._theta_dk / 2 / dt)
+    mag *= np.maximum(np.abs(H_j_next), np.abs(H_j_prev))
+    return mag
+
+def continuity_increment_j(model, dt):
+    inc = np.zeros(model.M)
+    H_j_next = model.H_j
+    H_j_prev = model.states['H_j']
+    inc += model.A_sj / dt
+    np.add.at(inc, model._J_uk, model._B_uk * model._dx_uk * model._theta_uk / 2 / dt)
+    np.add.at(inc, model._J_dk, model._B_dk * model._dx_dk * model._theta_dk / 2 / dt)
+    inc *= (H_j_next - H_j_prev)
+    return inc
+
+def momentum_error_uk(model, dt):
+    error = np.zeros(model.NK)
+    g = 9.81
+    Q_uk_next = model.Q_uk
+    h_uk_next = model._h_uk
+    H_juk_next = model.H_j[model._J_uk]
+    error += g * model._A_uk * model._kappa_uk * Q_uk_next
+    error += g * model._A_uk * model._lambda_uk * H_juk_next
+    error += g * model._A_uk * model._mu_uk
+    error -= g * model._A_uk * h_uk_next
+    # Friction and local losses
+    #error += 0.
+    return error
+
+def momentum_error_uk_2(model, dt):
+    error = np.zeros(model.NK)
+    g = 9.81
+    Q_1k_next = model.Q_ik[model._i_1k]
+    Q_uk_next = model.Q_uk
+    h_uk_next = model._h_uk
+    H_juk_next = model.H_j[model._J_uk]
+    b_uk = model._b_uk
+    c_uk = model._c_uk
+    P_uk = model._P_uk
+    A_uk = model._A_uk
+    LHS = b_uk * Q_uk_next + c_uk * Q_1k_next
+    RHS = P_uk + g * A_uk * (H_juk_next - h_uk_next)
+    error = LHS - RHS
+    return error
+
+# To be used with tsuperlink
+def momentum_error_uk_3(model, dt):
+    error = np.zeros(model.NK)
+    Q_uk_next = model.Q_uk
+    h_uk_next = model._h_uk
+    H_juk_next = model.H_j[model._J_uk]
+    a_uk = model._a_uk
+    b_uk = model._b_uk
+    c_uk = model._c_uk
+    P_uk = model._P_uk
+    LHS = a_uk * H_juk_next + b_uk * Q_uk_next + c_uk * h_uk_next
+    RHS = P_uk
+    error = LHS - RHS
+    return error
+
+def momentum_magnitude_uk(model, dt):
+    Q_uk_next = model.Q_uk
+    Q_uk_prev = model.states['Q_uk']
+    mag = np.maximum(np.abs(Q_uk_next), np.abs(Q_uk_prev))
+    return mag
+
+def momentum_error_dk(model, dt):
+    error = np.zeros(model.NK)
+    g = 9.81
+    Q_dk_next = model.Q_dk
+    h_dk_next = model._h_dk
+    H_jdk_next = model.H_j[model._J_dk]
+    error += g * model._A_dk * model._kappa_dk * Q_dk_next
+    error += g * model._A_dk * model._lambda_dk * H_jdk_next
+    error += g * model._A_dk * model._mu_dk
+    error -= g * model._A_dk * h_dk_next
+    # Friction and local losses
+    #error += 0.
+    return error
+
+def momentum_error_dk_2(model, dt):
+    error = np.zeros(model.NK)
+    g = 9.81
+    Q_dk_next = model.Q_dk
+    Q_nk_next = model.Q_ik[model._i_nk]
+    h_dk_next = model._h_dk
+    H_jdk_next = model.H_j[model._J_dk]
+    b_dk = model._b_dk
+    a_dk = model._a_dk
+    A_dk = model._A_dk
+    P_dk = model._P_dk
+    # Friction and local losses
+    LHS = b_dk * Q_dk_next + a_dk * Q_nk_next
+    RHS = P_dk + g * A_dk * (h_dk_next - H_jdk_next)
+    error = LHS - RHS
+    return error
+
+# To be used with tsuperlink
+def momentum_error_dk_3(model, dt):
+    error = np.zeros(model.NK)
+    Q_dk_next = model.Q_dk
+    h_dk_next = model._h_dk
+    H_jdk_next = model.H_j[model._J_dk]
+    a_dk = model._a_dk
+    b_dk = model._b_dk
+    c_dk = model._c_dk
+    P_dk = model._P_dk
+    LHS = a_dk * h_dk_next + b_dk * Q_dk_next + c_dk * H_jdk_next
+    RHS = P_dk
+    error = LHS - RHS
+    return error
+
+def momentum_magnitude_dk(model, dt):
+    Q_dk_next = model.Q_dk
+    Q_dk_prev = model.states['Q_dk']
+    mag = np.maximum(np.abs(Q_dk_next), np.abs(Q_dk_prev))
+    return mag
+
+def momentum_error_o(model, dt):
+    error = np.zeros(model.n_o)
+    raise NotImplementedError
+
+def momentum_error_w(model, dt):
+    error = np.zeros(model.n_w)
+    raise NotImplementedError
+
+def momentum_error_p(model, dt):
+    error = np.zeros(model.n_p)
+    raise NotImplementedError
+
+def continuity_error_Ik(model, dt):
+    """
+    err = [E_Ik h_Ik - Q_im1k + Q_ik] - [D_Ik] 
+    """
+    error = np.zeros(model._I.size)
+    h_Ik_next = model.h_Ik
+    h_Ik_prev = model.states['h_Ik']
+    _I_internal = (~model._I_start) & (~model._I_end)
+    error += model._E_Ik * h_Ik_next
+    error -= model._D_Ik
+    error[model._I_start] -= model._Q_uk
+    error[model._I_start] += model._Q_ik[model._i_1k]
+    error[model._I_end] -= model._Q_ik[model._i_nk]
+    error[model._I_end] += model._Q_dk
+    error[_I_internal] -= model._Q_ik[model.backward_I_i[_I_internal]]
+    error[_I_internal] += model._Q_ik[model.forward_I_i[_I_internal]]
+    return error
+
+def continuity_magnitude_Ik(model, dt):
+    h_Ik_next = model.h_Ik
+    h_Ik_prev = model.states['h_Ik']
+    mag = model._E_Ik * dt * np.maximum(np.abs(h_Ik_next), np.abs(h_Ik_prev))
+    return mag
+
+def momentum_error_ik(model, dt):
+    error = np.zeros(model._i.size)
+    _i_1k = model._i_1k
+    _i_nk = model._i_nk
+    _i_is_start = np.zeros(model._i.size, dtype=np.bool_)
+    _i_is_start[_i_1k] = True
+    _i_is_end = np.zeros(model._i.size, dtype=np.bool_)
+    _i_is_end[_i_nk] = True
+    _i_is_internal = (~_i_is_start) & (~_i_is_end)
+    _im1 = (model._i - 1)[_i_is_internal]
+    _ip1 = (model._i + 1)[_i_is_internal]
+    g = 9.81
+    Q_ik_next = model.Q_ik
+    Q_ik_prev = model.states['Q_ik']
+    h_Ik_next = model.h_Ik
+    error -= model._P_ik
+    error -= g * model._A_ik * (h_Ik_next[model._Ik] - h_Ik_next[model._Ip1k])
+    error += model._b_ik * Q_ik_next
+    error[_i_is_start] += model._a_ik[_i_is_start] * model.Q_uk
+    error[_i_is_end] += model._c_ik[_i_is_end] * model.Q_dk
+    # Addition 2025-06-05
+    if _i_is_internal.any():
+        error[_i_is_start] += model._c_ik[_i_is_start] * model.Q_ik[_i_1k + 1]
+        error[_i_is_end] += model._a_ik[_i_is_end] * model.Q_ik[_i_nk - 1]
+        error[_i_is_internal] += model._a_ik[_i_is_internal] * model.Q_ik[_im1]
+        error[_i_is_internal] += model._c_ik[_i_is_internal] * model.Q_ik[_ip1]
+    return error
+
+# To be used with tsuperlink
+def momentum_error_ik_2(model, dt):
+    error = np.zeros(model._i.size)
+    Q_ik_next = model.Q_ik
+    h_Ik_next = model.h_Ik
+    error -= model._P_ik
+    error += model._a_ik * h_Ik_next[model._Ik]
+    error += model._b_ik * Q_ik_next
+    error += model._c_ik * h_Ik_next[model._Ip1k]
+    return error
+
+def momentum_magnitude_ik(model, dt):
+    Q_ik_next = model.Q_ik
+    Q_ik_prev = model.states['Q_ik']
+    mag = np.maximum(np.abs(Q_ik_next), np.abs(Q_ik_prev))
+    return mag
+
+def compute_error(model, dt):
+    error_j = continuity_error_j(model, dt)
+    error_Ik = continuity_error_Ik(model, dt)
+    error_ik = momentum_error_ik(model, dt)
+    error = np.concatenate([error_j, error_Ik, error_ik])
+    return error
+
+def volume_j(model):
+    # Import instance variables
+    _functional = model._functional              # Superlinks with functional area curves
+    _tabular = model._tabular                    # Superlinks with tabular area curves
+    _storage_a = model._storage_a                # Coefficient of functional storage curve
+    _storage_b = model._storage_b                # Exponent of functional storage curve
+    _storage_c = model._storage_c                # Constant of functional storage curve
+    H_j = model.H_j                              # Head at superjunction j
+    _z_inv_j = model._z_inv_j                    # Invert elevation at superjunction j
+    min_depth = model.min_depth                  # Minimum depth allowed at superjunctions/nodes
+    _storage_hs = model._storage_hs
+    _storage_As = model._storage_As
+    _storage_Vs = model._storage_Vs
+    _storage_inds = model._storage_inds
+    _storage_lens = model._storage_lens
+    _storage_js = model._storage_js
+    _storage_codes = model._storage_codes
+    # Compute storage areas
+    _V_sj = np.zeros(model.M, dtype=np.float64)
+    _h_j = np.maximum(H_j - _z_inv_j, min_depth)
+    numba_compute_functional_storage_volumes(_h_j, _V_sj, _storage_a, _storage_b,
+                                                _storage_c, _functional)
+    if _tabular.any():
+        numba_compute_tabular_storage_volumes(_h_j, _V_sj, _storage_hs, _storage_As,
+                                                _storage_Vs, _storage_js, _storage_codes,
+                                                _storage_inds, _storage_lens)
+    return _V_sj
+
+def volume_uk(model):
+    V_uk = model._A_uk * model._dx_uk
+    return V_uk
+
+def volume_dk(model):
+    V_dk = model._A_dk * model._dx_dk
+    return V_dk
+
+def volume_Ik(model):
+    V_Ik = model._A_SIk * model.h_Ik
+    return V_Ik
+
+def volume_ik(model):
+    V_ik = model._A_ik * model._dx_ik
+    return V_ik
+
+def volume_flux_j(Q_in, Q_bc, dt):
+    V_in_exog_j = Q_in * dt
+    V_bc_j = Q_bc * dt
+    V_in_j = V_in_exog_j + V_bc_j
+    return V_in_j
+
+def volume_flux_Ik(Q_0Ik, dt):
+    V_in_Ik = Q_0Ik * dt
+    return V_in_Ik
+
+def total_volume(model):
+    V_sj = volume_j(model)
+    V_uk = volume_uk(model)
+    V_dk = volume_dk(model)
+    V_o = np.zeros(model.n_o, dtype=np.float64)
+    V_w = np.zeros(model.n_w, dtype=np.float64)
+    V_p = np.zeros(model.n_p, dtype=np.float64)
+    V_Ik = volume_Ik(model)
+    V_ik = volume_ik(model)
+    V_total = np.concatenate([V_sj, V_uk, V_dk, V_o, V_w, V_p, V_Ik, V_ik])
+    return V_total
+
+def max_if(arr):
+    if arr.size > 0:
+        return arr.max()
+    else:
+        return 0.
+
+def min_if(arr):
+    if arr.size > 0:
+        return arr.min()
+    else:
+        return 0.
\ No newline at end of file
diff --git a/pipedream_solver/ngeometry.py b/pipedream_solver/ngeometry.py
index 84ee754..edbd2b6 100644
--- a/pipedream_solver/ngeometry.py
+++ b/pipedream_solver/ngeometry.py
@@ -3,6 +3,8 @@
 from numba import njit
 from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, DictType, void
 
+MIN_DEPTH = 1e-5
+
 geom_code = {
     'circular' : 1,
     'rect_closed' : 2,
@@ -20,7 +22,7 @@
 
 @njit(float64(float64, float64, float64),
       cache=True)
-def Circular_A_ik(h_Ik, h_Ip1k, g1):
+def Circular_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -34,9 +36,10 @@ def Circular_A_ik(h_Ik, h_Ip1k, g1):
         Diameter of channel (meters)
     """
     d = g1
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    pslot = g2
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > d:
         y = d
     r = d / 2
@@ -51,7 +54,7 @@ def Circular_A_ik(h_Ik, h_Ip1k, g1):
 
 @njit(float64(float64, float64, float64),
       cache=True)
-def Circular_Pe_ik(h_Ik, h_Ip1k, g1):
+def Circular_Pe_ik(h_ik, g1, g2):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -65,9 +68,10 @@ def Circular_Pe_ik(h_Ik, h_Ip1k, g1):
         Diameter of channel (meters)
     """
     d = g1
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    pslot = g2
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > d:
         y = d
     r = d / 2
@@ -100,9 +104,9 @@ def Circular_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Circular_B_ik(h_Ik, h_Ip1k, g1, g2):
+def Circular_B_ik(h_ik, g1, g2):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -119,10 +123,9 @@ def Circular_B_ik(h_Ik, h_Ip1k, g1, g2):
     """
     d = g1
     pslot = g2
-    y = (h_Ik + h_Ip1k) / 2
-    # y[y < 0] = 0
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     r = d / 2
     phi = y / r
     if phi < 0:
@@ -134,13 +137,15 @@ def Circular_B_ik(h_Ik, h_Ip1k, g1, g2):
     if cond:
         B = 2 * r * np.sin(theta)
     else:
+        # TODO: Use absolute value instead of fraction?
         B = pslot * d
+    B = max(pslot * d, B)
     return B
 
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Rect_Closed_A_ik(h_Ik, h_Ip1k, g1, g2):
+def Rect_Closed_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -157,17 +162,17 @@ def Rect_Closed_A_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     A = y * b
     return A
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Rect_Closed_Pe_ik(h_Ik, h_Ip1k, g1, g2):
+def Rect_Closed_Pe_ik(h_ik, g1, g2):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -184,9 +189,9 @@ def Rect_Closed_Pe_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     Pe = b + 2 * y
@@ -212,9 +217,9 @@ def Rect_Closed_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64, float64),
+@njit(float64(float64, float64, float64, float64),
       cache=True)
-def Rect_Closed_B_ik(h_Ik, h_Ip1k, g1, g2, g3):
+def Rect_Closed_B_ik(h_ik, g1, g2, g3):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -234,9 +239,9 @@ def Rect_Closed_B_ik(h_Ik, h_Ip1k, g1, g2, g3):
     y_max = g1
     b = g2
     pslot = g3
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     cond = (y < y_max)
     if cond:
         B = b
@@ -245,9 +250,9 @@ def Rect_Closed_B_ik(h_Ik, h_Ip1k, g1, g2, g3):
     return B
 
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Rect_Open_A_ik(h_Ik, h_Ip1k, g1, g2):
+def Rect_Open_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -264,17 +269,17 @@ def Rect_Open_A_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     A = y * b
     return A
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Rect_Open_Pe_ik(h_Ik, h_Ip1k, g1, g2):
+def Rect_Open_Pe_ik(h_ik, g1, g2):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -291,9 +296,9 @@ def Rect_Open_Pe_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     Pe = b + 2 * y
@@ -319,9 +324,9 @@ def Rect_Open_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Rect_Open_B_ik(h_Ik, h_Ip1k, g1, g2):
+def Rect_Open_B_ik(h_ik, g1, g2):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -341,9 +346,9 @@ def Rect_Open_B_ik(h_Ik, h_Ip1k, g1, g2):
     return b
 
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Triangular_A_ik(h_Ik, h_Ip1k, g1, g2):
+def Triangular_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -360,17 +365,17 @@ def Triangular_A_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     m = g2
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     A = m * y**2
     return A
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Triangular_Pe_ik(h_Ik, h_Ip1k, g1, g2):
+def Triangular_Pe_ik(h_ik, g1, g2):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -387,9 +392,9 @@ def Triangular_Pe_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     m = g2
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     Pe = 2 * y * np.sqrt(1 + m**2)
@@ -415,9 +420,9 @@ def Triangular_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Triangular_B_ik(h_Ik, h_Ip1k, g1, g2):
+def Triangular_B_ik(h_ik, g1, g2):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -434,9 +439,9 @@ def Triangular_B_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     m = g2
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     cond = (y < y_max)
     if cond:
         B = 2 * m * y
@@ -445,9 +450,9 @@ def Triangular_B_ik(h_Ik, h_Ip1k, g1, g2):
     return B
 
 
-@njit(float64(float64, float64, float64, float64, float64),
+@njit(float64(float64, float64, float64, float64),
       cache=True)
-def Trapezoidal_A_ik(h_Ik, h_Ip1k, g1, g2, g3):
+def Trapezoidal_A_ik(h_ik, g1, g2, g3):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -467,17 +472,17 @@ def Trapezoidal_A_ik(h_Ik, h_Ip1k, g1, g2, g3):
     y_max = g1
     b = g2
     m = g3
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     A = y * (b + m * y)
     return A
 
-@njit(float64(float64, float64, float64, float64, float64),
+@njit(float64(float64, float64, float64, float64),
       cache=True)
-def Trapezoidal_Pe_ik(h_Ik, h_Ip1k, g1, g2, g3):
+def Trapezoidal_Pe_ik(h_ik, g1, g2, g3):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -497,9 +502,9 @@ def Trapezoidal_Pe_ik(h_Ik, h_Ip1k, g1, g2, g3):
     y_max = g1
     b = g2
     m = g3
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     if y > y_max:
         y = y_max
     Pe = b + 2 * y * np.sqrt(1 + m**2)
@@ -525,9 +530,9 @@ def Trapezoidal_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64, float64),
+@njit(float64(float64, float64, float64, float64),
       cache=True)
-def Trapezoidal_B_ik(h_Ik, h_Ip1k, g1, g2, g3):
+def Trapezoidal_B_ik(h_ik, g1, g2, g3):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -547,9 +552,9 @@ def Trapezoidal_B_ik(h_Ik, h_Ip1k, g1, g2, g3):
     y_max = g1
     b = g2
     m = g3
-    y = (h_Ik + h_Ip1k) / 2
-    if y < 0:
-        y = 0
+    y = h_ik
+    if y < MIN_DEPTH:
+        y = MIN_DEPTH
     cond = (y < y_max)
     if cond:
         B = b + 2 * m * y
@@ -559,7 +564,38 @@ def Trapezoidal_B_ik(h_Ik, h_Ip1k, g1, g2, g3):
 
 @njit(float64(float64, float64, float64, float64),
       cache=True)
-def Parabolic_A_ik(h_Ik, h_Ip1k, g1, g2):
+def Trapezoidal_dA_dh(h, g1, g2, g3):
+    h_max = g1
+    b = g2
+    m = g3
+    if h < 0:
+        h = 0
+    if h > h_max:
+        h = h_max
+    dA_dh = b + 2 * m * h
+    return dA_dh
+
+@njit(float64(float64, float64, float64, float64),
+      cache=True)
+def Trapezoidal_dPe_dh(h, g1, g2, g3):
+    h_max = g1
+    b = g2
+    m = g3
+    dPe_dh = 2 * np.sqrt(1 + m**2)
+    return dPe_dh
+
+@njit(float64(float64, float64, float64, float64),
+      cache=True)
+def Trapezoidal_dB_dh(h, g1, g2, g3):
+    h_max = g1
+    b = g2
+    m = g3
+    dB_dh = 2 * m
+    return dB_dh
+
+@njit(float64(float64, float64, float64),
+      cache=True)
+def Parabolic_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -576,7 +612,7 @@ def Parabolic_A_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0:
         y = 0
     if y > y_max:
@@ -584,9 +620,9 @@ def Parabolic_A_ik(h_Ik, h_Ip1k, g1, g2):
     A = 2 * b * y / 3
     return A
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Parabolic_Pe_ik(h_Ik, h_Ip1k, g1, g2):
+def Parabolic_Pe_ik(h_ik, g1, g2):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -603,7 +639,7 @@ def Parabolic_Pe_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y <= 0:
         y = eps
     if y > y_max:
@@ -632,9 +668,9 @@ def Parabolic_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Parabolic_B_ik(h_Ik, h_Ip1k, g1, g2):
+def Parabolic_B_ik(h_ik, g1, g2):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -651,15 +687,15 @@ def Parabolic_B_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0:
         y = 0
     B = b * np.sqrt(y / y_max)
     return B
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Elliptical_A_ik(h_Ik, h_Ip1k, g1, g2):
+def Elliptical_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -677,7 +713,7 @@ def Elliptical_A_ik(h_Ik, h_Ip1k, g1, g2):
     y_max = g1
     b = g1 / 2
     a = g2 / 2
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0:
         y = 0
     if y > y_max:
@@ -708,9 +744,9 @@ def Elliptical_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Elliptical_B_ik(h_Ik, h_Ip1k, g1, g2):
+def Elliptical_B_ik(h_ik, g1, g2):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -728,7 +764,7 @@ def Elliptical_B_ik(h_Ik, h_Ip1k, g1, g2):
     y_max = g1
     b = g1 / 2
     a = g2 / 2
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0:
         y = 0
     if y > y_max:
@@ -739,9 +775,9 @@ def Elliptical_B_ik(h_Ik, h_Ip1k, g1, g2):
     return B
 
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Wide_A_ik(h_Ik, h_Ip1k, g1, g2):
+def Wide_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -758,7 +794,7 @@ def Wide_A_ik(h_Ik, h_Ip1k, g1, g2):
     """
     y_max = g1
     b = g2
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0:
         y = 0
     if y > y_max:
@@ -766,9 +802,9 @@ def Wide_A_ik(h_Ik, h_Ip1k, g1, g2):
     A = y * b
     return A
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Wide_Pe_ik(h_Ik, h_Ip1k, g1, g2):
+def Wide_Pe_ik(h_ik, g1, g2):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -807,9 +843,9 @@ def Wide_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Wide_B_ik(h_Ik, h_Ip1k, g1, g2):
+def Wide_B_ik(h_ik, g1, g2):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -828,9 +864,9 @@ def Wide_B_ik(h_Ik, h_Ip1k, g1, g2):
     b = g2
     return b
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Force_Main_A_ik(h_Ik, h_Ip1k, g1, g2):
+def Force_Main_A_ik(h_ik, g1, g2):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -850,9 +886,9 @@ def Force_Main_A_ik(h_Ik, h_Ip1k, g1, g2):
     A = np.pi * r**2
     return A
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Force_Main_Pe_ik(h_Ik, h_Ip1k, g1, g2):
+def Force_Main_Pe_ik(h_ik, g1, g2):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -891,9 +927,9 @@ def Force_Main_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64),
+@njit(float64(float64, float64, float64),
       cache=True)
-def Force_Main_B_ik(h_Ik, h_Ip1k, g1, g2):
+def Force_Main_B_ik(h_ik, g1, g2):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -913,9 +949,9 @@ def Force_Main_B_ik(h_Ik, h_Ip1k, g1, g2):
     B = pslot * d
     return B
 
-@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64, float64),
+@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64),
       cache=True)
-def Floodplain_A_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7):
+def Floodplain_A_ik(h_ik, g1, g2, g3, g4, g5, g6, g7):
     """
     Compute cross-sectional area of flow for link i, superlink k.
 
@@ -950,7 +986,7 @@ def Floodplain_A_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7):
     m_top = g4
     m_base = g5
     m_middle = g6
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0.:
         y = 0.
     if y > h_max:
@@ -980,9 +1016,9 @@ def Floodplain_A_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7):
     return A
 
 
-@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64, float64),
+@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64),
       cache=True)
-def Floodplain_Pe_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7):
+def Floodplain_Pe_ik(h_ik, g1, g2, g3, g4, g5, g6, g7):
     """
     Compute perimeter of flow for link i, superlink k.
 
@@ -1013,7 +1049,7 @@ def Floodplain_Pe_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7):
     m_top = g4
     m_base = g5
     m_middle = g6
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0.:
         y = 0.
     if y > h_max:
@@ -1062,9 +1098,9 @@ def Floodplain_R_ik(A_ik, Pe_ik):
         R = 0
     return R
 
-@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64, float64),
+@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64),
       cache=True)
-def Floodplain_B_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7):
+def Floodplain_B_ik(h_ik, g1, g2, g3, g4, g5, g6, g7):
     """
     Compute top width of flow for link i, superlink k.
 
@@ -1095,7 +1131,7 @@ def Floodplain_B_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7):
     m_top = g4
     m_base = g5
     m_middle = g6
-    y = (h_Ik + h_Ip1k) / 2
+    y = h_ik
     if y < 0:
         y = 0.
     if y > h_max:
diff --git a/pipedream_solver/ninfiltration.py b/pipedream_solver/ninfiltration.py
index 908b2cf..d8ac37d 100644
--- a/pipedream_solver/ninfiltration.py
+++ b/pipedream_solver/ninfiltration.py
@@ -1,8 +1,9 @@
 import numpy as np
-from pipedream_solver.nutils import newton_raphson, bounded_newton_raphson, numba_any
-from pipedream_solver.infiltration import GreenAmpt
 from numba import njit
 import scipy.optimize
+from pipedream_solver.nutils import newton_raphson, bounded_newton_raphson, numba_any
+from pipedream_solver.infiltration import GreenAmpt
+from pipedream_solver._ninfiltration import *
 
 class nGreenAmpt(GreenAmpt):
     """
@@ -202,119 +203,3 @@ def step(self, dt, i):
         self.iter_count += 1
         self.t += dt
 
-@njit
-def run_green_ampt_newton(F_2, x0, F_1, dt, Ks, theta_d, psi_f, ia, max_iter=50,
-                          atol=1.48e-8, rtol=0.0, bounded=True):
-    """
-    Use Newton-Raphson iteration to find cumulative infiltration at next time step (F_2).
-
-    Inputs:
-    -------
-    F_2 : np.ndarray (float)
-        Cumulative infiltration at next time step (meters).
-    x0 : np.ndarray (float)
-        Initial guess for cumulative infiltration (meters)
-    F_1 : np.ndarray (float)
-        Cumulative infiltration at previous time step (meters)
-    dt : np.ndarray (float)
-        Time step (seconds)
-    Ks : np.ndarray (float)
-        Saturated hydraulic conductivity (m/s)
-    theta_d : np.ndarray (float)
-        Soil moisture deficit (-)
-    psi_f : np.ndarray (float)
-        Matric potential of the wetting front (m)
-    ia : np.ndarray (float)
-        Available rainfall depth (meters)
-    max_iter : int
-        Maximum number of Newton-Raphson iterations
-    atol : float
-        Allowable (absolute) error of the zero value
-    rtol : float
-        Allowable (relative) error of the zero value
-    bounded : bool
-        If True, use bounded Newton-Raphson iteration
-    """
-    n = F_2.size
-    for i in range(n):
-        x_0_i = x0[i]
-        F_1_i = F_1[i]
-        dt_i = dt[i]
-        nargs = np.zeros(5)
-        nargs[0] = F_1_i
-        nargs[1] = dt_i
-        nargs[2] = Ks[i]
-        nargs[3] = theta_d[i]
-        nargs[4] = psi_f[i]
-        if bounded:
-            min_F = 0
-            max_F = F_1_i + ia[i] * dt_i
-            F_est = bounded_newton_raphson(numba_integrated_green_ampt,
-                                        numba_derivative_green_ampt,
-                                        x_0_i, min_F, max_F,
-                                        nargs, max_iter=max_iter,
-                                        atol=atol, rtol=rtol)
-        else:
-            F_est = newton_raphson(numba_integrated_green_ampt,
-                                   numba_derivative_green_ampt,
-                                   x_0_i, nargs, max_iter=max_iter,
-                                   atol=atol, rtol=rtol)
-        F_2[i] = F_est
-    return F_2
-
-@njit
-def numba_integrated_green_ampt(F_2, args):
-    """
-    Solve integrated form of Green Ampt equation for cumulative infiltration.
-
-    Inputs:
-    -------
-    F_2: np.ndarray (float)
-        Cumulative infiltration at current timestep (m)
-    F_1: np.ndarray (float)
-        Cumulative infiltration at next timestep (m)
-    dt: float
-        Time step (seconds)
-    Ks: np.ndarray (float)
-        Saturated hydraulic conductivity (m/s)
-    theta_d: np.ndarray (float)
-        Soil moisture deficit
-    psi_s: np.ndarray (float)
-        Soil suction head (m)
-    """
-    F_1 = args[0]
-    dt = args[1]
-    Ks = args[2]
-    theta_d = args[3]
-    psi_s = args[4]
-    C = Ks * dt + F_1 - psi_s * theta_d * np.log(F_1 + np.abs(psi_s) * theta_d)
-    zero = C + psi_s * theta_d * np.log(F_2 + np.abs(psi_s) * theta_d) - F_2
-    return zero
-
-@njit
-def numba_derivative_green_ampt(F_2, args):
-    """
-    Derivative of Green Ampt equation for cumulative infiltration.
-
-    Inputs:
-    -------
-    F_2: np.ndarray (float)
-        Cumulative infiltration at current timestep (m)
-    F_1: np.ndarray (float)
-        Cumulative infiltration at next timestep (m)
-    dt: float
-        Time step (seconds)
-    Ks: np.ndarray (float)
-        Saturated hydraulic conductivity (m/s)
-    theta_d: np.ndarray (float)
-        Soil moisture deficit
-    psi_s: np.ndarray (float)
-        Soil suction head (m)
-    """
-    F_1 = args[0]
-    dt = args[1]
-    Ks = args[2]
-    theta_d = args[3]
-    psi_s = args[4]
-    zero = (psi_s * theta_d / (psi_s * theta_d + F_2)) - 1
-    return zero
diff --git a/pipedream_solver/nquality.py b/pipedream_solver/nquality.py
index 0818254..c68bde0 100644
--- a/pipedream_solver/nquality.py
+++ b/pipedream_solver/nquality.py
@@ -101,7 +101,7 @@ def __init__(self, hydraulics, superjunction_params, superlink_params,
         self._kI = self.hydraulics._kI
         self._K_j = superjunction_params['K'].values.astype(np.float64)
         self._c_j = superjunction_params['c_0'].values.astype(np.float64)
-        self.bc = superjunction_params['bc'].values.astype(np.bool8)
+        self.bc = superjunction_params['bc'].values.astype(np.bool_)
         if junction_params is not None:
             self._K_Ik = junction_params['K'].values.astype(np.float64)
             self._c_Ik = junction_params['c_0'].values.astype(np.float64)
diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py
index 107355e..6afdaf0 100644
--- a/pipedream_solver/nsuperlink.py
+++ b/pipedream_solver/nsuperlink.py
@@ -11,6 +11,7 @@
 import pipedream_solver.ngeometry
 import pipedream_solver.storage
 from pipedream_solver.superlink import SuperLink
+from pipedream_solver._nsuperlink import *
 
 class nSuperLink(SuperLink):
     """
@@ -352,6 +353,8 @@ def configure_hydraulic_geometry(self):
                 A = np.array([pipedream_solver.ngeometry.Transect_A_ik(w_i, x, y) for w_i in w])
                 B = np.array([pipedream_solver.ngeometry.Transect_B_ik(w_i, x, y) for w_i in w])
                 Pe = np.array([pipedream_solver.ngeometry.Transect_Pe_ik(w_i, x, y) for w_i in w])
+                # TODO: See if this is affected by change in definition of safe_divide_vec
+                # NOTE: R will always tend towards zero from l'hopital's rule
                 R = safe_divide_vec(A, Pe)
                 _transect_As.append(A)
                 _transect_Bs.append(B)
@@ -450,7 +453,7 @@ def configure_storages(self):
             for name, storage in storages.items():
                 A = storage['A']
                 h = storage['h']
-                V = scipy.integrate.cumtrapz(h, A, initial=0.)
+                V = scipy.integrate.cumtrapz(A, h, initial=0.)
                 _storage_As.append(A)
                 _storage_Vs.append(V)
                 _storage_hs.append(h)
@@ -479,6 +482,97 @@ def configure_storages(self):
         self._functional = _functional
         self._tabular = _tabular
 
+    def min_hydraulic_geometry(self):
+        min_depth = self.min_depth
+        _ik = self._ik                 # Link index
+        _Ik = self._Ik                 # Junction index
+        _Ip1k = self._Ip1k             # Index of next junction
+        _h_Ik_min = self._h_Ik_min
+        _A_ik_min = self._A_ik_min
+        _Pe_ik_min = self._Pe_ik_min
+        _R_ik_min = self._R_ik_min
+        _B_ik_min = self._B_ik_min
+        _A_uk_min = self._A_uk_min
+        _Pe_uk_min = self._Pe_uk_min
+        _R_uk_min = self._R_uk_min
+        _B_uk_min = self._B_uk_min
+        _A_dk_min = self._A_dk_min
+        _Pe_dk_min = self._Pe_dk_min
+        _R_dk_min = self._R_dk_min
+        _B_dk_min = self._B_dk_min
+        _z_inv_uk = self._z_inv_uk
+        _z_inv_dk = self._z_inv_dk
+        _g1_ik = self._g1_ik           # Geometry 1 of link ik (vertical)
+        _g2_ik = self._g2_ik           # Geometry 2 of link ik (horizontal)
+        _g3_ik = self._g3_ik           # Geometry 3 of link ik (other)
+        _g4_ik = self._g4_ik           # Geometry 4 of link ik (other)
+        _g5_ik = self._g5_ik           # Geometry 5 of link ik (other)
+        _g6_ik = self._g6_ik           # Geometry 6 of link ik (other)
+        _g7_ik = self._g7_ik           # Geometry 7 of link ik (other)
+        _geom_codes = self._geom_codes
+        _ellipse_ix = self._ellipse_ix
+        _is_irregular = self._is_irregular
+        _has_irregular = self._has_irregular
+        _transect_zs = self._transect_zs
+        _transect_As = self._transect_As
+        _transect_Bs = self._transect_Bs
+        _transect_Pes = self._transect_Pes
+        _transect_Rs = self._transect_Rs
+        _transect_codes = self._transect_codes
+        _transect_inds = self._transect_inds
+        _transect_lens = self._transect_lens
+        _z_inv_uk = self._z_inv_uk     # Invert offset of upstream end of superlink k
+        _J_uk = self._J_uk             # Index of junction upstream of superlink k
+        _z_inv_dk = self._z_inv_dk     # Invert offset of downstream end of superlink k
+        _J_dk = self._J_dk             # Index of junction downstream of superlink k
+        _uk_has_irregular = self._uk_has_irregular
+        _i_1k = self._i_1k
+        _I_1k = self._I_1k
+        _dk_has_irregular = self._dk_has_irregular
+        _i_nk = self._i_nk
+        _I_Np1k = self._I_Np1k
+        _theta_uk = np.ones(self._theta_uk.size, dtype=np.float64)
+        _theta_dk = np.ones(self._theta_dk.size, dtype=np.float64)
+        # Compute hydraulic geometry for regular geometries
+        # NOTE: Handle case for elliptical perimeter first
+        handle_elliptical_perimeter(_Pe_ik_min, _ellipse_ix, _Ik, _Ip1k, _h_Ik_min,
+                                    _g1_ik, _g2_ik)
+        # Compute hydraulic geometries for all other regular geometries
+        numba_hydraulic_geometry(_A_ik_min, _Pe_ik_min, _R_ik_min, _B_ik_min, _h_Ik_min,
+                                 _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
+                                 _geom_codes, _Ik, _ik)
+        # Compute hydraulic geometry for irregular geometries
+        if _has_irregular:
+            numba_transect_geometry(_A_ik_min, _Pe_ik_min, _R_ik_min, _B_ik_min, _h_Ik_min, _is_irregular,
+                                    _transect_zs, _transect_As, _transect_Bs, _transect_Pes,
+                                    _transect_Rs, _transect_codes, _transect_inds,
+                                    _transect_lens, _Ik, _ik)
+        # Compute hydraulic geometry for regular geometries
+        numba_boundary_geometry(_A_uk_min, _Pe_uk_min, _R_uk_min, _B_uk_min, _h_Ik_min, _z_inv_uk + min_depth, _z_inv_uk, _theta_uk,
+                                _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
+                                _geom_codes, _i_1k, _I_1k, _J_uk)
+        # Compute hydraulic geometry for irregular geometries
+        if _uk_has_irregular:
+            numba_boundary_transect(_A_uk_min, _Pe_uk_min, _R_uk_min, _B_uk_min, _h_Ik_min, _z_inv_uk + min_depth, _z_inv_uk, _theta_uk,
+                                    _is_irregular, _transect_zs, _transect_As, _transect_Bs,
+                                    _transect_Pes, _transect_Rs, _transect_codes, _transect_inds,
+                                    _transect_lens, _I_1k, _i_1k, _J_uk)
+        # Compute hydraulic geometry for regular geometries
+        numba_boundary_geometry(_A_dk_min, _Pe_dk_min, _R_dk_min, _B_dk_min, _h_Ik_min, _z_inv_dk + min_depth, _z_inv_dk, _theta_dk,
+                                _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
+                                _geom_codes, _i_nk, _I_Np1k, _J_dk)
+        # Compute hydraulic geometry for irregular geometries
+        if _dk_has_irregular:
+            numba_boundary_transect(_A_dk_min, _Pe_dk_min, _R_dk_min, _B_dk_min, _h_Ik_min, _z_inv_dk + min_depth, _z_inv_dk, _theta_dk,
+                                    _is_irregular, _transect_zs, _transect_As, _transect_Bs,
+                                    _transect_Pes, _transect_Rs, _transect_codes, _transect_inds,
+                                    _transect_lens, _I_Np1k, _i_nk, _J_dk)
+        # Export to instance variables
+        self._A_ik_min = _A_ik_min
+        self._Pe_ik_min = _Pe_ik_min
+        self._R_ik_min = _R_ik_min
+        self._B_ik_min = _B_ik_min
+
     def link_hydraulic_geometry(self):
         """
         Compute hydraulic geometry for each link.
@@ -707,26 +801,60 @@ def node_velocities(self):
         _u_ik = self._u_ik
         _u_Ik = self._u_Ik                   # Flow velocity at junction Ik
         _u_Ip1k = self._u_Ip1k               # Flow velocity at junction I + 1k
+        _Q_uk = self._Q_uk
+        _Q_dk = self._Q_dk
+        _A_uk = self._A_uk
+        _A_dk = self._A_dk
+        _u_uk = self._u_uk
+        _u_dk = self._u_dk
         _dx_ik = self._dx_ik                 # Length of link ik
+        _dx_uk = self._dx_uk
+        _dx_dk = self._dx_dk
+        _ki = self._ki
         _link_start = self._link_start
         _link_end = self._link_end
+        _A_ik_min = self._A_ik_min
+        _A_uk_min = self._A_uk_min
+        _A_dk_min = self._A_dk_min
         # Determine start and end nodes
         # Compute link velocities
-        numba_u_ik(_Q_ik, _A_ik, _u_ik)
+        _u_ik = numba_u_ik(_Q_ik, _A_ik, _A_ik_min, _u_ik)
+        # Compute boundary velocities
+        _u_uk = numba_u_ik(_Q_uk, _A_uk, _A_uk_min, _u_uk)
+        _u_dk = numba_u_ik(_Q_dk, _A_dk, _A_dk_min, _u_dk)
         # Compute velocities for start nodes (1 -> Nk)
-        numba_u_Ik(_dx_ik, _u_ik, _link_start, _u_Ik)
+        numba_u_Ik(_dx_ik, _u_ik, _dx_uk, _u_uk, _link_start, _ki, _u_Ik)
         # Compute velocities for end nodes (2 -> Nk+1)
-        numba_u_Ip1k(_dx_ik, _u_ik, _link_end, _u_Ip1k)
+        numba_u_Ip1k(_dx_ik, _u_ik, _dx_dk, _u_dk, _link_end, _ki, _u_Ip1k)
+        # Compute effective velocity at superjunctions
+        #M = self.M
+        #_J_uk = self._J_uk
+        #_J_dk = self._J_dk
+        #Qu_j = np.zeros(M)
+        #np.add.at(Qu_j, _J_uk, (_u_uk * np.abs(_Q_uk)))
+        #np.add.at(Qu_j, _J_dk, (_u_dk * np.abs(_Q_dk)))
+        #Q_j = np.zeros(M)
+        #np.add.at(Q_j, _J_dk, np.abs(_Q_dk))
+        #np.add.at(Q_j, _J_uk, np.abs(_Q_uk))
+        #_u_j = safe_divide_vec(Qu_j, Q_j)
         # Export to instance variables
         self._u_ik = _u_ik
+        self._u_uk = _u_uk
+        self._u_dk = _u_dk
         self._u_Ik = _u_Ik
         self._u_Ip1k = _u_Ip1k
+        #self._u_j = _u_j
 
     def link_coeffs(self, _dt=None, first_iter=True):
         """
         Compute link momentum coefficients: a_ik, b_ik, c_ik and P_ik.
         """
         # Import instance variables
+        _i_1k = self._i_1k
+        _i_nk = self._i_nk
+        _u_ik = self._u_ik
+        _u_uk = self._u_uk
+        _u_dk = self._u_dk
         _u_Ik = self._u_Ik         # Flow velocity at junction Ik
         _u_Ip1k = self._u_Ip1k     # Flow velocity at junction I + 1k
         _dx_ik = self._dx_ik       # Length of link ik
@@ -740,9 +868,35 @@ def link_coeffs(self, _dt=None, first_iter=True):
         _A_c_ik = self._A_c_ik     # Area of control structure at link ik
         _C_ik = self._C_ik         # Discharge coefficient of control structure at link ik
         _ctrl = self._ctrl         # Control structure exists at link ik (y/n)
-        inertial_damping = self.inertial_damping    # Use inertial damping (y/n)
         _sigma_ik = self._sigma_ik  # Inertial damping coefficient
+        NK = self.NK
         g = 9.81
+        # Upstream parameters
+        _link_start = self._link_start
+        _n_uk = self._n_uk
+        _Q_uk_next = self._Q_uk
+        _Q_uk_prev = self.states['Q_uk']
+        _A_uk = self._A_uk
+        _R_uk = self._R_uk
+        _dx_uk = self._dx_uk
+        _Sf_method_uk = self._Sf_method_uk
+        _C_uk = self._C_uk
+        _S_o_uk = self._S_o_uk
+        _theta_uk = self._theta_uk
+        _z_inv_uk = self._z_inv_uk
+        # Downstream parameters
+        _link_end = self._link_end
+        _n_dk = self._n_dk
+        _Q_dk_next = self._Q_dk
+        _Q_dk_prev = self.states['Q_dk']
+        _A_dk = self._A_dk
+        _R_dk = self._R_dk
+        _dx_dk = self._dx_dk
+        _Sf_method_dk = self._Sf_method_dk
+        _C_dk = self._C_dk
+        _S_o_dk = self._S_o_dk
+        _theta_dk = self._theta_dk
+        _z_inv_dk = self._z_inv_dk
         # If time step not specified, use instance time
         if _dt is None:
             _dt = self._dt
@@ -751,13 +905,55 @@ def link_coeffs(self, _dt=None, first_iter=True):
         _c_ik = numba_c_ik(_u_Ip1k, _sigma_ik)
         _b_ik = numba_b_ik(_dx_ik, _dt, _n_ik, _Q_ik_next, _A_ik, _R_ik, _A_c_ik,
                            _C_ik, _a_ik, _c_ik, _ctrl, _sigma_ik, _Sf_method_ik, g)
+        #_b_ik = numba_b_ik(_dx_ik, _dt, _n_ik, _Q_ik_next, _A_ik, _R_ik, _A_c_ik,
+        #                   _C_ik, _u_ik, _ctrl, _sigma_ik, _Sf_method_ik, g)
         _P_ik = numba_P_ik(_Q_ik_prev, _dx_ik, _dt, _A_ik, _S_o_ik,
                            _sigma_ik, g)
+        #_a_ik[_i_1k] *= _theta_uk
+        #_c_ik[_i_nk] *= _theta_dk
+        # Compute momentum coefficients for upstream boundary
+        _ctrl_uk = np.ones(NK, dtype=np.bool_)
+        _sigma_uk = _sigma_ik[_link_start]
+        _a_uk = np.zeros(NK, dtype=np.float64)
+        #_c_uk = numba_c_ik(_u_Ik[_link_start], _sigma_uk)
+        #_c_uk = numba_c_ik(_u_uk, _sigma_uk)
+        _c_uk = numba_c_ik(_u_ik[_link_start], _sigma_uk)    # This actually seems to be more stable
+        _b_uk = numba_b_ik(_dx_uk, _dt, _n_uk, _Q_uk_next, _A_uk, _R_uk, _A_uk,
+                           _C_uk, _a_uk, _c_uk, _ctrl_uk , _sigma_uk, _Sf_method_uk, g)
+        #_b_uk = numba_b_ik(_dx_uk, _dt, _n_uk, _Q_uk_next, _A_uk, _R_uk, _A_uk,
+        #                   _C_uk, _u_uk, _ctrl_uk , _sigma_uk, _Sf_method_uk, g)
+        _P_uk = numba_P_ik(_Q_uk_prev, _dx_uk, _dt, _A_uk, _S_o_uk, _sigma_uk, g)
+        _P_uk -= g * _A_uk * _theta_uk * _z_inv_uk
+        # Try to make depth critical
+        #_c_uk *= _theta_uk
+        # Compute momentum coefficients for downstream boundary
+        _ctrl_dk = np.ones(NK, dtype=np.bool_)
+        _sigma_dk = _sigma_ik[_link_end]
+        #_a_dk = numba_a_ik(_u_Ip1k[_link_end], _sigma_dk)
+        #_a_dk = numba_a_ik(_u_dk, _sigma_dk)
+        _a_dk = numba_a_ik(_u_ik[_link_end], _sigma_dk)    # This actually seems to be more stable
+        _c_dk = np.zeros(NK, dtype=np.float64)
+        _b_dk = numba_b_ik(_dx_dk, _dt, _n_dk, _Q_dk_next, _A_dk, _R_dk, _A_dk,
+                           _C_dk, _a_dk, _c_dk, _ctrl_dk , _sigma_dk, _Sf_method_dk, g)
+        #_b_dk = numba_b_ik(_dx_dk, _dt, _n_dk, _Q_dk_next, _A_dk, _R_dk, _A_dk,
+        #                   _C_dk, _u_dk, _ctrl_dk , _sigma_dk, _Sf_method_dk, g)
+        _P_dk = numba_P_ik(_Q_dk_prev, _dx_dk, _dt, _A_dk, _S_o_dk, _sigma_dk, g)
+        _P_dk += g * _A_dk * _theta_dk * _z_inv_dk
+        # Try to make depth critical
+        #_a_dk *= _theta_dk
         # Export to instance variables
         self._a_ik = _a_ik
         self._b_ik = _b_ik
         self._c_ik = _c_ik
         self._P_ik = _P_ik
+        self._a_uk = _a_uk
+        self._b_uk = _b_uk
+        self._c_uk = _c_uk
+        self._P_uk = _P_uk
+        self._a_dk = _a_dk
+        self._b_dk = _b_dk
+        self._c_dk = _c_dk
+        self._P_dk = _P_dk
 
     def node_coeffs(self, _Q_0Ik=None, _dt=None, first_iter=True):
         """
@@ -848,7 +1044,132 @@ def backward_recurrence(self):
         self._Y_Ik = _Y_Ik
         self._Z_Ik = _Z_Ik
 
+    def superlink_boundary_depth_coefficients(self, _dt=None):
+        _U_Ik = self._U_Ik                # Recurrence coefficient U_Ik
+        _V_Ik = self._V_Ik                # Recurrence coefficient V_Ik
+        _W_Ik = self._W_Ik                # Recurrence coefficient W_Ik
+        _X_Ik = self._X_Ik                # Recurrence coefficient X_Ik
+        _Y_Ik = self._Y_Ik                # Recurrence coefficient Y_Ik
+        _Z_Ik = self._Z_Ik                # Recurrence coefficient Z_Ik
+        _E_Ik = self._E_Ik
+        _D_Ik = self._D_Ik
+        _I_1k = self._I_1k
+        _I_Nk = self._I_Nk
+        _I_Np1k = self._I_Np1k
+        # Get boundary coefficients
+        U_Nk = _U_Ik[_I_Nk]
+        E_Np1k = _E_Ik[_I_Np1k]
+        Z_1k = _Z_Ik[_I_1k]
+        W_Nk = _W_Ik[_I_Nk]
+        X_1k = _X_Ik[_I_1k]
+        E_1k = _E_Ik[_I_1k]
+        Y_1k = _Y_Ik[_I_1k]
+        D_1k = _D_Ik[_I_1k]
+        D_Np1k = _D_Ik[_I_Np1k]
+        V_Nk = _V_Ik[_I_Nk]
+        # Formulate expressions
+        a = U_Nk - E_Np1k
+        b = Z_1k
+        c = W_Nk
+        d = X_1k + E_1k
+        e = Y_1k - D_1k
+        f = D_Np1k + V_Nk
+        denom = a*d - b*c
+        # Compute coefficients
+        _kappa_uk = a / denom
+        _lambda_uk = -b / denom
+        _mu_uk = (b*f - a*e) / denom
+        _kappa_dk = -c / denom
+        _lambda_dk = d / denom
+        _mu_dk = (c*e - f*d) / denom
+        # Store coefficients
+        self._kappa_uk = _kappa_uk
+        self._lambda_uk = _lambda_uk
+        self._mu_uk = _mu_uk
+        self._kappa_dk = _kappa_dk
+        self._lambda_dk = _lambda_dk
+        self._mu_dk = _mu_dk
+
+    def superlink_boundary_flow_coefficients(self, _dt=None):
+        _U_Ik = self._U_Ik                # Recurrence coefficient U_Ik
+        _V_Ik = self._V_Ik                # Recurrence coefficient V_Ik
+        _W_Ik = self._W_Ik                # Recurrence coefficient W_Ik
+        _X_Ik = self._X_Ik                # Recurrence coefficient X_Ik
+        _Y_Ik = self._Y_Ik                # Recurrence coefficient Y_Ik
+        _Z_Ik = self._Z_Ik                # Recurrence coefficient Z_Ik
+        _E_Ik = self._E_Ik
+        _D_Ik = self._D_Ik
+        _I_1k = self._I_1k
+        _I_Nk = self._I_Nk
+        _I_Np1k = self._I_Np1k
+        A_uk = self._A_uk
+        A_dk = self._A_dk
+        a_dk = self._a_dk
+        b_uk = self._b_uk
+        b_dk = self._b_dk
+        c_uk = self._c_uk
+        P_uk = self._P_uk
+        P_dk = self._P_dk
+        theta_uk = self._theta_uk
+        theta_dk = self._theta_dk
+        g = 9.81
+        # Get boundary coefficients
+        U_Nk = _U_Ik[_I_Nk]
+        Z_1k = _Z_Ik[_I_1k]
+        W_Nk = _W_Ik[_I_Nk]
+        X_1k = _X_Ik[_I_1k]
+        E_1k = _E_Ik[_I_1k]
+        E_Np1k = _E_Ik[_I_Np1k]
+        D_1k = _D_Ik[_I_1k]
+        D_Np1k = _D_Ik[_I_Np1k]
+        Y_1k = _Y_Ik[_I_1k]
+        V_Nk = _V_Ik[_I_Nk]
+        # Formulate expressions
+        a = U_Nk - E_Np1k
+        b = Z_1k
+        c = W_Nk
+        d = X_1k + E_1k
+        e = Y_1k - D_1k
+        f = D_Np1k + V_Nk
+        q = (-A_dk * g + a_dk * U_Nk)
+        r = (A_uk * g + c_uk * X_1k)
+        s = a * b_dk
+        t = d * b_uk
+        u = (a_dk + b_dk)
+        v = (b_uk + c_uk)
+        w = (A_uk * g - c_uk * E_1k)
+        x = (A_dk * g - a_dk * E_Np1k)
+        y = (c_uk * D_1k - P_uk)
+        z = (P_dk - b_dk * D_Np1k - u * V_Nk)
+        p = (d * P_uk + w * Y_1k - r * D_1k)
+        n = (a_dk * D_Np1k + P_dk)
+        m = (P_uk + b_uk * D_1k - v * Y_1k)
+        o = (a * P_dk - x * V_Nk + q * D_Np1k)
+        # Create inverse matrix
+        aa = A_uk * theta_uk * g * ((c * b * u) - d * (q + s))
+        bb = A_dk * theta_dk * g * (b * w) 
+        cc = A_uk * theta_uk * g * (c * x)
+        dd = A_dk * theta_dk * g * (a * (r + t) - (c * b * v))
+        ee = (-p * (q + s) - (b * c * u * y) - (b * w * z))
+        ff = (-o * (r + t) + (b * c * v * n) + (c * x * m))
+        denom =  (c * b * u * v) - (q + s) * (r + t)
+        # Compute coefficients
+        alpha_uk = aa / denom
+        beta_uk = bb / denom
+        chi_uk = ee / denom
+        alpha_dk = cc / denom
+        beta_dk = dd / denom
+        chi_dk = ff / denom
+        # Store coefficients
+        self._alpha_uk = alpha_uk
+        self._beta_uk = beta_uk
+        self._chi_uk = chi_uk
+        self._alpha_dk = alpha_dk
+        self._beta_dk = beta_dk
+        self._chi_dk = chi_dk
+
     def superlink_upstream_head_coefficients(self, _dt=None):
+        raise NotImplementedError('Deprecated')
         """
         Compute upstream head coefficients for superlinks: kappa_uk, lambda_uk, and mu_uk.
         """
@@ -894,16 +1215,17 @@ def superlink_upstream_head_coefficients(self, _dt=None):
             self._mu_uk = - _theta_uk * _z_inv_uk
         elif _bc_method == 'b':
             # Compute superlink upstream coefficients (momentum)
-            self._kappa_uk = kappa_uk(_Q_uk_next, _dx_uk, _A_uk, _C_uk,
+            self._kappa_uk = kappa_uk_old(_Q_uk_next, _dx_uk, _A_uk, _C_uk,
                                       _R_uk, _n_uk, _Sf_method_uk, _dt, g)
             self._lambda_uk = _theta_uk
-            self._mu_uk = mu_uk(_Q_uk_prev, _dx_uk, _A_uk, _theta_uk, _z_inv_uk,
+            self._mu_uk = mu_uk_old(_Q_uk_prev, _dx_uk, _A_uk, _theta_uk, _z_inv_uk,
                                 _S_o_uk, _dt, g)
         else:
             raise ValueError('Invalid BC method {}.'.format(_bc_method))
         self._theta_uk = _theta_uk
 
     def superlink_downstream_head_coefficients(self, _dt=None):
+        raise NotImplementedError('Deprecated')
         """
         Compute downstream head coefficients for superlinks: kappa_dk, lambda_dk, and mu_dk.
         """
@@ -948,15 +1270,16 @@ def superlink_downstream_head_coefficients(self, _dt=None):
             self._mu_dk = - _theta_dk * _z_inv_dk
         elif _bc_method == 'b':
             # Compute superlink upstream coefficients (momentum)
-            self._kappa_dk = kappa_dk(_Q_dk_next, _dx_dk, _A_dk, _C_dk,
+            self._kappa_dk = kappa_dk_old(_Q_dk_next, _dx_dk, _A_dk, _C_dk,
                                       _R_dk, _n_dk, _Sf_method_dk, _dt, g)
             self._lambda_dk = _theta_dk
-            self._mu_dk = mu_dk(_Q_dk_prev, _dx_dk, _A_dk, _theta_dk, _z_inv_dk, _S_o_dk, _dt, g)
+            self._mu_dk = mu_dk_old(_Q_dk_prev, _dx_dk, _A_dk, _theta_dk, _z_inv_dk, _S_o_dk, _dt, g)
         else:
             raise ValueError('Invalid BC method {}.'.format(_bc_method))
         self._theta_dk = _theta_dk
 
     def superlink_flow_coefficients(self):
+        raise NotImplementedError('Deprecated')
         """
         Compute superlink flow coefficients: alpha_uk, beta_uk, chi_uk,
         alpha_dk, beta_dk, chi_dk.
@@ -1003,25 +1326,25 @@ def superlink_flow_coefficients(self):
             _V_Nk = _V_Ik[_I_Nk] + _D_Ik[_I_Np1k]
             _W_Nk = _W_Ik[_I_Nk]
         # Compute D_k_star
-        _D_k_star = numba_D_k_star(_X_1k, _kappa_uk, _U_Nk,
+        _D_k_star = numba_D_k_star_old(_X_1k, _kappa_uk, _U_Nk,
                                    _kappa_dk, _Z_1k, _W_Nk)
         # Compute upstream superlink flow coefficients
-        _alpha_uk = numba_alpha_uk(_U_Nk, _kappa_dk, _X_1k,
+        _alpha_uk = numba_alpha_uk_old(_U_Nk, _kappa_dk, _X_1k,
                                    _Z_1k, _W_Nk, _D_k_star,
                                    _lambda_uk)
-        _beta_uk = numba_beta_uk(_U_Nk, _kappa_dk, _Z_1k,
+        _beta_uk = numba_beta_uk_old(_U_Nk, _kappa_dk, _Z_1k,
                                  _W_Nk, _D_k_star, _lambda_dk)
-        _chi_uk = numba_chi_uk(_U_Nk, _kappa_dk, _Y_1k,
+        _chi_uk = numba_chi_uk_old(_U_Nk, _kappa_dk, _Y_1k,
                                _X_1k, _mu_uk, _Z_1k,
                                _mu_dk, _V_Nk, _W_Nk,
                                _D_k_star)
         # Compute downstream superlink flow coefficients
-        _alpha_dk = numba_alpha_dk(_X_1k, _kappa_uk, _W_Nk,
+        _alpha_dk = numba_alpha_dk_old(_X_1k, _kappa_uk, _W_Nk,
                                    _D_k_star, _lambda_uk)
-        _beta_dk = numba_beta_dk(_X_1k, _kappa_uk, _U_Nk,
+        _beta_dk = numba_beta_dk_old(_X_1k, _kappa_uk, _U_Nk,
                                  _W_Nk, _Z_1k, _D_k_star,
                                  _lambda_dk)
-        _chi_dk = numba_chi_dk(_X_1k, _kappa_uk, _V_Nk,
+        _chi_dk = numba_chi_dk_old(_X_1k, _kappa_uk, _V_Nk,
                                _W_Nk, _mu_uk, _U_Nk,
                                _mu_dk, _Y_1k, _Z_1k,
                                _D_k_star)
@@ -1130,13 +1453,20 @@ def pump_flow_coefficients(self, u=None):
         self._beta_p = _beta_p
         self._chi_p = _chi_p
 
-    def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, implicit=True,
-                                first_time=False):
-        """
-        Construct sparse matrices A, O, W, P and b.
-        """
-        # Import instance variables
-        _k = self._k                     # Superlink indices
+    def create_superjunction_matrix(self, H_bc, Q_in, _dt):
+        J = self.J
+        bc = self.bc
+        _A_sj = self._A_sj
+        _dt = self._dt
+        M = self.M
+        numba_create_J_matrix(J, bc, _A_sj, _dt, M)
+        _G_jh = np.where(bc, 0., _A_sj / _dt)
+        _G_je = np.where(bc, H_bc, Q_in)
+        return J, _G_jh, _G_je
+
+    def create_superlink_matrix(self, _dt):
+        K = self.K
+        bc = self.bc
         _J_uk = self._J_uk               # Index of superjunction upstream of superlink k
         _J_dk = self._J_dk               # Index of superjunction downstream of superlink k
         _alpha_uk = self._alpha_uk       # Superlink flow coefficient
@@ -1145,10 +1475,6 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, impli
         _beta_dk = self._beta_dk         # Superlink flow coefficient
         _chi_uk = self._chi_uk           # Superlink flow coefficient
         _chi_dk = self._chi_dk           # Superlink flow coefficient
-        _alpha_ukm = self._alpha_ukm     # Summation of superlink flow coefficients
-        _beta_dkl = self._beta_dkl       # Summation of superlink flow coefficients
-        _chi_ukl = self._chi_ukl         # Summation of superlink flow coefficients
-        _chi_dkm = self._chi_dkm         # Summation of superlink flow coefficients
         _F_jj = self._F_jj
         _A_sj = self._A_sj               # Surface area of superjunction j
         _dx_uk = self._dx_uk
@@ -1158,53 +1484,147 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, impli
         _theta_uk = self._theta_uk
         _theta_dk = self._theta_dk
         NK = self.NK
-        n_o = self.n_o                   # Number of orifices in system
-        n_w = self.n_w                   # Number of weirs in system
-        n_p = self.n_p                   # Number of pumps in system
+        M = self.M
+        # Clear old data
+        _F_jj.fill(0.)
+        _G_jh = np.zeros(M, dtype=np.float64)
+        _G_je = np.zeros(M, dtype=np.float64)
+        numba_clear_off_diagonals(K, bc, _J_uk, _J_dk, NK)
+        # Compute top width contributed by attached superlinks
+        _xi_uk = xi_uk(_dx_uk, _B_uk, _theta_uk)
+        _xi_dk = xi_dk(_dx_dk, _B_dk, _theta_dk)
+        # Create A matrix
+        numba_create_K_matrix(K, _F_jj, bc, _J_uk, _J_dk, _alpha_uk,
+                              _alpha_dk, _beta_uk, _beta_dk, _xi_uk, _xi_dk,
+                              _A_sj, _dt, M, NK)
+        # Create RHS vector
+        numba_add_at(_G_jh, _J_uk, _xi_uk / _dt)
+        numba_add_at(_G_jh, _J_dk, _xi_dk / _dt)
+        numba_add_at(_G_je, _J_uk, -_chi_uk)
+        numba_add_at(_G_je, _J_dk, _chi_dk)
+        # Ensure RHS is set to zero for boundary nodes
+        _G_jh[bc] = 0.
+        _G_je[bc] = 0.
+        return K, _G_jh, _G_je
+
+    def create_orifice_matrix(self):
+        O = self.O
+        n_o = self.n_o
+        bc = self.bc
+        _J_uo = self._J_uo               # Index of superjunction upstream of orifice o
+        _J_do = self._J_do               # Index of superjunction upstream of orifice o
+        _alpha_o = self._alpha_o         # Orifice flow coefficient
+        _beta_o = self._beta_o           # Orifice flow coefficient
+        _chi_o = self._chi_o             # Orifice flow coefficient
+        _O_diag = self._O_diag           # Diagonal elements of matrix O
+        _alpha_uo = _alpha_o
+        _alpha_do = _alpha_o
+        _beta_uo = _beta_o
+        _beta_do = _beta_o
+        _chi_uo = _chi_o
+        _chi_do = _chi_o
+        M = self.M
+        # Clear arrays
+        _O_diag.fill(0.)
+        _G_jh = np.zeros(M, dtype=np.float64)
+        _G_je = np.zeros(M, dtype=np.float64)
+        numba_clear_off_diagonals(O, bc, _J_uo, _J_do, n_o)
+        # Set diagonal
+        numba_create_OWP_matrix(O, _O_diag, bc, _J_uo, _J_do, _alpha_uo,
+                                _alpha_do, _beta_uo, _beta_do, M, n_o)
+        # Set right-hand side
+        numba_add_at(_G_je, _J_uo, -_chi_uo)
+        numba_add_at(_G_je, _J_do, _chi_do)
+        # Ensure RHS is set to zero for boundary nodes
+        _G_jh[bc] = 0.
+        _G_je[bc] = 0.
+        return O, _G_jh, _G_je
+
+    def create_weir_matrix(self):
+        W = self.W
+        n_w = self.n_w
+        bc = self.bc
+        _J_uw = self._J_uw               # Index of superjunction upstream of weir w
+        _J_dw = self._J_dw               # Index of superjunction downstream of weir w
+        _alpha_w = self._alpha_w         # Weir flow coefficient
+        _beta_w = self._beta_w           # Weir flow coefficient
+        _chi_w = self._chi_w             # Weir flow coefficient
+        _W_diag = self._W_diag           # Diagonal elements of matrix W
+        # Rename indexers
+        _alpha_uw = _alpha_w
+        _alpha_dw = _alpha_w
+        _beta_uw = _beta_w
+        _beta_dw = _beta_w
+        _chi_uw = _chi_w
+        _chi_dw = _chi_w
+        M = self.M
+        # Clear arrays
+        _W_diag.fill(0.)
+        _G_jh = np.zeros(M, dtype=np.float64)
+        _G_je = np.zeros(M, dtype=np.float64)
+        numba_clear_off_diagonals(W, bc, _J_uw, _J_dw, n_w)
+        # Set diagonal
+        numba_create_OWP_matrix(W, _W_diag, bc, _J_uw, _J_dw, _alpha_uw,
+                                _alpha_dw, _beta_uw, _beta_dw, M, n_w)
+        # Set right-hand side
+        numba_add_at(_G_je, _J_uw, -_chi_uw)
+        numba_add_at(_G_je, _J_dw, _chi_dw)
+        # Ensure RHS is set to zero for boundary nodes
+        _G_jh[bc] = 0.
+        _G_je[bc] = 0.
+        return W, _G_jh, _G_je
+
+    def create_pump_matrix(self):
+        P = self.P
+        n_p = self.n_p
+        bc = self.bc
+        _J_up = self._J_up               # Index of superjunction upstream of pump p
+        _J_dp = self._J_dp               # Index of superjunction downstream of pump p
+        _alpha_p = self._alpha_p         # Pump flow coefficient
+        _beta_p = self._beta_p           # Pump flow coefficient
+        _chi_p = self._chi_p             # Pump flow coefficient
+        _P_diag = self._P_diag           # Diagonal elements of matrix P
+        _alpha_up = _alpha_p
+        _alpha_dp = _alpha_p
+        _beta_up = _beta_p
+        _beta_dp = _beta_p
+        _chi_up = _chi_p
+        _chi_dp = _chi_p
+        M = self.M
+        # Clear arrays
+        _P_diag.fill(0.)
+        _G_jh = np.zeros(M, dtype=np.float64)
+        _G_je = np.zeros(M, dtype=np.float64)
+        numba_clear_off_diagonals(P, bc, _J_up, _J_dp, n_p)
+        # Set diagonal
+        numba_create_OWP_matrix(P, _P_diag, bc, _J_up, _J_dp, _alpha_up,
+                                _alpha_dp, _beta_up, _beta_dp, M, n_p)
+        # Set right-hand side
+        numba_add_at(_G_je, _J_up, -_chi_up)
+        numba_add_at(_G_je, _J_dp, _chi_dp)
+        # Ensure RHS is set to zero for boundary nodes
+        _G_jh[bc] = 0.
+        _G_je[bc] = 0.
+        return P, _G_jh, _G_je
+
+
+    def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, implicit=True,
+                                first_time=False):
+        """
+        Construct sparse matrices A, O, W, P and b.
+        """
+        # Import instance variables
         A = self.A
-        if n_o:
-            O = self.O
-            _J_uo = self._J_uo               # Index of superjunction upstream of orifice o
-            _J_do = self._J_do               # Index of superjunction upstream of orifice o
-            _alpha_o = self._alpha_o         # Orifice flow coefficient
-            _beta_o = self._beta_o           # Orifice flow coefficient
-            _chi_o = self._chi_o             # Orifice flow coefficient
-            _alpha_uom = self._alpha_uom     # Summation of orifice flow coefficients
-            _beta_dol = self._beta_dol       # Summation of orifice flow coefficients
-            _chi_uol = self._chi_uol         # Summation of orifice flow coefficients
-            _chi_dom = self._chi_dom         # Summation of orifice flow coefficients
-            _O_diag = self._O_diag           # Diagonal elements of matrix O
-        if n_w:
-            W = self.W
-            _J_uw = self._J_uw               # Index of superjunction upstream of weir w
-            _J_dw = self._J_dw               # Index of superjunction downstream of weir w
-            _alpha_w = self._alpha_w         # Weir flow coefficient
-            _beta_w = self._beta_w           # Weir flow coefficient
-            _chi_w = self._chi_w             # Weir flow coefficient
-            _alpha_uwm = self._alpha_uwm     # Summation of weir flow coefficients
-            _beta_dwl = self._beta_dwl       # Summation of weir flow coefficients
-            _chi_uwl = self._chi_uwl         # Summation of weir flow coefficients
-            _chi_dwm = self._chi_dwm         # Summation of weir flow coefficients
-            _W_diag = self._W_diag           # Diagonal elements of matrix W
-        if n_p:
-            P = self.P
-            _J_up = self._J_up               # Index of superjunction upstream of pump p
-            _J_dp = self._J_dp               # Index of superjunction downstream of pump p
-            _alpha_p = self._alpha_p         # Pump flow coefficient
-            _beta_p = self._beta_p           # Pump flow coefficient
-            _chi_p = self._chi_p             # Pump flow coefficient
-            _alpha_upm = self._alpha_upm     # Summation of pump flow coefficients
-            _beta_dpl = self._beta_dpl       # Summation of pump flow coefficients
-            _chi_upl = self._chi_upl         # Summation of pump flow coefficients
-            _chi_dpm = self._chi_dpm         # Summation of pump flow coefficients
-            _P_diag = self._P_diag           # Diagonal elements of matrix P
+        b = self.b                       # Right-hand side vector
+        M = self.M
+        NK = self.NK
+        n_o = self.n_o
+        n_w = self.n_w
+        n_p = self.n_p
         _sparse = self._sparse           # Use sparse matrix data structures (y/n)
-        M = self.M                       # Number of superjunctions in system
         H_j_next = self.H_j                   # Head at superjunction j
         H_j_prev = self.states['H_j']
         bc = self.bc                     # Superjunction j has a fixed boundary condition (y/n)
-        D = self.D                       # Vector for storing chi coefficients
-        b = self.b                       # Right-hand side vector
         # If no time step specified, use instance time step
         if _dt is None:
             _dt = self._dt
@@ -1217,80 +1637,43 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, impli
         # If no control input signal specified assume zero input
         if u is None:
             u = 0
-        # Compute upstream/downstream link volume parameters
-        _xi_uk = xi_uk(_dx_uk, _B_uk, _theta_uk, _dt)
-        _xi_dk = xi_dk(_dx_dk, _B_dk, _theta_dk, _dt)
-        # Clear old data
-        _F_jj.fill(0)
-        D.fill(0)
-        numba_clear_off_diagonals(A, bc, _J_uk, _J_dk, NK)
-        # Create A matrix
-        numba_create_A_matrix(A, _F_jj, bc, _J_uk, _J_dk, _alpha_uk,
-                              _alpha_dk, _beta_uk, _beta_dk, _xi_uk, _xi_dk,
-                              _A_sj, _dt, M, NK)
-        # Create D vector
-        numba_add_at(D, _J_uk, -_chi_uk)
-        numba_add_at(D, _J_dk, _chi_dk)
-        numba_add_at(D, _J_uk, _xi_uk * H_j_prev[_J_uk])
-        numba_add_at(D, _J_dk, _xi_dk * H_j_prev[_J_dk])
-        # Compute control matrix
+        # Create matrices
+        A = np.zeros((M, M), dtype=np.float64)
+        b = np.zeros(M, dtype=np.float64)
+        G_jh = np.zeros(M, dtype=np.float64)
+        G_je = np.zeros(M, dtype=np.float64)
+        # Create component matrices
+        J, G_jhj, G_jej = self.create_superjunction_matrix(H_bc, _Q_0j, _dt)
+        A += J
+        G_jh += G_jhj
+        G_je += G_jej
+        if NK:
+            K, G_jhk, G_jek = self.create_superlink_matrix(_dt)
+            A += K
+            G_jh += G_jhk
+            G_je += G_jek
         if n_o:
-            _alpha_uo = _alpha_o
-            _alpha_do = _alpha_o
-            _beta_uo = _beta_o
-            _beta_do = _beta_o
-            _chi_uo = _chi_o
-            _chi_do = _chi_o
-            _O_diag.fill(0)
-            numba_clear_off_diagonals(O, bc, _J_uo, _J_do, n_o)
-            # Set diagonal
-            numba_create_OWP_matrix(O, _O_diag, bc, _J_uo, _J_do, _alpha_uo,
-                                    _alpha_do, _beta_uo, _beta_do, M, n_o)
-            # Set right-hand side
-            numba_add_at(D, _J_uo, -_chi_uo)
-            numba_add_at(D, _J_do, _chi_do)
+            O, G_jho, G_jeo = self.create_orifice_matrix()
+            A += O
+            G_jh += G_jho
+            G_je += G_jeo
         if n_w:
-            _alpha_uw = _alpha_w
-            _alpha_dw = _alpha_w
-            _beta_uw = _beta_w
-            _beta_dw = _beta_w
-            _chi_uw = _chi_w
-            _chi_dw = _chi_w
-            _W_diag.fill(0)
-            numba_clear_off_diagonals(W, bc, _J_uw, _J_dw, n_w)
-            # Set diagonal
-            numba_create_OWP_matrix(W, _W_diag, bc, _J_uw, _J_dw, _alpha_uw,
-                                    _alpha_dw, _beta_uw, _beta_dw, M, n_w)
-            # Set right-hand side
-            numba_add_at(D, _J_uw, -_chi_uw)
-            numba_add_at(D, _J_dw, _chi_dw)
+            W, G_jhw, G_jew = self.create_weir_matrix()
+            A += W
+            G_jh += G_jhw
+            G_je += G_jew
         if n_p:
-            _alpha_up = _alpha_p
-            _alpha_dp = _alpha_p
-            _beta_up = _beta_p
-            _beta_dp = _beta_p
-            _chi_up = _chi_p
-            _chi_dp = _chi_p
-            _P_diag.fill(0)
-            numba_clear_off_diagonals(P, bc, _J_up, _J_dp, n_p)
-            # Set diagonal
-            numba_create_OWP_matrix(P, _P_diag, bc, _J_up, _J_dp, _alpha_up,
-                                    _alpha_dp, _beta_up, _beta_dp, M, n_p)
-            # Set right-hand side
-            numba_add_at(D, _J_up, -_chi_up)
-            numba_add_at(D, _J_dp, _chi_dp)
-        b.fill(0)
+            P, G_jhp, G_jep = self.create_pump_matrix()
+            A += P
+            G_jh += G_jhp
+            G_je += G_jep
         # TODO: Which A_sj? Might need to apply product rule here.
-        b = (_A_sj * H_j_prev / _dt) + _Q_0j + D
+        b[:] = (G_jh * H_j_prev) + G_je
         # Ensure boundary condition is specified
         b[bc] = H_bc[bc]
         # Export instance variables
-        self.D = D
+        self.A = A
         self.b = b
-        # self._beta_dkl = _beta_dkl
-        # self._alpha_ukm = _alpha_ukm
-        # self._chi_ukl = _chi_ukl
-        # self._chi_dkm = _chi_dkm
         if first_time and _sparse:
             self.A = self.A.tocsr()
 
@@ -1312,31 +1695,16 @@ def solve_sparse_matrix(self, u=None, implicit=True):
         _sparse = self._sparse        # Use sparse data structures (y/n)
         min_depth = self.min_depth    # Minimum depth at superjunctions
         max_depth = self.max_depth    # Maximum depth at superjunctions
-        # Does the system have control assets?
-        has_control = n_o + n_w + n_p
-        # Get right-hand size
-        if has_control:
-            if implicit:
-                l = A + O + W + P
-                r = b
-            else:
-                # TODO: Broken
-                # l = A
-                # r = b + np.squeeze(B @ u)
-                raise NotImplementedError
-        else:
-            l = A
-            r = b
         if _sparse:
-            H_j_next = scipy.sparse.linalg.spsolve(l, r)
+            H_j_next = scipy.sparse.linalg.spsolve(A, b)
         else:
-            H_j_next = scipy.linalg.solve(l, r)
+            H_j_next = np.linalg.solve(A, b)
         assert np.isfinite(H_j_next).all()
         # Constrain heads based on allowed maximum/minimum depths
         # TODO: Not sure what's happening here
         # H_j_next = np.maximum(H_j_next, _z_inv_j + min_depth)
-        H_j_next = np.maximum(H_j_next, _z_inv_j)
-        H_j_next = np.minimum(H_j_next, _z_inv_j + max_depth)
+        #H_j_next = np.maximum(H_j_next, _z_inv_j)
+        #H_j_next = np.minimum(H_j_next, _z_inv_j + max_depth)
         # Export instance variables
         self.H_j = H_j_next
 
@@ -1357,28 +1725,17 @@ def solve_banded_matrix(self, u=None, implicit=True):
         max_depth = self.max_depth    # Maximum depth at superjunctions
         bandwidth = self.bandwidth
         M = self.M
-        # Does the system have control assets?
-        has_control = n_o + n_w + n_p
-        # Get right-hand size
-        if has_control:
-            if implicit:
-                l = A + O + W + P
-                r = b
-            else:
-                raise NotImplementedError
-        else:
-            l = A
-            r = b
-        AB = numba_create_banded(l, bandwidth, M)
-        H_j_next = scipy.linalg.solve_banded((bandwidth, bandwidth), AB, r,
+        AB = numba_create_banded(A, bandwidth, M)
+        H_j_next = scipy.linalg.solve_banded((bandwidth, bandwidth), AB, b,
                                              check_finite=False, overwrite_ab=True)
         assert np.isfinite(H_j_next).all()
         # Constrain heads based on allowed maximum/minimum depths
         # TODO: Not sure what's happening here
         # H_j_next = np.maximum(H_j_next, _z_inv_j + min_depth)
         # H_j_next = np.minimum(H_j_next, _z_inv_j + max_depth)
-        H_j_next = np.maximum(H_j_next, _z_inv_j)
-        H_j_next = np.minimum(H_j_next, _z_inv_j + max_depth)
+        # NOTE: Changed this while debugging
+        #H_j_next = np.maximum(H_j_next, _z_inv_j)
+        #H_j_next = np.minimum(H_j_next, _z_inv_j + max_depth)
         # Export instance variables
         self.H_j = H_j_next
 
@@ -1414,7 +1771,7 @@ def solve_internals_backwards(self, subcritical_only=False):
         # TODO: Temporary
         assert np.isfinite(_h_Ik).all()
         # Ensure non-negative depths?
-        _h_Ik[_h_Ik < min_depth] = min_depth
+        #_h_Ik[_h_Ik < min_depth] = min_depth
         # _h_Ik[_h_Ik > junction_max_depth] = junction_max_depth
         # _h_Ik[_h_Ik > max_depth] = max_depth
         # Export instance variables
@@ -1451,7 +1808,7 @@ def solve_internals_forwards(self, subcritical_only=False):
                               _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, nk, NK,
                               min_depth, max_depth_k, first_link_backwards=False)
         # Ensure non-negative depths?
-        _h_Ik[_h_Ik < min_depth] = min_depth
+        #_h_Ik[_h_Ik < min_depth] = min_depth
         # _h_Ik[_h_Ik > max_depth] = max_depth
         # Export instance variables
         self._h_Ik = _h_Ik
@@ -1500,7 +1857,7 @@ def solve_internals_lsq(self):
         _h_Ik[_is_start] = _h_uk
         _h_Ik[_is_end] = _h_dk
         # Set min depth
-        _h_Ik[_h_Ik < min_depth] = min_depth
+        #_h_Ik[_h_Ik < min_depth] = min_depth
         # Solve for flows using new depths
         Q_ik_b, Q_ik_f = self.superlink_flow_from_recurrence()
         _Q_ik = (Q_ik_b + Q_ik_f) / 2
@@ -1531,6 +1888,42 @@ def superlink_flow_from_recurrence(self):
                                      _h_uk, _Ik, _ki, n)
         return Q_ik_b, Q_ik_f
 
+    def solve_superlink_flows(self):
+        """
+        Solve for superlink boundary discharges given superjunction
+        heads at time t + dt.
+        """
+        # Import instance variables
+        _J_uk = self._J_uk            # Index of superjunction upstream of superlink k
+        _J_dk = self._J_dk            # Index of superjunction downstream of superlink k
+        _alpha_uk = self._alpha_uk    # Superlink flow coefficient
+        _alpha_dk = self._alpha_dk    # Superlink flow coefficient
+        _beta_uk = self._beta_uk      # Superlink flow coefficient
+        _beta_dk = self._beta_dk      # Superlink flow coefficient
+        _chi_uk = self._chi_uk        # Superlink flow coefficient
+        _chi_dk = self._chi_dk        # Superlink flow coefficient
+        H_j = self.H_j                # Head at superjunction j
+        _theta_uk = self._theta_uk
+        _theta_dk = self._theta_dk
+        _A_uk = self._A_uk
+        _A_dk = self._A_dk
+        _B_uk = self._B_uk
+        _B_dk = self._B_dk
+        g = 9.81
+        # Compute flow at next time step
+        _Q_uk_next = _alpha_uk * H_j[_J_uk] + _beta_uk * H_j[_J_dk] + _chi_uk
+        _Q_dk_next = _alpha_dk * H_j[_J_uk] + _beta_dk * H_j[_J_dk] + _chi_dk
+        # If overflow, need to use critical depth, otherwise singular
+        # TODO: Does this need to be accounted for in recurrence relations?
+        # TODO: Should calculate depth first, then flow?
+        #_Q_uk_crit = np.sqrt(g * _A_uk**3 / _B_uk)
+        #_Q_dk_crit = np.sqrt(g * _A_dk**3 / _B_dk)
+        #_Q_uk_next[_theta_uk == 0] = _Q_uk_crit[_theta_uk == 0]
+        #_Q_dk_next[_theta_dk == 0] = _Q_dk_crit[_theta_dk == 0]
+        # Export instance variables
+        self._Q_uk = _Q_uk_next
+        self._Q_dk = _Q_dk_next
+
     def solve_orifice_flows(self, dt, u=None):
         """
         Solve for orifice discharges given superjunction heads at time t + dt.
@@ -1557,7 +1950,8 @@ def solve_orifice_flows(self, dt, u=None):
         # TODO: Move this inside numba function
         upstream_ctrl = (H_j[_J_uo] > H_j[_J_do])
         _Qo_max = np.where(upstream_ctrl, _V_sj[_J_uo], _V_sj[_J_do]) / dt
-        _Qo_next = np.sign(_Qo_next) * np.minimum(np.abs(_Qo_next), _Qo_max)
+        # TODO: Check if flow limiter is causing issues
+        #_Qo_next = np.sign(_Qo_next) * np.minimum(np.abs(_Qo_next), _Qo_max)
         # Export instance variables
         self._Qo = _Qo_next
 
@@ -1679,7 +2073,7 @@ def reposition_junctions(self, reposition=None):
         _H_dk = H_j[_J_dk]
         # Handle which superlinks to reposition
         if reposition is None:
-            reposition = np.ones(NK, dtype=np.bool8)
+            reposition = np.ones(NK, dtype=np.bool_)
         # Reposition junctions
         numba_reposition_junctions(_x_Ik, _z_inv_Ik, _h_Ik, _dx_ik, _Q_ik, _H_dk,
                                     _b0, _zc, _xc, _m, _elem_pos, _i_1k, _I_1k,
@@ -1695,1517 +2089,4 @@ def handle_elliptical_perimeter(_Pe_ik, _ellipse_ix, _Ik, _Ip1k, _h_Ik, _g1_ik,
                                                             _g1_ik[_ik_g],
                                                             _g2_ik[_ik_g])
 
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            int64[:], int64[:], int64[:]),
-      cache=True)
-def numba_hydraulic_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik,
-                             _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
-                             _geom_codes, _Ik, _ik):
-    n = len(_ik)
-    for i in range(n):
-        I = _Ik[i]
-        Ip1 = I + 1
-        geom_code = _geom_codes[i]
-        h_I = _h_Ik[I]
-        h_Ip1 = _h_Ik[Ip1]
-        g1_i = _g1_ik[i]
-        g2_i = _g2_ik[i]
-        g3_i = _g3_ik[i]
-        g4_i = _g4_ik[i]
-        g5_i = _g5_ik[i]
-        g6_i = _g6_ik[i]
-        g7_i = _g7_ik[i]
-        if geom_code:
-            if geom_code == 1:
-                _A_ik[i] = pipedream_solver.ngeometry.Circular_A_ik(h_I, h_Ip1, g1_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Circular_Pe_ik(h_I, h_Ip1, g1_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Circular_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Circular_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 2:
-                _A_ik[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Rect_Closed_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Rect_Closed_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-            elif geom_code == 3:
-                _A_ik[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Rect_Open_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Rect_Open_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 4:
-                _A_ik[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Triangular_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Triangular_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Triangular_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 5:
-                _A_ik[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Trapezoidal_Pe_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Trapezoidal_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-            elif geom_code == 6:
-                _A_ik[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Parabolic_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Parabolic_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Parabolic_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 7:
-                # NOTE: Assumes that perimeter has already been calculated
-                _A_ik[i] = pipedream_solver.ngeometry.Elliptical_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Elliptical_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Elliptical_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 8:
-                _A_ik[i] = pipedream_solver.ngeometry.Wide_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Wide_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Wide_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Wide_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 9:
-                _A_ik[i] = pipedream_solver.ngeometry.Force_Main_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Force_Main_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Force_Main_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Force_Main_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 10:
-                _A_ik[i] = pipedream_solver.ngeometry.Floodplain_A_ik(h_I, h_Ip1, g1_i, g2_i,
-                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
-                _Pe_ik[i] = pipedream_solver.ngeometry.Floodplain_Pe_ik(h_I, h_Ip1, g1_i, g2_i,
-                                                                        g3_i, g4_i, g5_i, g6_i, g7_i)
-                _R_ik[i] = pipedream_solver.ngeometry.Floodplain_R_ik(_A_ik[i], _Pe_ik[i])
-                _B_ik[i] = pipedream_solver.ngeometry.Floodplain_B_ik(h_I, h_Ip1, g1_i, g2_i,
-                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
-    return 1
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            int64[:], int64[:], int64[:], int64[:]),
-      cache=True)
-def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, _theta_bk,
-                            _g1_ik, _g2_ik, _g3_ik, _g4_ik, _g5_ik, _g6_ik, _g7_ik,
-                            _geom_codes, _i_bk, _I_bk, _J_bk):
-    n = len(_i_bk)
-    for k in range(n):
-        i = _i_bk[k]
-        I = _I_bk[k]
-        j = _J_bk[k]
-        # TODO: This should incorporate theta
-        h_I = _h_Ik[I]
-        h_Ip1 = _theta_bk[k] * (_H_j[j] - _z_inv_bk[k])
-        geom_code = _geom_codes[i]
-        g1_i = _g1_ik[i]
-        g2_i = _g2_ik[i]
-        g3_i = _g3_ik[i]
-        g4_i = _g4_ik[i]
-        g5_i = _g5_ik[i]
-        g6_i = _g6_ik[i]
-        g7_i = _g7_ik[i]
-        if geom_code:
-            if geom_code == 1:
-                _A_bk[k] = pipedream_solver.ngeometry.Circular_A_ik(h_I, h_Ip1, g1_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Circular_Pe_ik(h_I, h_Ip1, g1_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Circular_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Circular_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 2:
-                _A_bk[k] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Rect_Closed_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Rect_Closed_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-            elif geom_code == 3:
-                _A_bk[k] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Rect_Open_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Rect_Open_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 4:
-                _A_bk[k] = pipedream_solver.ngeometry.Triangular_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Triangular_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Triangular_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Triangular_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 5:
-                _A_bk[k] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Trapezoidal_Pe_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Trapezoidal_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_I, h_Ip1, g1_i, g2_i, g3_i)
-            elif geom_code == 6:
-                _A_bk[k] = pipedream_solver.ngeometry.Parabolic_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Parabolic_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Parabolic_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Parabolic_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 7:
-                _A_bk[k] = pipedream_solver.ngeometry.Elliptical_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Elliptical_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Elliptical_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 8:
-                _A_bk[k] = pipedream_solver.ngeometry.Wide_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Wide_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Wide_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Wide_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 9:
-                _A_bk[k] = pipedream_solver.ngeometry.Force_Main_A_ik(h_I, h_Ip1, g1_i, g2_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Force_Main_Pe_ik(h_I, h_Ip1, g1_i, g2_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Force_Main_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Force_Main_B_ik(h_I, h_Ip1, g1_i, g2_i)
-            elif geom_code == 10:
-                _A_bk[k] = pipedream_solver.ngeometry.Floodplain_A_ik(h_I, h_Ip1, g1_i, g2_i,
-                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
-                _Pe_bk[k] = pipedream_solver.ngeometry.Floodplain_Pe_ik(h_I, h_Ip1, g1_i, g2_i,
-                                                                        g3_i, g4_i, g5_i, g6_i, g7_i)
-                _R_bk[k] = pipedream_solver.ngeometry.Floodplain_R_ik(_A_bk[k], _Pe_bk[k])
-                _B_bk[k] = pipedream_solver.ngeometry.Floodplain_B_ik(h_I, h_Ip1, g1_i, g2_i,
-                                                                      g3_i, g4_i, g5_i, g6_i, g7_i)
-    return 1
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            int64[:], int64),
-      cache=True)
-def numba_orifice_geometry(_Ao, h_eo, u_o, _g1_o, _g2_o, _g3_o, _geom_codes_o, n_o):
-    for i in range(n_o):
-        geom_code = _geom_codes_o[i]
-        g1 = _g1_o[i]
-        g2 = _g2_o[i]
-        g3 = _g3_o[i]
-        u = u_o[i]
-        h_e = h_eo[i]
-        if geom_code:
-            if geom_code == 1:
-                _Ao[i] = pipedream_solver.ngeometry.Circular_A_ik(h_e, h_e, g1 * u)
-            elif geom_code == 2:
-                _Ao[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_e, h_e, g1 * u, g2)
-            elif geom_code == 3:
-                _Ao[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_e, h_e, g1 * u, g2)
-            elif geom_code == 4:
-                _Ao[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_e, h_e, g1 * u, g2)
-            elif geom_code == 5:
-                _Ao[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_e, h_e, g1 * u, g2, g3)
-            elif geom_code == 6:
-                _Ao[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_e, h_e, g1 * u, g2)
-            elif geom_code == 7:
-                _Ao[i] = pipedream_solver.ngeometry.Elliptical_A_ik(h_e, h_e, g1 * u, g2)
-            elif geom_code == 8:
-                _Ao[i] = pipedream_solver.ngeometry.Wide_A_ik(h_e, h_e, g1 * u, g2)
-            elif geom_code == 9:
-                _Ao[i] = pipedream_solver.ngeometry.Force_Main_A_ik(h_e, h_e, g1 * u, g2)
-    return 1
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:],
-            int64[:], int64[:], int64[:], int64[:], int64[:]),
-      cache=True)
-def numba_transect_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik, _is_irregular, _transect_zs,
-                            _transect_As, _transect_Bs, _transect_Pes, _transect_Rs,
-                            _transect_codes, _transect_inds, _transect_lens, _Ik, _ik):
-    n = len(_ik)
-    for i in range(n):
-        is_irregular = _is_irregular[i]
-        if is_irregular:
-            I = _Ik[i]
-            Ip1 = I + 1
-            transect_code = _transect_codes[i]
-            h_I = _h_Ik[I]
-            h_Ip1 = _h_Ik[Ip1]
-            h_i = (h_I + h_Ip1) / 2
-            start = _transect_inds[transect_code]
-            size = _transect_lens[transect_code]
-            end = start + size
-            _z_range = _transect_zs[start:end]
-            _A_range = _transect_As[start:end]
-            _B_range = _transect_Bs[start:end]
-            _Pe_range = _transect_Pes[start:end]
-            _R_range = _transect_Rs[start:end]
-            _A_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _A_range, 1)
-            _B_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _B_range, 1)
-            _Pe_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _Pe_range, 1)
-            _R_ik[i] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _R_range, 1)
-    return 1
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], boolean[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], int64[:], int64[:],
-            int64[:], int64[:], int64[:], int64[:]),
-      cache=True)
-def numba_boundary_transect(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk,
-                            _theta_bk, _is_irregular, _transect_zs, _transect_As, _transect_Bs,
-                            _transect_Pes, _transect_Rs, _transect_codes, _transect_inds,
-                            _transect_lens, _I_bk, _i_bk, _J_bk):
-    n = len(_i_bk)
-    for k in range(n):
-        i = _i_bk[k]
-        I = _I_bk[k]
-        j = _J_bk[k]
-        is_irregular_bk = _is_irregular[i]
-        if is_irregular_bk:
-            h_I = _h_Ik[I]
-            # TODO: This should incorporate theta
-            h_Ip1 = _theta_bk[k] * (_H_j[j] - _z_inv_bk[k])
-            transect_code = _transect_codes[i]
-            h_i = (h_I + h_Ip1) / 2
-            start = _transect_inds[transect_code]
-            size = _transect_lens[transect_code]
-            end = start + size
-            _z_range = _transect_zs[start:end]
-            _A_range = _transect_As[start:end]
-            _B_range = _transect_Bs[start:end]
-            _Pe_range = _transect_Pes[start:end]
-            _R_range = _transect_Rs[start:end]
-            _A_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _A_range, 1)
-            _B_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _B_range, 1)
-            _Pe_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _Pe_range, 1)
-            _R_bk[k] = pipedream_solver.ngeometry.interpolate_geometry(h_i, _z_range, _R_range, 1)
-    return 1
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]),
-      cache=True)
-def numba_compute_functional_storage_areas(h, A, a, b, c, _functional):
-    M = h.size
-    for j in range(M):
-        if _functional[j]:
-            if h[j] < 0:
-                A[j] = 0
-            else:
-                A[j] = a[j] * (h[j]**b[j]) + c[j]
-    return A
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]),
-      cache=True)
-def numba_compute_functional_storage_volumes(h, V, a, b, c, _functional):
-    M = h.size
-    for j in range(M):
-        if _functional[j]:
-            if h[j] < 0:
-                V[j] = 0
-            else:
-                V[j] = (a[j] / (b[j] + 1)) * h[j] ** (b[j] + 1) + c[j] * h[j]
-    return V
-
-@njit
-def numba_compute_tabular_storage_areas(h_j, A_sj, hs, As, sjs, sts, inds, lens):
-    n = sjs.size
-    for i in range(n):
-        sj = sjs[i]
-        st = sts[i]
-        ind = inds[st]
-        size = lens[st]
-        h_range = hs[ind:ind+size]
-        A_range = As[ind:ind+size]
-        Amin = A_range.min()
-        Amax = A_range.max()
-        h_search = h_j[sj]
-        ix = np.searchsorted(h_range, h_search)
-        # NOTE: np.interp not supported in this version of numba
-        # A_result = np.interp(h_search, h_range, A_range)
-        # A_out[i] = A_result
-        if (ix == 0):
-            A_sj[sj] = Amin
-        elif (ix >= size):
-            A_sj[sj] = Amax
-        else:
-            dx_0 = h_search - h_range[ix - 1]
-            dx_1 = h_range[ix] - h_search
-            frac = dx_0 / (dx_0 + dx_1)
-            A_sj[sj] = (1 - frac) * A_range[ix - 1] + (frac) * A_range[ix]
-    return A_sj
-
-@njit
-def numba_compute_tabular_storage_volumes(h_j, V_sj, hs, As, Vs, sjs, sts, inds, lens):
-    n = sjs.size
-    for i in range(n):
-        sj = sjs[i]
-        st = sts[i]
-        ind = inds[st]
-        size = lens[st]
-        h_range = hs[ind:ind+size]
-        A_range = As[ind:ind+size]
-        V_range = Vs[ind:ind+size]
-        hmax = h_range.max()
-        Vmin = V_range.min()
-        Vmax = V_range.max()
-        Amax = A_range.max()
-        h_search = h_j[sj]
-        ix = np.searchsorted(h_range, h_search)
-        # NOTE: np.interp not supported in this version of numba
-        # A_result = np.interp(h_search, h_range, A_range)
-        # A_out[i] = A_result
-        if (ix == 0):
-            V_sj[sj] = Vmin
-        elif (ix >= size):
-            V_sj[sj] = Vmax + Amax * (h_search - hmax)
-        else:
-            dx_0 = h_search - h_range[ix - 1]
-            dx_1 = h_range[ix] - h_search
-            frac = dx_0 / (dx_0 + dx_1)
-            V_sj[sj] = (1 - frac) * V_range[ix - 1] + (frac) * V_range[ix]
-    return V_sj
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64))
-def friction_slope(Q_ik_t, dx_ik, A_ik, R_ik, n_ik, Sf_method_ik, g=9.81):
-    if A_ik > 0:
-        # Chezy-Manning eq.
-        if Sf_method_ik == 0:
-            t_1 = (g * n_ik**2 * np.abs(Q_ik_t) * dx_ik
-                   / A_ik / R_ik**(4/3))
-        # Hazen-Williams eq.
-        elif Sf_method_ik == 1:
-            t_1 = (1.354 * g * np.abs(Q_ik_t)**0.85 * dx_ik
-                   / A_ik**0.85 / n_ik**1.85 / R_ik**1.1655)
-        # Darcy-Weisbach eq.
-        elif Sf_method_ik == 2:
-            # kinematic viscosity(meter^2/sec), we can consider this is constant.
-            nu = 0.0000010034
-            Re = (np.abs(Q_ik_t) / A_ik) * 4 * R_ik / nu
-            f = 0.25 / (np.log10(n_ik / (3.7 * 4 * R_ik) + 5.74 / (Re**0.9)))**2
-            t_1 = (0.01274 * g * f * np.abs(Q_ik_t) * dx_ik
-                   / (A_ik * R_ik))
-        else:
-            raise ValueError('Invalid friction method.')
-        return t_1
-    else:
-        return 0.
-
-@njit(float64[:](float64[:], float64[:]),
-      cache=True)
-def numba_a_ik(u_Ik, sigma_ik):
-    """
-    Compute link coefficient 'a' for link i, superlink k.
-    """
-    return -np.maximum(u_Ik, 0) * sigma_ik
-
-@njit(float64[:](float64[:], float64[:]),
-      cache=True)
-def numba_c_ik(u_Ip1k, sigma_ik):
-    """
-    Compute link coefficient 'c' for link i, superlink k.
-    """
-    return -np.maximum(-u_Ip1k, 0) * sigma_ik
-
-@njit(float64[:](float64[:], float64, float64[:], float64[:], float64[:], float64[:],
-                 float64[:], float64[:], float64[:], float64[:], boolean[:], float64[:], int64[:], float64),
-      cache=True)
-def numba_b_ik(dx_ik, dt, n_ik, Q_ik_t, A_ik, R_ik,
-               A_c_ik, C_ik, a_ik, c_ik, ctrl, sigma_ik, Sf_method_ik, g=9.81):
-    """
-    Compute link coefficient 'b' for link i, superlink k.
-    """
-    # TODO: Clean up
-    t_0 = (dx_ik / dt) * sigma_ik
-    t_1 = np.zeros(Q_ik_t.size)
-    k = len(Sf_method_ik)
-    for n in range(k):
-        t_1[n] = friction_slope(Q_ik_t[n], dx_ik[n], A_ik[n], R_ik[n],
-                                n_ik[n], Sf_method_ik[n], g)
-    t_2 = np.zeros(ctrl.size)
-    cond = ctrl
-    t_2[cond] = C_ik[cond] * A_ik[cond] * np.abs(Q_ik_t[cond]) / A_c_ik[cond]**2
-    t_3 = a_ik
-    t_4 = c_ik
-    return t_0 + t_1 + t_2 - t_3 - t_4
-
-@njit(float64[:](float64[:], float64[:], float64, float64[:], float64[:], float64[:], float64),
-      cache=True)
-def numba_P_ik(Q_ik_t, dx_ik, dt, A_ik, S_o_ik, sigma_ik, g=9.81):
-    """
-    Compute link coefficient 'P' for link i, superlink k.
-    """
-    t_0 = (Q_ik_t * dx_ik / dt) * sigma_ik
-    t_1 = g * A_ik * S_o_ik * dx_ik
-    return t_0 + t_1
-
-@njit(float64(float64, float64, float64, float64, float64, float64),
-      cache=True)
-def E_Ik(B_ik, dx_ik, B_im1k, dx_im1k, A_SIk, dt):
-    """
-    Compute node coefficient 'E' for node I, superlink k.
-    """
-    t_0 = B_ik * dx_ik / 2
-    t_1 = B_im1k * dx_im1k / 2
-    t_2 = A_SIk
-    t_3 = dt
-    return (t_0 + t_1 + t_2) / t_3
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64),
-      cache=True)
-def D_Ik(Q_0IK, B_ik, dx_ik, B_im1k, dx_im1k, A_SIk, h_Ik_t, dt):
-    """
-    Compute node coefficient 'D' for node I, superlink k.
-    """
-    t_0 = Q_0IK
-    t_1 = B_ik * dx_ik / 2
-    t_2 = B_im1k * dx_im1k / 2
-    t_3 = A_SIk
-    t_4 = h_Ik_t / dt
-    return t_0 + ((t_1 + t_2 + t_3) * t_4)
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], int64[:], float64,
-            int64[:], int64[:], boolean[:], boolean[:]),
-      cache=True)
-def numba_node_coeffs(_D_Ik, _E_Ik, _Q_0Ik, _B_ik, _h_Ik, _dx_ik, _A_SIk,
-                      _B_uk, _B_dk, _dx_uk, _dx_dk, _kI, _dt,
-                      _forward_I_i, _backward_I_i, _is_start, _is_end):
-    N = _h_Ik.size
-    for I in range(N):
-        k = _kI[I]
-        if _is_start[I]:
-            i = _forward_I_i[I]
-            _E_Ik[I] = E_Ik(_B_ik[i], _dx_ik[i], _B_uk[k], _dx_uk[k], _A_SIk[I], _dt)
-            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_ik[i], _dx_ik[i], _B_uk[k], _dx_uk[k], _A_SIk[I],
-                            _h_Ik[I], _dt)
-        elif _is_end[I]:
-            im1 = _backward_I_i[I]
-            _E_Ik[I] = E_Ik(_B_dk[k], _dx_dk[k], _B_ik[im1], _dx_ik[im1],
-                            _A_SIk[I], _dt)
-            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_dk[k], _dx_dk[k], _B_ik[im1],
-                            _dx_ik[im1], _A_SIk[I], _h_Ik[I], _dt)
-        else:
-            i = _forward_I_i[I]
-            im1 = i - 1
-            _E_Ik[I] = E_Ik(_B_ik[i], _dx_ik[i], _B_ik[im1], _dx_ik[im1],
-                            _A_SIk[I], _dt)
-            _D_Ik[I] = D_Ik(_Q_0Ik[I], _B_ik[i], _dx_ik[i], _B_ik[im1],
-                            _dx_ik[im1], _A_SIk[I], _h_Ik[I], _dt)
-    return 1
-
-@njit(float64(float64, float64),
-      cache=True)
-def safe_divide(num, den):
-    if (den == 0):
-        return 0
-    else:
-        return num / den
-
-@njit(float64[:](float64[:], float64[:]),
-      cache=True)
-def safe_divide_vec(num, den):
-    result = np.zeros_like(num)
-    cond = (den != 0)
-    result[cond] = num[cond] / den[cond]
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def Q_i_f(h_Ip1k, h_1k, U_Ik, V_Ik, W_Ik):
-    t_0 = U_Ik * h_Ip1k
-    t_1 = V_Ik
-    t_2 = W_Ik * h_1k
-    return t_0 + t_1 + t_2
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def Q_i_b(h_Ik, h_Np1k, X_Ik, Y_Ik, Z_Ik):
-    t_0 = X_Ik * h_Ik
-    t_1 = Y_Ik
-    t_2 = Z_Ik * h_Np1k
-    return t_0 + t_1 + t_2
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def h_i_b(Q_ik, h_Np1k, X_Ik, Y_Ik, Z_Ik):
-    num = Q_ik - Y_Ik - Z_Ik * h_Np1k
-    den = X_Ik
-    result = safe_divide(num, den)
-    return result
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], int64[:], int64[:], int64[:], int64,
-            float64, float64[:], boolean),
-      cache=True)
-def numba_solve_internals(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik,
-                          _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, nk, NK,
-                          min_depth, max_depth_k, first_link_backwards=True):
-    for k in range(NK):
-        n = nk[k]
-        i_1 = _i_1k[k]
-        I_1 = _I_1k[k]
-        i_n = i_1 + n - 1
-        I_Np1 = I_1 + n
-        I_N = I_Np1 - 1
-        # Set boundary depths
-        _h_1k = _h_uk[k]
-        _h_Np1k = _h_dk[k]
-        _h_Ik[I_1] = _h_1k
-        _h_Ik[I_Np1] = _h_Np1k
-        # Set max depth
-        max_depth = max_depth_k[k]
-        # Compute internal depths and flows (except first link flow)
-        for j in range(n - 1):
-            I = I_N - j
-            Ip1 = I + 1
-            i = i_n - j
-            _Q_ik[i] = Q_i_f(_h_Ik[Ip1], _h_1k, _U_Ik[I], _V_Ik[I], _W_Ik[I])
-            _h_Ik[I] = h_i_b(_Q_ik[i], _h_Np1k, _X_Ik[I], _Y_Ik[I], _Z_Ik[I])
-            if _h_Ik[I] < min_depth:
-                _h_Ik[I] = min_depth
-            if _h_Ik[I] > max_depth:
-                _h_Ik[I] = max_depth
-        if first_link_backwards:
-            _Q_ik[i_1] = Q_i_b(_h_Ik[I_1], _h_Np1k, _X_Ik[I_1], _Y_Ik[I_1],
-                            _Z_Ik[I_1])
-        else:
-            # Not theoretically correct, but seems to be more stable sometimes
-            _Q_ik[i_1] = Q_i_f(_h_Ik[I_1 + 1], _h_1k, _U_Ik[I_1], _V_Ik[I_1],
-                            _W_Ik[I_1])
-    return 1
-
-@njit(float64[:](float64[:], int64, int64[:], int64[:], int64[:], int64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_solve_internals_ls(_h_Ik, NK, nk, _k_1k, _i_1k, _I_1k, _U, _X, _b):
-    for k in range(NK):
-        nlinks = nk[k]
-        lstart = _k_1k[k]
-        rstart = _i_1k[k]
-        jstart = _I_1k[k]
-        _Ak = np.zeros((nlinks, nlinks - 1))
-        for i in range(nlinks - 1):
-            _Ak[i, i] = _U[lstart + i]
-            _Ak[i + 1, i] = -_X[lstart + i]
-        _AkT = _Ak.T.copy()
-        _bk = _b[rstart:rstart+nlinks].copy()
-        _AA = _AkT @ _Ak
-        _Ab = _AkT @ _bk
-        # If want to prevent singular matrix, set ( diag == 0 ) = 1
-        for i in range(nlinks - 1):
-            if (_AA[i, i] == 0.0):
-                _AA[i, i] = 1.0
-        _h_inner = np.linalg.solve(_AA, _Ab)
-        _h_Ik[jstart+1:jstart+nlinks] = _h_inner
-    return _h_Ik
-
-@njit(float64[:](float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_u_ik(_Q_ik, _A_ik, _u_ik):
-    n = _u_ik.size
-    for i in range(n):
-        _Q_i = _Q_ik[i]
-        _A_i = _A_ik[i]
-        if _A_i:
-            _u_ik[i] = _Q_i / _A_i
-        else:
-            _u_ik[i] = 0
-    return _u_ik
-
-@njit(float64[:](float64[:], float64[:], boolean[:], float64[:]),
-      cache=True)
-def numba_u_Ik(_dx_ik, _u_ik, _link_start, _u_Ik):
-    n = _u_Ik.size
-    for i in range(n):
-        if _link_start[i]:
-            _u_Ik[i] = _u_ik[i]
-        else:
-            im1 = i - 1
-            num = _dx_ik[i] * _u_ik[im1] + _dx_ik[im1] * _u_ik[i]
-            den = _dx_ik[i] + _dx_ik[im1]
-            if den:
-                _u_Ik[i] = num / den
-            else:
-                _u_Ik[i] = 0
-    return _u_Ik
-
-@njit(float64[:](float64[:], float64[:], boolean[:], float64[:]),
-      cache=True)
-def numba_u_Ip1k(_dx_ik, _u_ik, _link_end, _u_Ip1k):
-    n = _u_Ip1k.size
-    for i in range(n):
-        if _link_end[i]:
-            _u_Ip1k[i] = _u_ik[i]
-        else:
-            ip1 = i + 1
-            num = _dx_ik[i] * _u_ik[ip1] + _dx_ik[ip1] * _u_ik[i]
-            den = _dx_ik[i] + _dx_ik[ip1]
-            if den:
-                _u_Ip1k[i] = num / den
-            else:
-                _u_Ip1k[i] = 0
-    return _u_Ip1k
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-                 int64[:], float64, float64), cache=True)
-def kappa_uk(Q_uk, dx_uk, A_uk, C_uk, R_uk, n_uk, Sf_method_uk, dt, g=9.81):
-    """
-    Compute boundary coefficient 'kappa' for upstream end of superlink k.
-    """
-    k = Q_uk.size
-    t_0 = - dx_uk / g / A_uk / dt
-    t_1 = np.zeros(k, dtype=np.float64)
-    for n in range(k):
-        t_1[n] = friction_slope(Q_uk[n], dx_uk[n], A_uk[n], R_uk[n],
-                                n_uk[n], Sf_method_uk[n], g)
-    t_2 = - C_uk * np.abs(Q_uk) / 2 / g / A_uk**2
-    return t_0 + t_1 + t_2
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-                 int64[:], float64, float64), cache=True)
-def kappa_dk(Q_dk, dx_dk, A_dk, C_dk, R_dk, n_dk, Sf_method_dk, dt, g=9.81):
-    """
-    Compute boundary coefficient 'kappa' for downstream end of superlink k.
-    """
-    k = Q_dk.size
-    t_0 = dx_dk / g / A_dk / dt
-    t_1 = np.zeros(k, dtype=np.float64)
-    for n in range(k):
-        t_1[n] = friction_slope(Q_dk[n], dx_dk[n], A_dk[n], R_dk[n],
-                                n_dk[n], Sf_method_dk[n], g)
-    t_2 = C_dk * np.abs(Q_dk) / 2 / g / A_dk**2
-    return t_0 + t_1 + t_2
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:],
-                 float64[:], float64, float64), cache=True)
-def mu_uk(Q_uk_t, dx_uk, A_uk, theta_uk, z_inv_uk, S_o_uk, dt, g=9.81):
-    """
-    Compute boundary coefficient 'mu' for upstream end of superlink k.
-    """
-    t_0 = Q_uk_t * dx_uk / g / A_uk / dt
-    t_1 = - theta_uk * z_inv_uk
-    t_2 = dx_uk * S_o_uk
-    return t_0 + t_1 + t_2
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:],
-                 float64[:], float64, float64), cache=True)
-def mu_dk(Q_dk_t, dx_dk, A_dk, theta_dk, z_inv_dk, S_o_dk, dt, g=9.81):
-    """
-    Compute boundary coefficient 'mu' for downstream end of superlink k.
-    """
-    t_0 = - Q_dk_t * dx_dk / g / A_dk / dt
-    t_1 = - theta_dk * z_inv_dk
-    t_2 = - dx_dk * S_o_dk
-    return t_0 + t_1 + t_2
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def U_1k(E_2k, c_1k, A_1k, T_1k, g=9.81):
-    """
-    Compute forward recurrence coefficient 'U' for node 1, superlink k.
-    """
-    num = E_2k * c_1k - g * A_1k
-    den = T_1k
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64, float64),
-      cache=True)
-def V_1k(P_1k, D_2k, c_1k, T_1k, a_1k=0.0, D_1k=0.0):
-    """
-    Compute forward recurrence coefficient 'V' for node 1, superlink k.
-    """
-    num = P_1k - D_2k * c_1k + D_1k * a_1k
-    den = T_1k
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def W_1k(A_1k, T_1k, a_1k=0.0, E_1k=0.0, g=9.81):
-    """
-    Compute forward recurrence coefficient 'W' for node 1, superlink k.
-    """
-    num = g * A_1k - E_1k * a_1k
-    den = T_1k
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64),
-      cache=True)
-def T_1k(a_1k, b_1k, c_1k):
-    """
-    Compute forward recurrence coefficient 'T' for link 1, superlink k.
-    """
-    return a_1k + b_1k + c_1k
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def U_Ik(E_Ip1k, c_ik, A_ik, T_ik, g=9.81):
-    """
-    Compute forward recurrence coefficient 'U' for node I, superlink k.
-    """
-    num = E_Ip1k * c_ik - g * A_ik
-    den = T_ik
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64, float64, float64, float64),
-      cache=True)
-def V_Ik(P_ik, a_ik, D_Ik, D_Ip1k, c_ik, A_ik, E_Ik, V_Im1k, U_Im1k, T_ik, g=9.81):
-    """
-    Compute forward recurrence coefficient 'V' for node I, superlink k.
-    """
-    t_0 = P_ik
-    t_1 = a_ik * D_Ik
-    t_2 = D_Ip1k * c_ik
-    t_3 = (g * A_ik - E_Ik * a_ik)
-    t_4 = V_Im1k + D_Ik
-    t_5 = U_Im1k - E_Ik
-    t_6 = T_ik
-    # TODO: There is still a divide by zero here
-    num = (t_0 + t_1 - t_2 - (t_3 * t_4 / t_5))
-    den = t_6
-    result = safe_divide(num, den)
-    return  result
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64),
-      cache=True)
-def W_Ik(A_ik, E_Ik, a_ik, W_Im1k, U_Im1k, T_ik, g=9.81):
-    """
-    Compute forward recurrence coefficient 'W' for node I, superlink k.
-    """
-    num = -(g * A_ik - E_Ik * a_ik) * W_Im1k
-    den = (U_Im1k - E_Ik) * T_ik
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64),
-      cache=True)
-def T_ik(a_ik, b_ik, c_ik, A_ik, E_Ik, U_Im1k, g=9.81):
-    """
-    Compute forward recurrence coefficient 'T' for link i, superlink k.
-    """
-    t_0 = a_ik + b_ik + c_ik
-    t_1 = g * A_ik - E_Ik * a_ik
-    t_2 = U_Im1k - E_Ik
-    result = t_0 - safe_divide(t_1, t_2)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def X_Nk(A_nk, E_Nk, a_nk, O_nk, g=9.81):
-    """
-    Compute backward recurrence coefficient 'X' for node N, superlink k.
-    """
-    num = g * A_nk - E_Nk * a_nk
-    den = O_nk
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64, float64),
-      cache=True)
-def Y_Nk(P_nk, D_Nk, a_nk, O_nk, c_nk=0.0, D_Np1k=0.0):
-    """
-    Compute backward recurrence coefficient 'Y' for node N, superlink k.
-    """
-    num = P_nk + D_Nk * a_nk - D_Np1k * c_nk
-    den = O_nk
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def Z_Nk(A_nk, O_nk, c_nk=0.0, E_Np1k=0.0, g=9.81):
-    """
-    Compute backward recurrence coefficient 'Z' for node N, superlink k.
-    """
-    num = E_Np1k * c_nk - g * A_nk
-    den = O_nk
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64),
-      cache=True)
-def O_nk(a_nk, b_nk, c_nk):
-    """
-    Compute backward recurrence coefficient 'O' for link n, superlink k.
-    """
-    return a_nk + b_nk + c_nk
-
-@njit(float64(float64, float64, float64, float64, float64),
-      cache=True)
-def X_Ik(A_ik, E_Ik, a_ik, O_ik, g=9.81):
-    """
-    Compute backward recurrence coefficient 'X' for node I, superlink k.
-    """
-    num = g * A_ik - E_Ik * a_ik
-    den = O_ik
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64, float64, float64, float64, float64),
-      cache=True)
-def Y_Ik(P_ik, a_ik, D_Ik, D_Ip1k, c_ik, A_ik, E_Ip1k, Y_Ip1k, X_Ip1k, O_ik, g=9.81):
-    """
-    Compute backward recurrence coefficient 'Y' for node I, superlink k.
-    """
-    t_0 = P_ik
-    t_1 = a_ik * D_Ik
-    t_2 = D_Ip1k * c_ik
-    t_3 = (g * A_ik - E_Ip1k * c_ik)
-    t_4 = D_Ip1k - Y_Ip1k
-    t_5 = X_Ip1k + E_Ip1k
-    t_6 = O_ik
-    # TODO: There is still a divide by zero here
-    num = (t_0 + t_1 - t_2 - (t_3 * t_4 / t_5))
-    den = t_6
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64),
-      cache=True)
-def Z_Ik(A_ik, E_Ip1k, c_ik, Z_Ip1k, X_Ip1k, O_ik, g=9.81):
-    """
-    Compute backward recurrence coefficient 'Z' for node I, superlink k.
-    """
-    num = (g * A_ik - E_Ip1k * c_ik) * Z_Ip1k
-    den = (X_Ip1k + E_Ip1k) * O_ik
-    result = safe_divide(num, den)
-    return result
-
-@njit(float64(float64, float64, float64, float64, float64, float64, float64),
-      cache=True)
-def O_ik(a_ik, b_ik, c_ik, A_ik, E_Ip1k, X_Ip1k, g=9.81):
-    """
-    Compute backward recurrence coefficient 'O' for link i, superlink k.
-    """
-    t_0 = a_ik + b_ik + c_ik
-    t_1 = g * A_ik - E_Ip1k * c_ik
-    t_2 = X_Ip1k + E_Ip1k
-    result = t_0 + safe_divide(t_1, t_2)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_D_k_star(X_1k, kappa_uk, U_Nk, kappa_dk, Z_1k, W_Nk):
-    """
-    Compute superlink boundary condition coefficient 'D_k_star'.
-    """
-    t_0 = (X_1k * kappa_uk - 1) * (U_Nk * kappa_dk - 1)
-    t_1 = (Z_1k * kappa_dk) * (W_Nk * kappa_uk)
-    result = t_0 - t_1
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_alpha_uk(U_Nk, kappa_dk, X_1k, Z_1k, W_Nk, D_k_star, lambda_uk):
-    """
-    Compute superlink boundary condition coefficient 'alpha' for upstream end
-    of superlink k.
-    """
-    num = lambda_uk * ((1 - U_Nk * kappa_dk) * X_1k + (Z_1k * kappa_dk * W_Nk))
-    den = D_k_star
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_beta_uk(U_Nk, kappa_dk, Z_1k, W_Nk, D_k_star, lambda_dk):
-    """
-    Compute superlink boundary condition coefficient 'beta' for upstream end
-    of superlink k.
-    """
-    num = lambda_dk * ((1 - U_Nk * kappa_dk) * Z_1k + (Z_1k * kappa_dk * U_Nk))
-    den = D_k_star
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-                 float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_chi_uk(U_Nk, kappa_dk, Y_1k, X_1k, mu_uk, Z_1k,
-                 mu_dk, V_Nk, W_Nk, D_k_star):
-    """
-    Compute superlink boundary condition coefficient 'chi' for upstream end
-    of superlink k.
-    """
-    t_0 = (1 - U_Nk * kappa_dk) * (Y_1k + X_1k * mu_uk + Z_1k * mu_dk)
-    t_1 = (Z_1k * kappa_dk) * (V_Nk + W_Nk * mu_uk + U_Nk * mu_dk)
-    num = t_0 + t_1
-    den = D_k_star
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_alpha_dk(X_1k, kappa_uk, W_Nk, D_k_star, lambda_uk):
-    """
-    Compute superlink boundary condition coefficient 'alpha' for downstream end
-    of superlink k.
-    """
-    num = lambda_uk * ((1 - X_1k * kappa_uk) * W_Nk + (W_Nk * kappa_uk * X_1k))
-    den = D_k_star
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_beta_dk(X_1k, kappa_uk, U_Nk, W_Nk, Z_1k, D_k_star, lambda_dk):
-    """
-    Compute superlink boundary condition coefficient 'beta' for downstream end
-    of superlink k.
-    """
-    num = lambda_dk * ((1 - X_1k * kappa_uk) * U_Nk + (W_Nk * kappa_uk * Z_1k))
-    den = D_k_star
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-                 float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def numba_chi_dk(X_1k, kappa_uk, V_Nk, W_Nk, mu_uk, U_Nk,
-                    mu_dk, Y_1k, Z_1k, D_k_star):
-    """
-    Compute superlink boundary condition coefficient 'chi' for downstream end
-    of superlink k.
-    """
-    t_0 = (1 - X_1k * kappa_uk) * (V_Nk + W_Nk * mu_uk + U_Nk * mu_dk)
-    t_1 = (W_Nk * kappa_uk) * (Y_1k + X_1k * mu_uk + Z_1k * mu_dk)
-    num = t_0 + t_1
-    den = D_k_star
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64),
-      cache=True)
-def gamma_o(Q_o_t, Ao, Co, g=9.81):
-    """
-    Compute flow coefficient 'gamma' for orifice o.
-    """
-    num = 2 * g * Co**2 * Ao**2
-    den = np.abs(Q_o_t)
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def gamma_w(Q_w_t, H_w_t, L_w, s_w, Cwr, Cwt):
-    """
-    Compute flow coefficient 'gamma' for weir w.
-    """
-    num = (Cwr * L_w * H_w_t + Cwt * s_w * H_w_t**2)**2
-    den = np.abs(Q_w_t)
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:]),
-      cache=True)
-def gamma_p(Q_p_t, b_p, c_p, u):
-    """
-    Compute flow coefficient 'gamma' for pump p.
-    """
-    num = u
-    den = b_p * np.abs(Q_p_t)**(c_p - 1)
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64),
-      cache=True)
-def gamma_uk(Q_uk_t, C_uk, A_uk, g=9.81):
-    """
-    Compute flow coefficient 'gamma' for upstream end of superlink k
-    """
-    num = -np.abs(Q_uk_t) * C_uk
-    den = 2 * (A_uk**2) * g
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64),
-      cache=True)
-def gamma_dk(Q_dk_t, C_dk, A_dk, g=9.81):
-    """
-    Compute flow coefficient 'gamma' for downstream end of superlink k
-    """
-    num = np.abs(Q_dk_t) * C_dk
-    den = 2 * (A_dk**2) * g
-    result = safe_divide_vec(num, den)
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64),
-      cache=True)
-def xi_uk(dx_uk, B_uk, theta_uk, dt):
-    num = dx_uk * B_uk * theta_uk
-    den = 2 * dt
-    result = num / den
-    return result
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64),
-      cache=True)
-def xi_dk(dx_dk, B_dk, theta_dk, dt):
-    num = dx_dk * B_dk * theta_dk
-    den = 2 * dt
-    result = num / den
-    return result
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:],
-             int64[:], int64[:]),
-      cache=True)
-def numba_orifice_flow_coefficients(_alpha_o, _beta_o, _chi_o, H_j, _Qo, u, _z_inv_j,
-                                    _z_o, _tau_o, _Co, _Ao, _y_max_o, _unidir_o,
-                                     _J_uo, _J_do):
-    g = 9.81
-    _H_uo = H_j[_J_uo]
-    _H_do = H_j[_J_do]
-    _z_inv_uo = _z_inv_j[_J_uo]
-    # Create indicator functions
-    _omega_o = np.zeros_like(_H_uo)
-    _omega_o[_H_uo >= _H_do] = 1.0
-    # Compute universal coefficients
-    _gamma_o = gamma_o(_Qo, _Ao, _Co, g)
-    # Create conditionals
-    cond_0 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
-                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u))
-    cond_1 = ((1 - _omega_o) * _H_uo + _omega_o * _H_do >
-                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u / 2))
-    cond_2 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
-                _z_o + _z_inv_uo)
-    cond_3 = (_H_do >= _H_uo) & _unidir_o
-    # Fill coefficient arrays
-    # Submerged on both sides
-    a = (cond_0 & cond_1)
-    _alpha_o[a] = _gamma_o[a]
-    _beta_o[a] = -_gamma_o[a]
-    _chi_o[a] = 0.0
-    # Submerged on one side
-    b = (cond_0 & ~cond_1)
-    _alpha_o[b] = _gamma_o[b] * _omega_o[b] * (-1)**(1 - _omega_o[b])
-    _beta_o[b] = _gamma_o[b] * (1 - _omega_o[b]) * (-1)**(1 - _omega_o[b])
-    _chi_o[b] = (_gamma_o[b] * (-1)**(1 - _omega_o[b])
-                                    * (- _z_inv_uo[b] - _z_o[b] -
-                                        _tau_o[b] * _y_max_o[b] * u[b] / 2))
-    # Weir flow
-    c = (~cond_0 & cond_2)
-    _alpha_o[c] = _gamma_o[c] * _omega_o[c] * (-1)**(1 - _omega_o[c])
-    _beta_o[c] = _gamma_o[c] * (1 - _omega_o[c]) * (-1)**(1 - _omega_o[c])
-    _chi_o[c] = (_gamma_o[c] * (-1)**(1 - _omega_o[c])
-                                    * (- _z_inv_uo[c] - _z_o[c]))
-    # No flow
-    d = (~cond_0 & ~cond_2) | cond_3
-    _alpha_o[d] = 0.0
-    _beta_o[d] = 0.0
-    _chi_o[d] = 0.0
-    return 1
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], boolean[:],
-            int64[:], int64[:], float64),
-      cache=True)
-def numba_solve_orifice_flows(H_j, u, _z_inv_j, _z_o,
-                              _tau_o, _y_max_o, _Co, _Ao, _unidir_o, _J_uo, _J_do, g=9.81):
-    # Specify orifice heads at previous timestep
-    _H_uo = H_j[_J_uo]
-    _H_do = H_j[_J_do]
-    _z_inv_uo = _z_inv_j[_J_uo]
-    # Create indicator functions
-    _omega_o = np.zeros_like(_H_uo)
-    _omega_o[_H_uo >= _H_do] = 1.0
-    # Create arrays to store flow coefficients for current time step
-    _alpha_o = np.zeros_like(_H_uo)
-    _beta_o = np.zeros_like(_H_uo)
-    _chi_o = np.zeros_like(_H_uo)
-    # Compute universal coefficients
-    _gamma_o = 2 * g * _Co**2 * _Ao**2
-    # Create conditionals
-    cond_0 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
-                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u))
-    cond_1 = ((1 - _omega_o) * _H_uo + _omega_o * _H_do >
-                _z_o + _z_inv_uo + (_tau_o * _y_max_o * u / 2))
-    cond_2 = (_omega_o * _H_uo + (1 - _omega_o) * _H_do >
-                _z_o + _z_inv_uo)
-    cond_3 = (_H_do >= _H_uo) & _unidir_o
-    # Fill coefficient arrays
-    # Submerged on both sides
-    a = (cond_0 & cond_1)
-    _alpha_o[a] = _gamma_o[a]
-    _beta_o[a] = -_gamma_o[a]
-    _chi_o[a] = 0.0
-    # Submerged on one side
-    b = (cond_0 & ~cond_1)
-    _alpha_o[b] = _gamma_o[b] * _omega_o[b] * (-1)**(1 - _omega_o[b])
-    _beta_o[b] = _gamma_o[b] * (1 - _omega_o[b]) * (-1)**(1 - _omega_o[b])
-    _chi_o[b] = (_gamma_o[b] * (-1)**(1 - _omega_o[b])
-                                    * (- _z_inv_uo[b] - _z_o[b]
-                                        - _tau_o[b] * _y_max_o[b] * u[b] / 2))
-    # Weir flow on one side
-    c = (~cond_0 & cond_2)
-    _alpha_o[c] = _gamma_o[c] * _omega_o[c] * (-1)**(1 - _omega_o[c])
-    _beta_o[c] = _gamma_o[c] * (1 - _omega_o[c]) * (-1)**(1 - _omega_o[c])
-    _chi_o[c] = (_gamma_o[c] * (-1)**(1 - _omega_o[c])
-                                    * (- _z_inv_uo[c] - _z_o[c]))
-    # No flow
-    d = (~cond_0 & ~cond_2) | cond_3
-    _alpha_o[d] = 0.0
-    _beta_o[d] = 0.0
-    _chi_o[d] = 0.0
-    # Compute flow
-    _Qo_next = (-1)**(1 - _omega_o) * np.sqrt(np.abs(
-               _alpha_o * _H_uo + _beta_o * _H_do + _chi_o))
-    # Export instance variables
-    return _Qo_next
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:]),
-      cache=True)
-def numba_weir_flow_coefficients(_Hw, _Qw, _alpha_w, _beta_w, _chi_w, H_j, _z_inv_j, _z_w,
-                                 _y_max_w, u, _L_w, _s_w, _Cwr, _Cwt, _J_uw, _J_dw):
-    # Specify weir heads at previous timestep
-    _H_uw = H_j[_J_uw]
-    _H_dw = H_j[_J_dw]
-    _z_inv_uw = _z_inv_j[_J_uw]
-    # Create indicator functions
-    _omega_w = np.zeros(_H_uw.size)
-    _omega_w[_H_uw >= _H_dw] = 1.0
-    # Create conditionals
-    cond_0 = (_omega_w * _H_uw + (1 - _omega_w) * _H_dw >
-                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
-    cond_1 = ((1 - _omega_w) * _H_uw + _omega_w * _H_dw >
-                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
-    # Effective heads
-    a = (cond_0 & cond_1)
-    b = (cond_0 & ~cond_1)
-    c = (~cond_0)
-    _Hw[a] = _H_uw[a] - _H_dw[a]
-    _Hw[b] = (_omega_w[b] * _H_uw[b] + (1 - _omega_w[b]) * _H_dw[b]
-                    + (-_z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
-    _Hw[c] = 0.0
-    _Hw = np.abs(_Hw)
-    # Compute universal coefficients
-    _gamma_w = gamma_w(_Qw, _Hw, _L_w, _s_w, _Cwr, _Cwt)
-    # Fill coefficient arrays
-    # Submerged on both sides
-    a = (cond_0 & cond_1)
-    _alpha_w[a] = _gamma_w[a]
-    _beta_w[a] = -_gamma_w[a]
-    _chi_w[a] = 0.0
-    # Submerged on one side
-    b = (cond_0 & ~cond_1)
-    _alpha_w[b] = _gamma_w[b] * _omega_w[b] * (-1)**(1 - _omega_w[b])
-    _beta_w[b] = _gamma_w[b] * (1 - _omega_w[b]) * (-1)**(1 - _omega_w[b])
-    _chi_w[b] = (_gamma_w[b] * (-1)**(1 - _omega_w[b]) *
-                                (- _z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
-    # No flow
-    c = (~cond_0)
-    _alpha_w[c] = 0.0
-    _beta_w[c] = 0.0
-    _chi_w[c] = 0.0
-    return 1
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], int64[:], int64[:]),
-      cache=True)
-def numba_solve_weir_flows(_Hw, _Qw, H_j, _z_inv_j, _z_w, _y_max_w, u, _L_w,
-                           _s_w, _Cwr, _Cwt, _J_uw, _J_dw):
-    _H_uw = H_j[_J_uw]
-    _H_dw = H_j[_J_dw]
-    _z_inv_uw = _z_inv_j[_J_uw]
-    # Create indicator functions
-    _omega_w = np.zeros(_H_uw.size)
-    _omega_w[_H_uw >= _H_dw] = 1.0
-    # Create conditionals
-    cond_0 = (_omega_w * _H_uw + (1 - _omega_w) * _H_dw >
-                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
-    cond_1 = ((1 - _omega_w) * _H_uw + _omega_w * _H_dw >
-                _z_w + _z_inv_uw + (1 - u) * _y_max_w)
-    # TODO: Is this being recalculated for a reason?
-    # Effective heads
-    a = (cond_0 & cond_1)
-    b = (cond_0 & ~cond_1)
-    c = (~cond_0)
-    _Hw[a] = _H_uw[a] - _H_dw[a]
-    _Hw[b] = (_omega_w[b] * _H_uw[b] + (1 - _omega_w[b]) * _H_dw[b]
-                    + (-_z_inv_uw[b] - _z_w[b] - (1 - u[b]) * _y_max_w[b]))
-    _Hw[c] = 0.0
-    _Hw = np.abs(_Hw)
-    # Compute universal coefficient
-    _gamma_ww = (_Cwr * _L_w * _Hw + _Cwt * _s_w * _Hw**2)**2
-    # Compute flow
-    _Qw_next = (-1)**(1 - _omega_w) * np.sqrt(_gamma_ww * _Hw)
-    return _Qw_next
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            int64[:], int64[:]),
-      cache=True)
-def numba_pump_flow_coefficients(_alpha_p, _beta_p, _chi_p, H_j, _z_inv_j, _Qp, u,
-                                 _z_p, _dHp_max, _dHp_min, _a_p, _b_p, _c_p,
-                                 _J_up, _J_dp):
-    # Get upstream and downstream heads and invert elevation
-    _H_up = H_j[_J_up]
-    _H_dp = H_j[_J_dp]
-    _z_inv_up = _z_inv_j[_J_up]
-    # Compute effective head
-    _dHp = _H_dp - _H_up
-    # Condition 0: Upstream head is above inlet height
-    cond_0 = _H_up > _z_inv_up + _z_p
-    # Condition 1: Head difference is within range of pump curve
-    cond_1 = (_dHp > _dHp_min) & (_dHp < _dHp_max)
-    _dHp[_dHp > _dHp_max] = _dHp_max[_dHp > _dHp_max]
-    _dHp[_dHp < _dHp_min] = _dHp_min[_dHp < _dHp_min]
-    # Compute universal coefficients
-    _gamma_p = gamma_p(_Qp, _b_p, _c_p, u)
-    # Fill coefficient arrays
-    # Head in pump curve range
-    a = (cond_0 & cond_1)
-    _alpha_p[a] = _gamma_p[a]
-    _beta_p[a] = -_gamma_p[a]
-    _chi_p[a] = _gamma_p[a] * _a_p[a]
-    # Head outside of pump curve range
-    b = (cond_0 & ~cond_1)
-    _alpha_p[b] = 0.0
-    _beta_p[b] = 0.0
-    _chi_p[b] = _gamma_p[b] * (_a_p[b] - _dHp[b])
-    # Depth below inlet
-    c = (~cond_0)
-    _alpha_p[c] = 0.0
-    _beta_p[c] = 0.0
-    _chi_p[c] = 0.0
-    return 1
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-                 float64[:], float64[:], float64[:], int64[:], int64[:]),
-      cache=True)
-def numba_solve_pump_flows(H_j, u, _z_inv_j, _z_p, _dHp_max, _dHp_min, _a_p, _b_p, _c_p,
-                           _J_up, _J_dp):
-    _H_up = H_j[_J_up]
-    _H_dp = H_j[_J_dp]
-    _z_inv_up = _z_inv_j[_J_up]
-    # Create conditionals
-    _dHp = _H_dp - _H_up
-    _dHp[_dHp > _dHp_max] = _dHp_max[_dHp > _dHp_max]
-    _dHp[_dHp < _dHp_min] = _dHp_min[_dHp < _dHp_min]
-    cond_0 = _H_up > _z_inv_up + _z_p
-    # Compute universal coefficients
-    _Qp_next = (u / _b_p * (_a_p - _dHp))**(1 / _c_p)
-    _Qp_next[~cond_0] = 0.0
-    return _Qp_next
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], int64, int64[:], int64[:], int64[:]),
-      cache=True)
-def numba_forward_recurrence(_T_ik, _U_Ik, _V_Ik, _W_Ik, _a_ik, _b_ik, _c_ik,
-                             _P_ik, _A_ik, _E_Ik, _D_Ik, NK, nk, _I_1k, _i_1k):
-    g = 9.81
-    for k in range(NK):
-        # Start at junction 1
-        _I_1 = _I_1k[k]
-        _i_1 = _i_1k[k]
-        _I_2 = _I_1 + 1
-        _i_2 = _i_1 + 1
-        nlinks = nk[k]
-        _T_ik[_i_1] = T_1k(_a_ik[_i_1], _b_ik[_i_1], _c_ik[_i_1])
-        _U_Ik[_I_1] = U_1k(_E_Ik[_I_2], _c_ik[_i_1], _A_ik[_i_1], _T_ik[_i_1], g)
-        _V_Ik[_I_1] = V_1k(_P_ik[_i_1], _D_Ik[_I_2], _c_ik[_i_1], _T_ik[_i_1],
-                            _a_ik[_i_1], _D_Ik[_I_1])
-        _W_Ik[_I_1] = W_1k(_A_ik[_i_1], _T_ik[_i_1], _a_ik[_i_1], _E_Ik[_I_1], g)
-        # Loop from junction 2 -> Nk
-        for i in range(nlinks - 1):
-            _i_next = _i_2 + i
-            _I_next = _I_2 + i
-            _Im1_next = _I_next - 1
-            _Ip1_next = _I_next + 1
-            _T_ik[_i_next] = T_ik(_a_ik[_i_next], _b_ik[_i_next], _c_ik[_i_next],
-                                  _A_ik[_i_next], _E_Ik[_I_next], _U_Ik[_Im1_next], g)
-            _U_Ik[_I_next] = U_Ik(_E_Ik[_Ip1_next], _c_ik[_i_next],
-                                  _A_ik[_i_next], _T_ik[_i_next], g)
-            _V_Ik[_I_next] = V_Ik(_P_ik[_i_next], _a_ik[_i_next], _D_Ik[_I_next],
-                                  _D_Ik[_Ip1_next], _c_ik[_i_next], _A_ik[_i_next],
-                                  _E_Ik[_I_next], _V_Ik[_Im1_next], _U_Ik[_Im1_next],
-                                  _T_ik[_i_next], g)
-            _W_Ik[_I_next] = W_Ik(_A_ik[_i_next], _E_Ik[_I_next], _a_ik[_i_next],
-                                  _W_Ik[_Im1_next], _U_Ik[_Im1_next], _T_ik[_i_next], g)
-    return 1
-
-@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-            float64[:], float64[:], float64[:], float64[:], int64, int64[:], int64[:], int64[:]),
-      cache=True)
-def numba_backward_recurrence(_O_ik, _X_Ik, _Y_Ik, _Z_Ik, _a_ik, _b_ik, _c_ik,
-                              _P_ik, _A_ik, _E_Ik, _D_Ik, NK, nk, _I_Nk, _i_nk):
-    g = 9.81
-    for k in range(NK):
-        _I_N = _I_Nk[k]
-        _i_n = _i_nk[k]
-        _I_Nm1 = _I_N - 1
-        _i_nm1 = _i_n - 1
-        _I_Np1 = _I_N + 1
-        nlinks = nk[k]
-        _O_ik[_i_n] = O_nk(_a_ik[_i_n], _b_ik[_i_n], _c_ik[_i_n])
-        _X_Ik[_I_N] = X_Nk(_A_ik[_i_n], _E_Ik[_I_N], _a_ik[_i_n], _O_ik[_i_n], g)
-        _Y_Ik[_I_N] = Y_Nk(_P_ik[_i_n], _D_Ik[_I_N], _a_ik[_i_n], _O_ik[_i_n],
-                            _c_ik[_i_n], _D_Ik[_I_Np1])
-        _Z_Ik[_I_N] = Z_Nk(_A_ik[_i_n], _O_ik[_i_n], _c_ik[_i_n], _E_Ik[_I_Np1], g)
-        for i in range(nlinks - 1):
-            _i_next = _i_nm1 - i
-            _I_next = _I_Nm1 - i
-            _Ip1_next = _I_next + 1
-            _O_ik[_i_next] = O_ik(_a_ik[_i_next], _b_ik[_i_next], _c_ik[_i_next],
-                                  _A_ik[_i_next], _E_Ik[_Ip1_next], _X_Ik[_Ip1_next], g)
-            _X_Ik[_I_next] = X_Ik(_A_ik[_i_next], _E_Ik[_I_next], _a_ik[_i_next],
-                                  _O_ik[_i_next], g)
-            _Y_Ik[_I_next] = Y_Ik(_P_ik[_i_next], _a_ik[_i_next], _D_Ik[_I_next],
-                                  _D_Ik[_Ip1_next], _c_ik[_i_next], _A_ik[_i_next],
-                                  _E_Ik[_Ip1_next], _Y_Ik[_Ip1_next], _X_Ik[_Ip1_next],
-                                  _O_ik[_i_next], g)
-            _Z_Ik[_I_next] = Z_Ik(_A_ik[_i_next], _E_Ik[_Ip1_next], _c_ik[_i_next],
-                                  _Z_Ik[_Ip1_next], _X_Ik[_Ip1_next], _O_ik[_i_next], g)
-    return 1
-
-@njit(float64[:,:](float64[:,:], int64, int64),
-      cache=True)
-def numba_create_banded(l, bandwidth, M):
-    AB = np.zeros((2*bandwidth + 1, M))
-    for i in range(M):
-        AB[bandwidth, i] = l[i, i]
-    for n in range(bandwidth):
-        for j in range(M - n - 1):
-            AB[bandwidth - n - 1, -j - 1] = l[-j - 2 - n, -j - 1]
-            AB[bandwidth + n + 1, j] = l[j + n + 1, j]
-    return AB
-
-@njit(void(float64[:], int64[:], float64[:]),
-      cache=True,
-      fastmath=True)
-def numba_add_at(a, indices, b):
-    n = len(indices)
-    for k in range(n):
-        i = indices[k]
-        a[i] += b[k]
-
-@njit(void(float64[:, :], boolean[:], int64[:], int64[:], int64),
-      cache=True)
-def numba_clear_off_diagonals(A, bc, _J_uk, _J_dk, NK):
-    for k in range(NK):
-        _J_u = _J_uk[k]
-        _J_d = _J_dk[k]
-        _bc_u = bc[_J_u]
-        _bc_d = bc[_J_d]
-        if not _bc_u:
-            A[_J_u, _J_d] = 0.0
-        if not _bc_d:
-            A[_J_d, _J_u] = 0.0
-
-@njit(void(float64[:, :], float64[:], boolean[:], int64[:], int64[:], float64[:],
-           float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-           float64, int64, int64),
-      cache=True,
-      fastmath=True)
-def numba_create_A_matrix(A, _F_jj, bc, _J_uk, _J_dk, _alpha_uk,
-                          _alpha_dk, _beta_uk, _beta_dk, _xi_uk, _xi_dk,
-                          _A_sj, _dt, M, NK):
-    numba_add_at(_F_jj, _J_uk, _alpha_uk)
-    numba_add_at(_F_jj, _J_dk, -_beta_dk)
-    numba_add_at(_F_jj, _J_uk, _xi_uk)
-    numba_add_at(_F_jj, _J_dk, _xi_dk)
-    _F_jj += (_A_sj / _dt)
-    # Set diagonal of A matrix
-    for i in range(M):
-        if bc[i]:
-            A[i,i] = 1.0
-        else:
-            A[i,i] = _F_jj[i]
-    for k in range(NK):
-        _J_u = _J_uk[k]
-        _J_d = _J_dk[k]
-        _bc_u = bc[_J_u]
-        _bc_d = bc[_J_d]
-        if not _bc_u:
-            A[_J_u, _J_d] += _beta_uk[k]
-        if not _bc_d:
-            A[_J_d, _J_u] -= _alpha_dk[k]
-
-@njit(void(float64[:, :], float64[:], boolean[:], int64[:], int64[:], float64[:],
-           float64[:], float64[:], float64[:], int64, int64),
-      cache=True,
-      fastmath=True)
-def numba_create_OWP_matrix(X, diag, bc, _J_uc, _J_dc, _alpha_uc,
-                            _alpha_dc, _beta_uc, _beta_dc, M, NC):
-    # Set diagonal
-    numba_add_at(diag, _J_uc, _alpha_uc)
-    numba_add_at(diag, _J_dc, -_beta_dc)
-    for i in range(M):
-        if bc[i]:
-            X[i,i] = 0.0
-        else:
-            X[i,i] = diag[i]
-    # Set off-diagonal
-    for c in range(NC):
-        _J_u = _J_uc[c]
-        _J_d = _J_dc[c]
-        _bc_u = bc[_J_u]
-        _bc_d = bc[_J_d]
-        if not _bc_u:
-            X[_J_u, _J_d] += _beta_uc[c]
-        if not _bc_d:
-            X[_J_d, _J_u] -= _alpha_dc[c]
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64),
-      cache=True)
-def numba_Q_i_next_b(X_Ik, h_Ik, Y_Ik, Z_Ik, h_Np1k, _Ik, _ki, n):
-    _Q_i = np.zeros(n)
-    for i in range(n):
-        I = _Ik[i]
-        k = _ki[i]
-        t_0 = X_Ik[I] * h_Ik[I]
-        t_1 = Y_Ik[I]
-        t_2 = Z_Ik[I] * h_Np1k[k]
-        _Q_i[i] = t_0 + t_1 + t_2
-    return _Q_i
-
-@njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64),
-      cache=True)
-def numba_Q_im1k_next_f(U_Ik, h_Ik, V_Ik, W_Ik, h_1k, _Ik, _ki, n):
-    _Q_i = np.zeros(n)
-    for i in range(n):
-        I = _Ik[i]
-        Ip1 = I + 1
-        k = _ki[i]
-        t_0 = U_Ik[I] * h_Ik[Ip1]
-        t_1 = V_Ik[I]
-        t_2 = W_Ik[I] * h_1k[k]
-        _Q_i[i] = t_0 + t_1 + t_2
-    return _Q_i
-
-@njit(void(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],
-           float64[:], float64[:], float64[:], float64[:], int64[:], int64[:], int64[:],
-           int64[:], int64[:], int64, boolean[:]),
-      cache=True)
-def numba_reposition_junctions(_x_Ik, _z_inv_Ik, _h_Ik, _dx_ik, _Q_ik, _H_dk,
-                               _b0, _zc, _xc, _m, _elem_pos, _i_1k, _I_1k,
-                               _I_Np1k, nk, NK, reposition):
-    for k in range(NK):
-        if reposition[k]:
-            _i_1 = _i_1k[k]
-            _I_1 = _I_1k[k]
-            _I_Np1 = _I_Np1k[k]
-            nlinks = nk[k]
-            njunctions = nlinks + 1
-            _i_end = _i_1 + nlinks
-            _I_end = _I_1 + njunctions
-            _H_d = _H_dk[k]
-            _z_inv_1 = _z_inv_Ik[_I_1]
-            _z_inv_Np1 = _z_inv_Ik[_I_Np1]
-            pos_prev = _elem_pos[k]
-            # Junction arrays for superlink k
-            _x_I = _x_Ik[_I_1:_I_end]
-            _z_inv_I = _z_inv_Ik[_I_1:_I_end]
-            _h_I = _h_Ik[_I_1:_I_end]
-            _dx_i = _dx_ik[_i_1:_i_end]
-            # Move junction if downstream head is within range
-            move_junction = (_H_d > _z_inv_Np1) & (_H_d < _z_inv_1)
-            if move_junction:
-                z_m = _H_d
-                _x0 = _x_I[_I_1]
-                x_m = (_H_d - _b0[k]) / _m[k] + _x0
-            else:
-                z_m = _zc[k]
-                x_m = _xc[k]
-            # Determine new x-position of junction
-            c = np.searchsorted(_x_I, x_m)
-            cm1 = c - 1
-            # Compute fractional x-position along superlink k
-            frac = (x_m - _x_I[cm1]) / (_x_I[c] - _x_I[cm1])
-            # Interpolate depth at new position
-            h_m = (1 - frac) * _h_I[cm1] + (frac) * _h_I[c]
-            # Link length ratio
-            r = _dx_i[pos_prev - 1] / (_dx_i[pos_prev - 1]
-                                    + _dx_i[pos_prev])
-            # Set new values
-            _x_I[pos_prev] = x_m
-            _z_inv_I[pos_prev] = z_m
-            _h_I[pos_prev] = h_m
-            Ix = np.argsort(_x_I)
-            _dx_i = np.diff(_x_I[Ix])
-            _x_Ik[_I_1:_I_end] = _x_I[Ix]
-            _z_inv_Ik[_I_1:_I_end] = _z_inv_I[Ix]
-            _h_Ik[_I_1:_I_end] = _h_I[Ix]
-            _dx_ik[_i_1:_i_end] = _dx_i
-            # Set position to new position
-            pos_change = np.argsort(Ix)
-            pos_next = pos_change[pos_prev]
-            _elem_pos[k] = pos_next
-            shifted = (pos_prev != pos_next)
-            # If position has shifted interpolate flow
-            if shifted:
-                ix = np.arange(nlinks)
-                ix[pos_prev] = pos_next
-                ix.sort()
-                _Q_i = _Q_ik[_i_1:_i_end]
-                _Q_i[pos_prev - 1] = (1 - r) * _Q_i[pos_prev - 1] + r * _Q_i[pos_prev]
-                _Q_ik[_i_1:_i_end] = _Q_i[ix]
 
diff --git a/pipedream_solver/simulation.py b/pipedream_solver/simulation.py
index a0af8b1..30ec168 100644
--- a/pipedream_solver/simulation.py
+++ b/pipedream_solver/simulation.py
@@ -573,6 +573,46 @@ def step(self, dt=None, subdivisions=1, retries=0, tol=1, norm=2,
         assert np.isfinite(self.model.H_j).all()
         self._iter_count += 1
 
+    def _interpolate_inputs(self, dt=None, **kwargs):
+        # Specify current timestamps
+        t_next = self.t + dt
+        # Import inputs
+        interpolation_method = self.interpolation
+        if not 'Q_in' in kwargs:
+            Q_in = self.Q_in
+            # Get superjunction runoff input
+            if Q_in is not None:
+                Q_in_index, Q_in_values = Q_in.index.values, Q_in.values
+                Q_in_next = interpolate_sample(t_next, Q_in_index, Q_in_values,
+                                               interpolation_method)
+            else:
+                Q_in_next = None
+        else:
+            Q_in_next = kwargs.pop('Q_in')
+        if not 'H_bc' in kwargs:
+            H_bc = self.H_bc
+            # Get head boundary conditions
+            if H_bc is not None:
+                H_bc_index, H_bc_values = H_bc.index.values, H_bc.values
+                H_bc_next = interpolate_sample(t_next, H_bc_index, H_bc_values,
+                                               interpolation_method)
+            else:
+                H_bc_next = None
+        else:
+            H_bc_next = kwargs.pop('H_bc')
+        if not 'Q_Ik' in kwargs:
+            Q_Ik = self.Q_Ik
+            # Get junction runoff input
+            if Q_Ik is not None:
+                Q_Ik_index, Q_Ik_values = Q_Ik.index.values, Q_Ik.values
+                Q_Ik_next = interpolate_sample(t_next, Q_Ik_index, Q_Ik_values,
+                                               interpolation_method)
+            else:
+                Q_Ik_next = None
+        else:
+            Q_Ik_next = kwargs.pop('Q_Ik')
+        return Q_in_next, H_bc_next, Q_Ik_next
+
     def _step(self, dt=None, **kwargs):
         # Specify current timestamps
         t_next = self.t + dt
diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py
index 2abcd7a..f9ab035 100644
--- a/pipedream_solver/superlink.py
+++ b/pipedream_solver/superlink.py
@@ -8,6 +8,9 @@
 import pipedream_solver.geometry
 import pipedream_solver.storage
 import pipedream_solver.visualization
+from pipedream_solver.callbacks import BaseCallback
+from pipedream_solver.diagnostics import ConvergenceTracker, VolumeTracker
+from pipedream_solver.diagnostics import ExperimentalConvergenceTracker
 
 class SuperLink():
     """
@@ -293,6 +296,7 @@ def __init__(self, superlinks, superjunctions,
         orifices = copy.deepcopy(orifices)
         weirs = copy.deepcopy(weirs)
         pumps = copy.deepcopy(pumps)
+        self.callbacks = {}
         # TODO: Ensure index is int
         # TODO: This needs to be done for orifices/weirs/pumps as well
         # Ensure nominal direction of superlinks is correct
@@ -420,6 +424,7 @@ def __init__(self, superlinks, superjunctions,
         # Dimensions
         self.nk = np.bincount(self._ki)
         # Create forward and backward indexers
+        # TODO: Check if these are correct for single link
         self.forward_I_I = np.copy(self._I)
         self.forward_I_I[self._Ik] = self._Ip1k
         self.backward_I_I = np.copy(self._I)
@@ -689,11 +694,15 @@ def __init__(self, superlinks, superjunctions,
         self.bc = self.superjunctions['bc'].values.astype(np.bool_)
         if sparse:
             self.B = scipy.sparse.lil_matrix((self.M, self.n_o))
+            self.J = scipy.sparse.lil_matrix((self.M, self.M))
+            self.K = scipy.sparse.lil_matrix((self.M, self.M))
             self.O = scipy.sparse.lil_matrix((self.M, self.M))
             self.W = scipy.sparse.lil_matrix((self.M, self.M))
             self.P = scipy.sparse.lil_matrix((self.M, self.M))
         else:
             self.B = np.zeros((self.M, self.n_o))
+            self.J = np.zeros((self.M, self.M))
+            self.K = np.zeros((self.M, self.M))
             self.O = np.zeros((self.M, self.M))
             self.W = np.zeros((self.M, self.M))
             self.P = np.zeros((self.M, self.M))
@@ -715,6 +724,8 @@ def __init__(self, superlinks, superjunctions,
         self._Q_dk = self._Q_ik[self._i_nk]
         self._h_uk = self._h_Ik[self._I_1k]
         self._h_dk = self._h_Ik[self._I_Np1k]
+        # Set boundary condition in/outflows
+        self._Q_bc = np.zeros(self.M, dtype=np.float64)
         # Other parameters
         self._O_diag = np.zeros(self.M)
         self._W_diag = np.zeros(self.M)
@@ -728,6 +739,9 @@ def __init__(self, superlinks, superjunctions,
         self._Pe_dk = np.copy(self._Pe_ik[self._i_nk])
         self._R_uk = np.copy(self._R_ik[self._i_1k])
         self._R_dk = np.copy(self._R_ik[self._i_nk])
+        # End velocities
+        self._u_uk = np.copy(self._u_ik[self._i_1k])
+        self._u_dk = np.copy(self._u_ik[self._i_nk])
         # End slopes
         cond_uk = (self._dx_uk > 0.)
         cond_dk = (self._dx_uk > 0.)
@@ -759,6 +773,22 @@ def __init__(self, superlinks, superjunctions,
         self._A_k = np.zeros(self.NK, dtype=np.float64)
         self._dt_ck = np.ones(self.NK, dtype=np.float64)
         self._Q_in = np.zeros(self.M, dtype=np.float64)
+        self._H_bc = np.zeros(self.M, dtype=np.float64)
+        # Set minimum hydraulic geometries
+        self._h_Ik_min = np.full(self._h_Ik.size, self.min_depth, dtype=np.float64)
+        self._A_ik_min = np.copy(self._A_ik)
+        self._B_ik_min = np.copy(self._B_ik)
+        self._Pe_ik_min = np.copy(self._Pe_ik)
+        self._R_ik_min = np.copy(self._R_ik)
+        self._A_uk_min = np.copy(self._A_uk)
+        self._B_uk_min = np.copy(self._B_uk)
+        self._Pe_uk_min = np.copy(self._Pe_uk)
+        self._R_uk_min = np.copy(self._R_uk)
+        self._A_dk_min = np.copy(self._A_dk)
+        self._B_dk_min = np.copy(self._B_dk)
+        self._Pe_dk_min = np.copy(self._Pe_dk)
+        self._R_dk_min = np.copy(self._R_dk)
+        self.min_hydraulic_geometry()
         # Initialize state dictionary
         self.states = {}
         # Iteration counter
@@ -767,6 +797,10 @@ def __init__(self, superlinks, superjunctions,
         # Compute bandwidth
         self._compute_bandwidth()
         # Initialize to stable state
+        self.convergence_tracker = ExperimentalConvergenceTracker(self)
+        self.bind_callback(self.convergence_tracker, 'convergence_tracker')
+        self.volume_tracker = VolumeTracker(self)
+        self.bind_callback(self.volume_tracker, 'volume_tracker')
         self.step(dt=1e-6, first_time=True)
         # Reset iteration counter
         self.iter_count = 0
@@ -924,6 +958,24 @@ def x_Ik(self):
     def x_Ik(self, value):
         self._x_Ik = np.asarray(value)
 
+    @property
+    def state_vector_prior(self):
+        vec = np.concatenate([self.states['H_j'], self.states['Q_uk'],
+                              self.states['Q_dk'], self.states['Q_o'],
+                              self.states['Q_w'], self.states['Q_p'], 
+                              self.states['h_Ik'], self.states['Q_ik']])
+        return vec
+
+    @property
+    def state_vector_next(self):
+        vec = np.concatenate([self.H_j, self.Q_uk, self.Q_dk, self.Q_o, self.Q_w, self.Q_p, 
+                              self.h_Ik, self.Q_ik])
+        return vec
+
+    @property
+    def state_vector(self):
+        return self.state_vector_next
+
     @property
     def adjacency_matrix(self, J_u=None, J_d=None, symmetric=True):
         M = self.M
@@ -1931,6 +1983,28 @@ def orifice_hydraulic_geometry(self, u=None):
         # Export to instance variables
         self._Ao = _Ao
 
+    def compute_boundary_indicator_variables(self):
+        H_j = self.H_j
+        _J_uk = self._J_uk
+        _J_dk = self._J_dk
+        _theta_uk = self._theta_uk
+        _theta_dk = self._theta_dk
+        _z_inv_uk = self._z_inv_uk
+        _z_inv_dk = self._z_inv_dk
+        # Compute upstream indicator variable
+        _H_juk = H_j[_J_uk]
+        upstream_depth_above_invert = _H_juk > _z_inv_uk
+        _theta_uk.fill(0.)
+        _theta_uk[upstream_depth_above_invert] = 1.
+        # Compute downstream indicator variable
+        _H_jdk = H_j[_J_dk]
+        downstream_depth_above_invert = _H_jdk > _z_inv_dk
+        _theta_dk.fill(0.)
+        _theta_dk[downstream_depth_above_invert] = 1.
+        # Store output
+        self._theta_uk = _theta_uk
+        self._theta_dk = _theta_dk
+
     def compute_storage_areas(self):
         """
         Compute surface area of superjunctions at current time step.
@@ -3141,8 +3215,6 @@ def solve_superlink_depths(self):
         superjunction heads at time t + dt.
         """
         # Import instance variables
-        _J_uk = self._J_uk              # Index of superjunction upstream of superlink k
-        _J_dk = self._J_dk              # Index of superjunction downstream of superlink k
         _kappa_uk = self._kappa_uk      # Superlink head coefficient
         _kappa_dk = self._kappa_dk      # Superlink head coefficient
         _lambda_uk = self._lambda_uk    # Superlink head coefficient
@@ -3151,11 +3223,10 @@ def solve_superlink_depths(self):
         _mu_dk = self._mu_dk            # Superlink head coefficient
         _Q_uk = self._Q_uk              # Flow rate at upstream end of superlink k
         _Q_dk = self._Q_dk              # Flow rate at downstream end of superlink k
-        H_j = self.H_j                  # Head at superjunction j
         min_depth = self.min_depth      # Minimum allowable depth at boundaries
         # Compute flow at next time step
-        _h_uk_next = _kappa_uk * _Q_uk + _lambda_uk * H_j[_J_uk] + _mu_uk
-        _h_dk_next = _kappa_dk * _Q_dk + _lambda_dk * H_j[_J_dk] + _mu_dk
+        _h_uk_next = _kappa_uk * _Q_uk + _lambda_uk * _Q_dk + _mu_uk
+        _h_dk_next = _kappa_dk * _Q_uk + _lambda_dk * _Q_dk + _mu_dk
         # Set minimum values
         # TODO: Is this causing the difference between normal/numba versions?
         # _h_uk_next[_h_uk_next < min_depth] = min_depth
@@ -3533,6 +3604,46 @@ def solve_internals_lsq(self):
         self._Q_ik = _Q_ik
         self._h_Ik = _h_Ik
 
+    def compute_boundary_flows(self, dt=None):
+        # Import instance variables
+        bc = self.bc
+        if not bc.any():
+            return None
+        _Q_bc = self._Q_bc
+        _H_j_next = self.H_j
+        _H_j_prev = self.states['H_j']
+        _Q_in = self._Q_in
+        _A_sj = self._A_sj
+        _J_uk = self._J_uk
+        _J_dk = self._J_dk
+        _J_uo = self._J_uo
+        _J_do = self._J_do
+        _J_uw = self._J_uw
+        _J_dw = self._J_dw
+        _J_up = self._J_up
+        _J_dp = self._J_dp
+        _Q_uk = self._Q_uk
+        _Q_dk = self._Q_dk
+        _Q_o = self._Qo
+        _Q_w = self._Qw
+        _Q_p = self._Qp
+        if dt is None:
+            dt = self._dt
+        # Compute boundary flows
+        _Q_bc.fill(0.)
+        _Q_bc -= _Q_in
+        _Q_bc += (_H_j_next - _H_j_prev) * _A_sj / dt
+        _Q_bc[self._J_uk] += _Q_uk
+        _Q_bc[self._J_dk] -= _Q_dk
+        _Q_bc[self._J_uo] += _Q_o
+        _Q_bc[self._J_do] -= _Q_o
+        _Q_bc[self._J_uw] += _Q_w
+        _Q_bc[self._J_dw] -= _Q_w
+        _Q_bc[self._J_up] += _Q_p
+        _Q_bc[self._J_dp] -= _Q_p
+        _Q_bc[~bc] = 0.
+        self._Q_bc = _Q_bc
+
     def exit_conditions(self):
         """
         Determine which superlinks have exit depths below the pipe crown elevation.
@@ -3890,6 +4001,7 @@ def state_space_system(self, _dt=None):
         n_w = self.n_w                # Number of weirs
         n_p = self.n_p                # Number of pumps
         Q_in = self._Q_in
+        H_bc = self._H_bc
         # If no time step specified, use instance time step
         if _dt is None:
             _dt = self._dt
@@ -3900,8 +4012,26 @@ def state_space_system(self, _dt=None):
         else:
             A_1 = A
         # Get A_2
-        A_2 = np.diag(np.where(~bc, _A_sj / _dt, 0))
-        return A_1, A_2, D, H_j_next, H_j_prev, Q_in
+        A_2_diag = _A_sj / _dt
+        np.add.at(A_2_diag, self._J_uk, self._B_uk * self._dx_uk * self._theta_uk / 2 / _dt)
+        np.add.at(A_2_diag, self._J_dk, self._B_dk * self._dx_dk * self._theta_dk / 2 / _dt) 
+        A_2 = np.diag(np.where(~bc, A_2_diag, 0.))
+        D = np.where(~bc, D, 0.)
+        B = np.diag(np.where(~bc, 1., 0.))
+        H_bc = np.where(bc, H_bc, 0.)
+        return A_1, A_2, B, D, H_j_next, H_j_prev, Q_in, H_bc
+
+    def return_state(self):
+        states = {}
+        states['H_j'] = np.copy(self.H_j)
+        states['Q_uk'] = np.copy(self.Q_uk)
+        states['Q_dk'] = np.copy(self.Q_dk)
+        states['Q_o'] = np.copy(self.Q_o)
+        states['Q_w'] = np.copy(self.Q_w)
+        states['Q_p'] = np.copy(self.Q_p)
+        states['h_Ik'] = np.copy(self.h_Ik)
+        states['Q_ik'] = np.copy(self.Q_ik)
+        return states
 
     def save_state(self):
         """
@@ -3914,18 +4044,17 @@ def save_state(self):
         self.states['Q_uk'] = np.copy(self.Q_uk)
         self.states['Q_dk'] = np.copy(self.Q_dk)
         self.states['x_Ik'] = np.copy(self.x_Ik)
-        if self.n_o:
-            self.states['Q_o'] = np.copy(self.Q_o)
-            # TODO: Need to add orifice area here
-        if self.n_w:
-            self.states['Q_w'] = np.copy(self.Q_w)
-        if self.n_p:
-            self.states['Q_p'] = np.copy(self.Q_p)
+        self.states['Q_o'] = np.copy(self.Q_o)
+        # TODO: Need to add orifice area here
+        self.states['Q_w'] = np.copy(self.Q_w)
+        self.states['Q_p'] = np.copy(self.Q_p)
         self.states['A_ik'] = np.copy(self.A_ik)
         self.states['A_uk'] = np.copy(self.A_uk)
         self.states['A_dk'] = np.copy(self.A_dk)
         self.states['A_sj'] = np.copy(self.A_sj)
         self.states['V_j'] = np.copy(self.V_j)
+        for _, callback in self.callbacks.items():
+            callback.__on_save_state__()
 
     def load_state(self, states={}, exclude_states=set(),
                    compute_hydraulic_geometries=True):
@@ -3951,6 +4080,8 @@ def load_state(self, states={}, exclude_states=set(),
             self.downstream_hydraulic_geometry()
             self.compute_storage_areas()
             self.node_velocities()
+        for _, callback in self.callbacks.items():
+            callback.__on_load_state__()
 
     def spinup(self, n_steps=100, dt=10, Q_in=None, Q_0Ik=None, reposition_junctions=False,
                reset_counters=True, **kwargs):
@@ -4012,20 +4143,37 @@ def plot_network_3d(self, ax=None, superjunction_signal=None, junction_signal=No
                                  weir_kwargs=weir_kwargs,
                                  pump_kwargs=pump_kwargs))
 
+    def bind_callback(self, callback, key):
+        assert isinstance(callback, BaseCallback)
+        self.callbacks[key] = callback
+
+    def unbind_callback(self, key):
+        return self.callbacks.pop(key)
+
     def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, dt=None,
              first_time=False, implicit=True, banded=False, first_iter=True):
         if first_iter:
             self.save_state()
         if dt is None:
             dt = self._dt
-        self._Q_in = Q_in
-        self._Q_0Ik = Q_0Ik
+        else:
+            self._dt = dt
+        if Q_in is None:
+            self._Q_in = np.zeros(self.M, dtype=np.float64)
+        else:
+            self._Q_in = Q_in
+        if Q_0Ik is None:
+            self._Q_0Ik = np.zeros(self._I.size, dtype=np.float64)
+        else:
+            self._Q_0Ik = Q_0Ik
+        self._H_bc = H_bc
         if not implicit:
             raise NotImplementedError
         # Compute all hydraulic geometries
         self.link_hydraulic_geometry()
         self.upstream_hydraulic_geometry()
         self.downstream_hydraulic_geometry()
+        self.compute_boundary_indicator_variables()
         self.compute_storage_areas()
         self.compute_storage_volumes()
         self.node_velocities()
@@ -4037,9 +4185,11 @@ def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=
         self.node_coeffs(_Q_0Ik=Q_0Ik, _dt=dt, first_iter=first_iter)
         self.forward_recurrence()
         self.backward_recurrence()
-        self.superlink_upstream_head_coefficients(_dt=dt)
-        self.superlink_downstream_head_coefficients(_dt=dt)
-        self.superlink_flow_coefficients()
+        #self.superlink_upstream_head_coefficients(_dt=dt)
+        #self.superlink_downstream_head_coefficients(_dt=dt)
+        #self.superlink_flow_coefficients()
+        self.superlink_boundary_depth_coefficients(_dt=dt)
+        self.superlink_boundary_flow_coefficients(_dt=dt)
         if self.orifices is not None:
             self.orifice_flow_coefficients(u=u_o)
         if self.weirs is not None:
@@ -4079,12 +4229,13 @@ def _solve_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=
             self.solve_internals_nnls()
         elif _method == 'lsq':
             self.solve_internals_lsq()
+        self.compute_boundary_flows(dt=dt)
         self.iter_count += 1
         self.t += dt
 
     def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, dt=None,
              first_time=False, implicit=True, banded=None, first_iter=True,
-             num_iter=1, head_tol=0.0015):
+             num_iter=1, rtol=1e-5, atol=1e-5, head_tol=0.0015):
         """
         Advance model forward to next time step, computing hydraulic states.
 
@@ -4117,34 +4268,29 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d
         implicit : bool
             (Deprecated)
         """
+        for _, callback in self.callbacks.items():
+            callback.__on_step_start__(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                                       first_time=first_time, implicit=implicit, banded=banded,
+                                       first_iter=first_iter, num_iter=num_iter, rtol=rtol, atol=atol,
+                                       head_tol=head_tol)
         if banded is None:
             banded = self.banded
-        self._setup_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
-                         first_time=first_time, implicit=implicit, banded=banded,
-                         first_iter=first_iter)
-        self._solve_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
-                         first_time=first_time, implicit=implicit, banded=banded,
-                         first_iter=first_iter)
-        # Perform fixed-point iteration until convergence
+        try:
+            self._setup_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                            first_time=first_time, implicit=implicit, banded=banded,
+                            first_iter=first_iter)
+            self._solve_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+                            first_time=first_time, implicit=implicit, banded=banded,
+                            first_iter=first_iter)
+        except:
+            #self.load_state()
+            raise
+
+        first_iter = False
         num_iter -= 1
-        iter_elapsed = 1
-        if (num_iter > 0):
-            H_j_prev = self.states['H_j']
-            H_j_next = np.copy(self.H_j)
-            residual = np.abs(H_j_next - H_j_prev)
-            if not (residual < head_tol).all():
-                for _ in range(num_iter):
-                    self.iter_count -= 1
-                    self.t -= dt
-                    self._setup_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
-                                     first_time=first_time, implicit=implicit, banded=banded,
-                                     first_iter=False)
-                    self._solve_step(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
+        self.iter_elapsed = 1
+        for _, callback in self.callbacks.items():
+            callback.__on_step_end__(H_bc=H_bc, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o, u_w=u_w, u_p=u_p, dt=dt,
                                      first_time=first_time, implicit=implicit, banded=banded,
-                                     first_iter=False)
-                    iter_elapsed += 1
-                    residual = np.abs(H_j_next - self.H_j)
-                    if (residual < head_tol).all():
-                        break
-                    H_j_next = np.copy(self.H_j)
-        self.iter_elapsed = iter_elapsed
+                                     first_iter=first_iter, num_iter=num_iter, rtol=rtol, atol=atol,
+                                     head_tol=head_tol)
diff --git a/pipedream_solver/tsuperlink.py b/pipedream_solver/tsuperlink.py
new file mode 100644
index 0000000..a64537f
--- /dev/null
+++ b/pipedream_solver/tsuperlink.py
@@ -0,0 +1,296 @@
+import numpy as np
+
+from pipedream_solver.nsuperlink import nSuperLink
+from pipedream_solver._tsuperlink import tnumba_a_ik, tnumba_b_ik, tnumba_c_ik, tnumba_P_ik, tnumba_forward_recurrence, tnumba_backward_recurrence, tnumba_solve_internals
+
+class tSuperLink(nSuperLink):
+    def __init__(self, superlinks, superjunctions,
+                 links=None, junctions=None,
+                 transects={}, storages={},
+                 orifices=None, weirs=None, pumps=None,
+                 dt=60, sparse=False, min_depth=1e-5, method='b',
+                 inertial_damping=False, bc_method='z',
+                 exit_hydraulics=False, auto_permute=False,
+                 end_length=None, end_method='b', internal_links=4, mobile_elements=False):
+        super().__init__(superlinks, superjunctions,
+                         links, junctions, transects, storages,
+                         orifices, weirs, pumps, dt, sparse,
+                         min_depth, method, inertial_damping,
+                         bc_method, exit_hydraulics, auto_permute,
+                         end_length, end_method, internal_links, mobile_elements)
+
+    def link_coeffs(self, _dt=None, first_iter=True):
+        """
+        Compute link momentum coefficients: a_ik, b_ik, c_ik and P_ik.
+        """
+        # Import instance variables
+        _u_ik = self._u_ik
+        _u_uk = self._u_uk
+        _u_dk = self._u_dk
+        _dx_ik = self._dx_ik       # Length of link ik
+        _Sf_method_ik = self._Sf_method_ik
+        _n_ik = self._n_ik         # Manning's roughness of link ik
+        _Q_ik_prev = np.copy(self.states['Q_ik'])
+        _h_Ik_prev = np.copy(self.states['h_Ik'][self._Ik])
+        _h_Ip1k_prev = np.copy(self.states['h_Ik'][self._Ip1k])
+        _Q_ik_next = self._Q_ik         # Flow rate at link ik
+        _A_ik = self._A_ik         # Flow area at link ik
+        _B_ik = self._B_ik
+        _R_ik = self._R_ik         # Hydraulic radius at link ik
+        _S_o_ik = self._S_o_ik     # Channel bottom slope at link ik
+        _C_ik = self._C_ik         # Discharge coefficient of control structure at link ik
+        # Upstream parameters
+        _n_uk = self._n_uk
+        _Q_uk_next = self._Q_uk
+        _Q_uk_prev = self.states['Q_uk']
+        _A_uk = self._A_uk
+        _B_uk = self._B_uk
+        _R_uk = self._R_uk
+        _dx_uk = self._dx_uk
+        _Sf_method_uk = self._Sf_method_uk
+        _C_uk = self._C_uk
+        _S_o_uk = self._S_o_uk
+        _theta_uk = self._theta_uk
+        _z_inv_uk = self._z_inv_uk
+        _J_uk = self._J_uk
+        _I_1k = self._I_1k
+        _h_juk_prev = _theta_uk * (self.states['H_j'][_J_uk] - _z_inv_uk)
+        _h_1k_prev = np.copy(self.states['h_Ik'][_I_1k])
+        # Downstream parameters
+        _n_dk = self._n_dk
+        _Q_dk_next = self._Q_dk
+        _Q_dk_prev = self.states['Q_dk']
+        _A_dk = self._A_dk
+        _B_dk = self._B_dk
+        _R_dk = self._R_dk
+        _dx_dk = self._dx_dk
+        _Sf_method_dk = self._Sf_method_dk
+        _C_dk = self._C_dk
+        _S_o_dk = self._S_o_dk
+        _theta_dk = self._theta_dk
+        _z_inv_dk = self._z_inv_dk
+        _J_dk = self._J_dk
+        _I_Np1k = self._I_Np1k
+        _h_jdk_prev = _theta_dk * (self.states['H_j'][_J_dk] - _z_inv_dk)
+        _h_Np1k_prev = np.copy(self.states['h_Ik'][_I_Np1k])
+        # If time step not specified, use instance time
+        if _dt is None:
+            _dt = self._dt
+        g = 9.81
+        # Compute link coefficients
+        _a_ik = tnumba_a_ik(_u_ik, _B_ik, _A_ik, _dx_ik, _dt, g)
+        _c_ik = tnumba_c_ik(_u_ik, _B_ik, _A_ik, _dx_ik, _dt, g)
+        _b_ik = tnumba_b_ik(_dx_ik, _dt, _n_ik, _Q_ik_next, _A_ik, _R_ik, _C_ik, _Sf_method_ik, g)
+        _P_ik = tnumba_P_ik(_Q_ik_prev, _dx_ik, _dt, _u_ik, _B_ik, _A_ik, _h_Ik_prev, _h_Ip1k_prev, _S_o_ik, g)
+        # Compute momentum coefficients for upstream boundary
+        _a_uk = _theta_uk * tnumba_a_ik(_u_uk, _B_uk, _A_uk, _dx_uk, _dt, g)
+        _c_uk = tnumba_c_ik(_u_uk, _B_uk, _A_uk, _dx_uk, _dt, g)
+        _b_uk = tnumba_b_ik(_dx_uk, _dt, _n_uk, _Q_uk_next, _A_uk, _R_uk, _C_uk, _Sf_method_uk, g)
+        _P_uk = tnumba_P_ik(_Q_uk_prev, _dx_uk, _dt, _u_uk, _B_uk, _A_uk, _h_juk_prev, _h_1k_prev, _S_o_uk, g)
+        _P_uk += _theta_uk * _z_inv_uk * _a_uk 
+        # Compute momentum coefficients for downstream boundary
+        _a_dk = tnumba_a_ik(_u_dk, _B_dk, _A_dk, _dx_dk, _dt, g)
+        _c_dk = _theta_dk * tnumba_c_ik(_u_dk, _B_dk, _A_dk, _dx_dk, _dt, g)
+        _b_dk = tnumba_b_ik(_dx_dk, _dt, _n_dk, _Q_dk_next, _A_dk, _R_dk, _C_dk, _Sf_method_dk, g)
+        _P_dk = tnumba_P_ik(_Q_dk_prev, _dx_dk, _dt, _u_dk, _B_dk, _A_dk, _h_Np1k_prev, _h_jdk_prev, _S_o_dk, g)
+        _P_dk += _theta_dk * _z_inv_dk * _c_dk 
+        # Export to instance variables
+        self._a_ik = _a_ik
+        self._b_ik = _b_ik
+        self._c_ik = _c_ik
+        self._P_ik = _P_ik
+        self._a_uk = _a_uk
+        self._b_uk = _b_uk
+        self._c_uk = _c_uk
+        self._P_uk = _P_uk
+        self._a_dk = _a_dk
+        self._b_dk = _b_dk
+        self._c_dk = _c_dk
+        self._P_dk = _P_dk
+
+    def forward_recurrence(self):
+        """
+        Compute forward recurrence coefficients: T_ik, U_Ik, V_Ik, and W_Ik.
+        """
+        # Import instance variables
+        _I_1k = self._I_1k                # Index of first junction in each superlink
+        _i_1k = self._i_1k                # Index of first link in each superlink
+        _E_Ik = self._E_Ik                # Continuity coefficient E_Ik
+        _D_Ik = self._D_Ik                # Continuity coefficient D_Ik
+        _a_ik = self._a_ik                # Momentum coefficient a_ik
+        _b_ik = self._b_ik                # Momentum coefficient b_ik
+        _c_ik = self._c_ik                # Momentum coefficient c_ik
+        _P_ik = self._P_ik                # Momentum coefficient P_ik
+        _T_ik = self._T_ik                # Recurrence coefficient T_ik
+        _U_Ik = self._U_Ik                # Recurrence coefficient U_Ik
+        _V_Ik = self._V_Ik                # Recurrence coefficient V_Ik
+        _W_Ik = self._W_Ik                # Recurrence coefficient W_Ik
+        NK = self.NK
+        nk = self.nk
+        tnumba_forward_recurrence(_T_ik, _U_Ik, _V_Ik, _W_Ik, _a_ik, _b_ik, _c_ik,
+                                 _P_ik, _E_Ik, _D_Ik, NK, nk, _I_1k, _i_1k)
+        # Export instance variables
+        self._T_ik = _T_ik
+        self._U_Ik = _U_Ik
+        self._V_Ik = _V_Ik
+        self._W_Ik = _W_Ik
+
+    def backward_recurrence(self):
+        """
+        Compute backward recurrence coefficients: O_ik, X_Ik, Y_Ik, and Z_Ik.
+        """
+        _I_Nk = self._I_Nk                # Index of penultimate junction in each superlink
+        _i_nk = self._i_nk                # Index of last link in each superlink
+        _E_Ik = self._E_Ik                # Continuity coefficient E_Ik
+        _D_Ik = self._D_Ik                # Continuity coefficient D_Ik
+        _a_ik = self._a_ik                # Momentum coefficient a_ik
+        _b_ik = self._b_ik                # Momentum coefficient b_ik
+        _c_ik = self._c_ik                # Momentum coefficient c_ik
+        _P_ik = self._P_ik                # Momentum coefficient P_ik
+        _O_ik = self._O_ik                # Recurrence coefficient O_ik
+        _X_Ik = self._X_Ik                # Recurrence coefficient X_Ik
+        _Y_Ik = self._Y_Ik                # Recurrence coefficient Y_Ik
+        _Z_Ik = self._Z_Ik                # Recurrence coefficient Z_Ik
+        NK = self.NK
+        nk = self.nk
+        tnumba_backward_recurrence(_O_ik, _X_Ik, _Y_Ik, _Z_Ik, _a_ik, _b_ik, _c_ik,
+                                    _P_ik, _E_Ik, _D_Ik, NK, nk, _I_Nk, _i_nk)
+        # Export instance variables
+        self._O_ik = _O_ik
+        self._X_Ik = _X_Ik
+        self._Y_Ik = _Y_Ik
+        self._Z_Ik = _Z_Ik
+
+    def superlink_boundary_flow_coefficients(self, _dt=None):
+        _U_Ik = self._U_Ik                # Recurrence coefficient U_Ik
+        _V_Ik = self._V_Ik                # Recurrence coefficient V_Ik
+        _W_Ik = self._W_Ik                # Recurrence coefficient W_Ik
+        _X_Ik = self._X_Ik                # Recurrence coefficient X_Ik
+        _Y_Ik = self._Y_Ik                # Recurrence coefficient Y_Ik
+        _Z_Ik = self._Z_Ik                # Recurrence coefficient Z_Ik
+        _E_Ik = self._E_Ik
+        _D_Ik = self._D_Ik
+        _I_1k = self._I_1k
+        _I_Nk = self._I_Nk
+        _I_Np1k = self._I_Np1k
+        a_uk = self._a_uk
+        a_dk = self._a_dk
+        b_uk = self._b_uk
+        b_dk = self._b_dk
+        c_uk = self._c_uk
+        c_dk = self._c_dk
+        P_uk = self._P_uk
+        P_dk = self._P_dk
+        # Get boundary coefficients
+        U_Nk = _U_Ik[_I_Nk]
+        Z_1k = _Z_Ik[_I_1k]
+        W_Nk = _W_Ik[_I_Nk]
+        X_1k = _X_Ik[_I_1k]
+        E_1k = _E_Ik[_I_1k]
+        E_Np1k = _E_Ik[_I_Np1k]
+        D_1k = _D_Ik[_I_1k]
+        D_Np1k = _D_Ik[_I_Np1k]
+        Y_1k = _Y_Ik[_I_1k]
+        V_Nk = _V_Ik[_I_Nk]
+        # Formulate expressions
+        a = U_Nk - E_Np1k
+        b = Z_1k
+        c = W_Nk
+        d = X_1k + E_1k
+        e = Y_1k - D_1k
+        f = D_Np1k + V_Nk
+        p = a*d - b*c
+        q = -p
+        r = (a_dk * d - b_dk * q)
+        s = (b_uk * q - c_uk * a)
+        # Create inverse matrix
+        aa = -a_uk * q * r
+        bb = -(b * c_uk * c_dk * q)
+        cc = -(c * a_uk * a_dk * q)
+        dd = c_dk * q * s
+        ee = P_uk * (b_dk * q - a_dk * d) + c_uk * (b_dk * (b * f - a * e) - a_dk * e - P_dk * b)
+        ff = P_dk * (b_uk * q - c_uk * a) + a_dk * (b_uk * (c * e -d * f) - c_uk * f - P_uk * c)
+        denom =  (c * b * a_dk * c_uk) + (r * s)
+        exog_denom = b_dk * b_uk * q - a_dk * b_uk * d - c_uk * b_dk * a - c_uk * a_dk
+        # Compute coefficients
+        alpha_uk = aa / denom
+        beta_uk = bb / denom
+        chi_uk = ee / exog_denom
+        alpha_dk = cc / denom
+        beta_dk = dd / denom
+        chi_dk = ff / exog_denom
+        # Store coefficients
+        self._alpha_uk = alpha_uk
+        self._beta_uk = beta_uk
+        self._chi_uk = chi_uk
+        self._alpha_dk = alpha_dk
+        self._beta_dk = beta_dk
+        self._chi_dk = chi_dk
+
+    def solve_internals_backwards(self, subcritical_only=False):
+        """
+        Solve for internal states of each superlink in the backward direction.
+        """
+        # Import instance variables
+        _I_1k = self._I_1k                  # Index of first junction in superlink k
+        _I_Nk = self._I_Nk
+        _i_1k = self._i_1k                  # Index of first link in superlink k
+        nk = self.nk
+        NK = self.NK
+        _h_Ik = self._h_Ik                  # Depth at junction Ik
+        _Q_ik = self._Q_ik                  # Flow rate at link ik
+        _D_Ik = self._D_Ik                  # Continuity coefficient
+        _E_Ik = self._E_Ik                  # Continuity coefficient
+        _U_Ik = self._U_Ik                  # Forward recurrence coefficient
+        _V_Ik = self._V_Ik                  # Forward recurrence coefficient
+        _W_Ik = self._W_Ik                  # Forward recurrence coefficient
+        _X_Ik = self._X_Ik                  # Backward recurrence coefficient
+        _Y_Ik = self._Y_Ik                  # Backward recurrence coefficient
+        _Z_Ik = self._Z_Ik                  # Backward recurrence coefficient
+        _Q_uk = self._Q_uk                  # Flow rate at upstream end of superlink k
+        _Q_dk = self._Q_dk                  # Flow rate at downstream end of superlink k
+        _h_uk = self._h_uk                  # Depth at upstream end of superlink k
+        _h_dk = self._h_dk                  # Depth at downstream end of superlink k
+        min_depth = self.min_depth          # Minimum allowable water depth
+        max_depth_k = self.max_depth_k
+        # Solve internals
+        tnumba_solve_internals(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik,
+                              _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, _I_Nk, nk, NK)
+        # TODO: Temporary
+        assert np.isfinite(_h_Ik).all()
+        # Export instance variables
+        self._h_Ik = _h_Ik
+        self._Q_ik = _Q_ik
+
+    def solve_internals_forwards(self, subcritical_only=False):
+        """
+        Solve for internal states of each superlink in the backward direction.
+        """
+        # Import instance variables
+        _I_1k = self._I_1k                  # Index of first junction in superlink k
+        _i_1k = self._i_1k                  # Index of first link in superlink k
+        _I_Nk = self._I_Nk
+        nk = self.nk
+        NK = self.NK
+        _h_Ik = self._h_Ik                  # Depth at junction Ik
+        _Q_ik = self._Q_ik                  # Flow rate at link ik
+        _D_Ik = self._D_Ik                  # Continuity coefficient
+        _E_Ik = self._E_Ik                  # Continuity coefficient
+        _U_Ik = self._U_Ik                  # Forward recurrence coefficient
+        _V_Ik = self._V_Ik                  # Forward recurrence coefficient
+        _W_Ik = self._W_Ik                  # Forward recurrence coefficient
+        _X_Ik = self._X_Ik                  # Backward recurrence coefficient
+        _Y_Ik = self._Y_Ik                  # Backward recurrence coefficient
+        _Z_Ik = self._Z_Ik                  # Backward recurrence coefficient
+        _Q_uk = self._Q_uk                  # Flow rate at upstream end of superlink k
+        _Q_dk = self._Q_dk                  # Flow rate at downstream end of superlink k
+        _h_uk = self._h_uk                  # Depth at upstream end of superlink k
+        _h_dk = self._h_dk                  # Depth at downstream end of superlink k
+        min_depth = self.min_depth          # Minimum allowable water depth
+        max_depth_k = self.max_depth_k
+        # Solve internals
+        tnumba_solve_internals(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik,
+                              _X_Ik, _Y_Ik, _Z_Ik, _i_1k, _I_1k, _I_Nk, nk, NK)
+        # Export instance variables
+        self._h_Ik = _h_Ik
+        self._Q_ik = _Q_ik
\ No newline at end of file
diff --git a/test/test_pipedream.py b/test/coverage_test.py
similarity index 59%
rename from test/test_pipedream.py
rename to test/coverage_test.py
index 099eaa4..043117e 100644
--- a/test/test_pipedream.py
+++ b/test/coverage_test.py
@@ -17,9 +17,7 @@
 
 internal_links = 24
 
-hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                      hillslope_superjunctions,
-                                      internal_links=internal_links)
+hillslope_superlink_wq_params[['No_of_obs_point', 'KF_obs_point', 'KF_process_sigma', 'KF_measure_sigma']] = [1, 1, 1, 1]
 
 hillslope_nsuperlink_model = nSuperLink(hillslope_superlinks,
                                        hillslope_superjunctions,
@@ -28,46 +26,34 @@
 hillslope_greenampt_model = GreenAmpt(hillslope_soil_params)
 hillslope_ngreenampt_model = nGreenAmpt(hillslope_soil_params)
 
-hillslope_water_quality_model = QualityBuilder(hillslope_nsuperlink_model,
-                                               superjunction_params=hillslope_superjunction_wq_params,
-                                               superlink_params=hillslope_superlink_wq_params)
+# hillslope_water_quality_model = QualityBuilder(hillslope_nsuperlink_model,
+#                                                superjunction_params=hillslope_superjunction_wq_params,
+#                                                superlink_params=hillslope_superlink_wq_params)
 
 initial_nsuperlink_states = copy.deepcopy(hillslope_nsuperlink_model.states)
 
-def test_superlink_step():
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                        hillslope_superjunctions,
-                                        internal_links=4)
-    dt = 10
-    Q_in = 1e-2 * np.asarray([1., 0.])
-    Q_0Ik = 1e-3 * np.ones(hillslope_superlink_model.NIk)
-    hillslope_superlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
-    hillslope_superlink_model.reposition_junctions()
-
 def test_nsuperlink_step():
+    hillslope_nsuperlink_model = nSuperLink(hillslope_superlinks,
+                                        hillslope_superjunctions,
+                                           internal_links=4,
+                                           mobile_elements=True)
     dt = 10
     Q_in = 1e-2 * np.asarray([1., 0.])
     Q_0Ik = 1e-3 * np.ones(hillslope_nsuperlink_model.NIk)
     hillslope_nsuperlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
     hillslope_nsuperlink_model.reposition_junctions()
 
-def test_superlink_spinup():
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                        hillslope_superjunctions,
-                                        internal_links=4)
-    hillslope_superlink_model.spinup(n_steps=100)
-
 def test_superlink_convergence():
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
+    hillslope_superlink_model = nSuperLink(hillslope_superlinks,
                                         hillslope_superjunctions,
                                         internal_links=4)
     dt = 10
     Q_in = 1e-2 * np.asarray([1., 0.])
-    Q_0Ik = 1e-3 * np.ones(hillslope_superlink_model.NIk)
+    Q_0Ik = 1e-3 * np.ones(hillslope_nsuperlink_model.NIk)
     hillslope_superlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik, num_iter=8)
 
 def test_superlink_banded_step():
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
+    hillslope_superlink_model = nSuperLink(hillslope_superlinks,
                                           hillslope_superjunctions,
                                           internal_links=4, auto_permute=True)
     dt = 10
@@ -75,28 +61,11 @@ def test_superlink_banded_step():
     Q_0Ik = 1e-3 * np.ones(hillslope_superlink_model.NIk)
     hillslope_superlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
 
-def test_superlink_recurrence_method():
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                          hillslope_superjunctions,
-                                          internal_links=4, method='f')
-    dt = 10
-    Q_in = 1e-2 * np.asarray([1., 0.])
-    Q_0Ik = 1e-3 * np.ones(hillslope_superlink_model.NIk)
-    hillslope_superlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                          hillslope_superjunctions,
-                                          internal_links=4, method='nnls')
-    hillslope_superlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                          hillslope_superjunctions,
-                                          internal_links=4, method='lsq')
-    hillslope_superlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
-
 def test_simulation_manager():
     dt = 10
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                          hillslope_superjunctions,
-                                          internal_links=24)
+    hillslope_superlink_model = nSuperLink(hillslope_superlinks,
+                                           hillslope_superjunctions,
+                                           internal_links=24)
     Q_in = pd.DataFrame.from_dict(
         {
             0 :  np.zeros(hillslope_superlink_model.M),
@@ -106,7 +75,6 @@ def test_simulation_manager():
             18001 : np.zeros(hillslope_superlink_model.M),
             28000 : np.zeros(hillslope_superlink_model.M)
         }, orient='index')
-
     Q_Ik = pd.DataFrame.from_dict(
         {
             0 :  np.zeros(hillslope_superlink_model.NIk),
@@ -130,9 +98,9 @@ def test_simulation_manager():
 
 def test_adaptive_timestep():
     dt = 10
-    hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                          hillslope_superjunctions,
-                                          internal_links=24)
+    hillslope_superlink_model = nSuperLink(hillslope_superlinks,
+                                           hillslope_superjunctions,
+                                           internal_links=24)
     Q_in = pd.DataFrame.from_dict(
         {
             0 :  np.zeros(hillslope_superlink_model.M),
@@ -142,7 +110,6 @@ def test_adaptive_timestep():
             18001 : np.zeros(hillslope_superlink_model.M),
             28000 : np.zeros(hillslope_superlink_model.M)
         }, orient='index')
-
     Q_Ik = pd.DataFrame.from_dict(
         {
             0 :  np.zeros(hillslope_superlink_model.NIk),
@@ -176,9 +143,9 @@ def test_superlink_geometry():
     hillslope_superlinks['g3'] = 1
     for geom in geoms:
         hillslope_superlinks['shape'] = geom
-        hillslope_superlink_model = SuperLink(hillslope_superlinks,
-                                            hillslope_superjunctions,
-                                            internal_links=4)
+        hillslope_superlink_model = nSuperLink(hillslope_superlinks,
+                                               hillslope_superjunctions,
+                                               internal_links=4)
         dt = 10
         Q_in = 1e-2 * np.asarray([1., 0.])
         Q_0Ik = 1e-3 * np.ones(hillslope_superlink_model.NIk)
@@ -188,23 +155,14 @@ def test_superlink_geometry():
     hillslope_superlinks['g2'] = 5
 
 def test_plot_profile():
-    hillslope_superlink_model.plot_profile([0, 1], width=100)
     hillslope_nsuperlink_model.plot_profile([0, 1], width=100)
 
 def test_plot_network_2d():
-    hillslope_superlink_model.plot_network_2d(junction_kwargs={'s' : 4})
     hillslope_nsuperlink_model.plot_network_2d(junction_kwargs={'s' : 4})
 
-def test_greenampt_step():
-    dt = 120
-    i = 50 / 1000 / 3600 * np.ones(hillslope_superlink_model.NIk)
-    infiltration_model = hillslope_greenampt_model
-    for _ in range(100):
-        infiltration_model.step(dt=dt, i=i)
-
 def test_ngreenampt_step():
     dt = 10
-    i = 50 / 1000 / 3600 * np.ones(hillslope_superlink_model.NIk)
+    i = 50 / 1000 / 3600 * np.ones(hillslope_nsuperlink_model.NIk)
     hydraulic_model = hillslope_nsuperlink_model
     infiltration_model = hillslope_ngreenampt_model
     hydraulic_model.load_state(initial_nsuperlink_states)
@@ -214,42 +172,14 @@ def test_ngreenampt_step():
         Q_0Ik = infiltration_model.Q
         hydraulic_model.step(dt=dt, Q_0Ik=Q_0Ik)
 
-def test_water_quality_step():
-    dt = 10
-    Q_in = 1e-2 * np.asarray([1., 0.])
-    Q_0Ik = 1e-3 * np.ones(hillslope_nsuperlink_model.NIk)
-    c_0j = 10. * np.asarray([1., 0.])
-    for _ in range(100):
-        hillslope_nsuperlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
-        hillslope_water_quality_model.step(dt=dt, c_0j=c_0j)
-
-def test_orifice():
-    superjunctions = copy.deepcopy(hillslope_superjunctions)
-    superlinks = hillslope_superlinks
-    superjunctions.loc[2] = superjunctions.loc[0]
-    superjunctions.loc[2, ['name', 'id']] = 2
-    superjunctions.loc[2, 'h_0'] = 2.0
-    orifices = {
-        'id' : 0,
-        'sj_0' : 2,
-        'sj_1' : 0,
-        'A' : 0.3048**2,
-        'orientation' : 'side',
-        'z_o' : 0,
-        'y_max' : 0.3048,
-        'C' : 0.67}
-    orifices = pd.DataFrame(orifices, index=[0])
-    hydraulic_model = SuperLink(superlinks, superjunctions, orifices=orifices,
-                                internal_links=internal_links)
-    dt = 10
-    Q_in = 1e-2 * np.asarray([0., 0., 0.])
-    Q_0Ik = 1e-3 * np.ones(hydraulic_model.NIk)
-    for _ in range(100):
-        if (hydraulic_model.t > 200):
-            u_o = 0.5 * np.ones(1)
-        else:
-            u_o = np.zeros(1)
-        hydraulic_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o)
+# def test_water_quality_step():
+#     dt = 10
+#     Q_in = 1e-2 * np.asarray([1., 0.])
+#     Q_0Ik = 1e-3 * np.ones(hillslope_nsuperlink_model.NIk)
+#     c_0j = 10. * np.asarray([1., 0.])
+#     for _ in range(100):
+#         hillslope_nsuperlink_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik)
+#         hillslope_water_quality_model.step(dt=dt, c_0j=c_0j)
 
 def test_norifice():
     superjunctions = copy.deepcopy(hillslope_superjunctions)
@@ -279,36 +209,6 @@ def test_norifice():
             u_o = np.zeros(1)
         hydraulic_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik, u_o=u_o)
 
-def test_weir():
-    superjunctions = copy.deepcopy(hillslope_superjunctions)
-    superlinks = hillslope_superlinks
-    superjunctions.loc[2] = superjunctions.loc[0]
-    superjunctions.loc[2, ['name', 'id']] = 2
-    superjunctions.loc[2, 'h_0'] = 2.0
-    weirs = {
-        'id' : 0,
-        'sj_0' : 2,
-        'sj_1' : 0,
-        'z_w' : 0,
-        'y_max' : 0.3048,
-        'Cr' : 0.67,
-        'Ct' : 0.67,
-        'L' : 0.3048,
-        's' : 0.01
-    }
-    weirs = pd.DataFrame(weirs, index=[0])
-    hydraulic_model = SuperLink(superlinks, superjunctions, weirs=weirs,
-                                internal_links=internal_links)
-    dt = 10
-    Q_in = 1e-2 * np.asarray([0., 0., 0.])
-    Q_0Ik = 1e-3 * np.ones(hydraulic_model.NIk)
-    for _ in range(100):
-        if (hydraulic_model.t > 200):
-            u_w = 0.5 * np.ones(1)
-        else:
-            u_w = np.zeros(1)
-        hydraulic_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik, u_w=u_w)
-
 def test_nweir():
     superjunctions = copy.deepcopy(hillslope_superjunctions)
     superlinks = hillslope_superlinks
@@ -339,35 +239,6 @@ def test_nweir():
             u_w = np.zeros(1)
         hydraulic_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik, u_w=u_w)
 
-def test_pump():
-    superjunctions = copy.deepcopy(hillslope_superjunctions)
-    superlinks = hillslope_superlinks
-    superjunctions.loc[2] = superjunctions.loc[0]
-    superjunctions.loc[2, ['name', 'id']] = 2
-    superjunctions.loc[2, 'h_0'] = 2.0
-    pumps = {
-        'id' : 0,
-        'sj_0' : 0,
-        'sj_1' : 2,
-        'z_p' : 0,
-        'a_q' : 2.0,
-        'a_h' : 0.1,
-        'dH_min' : 0.5,
-        'dH_max' : 2.0
-    }
-    pumps = pd.DataFrame(pumps, index=[0])
-    hydraulic_model = SuperLink(superlinks, superjunctions, pumps=pumps,
-                                internal_links=internal_links)
-    dt = 10
-    Q_in = 1e-2 * np.asarray([0., 0., 0.])
-    Q_0Ik = 1e-3 * np.ones(hydraulic_model.NIk)
-    for _ in range(100):
-        if (hydraulic_model.t > 200):
-            u_p = 0.5 * np.ones(1)
-        else:
-            u_p = np.zeros(1)
-        hydraulic_model.step(dt=dt, Q_in=Q_in, Q_0Ik=Q_0Ik, u_p=u_p)
-
 def test_npump():
     superjunctions = copy.deepcopy(hillslope_superjunctions)
     superlinks = hillslope_superlinks
@@ -379,8 +250,9 @@ def test_npump():
         'sj_0' : 0,
         'sj_1' : 2,
         'z_p' : 0,
-        'a_q' : 2.0,
-        'a_h' : 0.1,
+        'a_p' : 1.0,
+        'b_p' : 1.0,
+        'c_p' : 1.0,
         'dH_min' : 0.5,
         'dH_max' : 2.0
     }