split and optimization of the new solve1pixel function for the mct br…

…anch

src/lisflood/hydrological_modules/kinematic_wave_parallel_tools.py

@@ -60,85 +60,103 @@ def kinematicRouting(discharge_avg, discharge, lateral_inflow, constant, upstrea
         # Iterate through each pixel in the current order in parallel
         for index in prange(first, last):  # this is a parallel loop
             # Solve the kinematic wave for the current pixel
-            solve1Pixel(ordered_pixels[index], discharge_avg, discharge, lateral_inflow, constant, upstream_lookup,\
-                        num_upstream_pixels, a_dx_div_dt, b_a_dx_div_dt, inv_time_delta, beta, inv_beta, b_minus_1, a_dx)
+            pix = ordered_pixels[index]
+            discharge_before = discharge[pix]
+            if solve1Pixel(pix, discharge, constant[pix], upstream_lookup,\
+                        num_upstream_pixels, a_dx_div_dt[pix], b_a_dx_div_dt[pix], beta, inv_beta, b_minus_1):                        
+                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_avg, discharge, lateral_inflow, constant,\
+def solve1Pixel(pix, discharge, constant,\
                       upstream_lookup, num_upstream_pixels, a_dx_div_dt,\
-                      b_a_dx_div_dt, inv_time_delta, beta, inv_beta, b_minus_1, a_dx):
+                      b_a_dx_div_dt, beta, inv_beta, b_minus_1):
     """
     Te Chow et al. 1988 - Applied Hydrology - Chapter 9.6
     :param pix:
-    :param discharge_avg: average outflow discharge
     :param discharge: instantaneous outflow discharge
-    :param lateral_inflow:
     :param constant:
     :param upstream_lookup:
     :param num_upstream_pixels:
     :param a_dx_div_dt:
     :param b_a_dx_div_dt:
-    :param inv_time_delta: 1/DtRouting
     :param beta:
     :param inv_beta:
     :param b_minus_1:
-    :param a_dx: ChannelAlpha * ChanLength
     :return:
     """
     count = 0
     previous_estimate = -1.0
     upstream_inflow = 0.0
-    upstream_inflow_avg = 0.0
-
-    # volume of water in channel at beginning of computation step
-    channel_volume_start = a_dx * discharge[pix]**beta
-
     # Inflow from upstream pixels
     for ups_ix in range(num_upstream_pixels[pix]):
         upstream_inflow += discharge[upstream_lookup[pix,ups_ix]]
-        upstream_inflow_avg += discharge_avg[upstream_lookup[pix, ups_ix]]
-
-    const_plus_ups_infl = upstream_inflow + constant[pix] # upstream_inflow + alpha*dx/dt*Qold**beta + dx*specific_lateral_inflow
+    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:
         discharge[pix] = 0
-        return
+        return False
+
     # Initial discharge guess using analytically derived boundary values
-    a_cpui_pow_b_m_1 = b_a_dx_div_dt[pix] * const_plus_ups_infl**b_minus_1
+    a_cpui_pow_b_m_1 = b_a_dx_div_dt * const_plus_ups_infl**b_minus_1
     if a_cpui_pow_b_m_1 <= 1:
         secant_bound = const_plus_ups_infl / (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[pix])**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[pix], beta)
+    error = closureError(discharge[pix], const_plus_ups_infl, a_dx_div_dt, beta)
     # 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[pix] * discharge[pix]**b_minus_1)
+        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[pix], beta)
+        error = closureError(discharge[pix], const_plus_ups_infl, a_dx_div_dt, beta)
         count += 1
     # If iterations converge to NEWTON_TOL, set value to 0
     if discharge[pix] == NEWTON_TOL:
         discharge[pix] = 0
+
+    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):
+    """
+    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 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:
+    """
+    upstream_inflow_avg = 0.0
+
+    # Inflow from upstream pixels
+    for ups_ix in range(num_upstream_pixels[pix]):
+        upstream_inflow_avg += discharge_avg[upstream_lookup[pix, ups_ix]]
+
+    # volume of water in channel at beginning of computation step
+    channel_volume_start = discharge_before**beta
 
     # volume of water in channel at end of computation step
-    channel_volume_end = a_dx * discharge[pix]**beta
+    channel_volume_end = discharge_after**beta
 
     # integration on control volume to calc average outflow (channel water mass balance)
-    discharge_avg[pix] = upstream_inflow_avg + lateral_inflow[pix] + (channel_volume_start[pix] - channel_volume_end[pix]) * inv_time_delta
+    discharge_avg[pix] = upstream_inflow_avg + lateral_inflow + a_dx*(channel_volume_start - channel_volume_end) * inv_time_delta
 
     # avoid negative average discharge
     if discharge_avg[pix] < 0:
         discharge_avg[pix] = 0
 
-    # to simulate inf or nan: discharge[pix] = 1.0/0.0
-    # with gil:
-    #    got_valid_value = np.isfinite(discharge[pix])
-    #    if got_valid_value==False:
-    #          print(str(discharge[pix]) + ' found at pix:' + str(pix))
-
 @njit(nogil=True, fastmath=False, cache=True)
 def closureError(discharge, upper_bound, a_dx_div_dt, beta):
     """"""