minor optimisation

src/lisflood/hydrological_modules/kinematic_wave_parallel_tools.py

@@ -37,13 +37,15 @@ def kinematicRouting(discharge_avg, discharge, lateral_inflow, constant, upstrea
                      b_minus_1, a_dx_div_dt, b_a_dx_div_dt, a_dx):
     """
     This function performs the kinematic wave routing algorithm to simulate the movement of water through a network of interconnected channels.
+    :param discharge_avg:
     :param discharge:
     :param lateral_inflow:
     :param constant:
     :param upstream_lookup:
     :param num_upstream_pixels:
     :param ordered_pixels:
     :param start_stop:
+    :param inv_time_delta: 1/DtRouting
     :param beta:
     :param inv_beta:
     :param b_minus_1:
@@ -67,10 +69,12 @@ def kinematicRouting(discharge_avg, discharge, lateral_inflow, constant, upstrea
                 solve1PixelAvg(pix, discharge_avg, discharge_before, discharge[pix], lateral_inflow[pix], upstream_lookup,\
                             num_upstream_pixels, inv_time_delta, beta, a_dx[pix])
 
+
 @njit(nogil=True, fastmath=False, cache=True)
-def solve1Pixel(pix, discharge, constant,\
-                      upstream_lookup, num_upstream_pixels, a_dx_div_dt,\
-                      b_a_dx_div_dt, beta, inv_beta, b_minus_1):
+def solve1Pixel(
+    pix, discharge, constant,
+    upstream_lookup, num_upstream_pixels, a_dx_div_dt,
+    b_a_dx_div_dt, beta, inv_beta, b_minus_1):
     """
     Te Chow et al. 1988 - Applied Hydrology - Chapter 9.6
     :param pix:
@@ -85,56 +89,63 @@ def solve1Pixel(pix, discharge, constant,\
     :param b_minus_1:
     :return:
     """
-    count = 0
-    previous_estimate = -1.0
-    upstream_inflow = 0.0
+
     # Inflow from upstream pixels
+    upstream_inflow = constant  # upstream_inflow + alpha*dx/dt*Qold**beta + dx*specific_lateral_inflow
     for ups_ix in range(num_upstream_pixels[pix]):
         upstream_inflow += discharge[upstream_lookup[pix,ups_ix]]
-    const_plus_ups_infl = upstream_inflow + constant # upstream_inflow + alpha*dx/dt*Qold**beta + dx*specific_lateral_inflow
+
     # If old discharge, upstream inflow and lateral inflow are below accuracy: set discharge to 0 and exit
-    if const_plus_ups_infl <= NEWTON_TOL:
+    if upstream_inflow <= NEWTON_TOL:
         discharge[pix] = 0
         return False
 
     # Initial discharge guess using analytically derived boundary values
-    a_cpui_pow_b_m_1 = b_a_dx_div_dt * const_plus_ups_infl**b_minus_1
+    a_cpui_pow_b_m_1 = b_a_dx_div_dt * upstream_inflow**b_minus_1
     if a_cpui_pow_b_m_1 <= 1:
-        secant_bound = const_plus_ups_infl / (1 + a_cpui_pow_b_m_1)
+        secant_bound = upstream_inflow / (1 + a_cpui_pow_b_m_1)
     else:
-        secant_bound = const_plus_ups_infl / (1 + a_cpui_pow_b_m_1**inv_beta)
-    other_bound = ((const_plus_ups_infl - secant_bound) / a_dx_div_dt)**inv_beta
-    discharge[pix] = (secant_bound + other_bound) / 2
-    error = closureError(discharge[pix], const_plus_ups_infl, a_dx_div_dt, beta)
+        secant_bound = upstream_inflow / (1 + a_cpui_pow_b_m_1**inv_beta)
+    other_bound = ((upstream_inflow - secant_bound) / a_dx_div_dt)**inv_beta
+    dis = (secant_bound + other_bound) / 2
+    error = dis + a_dx_div_dt * dis**beta - upstream_inflow
+
     # Iterations
-    while fabs(error) > NEWTON_TOL and discharge[pix] != previous_estimate and count < MAX_ITERS: # is previous_estimate useful?
-        previous_estimate = discharge[pix]
-        discharge[pix] -= error / (1 + b_a_dx_div_dt * discharge[pix]**b_minus_1)
-        discharge[pix] = max(discharge[pix], NEWTON_TOL)
-        error = closureError(discharge[pix], const_plus_ups_infl, a_dx_div_dt, beta)
+    count = 0
+    previous_estimate = -1.0
+    while count < MAX_ITERS:
+        if fabs(error) <= NEWTON_TOL or dis == previous_estimate:
+            break
+        previous_estimate = dis
+        factor = (1 + b_a_dx_div_dt * dis**b_minus_1)
+        dis = max(dis - error / factor, NEWTON_TOL)
+        error = dis + a_dx_div_dt * dis**beta - upstream_inflow
         count += 1
+
     # If iterations converge to NEWTON_TOL, set value to 0
-    if discharge[pix] == NEWTON_TOL:
+    if dis == NEWTON_TOL:
         discharge[pix] = 0
-
-    return True
+        return False
+
+    discharge[pix] = dis
+    return True    
+
 
 @njit(nogil=True, fastmath=False, cache=True)
-def solve1PixelAvg(pix, discharge_avg, discharge_before, discharge_after, lateral_inflow, \
-                      upstream_lookup, num_upstream_pixels, \
-                      inv_time_delta, beta, a_dx):
+def solve1PixelAvg(
+    pix, discharge_avg, discharge_before, discharge_after, lateral_inflow,
+    upstream_lookup, num_upstream_pixels,
+    inv_time_delta, beta, a_dx):
     """
-    Te Chow et al. 1988 - Applied Hydrology - Chapter 9.6
     :param pix:
     :param discharge_avg: average outflow discharge
     :param discharge_before: instantaneous outflow discharge before solve1Pixel computation
-    :param discharge_before: instantaneous outflow discharge after solve1Pixelcomputation
+    :param discharge_after: instantaneous outflow discharge after solve1Pixel computation
     :param lateral_inflow:
     :param upstream_lookup:
     :param num_upstream_pixels:
     :param inv_time_delta: 1/DtRouting
     :param beta:
-    :param inv_beta:
     :param a_dx: ChannelAlpha * ChanLength
     :return:
     """
@@ -157,12 +168,6 @@ def solve1PixelAvg(pix, discharge_avg, discharge_before, discharge_after, latera
     if discharge_avg[pix] < 0:
         discharge_avg[pix] = 0
 
-@njit(nogil=True, fastmath=False, cache=True)
-def closureError(discharge, upper_bound, a_dx_div_dt, beta):
-    """"""
-    return discharge + a_dx_div_dt * discharge**beta - upper_bound
-
-
 
 # -------------------------------------------------------------------------------------------------
 # FLOW DIRECTION MATRIX PRE-PROCESSING FUNCTIONS
@@ -208,6 +213,7 @@ def updateUpstreamRouted(pixels_routed, downstream_lookup,\
             # Mark the upstream pixel as routed
             upstream_routed[downs_pix,ups_index] = True
 
+
 @njit(nogil=True, fastmath=False, cache=True)
 def upDownLookups(flow_d8, land_mask, land_points, num_pixs, ix_adds):
     """
@@ -278,7 +284,6 @@ def upDownLookups(flow_d8, land_mask, land_points, num_pixs, ix_adds):
     return np.asarray(downstream_lookup), np.asarray(upstream_lookup)
 
 
-
 # -------------------------------------------------------------------------------------------------
 # MISCELLANOUS TOOLS
 # -------------------------------------------------------------------------------------------------
@@ -318,4 +323,3 @@ def immediateUpstreamInflow(discharge, upstream_lookup, num_upstream_pixels):
 
     # Return the inflow array as a NumPy array
     return np.asarray(inflow)
-