From c51df7e3c72c69834255b472b65d6688303269ee Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Tue, 26 Mar 2024 21:38:52 -0500 Subject: [PATCH 01/27] Draft of coupling class --- pipedream_solver/coupling.py | 32 ++++++++++++++++++++++++++++++++ pipedream_solver/nquality.py | 24 ++++++++++++++++++++++-- 2 files changed, 54 insertions(+), 2 deletions(-) create mode 100644 pipedream_solver/coupling.py 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/nquality.py b/pipedream_solver/nquality.py index 0818254..3a4ad90 100644 --- a/pipedream_solver/nquality.py +++ b/pipedream_solver/nquality.py @@ -223,7 +223,7 @@ def __init__(self, hydraulics, superjunction_params, superlink_params, self.N_process_sigma = superlink_params['KF_process_sigma'].values[0].astype(np.float64) self.N_measure_sigma = superlink_params['KF_measure_sigma'].values[0].astype(np.float64) self.step(dt=1e-6) - + # TODO: It might be safer to have these as @properties def import_hydraulic_states(self, _dt): self._H_j_next = self.hydraulics.H_j @@ -277,7 +277,27 @@ def import_hydraulic_states(self, _dt): self.hydraulics._geom_codes, self.hydraulics._g1_ik, self._H_j_next, self._z_inv_j, self._H_j_prev, self._Q_dk_next, self._B_dk, self._dx_dk, self._Q_uk_up_next, self._Q_uk_dn_next, self._Q_dk_up_next, self._Q_dk_dn_next, self._A_uk_next, self._A_dk_next, self._A_uk_prev, self._A_dk_prev) - + + @property + def c(self): + # NOTE: more performant to write in-place + result = np.concatenate([self._c_j, self._c_Ik, self._c_ik, + self._c_uk, self._c_dk]) + return result + + @c.setter + def c(self, value): + n1 = self._c_j.size + n2 = n1 + self._c_Ik.size + n3 = n2 + self._c_ik.size + n4 = n3 + self._c_uk.size + n5 = n4 + self._c_dk.size + self._c_j[:] = np.asarray(value[:n1]) + self._c_Ik[:] = np.asarray(value[n1:n2]) + self._c_ik[:] = np.asarray(value[n2:n3]) + self._c_uk[:] = np.asarray(value[n3:n4]) + self._c_dk[:] = np.asarray(value[n4:n5]) + @property def c_j(self): return self._c_j From 04057af7cd5b0cda56a9365828d39550f6f4eac9 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Thu, 30 May 2024 23:33:18 -0500 Subject: [PATCH 02/27] Re-create tests --- pipedream_solver/nsuperlink.py | 2 +- test/test_pipedream.py | 192 ++++++--------------------------- 2 files changed, 33 insertions(+), 161 deletions(-) diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index 6dae001..476bb12 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -1558,7 +1558,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, diff --git a/test/test_pipedream.py b/test/test_pipedream.py index 099eaa4..043117e 100644 --- a/test/test_pipedream.py +++ b/test/test_pipedream.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 } From 3771ba6014586976e2f13e22fdc447d7be35dd64 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 31 May 2024 00:16:37 -0500 Subject: [PATCH 03/27] Separate numba functions for coverage tests --- .coveragerc | 2 +- pipedream_solver/_ninfiltration.py | 121 ++ pipedream_solver/_nsuperlink.py | 1446 ++++++++++++++++++ pipedream_solver/ninfiltration.py | 121 +- pipedream_solver/nsuperlink.py | 1443 +---------------- test/{test_pipedream.py => coverage_test.py} | 0 6 files changed, 1572 insertions(+), 1561 deletions(-) create mode 100644 pipedream_solver/_ninfiltration.py create mode 100644 pipedream_solver/_nsuperlink.py rename test/{test_pipedream.py => coverage_test.py} (100%) 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..491c38e --- /dev/null +++ b/pipedream_solver/_nsuperlink.py @@ -0,0 +1,1446 @@ +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 + +@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[:], + 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, + _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: does not handle "max" mode + h_I = _h_Ik[I] + h_Ip1 = _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(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/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/nsuperlink.py b/pipedream_solver/nsuperlink.py index 476bb12..fb8f958 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): """ @@ -1574,1445 +1575,3 @@ 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[:], - 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, - _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: does not handle "max" mode - h_I = _h_Ik[I] - h_Ip1 = _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(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/test/test_pipedream.py b/test/coverage_test.py similarity index 100% rename from test/test_pipedream.py rename to test/coverage_test.py From 832c0b95f0d56da3b455286d3b282eefec6e2c17 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Sat, 1 Jun 2024 16:35:35 -0500 Subject: [PATCH 04/27] Simplify geometry function inputs --- pipedream_solver/_nsuperlink.py | 143 ++++++++++++----------- pipedream_solver/ngeometry.py | 199 ++++++++++++++++++-------------- 2 files changed, 188 insertions(+), 154 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 491c38e..7405940 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -17,6 +17,7 @@ def numba_hydraulic_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik, 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] @@ -26,55 +27,58 @@ def numba_hydraulic_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik, 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) + _A_ik[i] = pipedream_solver.ngeometry.Circular_A_ik(h_i, g1_i) + _Pe_ik[i] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, 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) + _B_ik[i] = pipedream_solver.ngeometry.Circular_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i, g3_i) + _B_ik[i] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_ik[i] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_ik[i] = pipedream_solver.ngeometry.Triangular_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i, g3_i) + _B_ik[i] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_ik[i] = pipedream_solver.ngeometry.Parabolic_B_ik(h_i, 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) + raise NotImplementedError 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) + _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, h_Ip1, g1_i, g2_i) + _B_ik[i] = pipedream_solver.ngeometry.Wide_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_ik[i] = pipedream_solver.ngeometry.Force_Main_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i, g3_i, g4_i, g5_i, g6_i, g7_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[:], @@ -92,6 +96,7 @@ def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, # TODO: does not handle "max" mode h_I = _h_Ik[I] h_Ip1 = _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] @@ -102,54 +107,52 @@ def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, 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) + _A_bk[k] = pipedream_solver.ngeometry.Circular_A_ik(h_i, g1_i) + _Pe_bk[k] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, 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) + _B_bk[k] = pipedream_solver.ngeometry.Circular_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i, g3_i) + _B_bk[k] = pipedream_solver.ngeometry.Rect_Closed_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_bk[k] = pipedream_solver.ngeometry.Rect_Open_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_bk[k] = pipedream_solver.ngeometry.Triangular_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i, g3_i) + _B_bk[k] = pipedream_solver.ngeometry.Trapezoidal_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_bk[k] = pipedream_solver.ngeometry.Parabolic_B_ik(h_i, 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) + raise NotImplementedError 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) + _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, h_Ip1, g1_i, g2_i) + _B_bk[k] = pipedream_solver.ngeometry.Wide_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i) + _B_bk[k] = pipedream_solver.ngeometry.Force_Main_B_ik(h_i, 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) + _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, h_Ip1, g1_i, g2_i, g3_i, g4_i, g5_i, g6_i, g7_i) + _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[:], @@ -165,23 +168,23 @@ def numba_orifice_geometry(_Ao, h_eo, u_o, _g1_o, _g2_o, _g3_o, _geom_codes_o, n 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) + _Ao[i] = pipedream_solver.ngeometry.Circular_A_ik(h_e, g1 * u) elif geom_code == 2: - _Ao[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(h_e, h_e, g1 * u, g2) + _Ao[i] = pipedream_solver.ngeometry.Rect_Closed_A_ik(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) + _Ao[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_e, g1 * u, g2) elif geom_code == 4: - _Ao[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_e, h_e, g1 * u, g2) + _Ao[i] = pipedream_solver.ngeometry.Triangular_A_ik(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) + _Ao[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(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) + _Ao[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_e, g1 * u, g2) elif geom_code == 7: - _Ao[i] = pipedream_solver.ngeometry.Elliptical_A_ik(h_e, h_e, g1 * u, g2) + raise NotImplementedError elif geom_code == 8: - _Ao[i] = pipedream_solver.ngeometry.Wide_A_ik(h_e, h_e, g1 * u, g2) + _Ao[i] = pipedream_solver.ngeometry.Wide_A_ik(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) + _Ao[i] = pipedream_solver.ngeometry.Force_Main_A_ik(h_e, g1 * u, g2) return 1 @njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]), diff --git a/pipedream_solver/ngeometry.py b/pipedream_solver/ngeometry.py index e767879..514d928 100644 --- a/pipedream_solver/ngeometry.py +++ b/pipedream_solver/ngeometry.py @@ -18,9 +18,9 @@ eps = np.finfo(float).eps -@njit(float64(float64, float64, float64), +@njit(float64(float64, float64), cache=True) -def Circular_A_ik(h_Ik, h_Ip1k, g1): +def Circular_A_ik(h_ik, g1): """ Compute cross-sectional area of flow for link i, superlink k. @@ -34,7 +34,7 @@ def Circular_A_ik(h_Ik, h_Ip1k, g1): Diameter of channel (meters) """ d = g1 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 if y > d: @@ -49,9 +49,9 @@ def Circular_A_ik(h_Ik, h_Ip1k, g1): A = r**2 * (theta - np.cos(theta) * np.sin(theta)) return A -@njit(float64(float64, float64, float64), +@njit(float64(float64, float64), cache=True) -def Circular_Pe_ik(h_Ik, h_Ip1k, g1): +def Circular_Pe_ik(h_ik, g1): """ Compute perimeter of flow for link i, superlink k. @@ -65,7 +65,7 @@ def Circular_Pe_ik(h_Ik, h_Ip1k, g1): Diameter of channel (meters) """ d = g1 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 if y > d: @@ -100,9 +100,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,8 +119,7 @@ def Circular_B_ik(h_Ik, h_Ip1k, g1, g2): """ d = g1 pslot = g2 - y = (h_Ik + h_Ip1k) / 2 - # y[y < 0] = 0 + y = h_ik if y < 0: y = 0 r = d / 2 @@ -134,13 +133,14 @@ 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 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,7 +157,7 @@ def Rect_Closed_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: @@ -165,9 +165,9 @@ def Rect_Closed_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 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,7 +184,7 @@ def Rect_Closed_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 = 0 if y > y_max: @@ -212,9 +212,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,7 +234,7 @@ 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 + y = h_ik if y < 0: y = 0 cond = (y < y_max) @@ -245,9 +245,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,7 +264,7 @@ def Rect_Open_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: @@ -272,9 +272,9 @@ def Rect_Open_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 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,7 +291,7 @@ def Rect_Open_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 = 0 if y > y_max: @@ -319,9 +319,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 +341,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,7 +360,7 @@ def Triangular_A_ik(h_Ik, h_Ip1k, g1, g2): """ y_max = g1 m = g2 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 if y > y_max: @@ -368,9 +368,9 @@ def Triangular_A_ik(h_Ik, h_Ip1k, g1, g2): 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,7 +387,7 @@ def Triangular_Pe_ik(h_Ik, h_Ip1k, g1, g2): """ y_max = g1 m = g2 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 if y > y_max: @@ -415,9 +415,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,7 +434,7 @@ def Triangular_B_ik(h_Ik, h_Ip1k, g1, g2): """ y_max = g1 m = g2 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 cond = (y < y_max) @@ -445,9 +445,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,7 +467,7 @@ def Trapezoidal_A_ik(h_Ik, h_Ip1k, g1, g2, g3): y_max = g1 b = g2 m = g3 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 if y > y_max: @@ -475,9 +475,9 @@ def Trapezoidal_A_ik(h_Ik, h_Ip1k, g1, g2, g3): 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,7 +497,7 @@ def Trapezoidal_Pe_ik(h_Ik, h_Ip1k, g1, g2, g3): y_max = g1 b = g2 m = g3 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 if y > y_max: @@ -525,9 +525,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,7 +547,7 @@ def Trapezoidal_B_ik(h_Ik, h_Ip1k, g1, g2, g3): y_max = g1 b = g2 m = g3 - y = (h_Ik + h_Ip1k) / 2 + y = h_ik if y < 0: y = 0 cond = (y < y_max) @@ -559,7 +559,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 +607,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 +615,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 +634,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 +663,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 +682,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 +708,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 +739,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 +759,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 +770,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 +789,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 +797,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 +838,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 +859,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 +881,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 +922,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 +944,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. @@ -942,7 +973,7 @@ def Floodplain_A_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7): g6: np.ndarray Inverse slope of middle channel sides (run/rise) g7: np.ndarray - bottom width of channel (meters) + bottom width of channel (meters) """ h_max = g1 h_low = g2 @@ -950,7 +981,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 +1011,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. @@ -1005,7 +1036,7 @@ def Floodplain_Pe_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7): g6: np.ndarray Inverse slope of middle channel sides (run/rise) g7: np.ndarray - bottom width of channel (meters) + bottom width of channel (meters) """ h_max = g1 h_low = g2 @@ -1013,7 +1044,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 +1093,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. @@ -1087,7 +1118,7 @@ def Floodplain_B_ik(h_Ik, h_Ip1k, g1, g2, g3, g4, g5, g6, g7): g6: np.ndarray Inverse slope of middle channel sides (run/rise) g7: np.ndarray - bottom width of channel (meters) + bottom width of channel (meters) """ h_max = g1 h_low = g2 @@ -1095,7 +1126,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: From dacd6581ae7bd9a036fd8eba03f04d33073d3c80 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Sun, 2 Jun 2024 02:08:05 -0500 Subject: [PATCH 05/27] Fix merge conflict --- pipedream_solver/_nsuperlink.py | 92 ++++++++++++++++++++++++++++----- 1 file changed, 79 insertions(+), 13 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 7405940..8996ce8 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -70,22 +70,19 @@ def numba_hydraulic_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik, _B_ik[i] = pipedream_solver.ngeometry.Force_Main_B_ik(h_i, g1_i, g2_i) elif geom_code == 10: _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) + 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) + 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) + g3_i, g4_i, g5_i, g6_i, g7_i) return 1 -@njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], +@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, +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) @@ -93,9 +90,9 @@ def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, i = _i_bk[k] I = _I_bk[k] j = _J_bk[k] - # TODO: does not handle "max" mode + # TODO: This should incorporate theta h_I = _h_Ik[I] - h_Ip1 = _H_j[j] - _z_inv_bk[k] + 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] @@ -149,10 +146,13 @@ def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, _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 == 10: - _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) + _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) + _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[:], @@ -187,6 +187,72 @@ def numba_orifice_geometry(_Ao, h_eo, u_o, _g1_o, _g2_o, _g3_o, _geom_codes_o, n _Ao[i] = pipedream_solver.ngeometry.Force_Main_A_ik(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): From 73006f47c7d4293a2e376da1a8d53301413ac503 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 9 Aug 2024 01:54:12 -0500 Subject: [PATCH 06/27] Refactoring of geometry and state space representation --- pipedream_solver/_nsuperlink.py | 82 +++++++++++++++++++-------------- pipedream_solver/ngeometry.py | 10 ++-- pipedream_solver/superlink.py | 13 +++++- 3 files changed, 65 insertions(+), 40 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 8996ce8..ba0960a 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -3,6 +3,17 @@ from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, DictType, void import pipedream_solver.ngeometry +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[:]), @@ -26,49 +37,49 @@ def numba_hydraulic_geometry(_A_ik, _Pe_ik, _R_ik, _B_ik, _h_Ik, 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, g1_i) - _Pe_ik[i] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, g1_i) + 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 == 2: + 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 == 3: + 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 == 4: + 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 == 5: + 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 == 6: + 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 == 7: + elif geom_code == ELLIPTICAL: raise NotImplementedError - elif geom_code == 8: + 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 == 9: + 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 == 10: + 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, @@ -103,49 +114,49 @@ def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, 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, g1_i) - _Pe_bk[k] = pipedream_solver.ngeometry.Circular_Pe_ik(h_i, g1_i) + 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 == 2: + 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 == 3: + 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 == 4: + 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 == 5: + 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 == 6: + 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 == 7: + elif geom_code == ELLIPTICAL: raise NotImplementedError - elif geom_code == 8: + 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 == 9: + 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 == 10: + 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, @@ -167,24 +178,27 @@ def numba_orifice_geometry(_Ao, h_eo, u_o, _g1_o, _g2_o, _g3_o, _geom_codes_o, n 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, g1 * u) - elif geom_code == 2: + 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 == 3: + elif geom_code == RECT_OPEN: _Ao[i] = pipedream_solver.ngeometry.Rect_Open_A_ik(h_e, g1 * u, g2) - elif geom_code == 4: + elif geom_code == TRIANGULAR: _Ao[i] = pipedream_solver.ngeometry.Triangular_A_ik(h_e, g1 * u, g2) - elif geom_code == 5: + elif geom_code == TRAPEZOIDAL: _Ao[i] = pipedream_solver.ngeometry.Trapezoidal_A_ik(h_e, g1 * u, g2, g3) - elif geom_code == 6: + elif geom_code == PARABOLIC: _Ao[i] = pipedream_solver.ngeometry.Parabolic_A_ik(h_e, g1 * u, g2) - elif geom_code == 7: + elif geom_code == ELLIPTICAL: raise NotImplementedError - elif geom_code == 8: + elif geom_code == WIDE: _Ao[i] = pipedream_solver.ngeometry.Wide_A_ik(h_e, g1 * u, g2) - elif geom_code == 9: + 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[:], diff --git a/pipedream_solver/ngeometry.py b/pipedream_solver/ngeometry.py index d0250a1..ec421a5 100644 --- a/pipedream_solver/ngeometry.py +++ b/pipedream_solver/ngeometry.py @@ -18,9 +18,9 @@ eps = np.finfo(float).eps -@njit(float64(float64, float64), +@njit(float64(float64, float64, float64), cache=True) -def Circular_A_ik(h_ik, g1): +def Circular_A_ik(h_ik, g1, g2): """ Compute cross-sectional area of flow for link i, superlink k. @@ -34,6 +34,7 @@ def Circular_A_ik(h_ik, g1): Diameter of channel (meters) """ d = g1 + pslot = g2 y = h_ik if y < 0: y = 0 @@ -49,9 +50,9 @@ def Circular_A_ik(h_ik, g1): A = r**2 * (theta - np.cos(theta) * np.sin(theta)) return A -@njit(float64(float64, float64), +@njit(float64(float64, float64, float64), cache=True) -def Circular_Pe_ik(h_ik, g1): +def Circular_Pe_ik(h_ik, g1, g2): """ Compute perimeter of flow for link i, superlink k. @@ -65,6 +66,7 @@ def Circular_Pe_ik(h_ik, g1): Diameter of channel (meters) """ d = g1 + pslot = g2 y = h_ik if y < 0: y = 0 diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 2abcd7a..45c1c25 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -759,6 +759,7 @@ 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) # Initialize state dictionary self.states = {} # Iteration counter @@ -3890,6 +3891,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 +3902,14 @@ 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 save_state(self): """ @@ -4018,6 +4026,7 @@ def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p= self.save_state() if dt is None: dt = self._dt + self._H_bc = H_bc self._Q_in = Q_in self._Q_0Ik = Q_0Ik if not implicit: From d05739a7a2de6dbe07bff7febf2adc10fbf8c62d Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Wed, 14 Aug 2024 23:04:40 -0500 Subject: [PATCH 07/27] Add callback functionality --- pipedream_solver/callbacks.py | 22 +++++ pipedream_solver/convergence.py | 70 ++++++++++++++ pipedream_solver/diagnostics.py | 157 ++++++++++++++++++++++++++++++++ pipedream_solver/nsuperlink.py | 3 +- pipedream_solver/superlink.py | 22 +++++ 5 files changed, 273 insertions(+), 1 deletion(-) create mode 100644 pipedream_solver/callbacks.py create mode 100644 pipedream_solver/convergence.py create mode 100644 pipedream_solver/diagnostics.py 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/diagnostics.py b/pipedream_solver/diagnostics.py new file mode 100644 index 0000000..b81bd10 --- /dev/null +++ b/pipedream_solver/diagnostics.py @@ -0,0 +1,157 @@ +import numpy as np +from pipedream_solver._nsuperlink import numba_compute_functional_storage_volumes, numba_compute_tabular_storage_volumes +from pipedream_solver.callbacks import BaseCallback + +class ErrorTracker(BaseCallback): + def __init__(self, model): + self.model = model + self.error = None + + def __on_step_end__(self): + model = self.model + dt = model._dt + error = compute_error(model, dt) + self.error = error + + +def continuity_error_j(model, dt): + 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 + 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[bc] = 0. + return error + +def momentum_error_uk(model, dt): + error = np.zeros(model.NK) + raise NotImplementedError + +def momentum_error_dk(model, dt): + error = np.zeros(model.NK) + raise NotImplementedError + +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): + 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 momentum_error_ik(model, dt): + error = np.zeros(model._i.size) + _i_is_start = np.zeros(model._i.size, dtype=np.bool_) + _i_is_start[model._i_1k] = True + _i_is_end = np.zeros(model._i.size, dtype=np.bool_) + _i_is_end[model._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 + 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 + +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 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 diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index a0baf67..e4dc3c6 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -1558,7 +1558,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 diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 45c1c25..c4e294f 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -8,6 +8,7 @@ import pipedream_solver.geometry import pipedream_solver.storage import pipedream_solver.visualization +from pipedream_solver.callbacks import BaseCallback class SuperLink(): """ @@ -293,6 +294,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 +422,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) @@ -3934,6 +3937,8 @@ def save_state(self): 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): @@ -3959,6 +3964,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): @@ -4020,12 +4027,21 @@ 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 + else: + self._dt = dt self._H_bc = H_bc self._Q_in = Q_in self._Q_0Ik = Q_0Ik @@ -4126,6 +4142,8 @@ 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__() 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, @@ -4157,3 +4175,7 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d break H_j_next = np.copy(self.H_j) self.iter_elapsed = iter_elapsed + 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) From dba305da3e91790a5da5ef457ddbae8d857f2617 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Tue, 21 Jan 2025 13:16:00 -0600 Subject: [PATCH 08/27] Ensure storage areas are not negative --- pipedream_solver/_nsuperlink.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index ba0960a..1256607 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -276,7 +276,7 @@ def numba_compute_functional_storage_areas(h, A, a, b, c, _functional): if h[j] < 0: A[j] = 0 else: - A[j] = a[j] * (h[j]**b[j]) + c[j] + A[j] = max(a[j] * (h[j]**b[j]) + c[j], MIN_SJ_AREA) return A @njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], boolean[:]), @@ -288,6 +288,7 @@ def numba_compute_functional_storage_volumes(h, V, a, b, c, _functional): 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 From a2ea123c784c5e5cf276a3d19d4031dd9b0663fc Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Tue, 21 Jan 2025 13:17:52 -0600 Subject: [PATCH 09/27] Ensure storage areas are never negative (include constant) --- pipedream_solver/_nsuperlink.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 1256607..0541123 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -3,6 +3,8 @@ from numba.types import float64, int64, uint32, uint16, uint8, boolean, UniTuple, Tuple, List, DictType, void import pipedream_solver.ngeometry +MIN_SJ_AREA = 1e-8 + CIRCULAR = 1 RECT_CLOSED = 2 RECT_OPEN = 3 From dab8841ab1bb134de72bbfdf6569a4c90f66508d Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Thu, 6 Feb 2025 19:59:46 -0600 Subject: [PATCH 10/27] Fix Darcy-Weisbach friction formulation for low-flow conditions --- pipedream_solver/_nsuperlink.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 0541123..d6be5cc 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -369,7 +369,14 @@ def friction_slope(Q_ik_t, dx_ik, A_ik, R_ik, n_ik, Sf_method_ik, g=9.81): # 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 + 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: From b4bcf2203eb717e875d28aca9e133c90cce09e05 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 28 Feb 2025 15:57:44 -0600 Subject: [PATCH 11/27] Fix tabular volume calculation --- pipedream_solver/nsuperlink.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index e4dc3c6..7d0a46d 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -451,7 +451,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) From 7afbe0d21527e2e620af27807fb1fa64cc982da0 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Sun, 6 Jul 2025 04:22:44 -0500 Subject: [PATCH 12/27] Fix error in boundary friction slope --- pipedream_solver/_nsuperlink.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index d6be5cc..ae73e3f 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -650,7 +650,7 @@ def kappa_uk(Q_uk, dx_uk, A_uk, C_uk, R_uk, n_uk, Sf_method_uk, dt, g=9.81): 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], + 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 From 7af88fe29d93e8bcad9fb4bd0af3be9aa9ecec2b Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Sun, 6 Jul 2025 04:23:23 -0500 Subject: [PATCH 13/27] Add diagnostics --- pipedream_solver/diagnostics.py | 188 ++++++++++++++++++++++++++++++-- pipedream_solver/nsuperlink.py | 6 +- pipedream_solver/superlink.py | 105 +++++++++++++----- 3 files changed, 259 insertions(+), 40 deletions(-) diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index b81bd10..8ea512a 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -5,13 +5,116 @@ class ErrorTracker(BaseCallback): def __init__(self, model): self.model = model - self.error = None + 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) def __on_step_end__(self): model = self.model + # TODO: This could cause problems, need to make sure this stays updated at each step dt = model._dt - error = compute_error(model, dt) - self.error = error + 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(model, dt) + self.momentum_error_dk = momentum_error_dk(model, dt) + #self.momentum_error_o = + #self.momentum_error_w = + #self.momentum_error_p = + #self.error = error + +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) + + def __on_step_end__(self): + 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): + self.model = model + self.prior_guess = 0. + self.next_guess = 0. + self.rtol = 1e-5 + self.atol = 1e-8 + + def _convergence_met(self, prior_guess, next_guess, rtol=None, atol=None): + if rtol is None: + rtol = self.rtol + if atol is None: + atol = self.atol + e = np.abs(next_guess - prior_guess) + ewt = rtol * np.abs(next_guess) + atol + valid = (e > 0.) & (ewt > 0.) + condition = (np.log(e[valid]) - np.log(ewt[valid])).max() <= 0. + return condition + + def _compute_prior_guess(self): + #return np.copy(self.model.H_j) + return self.model.state_vector + + def _compute_next_guess(self): + #return np.copy(self.model.H_j) + return self.model.state_vector + + 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.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, rtol=rtol, atol=atol) + 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.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) + self.next_guess = self._compute_next_guess() + iter_elapsed += 1 + convergence_met = self._convergence_met(self.prior_guess, self.next_guess, rtol=rtol, atol=atol) + if convergence_met: + break + self.model.iter_elapsed = iter_elapsed def continuity_error_j(model, dt): @@ -20,6 +123,7 @@ def continuity_error_j(model, dt): 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) @@ -33,16 +137,69 @@ def continuity_error_j(model, dt): np.add.at(error, model._J_up, model._Qp) np.subtract.at(error, model._J_dp, model._Qp) error -= Q_in - error[bc] = 0. + error -= Q_bc + #error[bc] = 0. return error def momentum_error_uk(model, dt): error = np.zeros(model.NK) - raise NotImplementedError + 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_uk_next = model.Q_uk + Q_uk_prev = model.states['Q_uk'] + h_uk_next = model._h_uk + H_juk_next = model.H_j[model._J_uk] + z_inv_uk = model._z_inv_uk + theta_uk = model._theta_uk + error += (Q_uk_next - Q_uk_prev) * model._dx_uk / dt + error += g * model._A_uk * (h_uk_next - theta_uk * (H_juk_next - z_inv_uk)) + error -= g * model._A_uk * model._dx_uk * model._S_o_uk + # Friction and local losses + #error += 0. + return error def momentum_error_dk(model, dt): error = np.zeros(model.NK) - raise NotImplementedError + 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_dk_prev = model.states['Q_dk'] + h_dk_next = model._h_dk + H_jdk_next = model.H_j[model._J_dk] + z_inv_dk = model._z_inv_dk + theta_dk = model._theta_dk + error += (Q_dk_next - Q_dk_prev) * model._dx_dk / dt + error += g * model._A_dk * (theta_dk * (H_jdk_next - z_inv_dk) - h_dk_next) + error -= g * model._A_dk * model._dx_dk * model._S_o_dk + # Friction and local losses + #error += 0. + return error def momentum_error_o(model, dt): error = np.zeros(model.n_o) @@ -73,10 +230,12 @@ def continuity_error_Ik(model, dt): 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[model._i_1k] = True + _i_is_start[_i_1k] = True _i_is_end = np.zeros(model._i.size, dtype=np.bool_) - _i_is_end[model._i_nk] = True + _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] @@ -89,6 +248,9 @@ def momentum_error_ik(model, dt): 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 + 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 @@ -144,6 +306,16 @@ 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) diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index 7d0a46d..bc32e9d 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -1131,7 +1131,7 @@ 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, + def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, _Q_bc=None, u=None, _dt=None, implicit=True, first_time=False): """ Construct sparse matrices A, O, W, P and b. @@ -1215,6 +1215,8 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, impli # If no flow input specified, assume zero external inflow if _Q_0j is None: _Q_0j = 0 + if _Q_bc is None: + _Q_bc = 0 # If no control input signal specified assume zero input if u is None: u = 0 @@ -1282,7 +1284,7 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, impli numba_add_at(D, _J_dp, _chi_dp) b.fill(0) # TODO: Which A_sj? Might need to apply product rule here. - b = (_A_sj * H_j_prev / _dt) + _Q_0j + D + b = (_A_sj * H_j_prev / _dt) + _Q_0j + _Q_bc + D # Ensure boundary condition is specified b[bc] = H_bc[bc] # Export instance variables diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index c4e294f..8d73b9f 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -9,6 +9,7 @@ import pipedream_solver.storage import pipedream_solver.visualization from pipedream_solver.callbacks import BaseCallback +from pipedream_solver.diagnostics import ConvergenceTracker class SuperLink(): """ @@ -718,6 +719,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) @@ -771,6 +774,8 @@ def __init__(self, superlinks, superjunctions, # Compute bandwidth self._compute_bandwidth() # Initialize to stable state + self.convergence_tracker = ConvergenceTracker(self) + self.bind_callback(self.convergence_tracker, 'convergence_tracker') self.step(dt=1e-6, first_time=True) # Reset iteration counter self.iter_count = 0 @@ -928,6 +933,12 @@ def x_Ik(self): def x_Ik(self, value): self._x_Ik = np.asarray(value) + @property + def state_vector(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 adjacency_matrix(self, J_u=None, J_d=None, symmetric=True): M = self.M @@ -2673,7 +2684,7 @@ 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, + def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, _Q_bc=None, u=None, _dt=None, implicit=True, first_time=False): """ Construct sparse matrices A, O, W, P and b. @@ -3537,6 +3548,45 @@ 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. + def exit_conditions(self): """ Determine which superlinks have exit depths below the pipe crown elevation. @@ -4034,7 +4084,7 @@ def bind_callback(self, callback, key): 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, + def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, Q_bc=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() @@ -4042,9 +4092,19 @@ def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p= dt = self._dt 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_bc is None: + self._Q_bc = np.zeros(self.M, dtype=np.float64) + else: + self._Q_bc = Q_bc + 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 - self._Q_in = Q_in - self._Q_0Ik = Q_0Ik if not implicit: raise NotImplementedError # Compute all hydraulic geometries @@ -4071,7 +4131,7 @@ def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p= self.weir_flow_coefficients(u=u_w) if self.pumps is not None: self.pump_flow_coefficients(u=u_p) - self.sparse_matrix_equations(H_bc=H_bc, _Q_0j=Q_in, + self.sparse_matrix_equations(H_bc=H_bc, _Q_0j=Q_in, _Q_bc=Q_bc, first_time=first_time, _dt=dt, implicit=implicit) @@ -4104,12 +4164,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. @@ -4143,7 +4204,10 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d (Deprecated) """ for _, callback in self.callbacks.items(): - callback.__on_step_start__() + 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, @@ -4152,30 +4216,11 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d 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 + 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, - 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 + 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) + first_iter=first_iter, num_iter=num_iter, rtol=rtol, atol=atol, + head_tol=head_tol) From 04b94a9f7828cb349162ca2d274fbfbb6a8c3781 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Mon, 7 Jul 2025 11:36:51 -0500 Subject: [PATCH 14/27] Default to old convergence behavior --- pipedream_solver/diagnostics.py | 57 +++++++++++++++++++++++++++++---- pipedream_solver/simulation.py | 7 +++- pipedream_solver/superlink.py | 4 +-- 3 files changed, 59 insertions(+), 9 deletions(-) diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index 8ea512a..cd05a9d 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -14,7 +14,7 @@ def __init__(self, model): self.momentum_error_w = np.zeros(model.n_w) self.momentum_error_p = np.zeros(model.n_p) - def __on_step_end__(self): + def __on_step_end__(self, *args, **kwargs): model = self.model # TODO: This could cause problems, need to make sure this stays updated at each step dt = model._dt @@ -41,7 +41,7 @@ def __init__(self, model): self.cumulative_vol_flux_j = np.zeros(model.M) self.cumulative_vol_flux_Ik = np.zeros(model._I.size) - def __on_step_end__(self): + def __on_step_end__(self, *args, **kwargs): model = self.model dt = model._dt Q_in = model._Q_in @@ -58,12 +58,12 @@ def __on_step_end__(self): self.cumulative_vol_flux_Ik += self.volume_flux_Ik class ConvergenceTracker(BaseCallback): - def __init__(self, model): + def __init__(self, model, rtol=1e-5, atol=1e-8): self.model = model self.prior_guess = 0. self.next_guess = 0. - self.rtol = 1e-5 - self.atol = 1e-8 + self.rtol = rtol + self.atol = atol def _convergence_met(self, prior_guess, next_guess, rtol=None, atol=None): if rtol is None: @@ -71,7 +71,7 @@ def _convergence_met(self, prior_guess, next_guess, rtol=None, atol=None): if atol is None: atol = self.atol e = np.abs(next_guess - prior_guess) - ewt = rtol * np.abs(next_guess) + atol + ewt = np.maximum(rtol * np.maximum(np.abs(prior_guess), np.abs(next_guess)), atol) valid = (e > 0.) & (ewt > 0.) condition = (np.log(e[valid]) - np.log(ewt[valid])).max() <= 0. return condition @@ -116,6 +116,51 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, break self.model.iter_elapsed = iter_elapsed +class LegacyConvergenceTracker(ConvergenceTracker): + def __init__(self, model): + self.model = model + self.prior_guess = 0. + self.next_guess = 0. + + 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 np.copy(self.model.H_j) + return self.model.H_j + + def _compute_next_guess(self): + #return np.copy(self.model.H_j) + return self.model.H_j + + 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.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) + self.next_guess = self._compute_next_guess() + iter_elapsed += 1 + convergence_met = self._convergence_met(self.prior_guess, self.next_guess, head_tol=head_tol) + if convergence_met: + break + self.model.iter_elapsed = iter_elapsed def continuity_error_j(model, dt): error = np.zeros(model.M) diff --git a/pipedream_solver/simulation.py b/pipedream_solver/simulation.py index a0af8b1..d36d85f 100644 --- a/pipedream_solver/simulation.py +++ b/pipedream_solver/simulation.py @@ -573,7 +573,7 @@ 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 _step(self, dt=None, **kwargs): + def _interpolate_inputs(self, dt=None, **kwargs): # Specify current timestamps t_next = self.t + dt # Import inputs @@ -611,11 +611,16 @@ def _step(self, dt=None, **kwargs): 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): # Infer if system is banded if not 'banded' in kwargs: banded = self.model.banded else: banded = kwargs.pop('banded') + # Get next sample + Q_in_next, H_bc_next, Q_Ik_next = self._interpolate_inputs(dt=dt, **kwargs) # Step model forward with stepsize dt self.Q_in_next = Q_in_next self.H_bc_next = H_bc_next diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 8d73b9f..fc54441 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -9,7 +9,7 @@ import pipedream_solver.storage import pipedream_solver.visualization from pipedream_solver.callbacks import BaseCallback -from pipedream_solver.diagnostics import ConvergenceTracker +from pipedream_solver.diagnostics import LegacyConvergenceTracker class SuperLink(): """ @@ -774,7 +774,7 @@ def __init__(self, superlinks, superjunctions, # Compute bandwidth self._compute_bandwidth() # Initialize to stable state - self.convergence_tracker = ConvergenceTracker(self) + self.convergence_tracker = LegacyConvergenceTracker(self) self.bind_callback(self.convergence_tracker, 'convergence_tracker') self.step(dt=1e-6, first_time=True) # Reset iteration counter From 65bdfdb63e3099303d0de2c972c7bf99ab5d3b60 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Wed, 9 Jul 2025 03:01:01 -0500 Subject: [PATCH 15/27] Add more diagnostic tooling; work on min depth constraint --- pipedream_solver/_nsuperlink.py | 34 +++++---- pipedream_solver/diagnostics.py | 131 ++++++++++++++++++++++++++++---- pipedream_solver/ngeometry.py | 58 +++++++------- pipedream_solver/nsuperlink.py | 11 +-- pipedream_solver/simulation.py | 39 +++++++++- pipedream_solver/superlink.py | 22 ++++-- 6 files changed, 223 insertions(+), 72 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index ae73e3f..384f474 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -3,6 +3,8 @@ 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 @@ -410,7 +412,7 @@ def numba_b_ik(dx_ik, dt, n_ik, Q_ik_t, A_ik, R_ik, Compute link coefficient 'b' for link i, superlink k. """ # TODO: Clean up - t_0 = (dx_ik / dt) * sigma_ik + t_0 = (dx_ik / dt) t_1 = np.zeros(Q_ik_t.size) k = len(Sf_method_ik) for n in range(k): @@ -419,8 +421,8 @@ def numba_b_ik(dx_ik, dt, n_ik, Q_ik_t, A_ik, R_ik, 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 + 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), @@ -491,18 +493,20 @@ def numba_node_coeffs(_D_Ik, _E_Ik, _Q_0Ik, _B_ik, _h_Ik, _dx_ik, _A_SIk, @njit(float64(float64, float64), cache=True) def safe_divide(num, den): - if (den == 0): - return 0 - else: - return num / den + #if (den == 0): + # return 0 + #else: + # return num / den + return (num + SMALLEST_NORMAL) / (den + SMALLEST_NORMAL) @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 + #result = np.zeros_like(num) + #cond = (den != 0) + #result[cond] = num[cond] / den[cond] + #return result + return (num + SMALLEST_NORMAL) / (den + SMALLEST_NORMAL) @njit(float64(float64, float64, float64, float64, float64), cache=True) @@ -556,10 +560,10 @@ def numba_solve_internals(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik, 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 _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]) diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index cd05a9d..dbdb1ab 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -5,6 +5,7 @@ class ErrorTracker(BaseCallback): def __init__(self, model): 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) @@ -13,8 +14,40 @@ def __init__(self, model): 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) - def __on_step_end__(self, *args, **kwargs): + @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]) + + 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 @@ -26,7 +59,21 @@ def __on_step_end__(self, *args, **kwargs): #self.momentum_error_o = #self.momentum_error_w = #self.momentum_error_p = - #self.error = error + + def _compute_magnitudes(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_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) + #self._record_error(*args, **kwargs) class VolumeTracker(BaseCallback): def __init__(self, model): @@ -103,12 +150,16 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, self.model.iter_count -= 1 self.model.t -= dt self.prior_guess = self._compute_prior_guess() - 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) + 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() iter_elapsed += 1 convergence_met = self._convergence_met(self.prior_guess, self.next_guess, rtol=rtol, atol=atol) @@ -149,12 +200,16 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, self.model.iter_count -= 1 self.model.t -= dt self.prior_guess = self._compute_prior_guess() - 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) + 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() iter_elapsed += 1 convergence_met = self._convergence_met(self.prior_guess, self.next_guess, head_tol=head_tol) @@ -182,10 +237,30 @@ def continuity_error_j(model, dt): 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. + #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 @@ -216,6 +291,12 @@ def momentum_error_uk_2(model, dt): #error += 0. 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 @@ -246,6 +327,12 @@ def momentum_error_dk_2(model, dt): #error += 0. 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 @@ -273,6 +360,12 @@ def continuity_error_Ik(model, dt): 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 @@ -300,6 +393,12 @@ def momentum_error_ik(model, dt): error[_i_is_internal] += model._c_ik[_i_is_internal] * model.Q_ik[_ip1] 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) diff --git a/pipedream_solver/ngeometry.py b/pipedream_solver/ngeometry.py index ec421a5..b5aca08 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, @@ -36,8 +38,8 @@ def Circular_A_ik(h_ik, g1, g2): d = g1 pslot = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > d: y = d r = d / 2 @@ -68,8 +70,8 @@ def Circular_Pe_ik(h_ik, g1, g2): d = g1 pslot = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > d: y = d r = d / 2 @@ -122,8 +124,8 @@ def Circular_B_ik(h_ik, g1, g2): d = g1 pslot = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH r = d / 2 phi = y / r if phi < 0: @@ -160,8 +162,8 @@ def Rect_Closed_A_ik(h_ik, g1, g2): y_max = g1 b = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max A = y * b @@ -187,8 +189,8 @@ def Rect_Closed_Pe_ik(h_ik, g1, g2): y_max = g1 b = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max Pe = b + 2 * y @@ -237,8 +239,8 @@ def Rect_Closed_B_ik(h_ik, g1, g2, g3): b = g2 pslot = g3 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH cond = (y < y_max) if cond: B = b @@ -267,8 +269,8 @@ def Rect_Open_A_ik(h_ik, g1, g2): y_max = g1 b = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max A = y * b @@ -294,8 +296,8 @@ def Rect_Open_Pe_ik(h_ik, g1, g2): y_max = g1 b = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max Pe = b + 2 * y @@ -363,8 +365,8 @@ def Triangular_A_ik(h_ik, g1, g2): y_max = g1 m = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max A = m * y**2 @@ -390,8 +392,8 @@ def Triangular_Pe_ik(h_ik, g1, g2): y_max = g1 m = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max Pe = 2 * y * np.sqrt(1 + m**2) @@ -437,8 +439,8 @@ def Triangular_B_ik(h_ik, g1, g2): y_max = g1 m = g2 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH cond = (y < y_max) if cond: B = 2 * m * y @@ -470,8 +472,8 @@ def Trapezoidal_A_ik(h_ik, g1, g2, g3): b = g2 m = g3 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max A = y * (b + m * y) @@ -500,8 +502,8 @@ def Trapezoidal_Pe_ik(h_ik, g1, g2, g3): b = g2 m = g3 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH if y > y_max: y = y_max Pe = b + 2 * y * np.sqrt(1 + m**2) @@ -550,8 +552,8 @@ def Trapezoidal_B_ik(h_ik, g1, g2, g3): b = g2 m = g3 y = h_ik - if y < 0: - y = 0 + if y < MIN_DEPTH: + y = MIN_DEPTH cond = (y < y_max) if cond: B = b + 2 * m * y diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index bc32e9d..1c02f30 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -1380,8 +1380,9 @@ def solve_banded_matrix(self, u=None, implicit=True): # 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 @@ -1417,7 +1418,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 @@ -1454,7 +1455,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 @@ -1503,7 +1504,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 diff --git a/pipedream_solver/simulation.py b/pipedream_solver/simulation.py index d36d85f..30ec168 100644 --- a/pipedream_solver/simulation.py +++ b/pipedream_solver/simulation.py @@ -614,13 +614,48 @@ def _interpolate_inputs(self, dt=None, **kwargs): return Q_in_next, H_bc_next, Q_Ik_next def _step(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') # Infer if system is banded if not 'banded' in kwargs: banded = self.model.banded else: banded = kwargs.pop('banded') - # Get next sample - Q_in_next, H_bc_next, Q_Ik_next = self._interpolate_inputs(dt=dt, **kwargs) # Step model forward with stepsize dt self.Q_in_next = Q_in_next self.H_bc_next = H_bc_next diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index fc54441..02b7d29 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -4210,12 +4210,22 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d 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) + 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 + #new_dt = dt / 2 + #self.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=new_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) + first_iter = False num_iter -= 1 self.iter_elapsed = 1 From 6793902ea80552ad726a4892665b7ef1293b23b6 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Mon, 14 Jul 2025 18:01:02 -0500 Subject: [PATCH 16/27] Add boundary velocities --- pipedream_solver/_nsuperlink.py | 104 +++++++++++++++++++++++++++++--- pipedream_solver/diagnostics.py | 44 +++++++++++++- pipedream_solver/nsuperlink.py | 24 ++++++-- pipedream_solver/superlink.py | 44 ++++++++------ 4 files changed, 180 insertions(+), 36 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 384f474..7785309 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -573,6 +573,56 @@ def numba_solve_internals(_h_Ik, _Q_ik, _h_uk, _h_dk, _U_Ik, _V_Ik, _W_Ik, _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): @@ -610,13 +660,45 @@ def numba_u_ik(_Q_ik, _A_ik, _u_ik): _u_ik[i] = 0 return _u_ik -@njit(float64[:](float64[:], float64[:], boolean[:], float64[:]), +@njit(float64[:](float64[:], float64[:], float64[:]), + cache=True) +def numba_u_uk(_Q_uk, _A_uk, _u_uk): + n = _u_uk.size + for i in range(n): + _Q_u = _Q_uk[i] + _A_u = _A_uk[i] + if _A_u: + _u_uk[i] = _Q_u / _A_u + else: + _u_uk[i] = 0 + return _u_uk + +@njit(float64[:](float64[:], float64[:], float64[:]), + cache=True) +def numba_u_dk(_Q_dk, _A_dk, _u_dk): + n = _u_dk.size + for i in range(n): + _Q_d = _Q_dk[i] + _A_d = _A_dk[i] + if _A_d: + _u_dk[i] = _Q_d / _A_d + else: + _u_dk[i] = 0 + return _u_dk + +@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]), cache=True) -def numba_u_Ik(_dx_ik, _u_ik, _link_start, _u_Ik): +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_ik[i] + num = _dx_ik[i] * _u_uk[k] + _dx_uk[k] * _u_ik[i] + den = _dx_ik[i] + _dx_uk[k] + if den: + _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] @@ -624,16 +706,22 @@ def numba_u_Ik(_dx_ik, _u_ik, _link_start, _u_Ik): if den: _u_Ik[i] = num / den else: - _u_Ik[i] = 0 + _u_Ik[i] = 0. return _u_Ik -@njit(float64[:](float64[:], float64[:], boolean[:], float64[:]), +@njit(float64[:](float64[:], float64[:], float64[:], float64[:], boolean[:], int64[:], float64[:]), cache=True) -def numba_u_Ip1k(_dx_ik, _u_ik, _link_end, _u_Ip1k): +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_ik[i] + num = _dx_ik[i] * _u_dk[k] + _dx_dk[k] * _u_ik[i] + den = _dx_ik[i] + _dx_dk[k] + if den: + _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] @@ -641,7 +729,7 @@ def numba_u_Ip1k(_dx_ik, _u_ik, _link_end, _u_Ip1k): if den: _u_Ip1k[i] = num / den else: - _u_Ip1k[i] = 0 + _u_Ip1k[i] = 0. return _u_Ip1k @njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index dbdb1ab..3f66da4 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -3,7 +3,10 @@ from pipedream_solver.callbacks import BaseCallback class ErrorTracker(BaseCallback): - def __init__(self, model): + """ + 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) @@ -22,6 +25,8 @@ def __init__(self, model): 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): @@ -43,6 +48,23 @@ def magnitude(self): 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 @@ -73,7 +95,18 @@ def _compute_magnitudes(self, *args, **kwargs): def __on_step_end__(self, *args, **kwargs): self._compute_error(*args, **kwargs) self._compute_magnitudes(*args, **kwargs) - #self._record_error(*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): @@ -104,6 +137,7 @@ def __on_step_end__(self, *args, **kwargs): 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, rtol=1e-5, atol=1e-8): self.model = model @@ -218,6 +252,9 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, 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'] @@ -346,6 +383,9 @@ def momentum_error_p(model, dt): 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'] diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index 1c02f30..4946f82 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -708,18 +708,32 @@ 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 # Determine start and end nodes # Compute link velocities numba_u_ik(_Q_ik, _A_ik, _u_ik) + # Compute boundary velocities + numba_u_uk(_Q_uk, _A_uk, _u_uk) + numba_u_dk(_Q_dk, _A_dk, _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) # 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 @@ -1131,7 +1145,7 @@ 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, _Q_bc=None, u=None, _dt=None, implicit=True, + 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. @@ -1215,8 +1229,6 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, _Q_bc=None, u=None, _dt # If no flow input specified, assume zero external inflow if _Q_0j is None: _Q_0j = 0 - if _Q_bc is None: - _Q_bc = 0 # If no control input signal specified assume zero input if u is None: u = 0 @@ -1284,7 +1296,7 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, _Q_bc=None, u=None, _dt numba_add_at(D, _J_dp, _chi_dp) b.fill(0) # TODO: Which A_sj? Might need to apply product rule here. - b = (_A_sj * H_j_prev / _dt) + _Q_0j + _Q_bc + D + b = (_A_sj * H_j_prev / _dt) + _Q_0j + D # Ensure boundary condition is specified b[bc] = H_bc[bc] # Export instance variables diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 02b7d29..5bd18fa 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -734,6 +734,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.) @@ -934,11 +937,23 @@ def x_Ik(self, value): self._x_Ik = np.asarray(value) @property - def state_vector(self): + 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 @@ -2684,7 +2699,7 @@ 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, _Q_bc=None, u=None, _dt=None, implicit=True, + 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. @@ -3586,6 +3601,7 @@ def compute_boundary_flows(self, dt=None): _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): """ @@ -3975,13 +3991,10 @@ 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) @@ -4084,7 +4097,7 @@ def bind_callback(self, callback, key): def unbind_callback(self, key): return self.callbacks.pop(key) - def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, Q_bc=None, u_o=None, u_w=None, u_p=None, dt=None, + 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() @@ -4096,10 +4109,6 @@ def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, Q_bc=None, u_o=None, u_w self._Q_in = np.zeros(self.M, dtype=np.float64) else: self._Q_in = Q_in - if Q_bc is None: - self._Q_bc = np.zeros(self.M, dtype=np.float64) - else: - self._Q_bc = Q_bc if Q_0Ik is None: self._Q_0Ik = np.zeros(self._I.size, dtype=np.float64) else: @@ -4131,7 +4140,7 @@ def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, Q_bc=None, u_o=None, u_w self.weir_flow_coefficients(u=u_w) if self.pumps is not None: self.pump_flow_coefficients(u=u_p) - self.sparse_matrix_equations(H_bc=H_bc, _Q_0j=Q_in, _Q_bc=Q_bc, + self.sparse_matrix_equations(H_bc=H_bc, _Q_0j=Q_in, first_time=first_time, _dt=dt, implicit=implicit) @@ -4220,11 +4229,6 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d except: self.load_state() raise - #new_dt = dt / 2 - #self.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=new_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) first_iter = False num_iter -= 1 From e438f5adc77723258dd849e216c4347b4566fd37 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Mon, 14 Jul 2025 20:48:41 -0500 Subject: [PATCH 17/27] Refactor matrix creation --- pipedream_solver/_nsuperlink.py | 57 +++--- pipedream_solver/nsuperlink.py | 333 ++++++++++++++++++-------------- pipedream_solver/superlink.py | 4 + 3 files changed, 219 insertions(+), 175 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 7785309..6aa6e73 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -1118,20 +1118,16 @@ def gamma_dk(Q_dk_t, C_dk, A_dk, g=9.81): result = safe_divide_vec(num, den) return result -@njit(float64[:](float64[:], float64[:], float64[:], float64), +@njit(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 +def xi_uk(dx_uk, B_uk, theta_uk): + result = (dx_uk * B_uk * theta_uk) / 2 return result -@njit(float64[:](float64[:], float64[:], float64[:], float64), +@njit(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 +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[:], @@ -1460,8 +1456,7 @@ def numba_create_banded(l, bandwidth, M): return AB @njit(void(float64[:], int64[:], float64[:]), - cache=True, - fastmath=True) + cache=True) def numba_add_at(a, indices, b): n = len(indices) for k in range(n): @@ -1481,39 +1476,51 @@ def numba_clear_off_diagonals(A, bc, _J_uk, _J_dk, NK): 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, - fastmath=True) -def numba_create_A_matrix(A, _F_jj, bc, _J_uk, _J_dk, _alpha_uk, + 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) - numba_add_at(_F_jj, _J_dk, _xi_dk) - _F_jj += (_A_sj / _dt) - # Set diagonal of A matrix + 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]: - A[i,i] = 1.0 + K[i,i] = 0.0 else: - A[i,i] = _F_jj[i] + 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: - A[_J_u, _J_d] += _beta_uk[k] + K[_J_u, _J_d] += _beta_uk[k] if not _bc_d: - A[_J_d, _J_u] -= _alpha_dk[k] + K[_J_d, _J_u] -= _alpha_dk[k] @njit(void(float64[:, :], float64[:], boolean[:], int64[:], int64[:], float64[:], float64[:], float64[:], float64[:], int64, int64), - cache=True, - fastmath=True) + 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 diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index 4946f82..1c9ddac 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -1145,13 +1145,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 @@ -1160,10 +1167,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 @@ -1173,53 +1176,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 @@ -1232,80 +1329,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.fill(0.) + b.fill(0.) + 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() @@ -1327,25 +1387,10 @@ 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 = scipy.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 @@ -1372,20 +1417,8 @@ 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 diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 5bd18fa..8c87f26 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -693,11 +693,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)) From c5ea1064fe70865c28e18561d81ae8df7a0ccd1d Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Mon, 14 Jul 2025 22:41:02 -0500 Subject: [PATCH 18/27] Add computation of momentum at boundaries --- pipedream_solver/_nsuperlink.py | 2 +- pipedream_solver/nsuperlink.py | 54 ++++++++++++++++++++++++++++++++- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 6aa6e73..0748017 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -420,7 +420,7 @@ def numba_b_ik(dx_ik, dt, n_ik, Q_ik_t, A_ik, R_ik, 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_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 diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index 1c9ddac..b9acbb5 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -755,9 +755,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 @@ -768,11 +794,37 @@ def link_coeffs(self, _dt=None, first_iter=True): _C_ik, _a_ik, _c_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) + # 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) + _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) + _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 + # 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) + _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) + _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 # 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): """ From ec9f4f5cb7014d99f44add6d6d091e543b5c44af Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Tue, 15 Jul 2025 03:13:53 -0500 Subject: [PATCH 19/27] Downstream boundary problem fixed --- pipedream_solver/_nsuperlink.py | 22 ++--- pipedream_solver/diagnostics.py | 35 +++---- pipedream_solver/nsuperlink.py | 159 ++++++++++++++++++++++++++++---- pipedream_solver/superlink.py | 15 ++- 4 files changed, 179 insertions(+), 52 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 0748017..80bca61 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -734,7 +734,7 @@ def numba_u_Ip1k(_dx_ik, _u_ik, _dx_dk, _u_dk, _link_end, _ki, _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): +def kappa_uk_old(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. """ @@ -749,7 +749,7 @@ def kappa_uk(Q_uk, dx_uk, A_uk, C_uk, R_uk, n_uk, Sf_method_uk, dt, g=9.81): @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): +def kappa_dk_old(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. """ @@ -764,7 +764,7 @@ def kappa_dk(Q_dk, dx_dk, A_dk, C_dk, R_dk, n_dk, Sf_method_dk, dt, g=9.81): @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): +def mu_uk_old(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. """ @@ -775,7 +775,7 @@ def mu_uk(Q_uk_t, dx_uk, A_uk, theta_uk, z_inv_uk, S_o_uk, dt, g=9.81): @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): +def mu_dk_old(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. """ @@ -974,7 +974,7 @@ def O_ik(a_ik, b_ik, c_ik, A_ik, E_Ip1k, X_Ip1k, g=9.81): @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): +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'. """ @@ -985,7 +985,7 @@ def numba_D_k_star(X_1k, kappa_uk, U_Nk, kappa_dk, Z_1k, W_Nk): @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): +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. @@ -997,7 +997,7 @@ def numba_alpha_uk(U_Nk, kappa_dk, X_1k, Z_1k, W_Nk, D_k_star, lambda_uk): @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): +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. @@ -1010,7 +1010,7 @@ def numba_beta_uk(U_Nk, kappa_dk, Z_1k, W_Nk, D_k_star, lambda_dk): @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, +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 @@ -1025,7 +1025,7 @@ def numba_chi_uk(U_Nk, kappa_dk, Y_1k, X_1k, mu_uk, Z_1k, @njit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:]), cache=True) -def numba_alpha_dk(X_1k, kappa_uk, W_Nk, D_k_star, lambda_uk): +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. @@ -1037,7 +1037,7 @@ def numba_alpha_dk(X_1k, kappa_uk, W_Nk, D_k_star, lambda_uk): @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): +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. @@ -1050,7 +1050,7 @@ def numba_beta_dk(X_1k, kappa_uk, U_Nk, W_Nk, Z_1k, D_k_star, lambda_dk): @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, +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 diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index 3f66da4..08a7339 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -76,8 +76,8 @@ def _compute_error(self, *args, **kwargs): 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(model, dt) - self.momentum_error_dk = momentum_error_dk(model, dt) + self.momentum_error_uk = momentum_error_uk_2(model, dt) + self.momentum_error_dk = momentum_error_dk_2(model, dt) #self.momentum_error_o = #self.momentum_error_w = #self.momentum_error_p = @@ -315,17 +315,17 @@ def momentum_error_uk(model, dt): 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 - Q_uk_prev = model.states['Q_uk'] h_uk_next = model._h_uk H_juk_next = model.H_j[model._J_uk] - z_inv_uk = model._z_inv_uk - theta_uk = model._theta_uk - error += (Q_uk_next - Q_uk_prev) * model._dx_uk / dt - error += g * model._A_uk * (h_uk_next - theta_uk * (H_juk_next - z_inv_uk)) - error -= g * model._A_uk * model._dx_uk * model._S_o_uk - # Friction and local losses - #error += 0. + 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 def momentum_magnitude_uk(model, dt): @@ -352,16 +352,17 @@ def momentum_error_dk_2(model, dt): error = np.zeros(model.NK) g = 9.81 Q_dk_next = model.Q_dk - Q_dk_prev = model.states['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] - z_inv_dk = model._z_inv_dk - theta_dk = model._theta_dk - error += (Q_dk_next - Q_dk_prev) * model._dx_dk / dt - error += g * model._A_dk * (theta_dk * (H_jdk_next - z_inv_dk) - h_dk_next) - error -= g * model._A_dk * model._dx_dk * model._S_o_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 - #error += 0. + 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 def momentum_magnitude_dk(model, dt): diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index b9acbb5..0eed195 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -915,7 +915,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. """ @@ -961,16 +1086,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. """ @@ -1015,15 +1141,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. @@ -1070,25 +1197,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) @@ -1382,8 +1509,8 @@ def sparse_matrix_equations(self, H_bc=None, _Q_0j=None, u=None, _dt=None, impli if u is None: u = 0 # Create matrices - A.fill(0.) - b.fill(0.) + 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 @@ -1442,13 +1569,13 @@ def solve_sparse_matrix(self, u=None, implicit=True): if _sparse: H_j_next = scipy.sparse.linalg.spsolve(A, b) else: - H_j_next = scipy.linalg.solve(A, b) + 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 diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 8c87f26..934eb22 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -3175,8 +3175,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 @@ -3185,11 +3183,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 @@ -4135,9 +4132,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: From 385140636f63900ee2bba37fb4fca1668c6a80da Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 18 Jul 2025 21:56:54 -0500 Subject: [PATCH 20/27] More stable --- pipedream_solver/_nsuperlink.py | 169 +++++++++++--------------------- pipedream_solver/diagnostics.py | 9 +- pipedream_solver/ngeometry.py | 1 + pipedream_solver/nsuperlink.py | 81 ++++++++++++++- pipedream_solver/superlink.py | 23 +++++ 5 files changed, 164 insertions(+), 119 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 80bca61..05d13d9 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -105,7 +105,7 @@ def numba_boundary_geometry(_A_bk, _Pe_bk, _R_bk, _B_bk, _h_Ik, _H_j, _z_inv_bk, i = _i_bk[k] I = _I_bk[k] j = _J_bk[k] - # TODO: This should incorporate theta + # 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 @@ -425,6 +425,27 @@ def numba_b_ik(dx_ik, dt, n_ik, Q_ik_t, A_ik, R_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): @@ -499,14 +520,15 @@ def safe_divide(num, den): # 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) - #cond = (den != 0) - #result[cond] = num[cond] / den[cond] - #return result - return (num + SMALLEST_NORMAL) / (den + SMALLEST_NORMAL) + 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) @@ -654,38 +676,12 @@ def numba_u_ik(_Q_ik, _A_ik, _u_ik): for i in range(n): _Q_i = _Q_ik[i] _A_i = _A_ik[i] - if _A_i: + if _A_i > SMALLEST_NORMAL: _u_ik[i] = _Q_i / _A_i else: - _u_ik[i] = 0 + _u_ik[i] = 0. return _u_ik -@njit(float64[:](float64[:], float64[:], float64[:]), - cache=True) -def numba_u_uk(_Q_uk, _A_uk, _u_uk): - n = _u_uk.size - for i in range(n): - _Q_u = _Q_uk[i] - _A_u = _A_uk[i] - if _A_u: - _u_uk[i] = _Q_u / _A_u - else: - _u_uk[i] = 0 - return _u_uk - -@njit(float64[:](float64[:], float64[:], float64[:]), - cache=True) -def numba_u_dk(_Q_dk, _A_dk, _u_dk): - n = _u_dk.size - for i in range(n): - _Q_d = _Q_dk[i] - _A_d = _A_dk[i] - if _A_d: - _u_dk[i] = _Q_d / _A_d - else: - _u_dk[i] = 0 - return _u_dk - @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): @@ -695,7 +691,7 @@ def numba_u_Ik(_dx_ik, _u_ik, _dx_uk, _u_uk, _link_start, _ki, _u_Ik): 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: + if den > SMALLEST_NORMAL: _u_Ik[i] = num / den else: _u_Ik[i] = 0. @@ -703,7 +699,7 @@ def numba_u_Ik(_dx_ik, _u_ik, _dx_uk, _u_uk, _link_start, _ki, _u_Ik): 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: + if den > SMALLEST_NORMAL: _u_Ik[i] = num / den else: _u_Ik[i] = 0. @@ -718,7 +714,7 @@ def numba_u_Ip1k(_dx_ik, _u_ik, _dx_dk, _u_dk, _link_end, _ki, _u_Ip1k): 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: + if den > SMALLEST_NORMAL: _u_Ip1k[i] = num / den else: _u_Ip1k[i] = 0. @@ -726,63 +722,38 @@ def numba_u_Ip1k(_dx_ik, _u_ik, _dx_dk, _u_dk, _link_end, _ki, _u_Ip1k): 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: + if den > SMALLEST_NORMAL: _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_old(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[:], 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[:], float64[:], - int64[:], float64, float64), cache=True) -def kappa_dk_old(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_old(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_old(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) @@ -1096,28 +1067,6 @@ def gamma_p(Q_p_t, b_p, c_p, u): 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[:]), cache=True) def xi_uk(dx_uk, B_uk, theta_uk): diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index 08a7339..d5d468a 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -428,10 +428,11 @@ def momentum_error_ik(model, dt): 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 - 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] + 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 def momentum_magnitude_ik(model, dt): diff --git a/pipedream_solver/ngeometry.py b/pipedream_solver/ngeometry.py index b5aca08..edbd2b6 100644 --- a/pipedream_solver/ngeometry.py +++ b/pipedream_solver/ngeometry.py @@ -139,6 +139,7 @@ def Circular_B_ik(h_ik, g1, g2): else: # TODO: Use absolute value instead of fraction? B = pslot * d + B = max(pslot * d, B) return B diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index 0eed195..bf46fb0 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -353,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) @@ -722,26 +724,43 @@ def node_velocities(self): _link_end = self._link_end # 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, _u_ik) # Compute boundary velocities - numba_u_uk(_Q_uk, _A_uk, _u_uk) - numba_u_dk(_Q_dk, _A_dk, _u_dk) + _u_uk = numba_u_ik(_Q_uk, _A_uk, _u_uk) + _u_dk = numba_u_ik(_Q_dk, _A_dk, _u_dk) # Compute velocities for start nodes (1 -> Nk) 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, _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 @@ -792,26 +811,42 @@ 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_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_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 @@ -1759,6 +1794,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. diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 934eb22..a5572e0 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -1965,6 +1965,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. @@ -4121,6 +4143,7 @@ def _setup_step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p= 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() From 5efbf7b3e7c995313595f28dba230631e7985d79 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Tue, 29 Jul 2025 05:22:46 -0500 Subject: [PATCH 21/27] Add quotient rule formulation of SVE --- pipedream_solver/_tsuperlink.py | 1407 +++++++++++++++++++++++++++++++ pipedream_solver/diagnostics.py | 68 +- pipedream_solver/superlink.py | 12 + pipedream_solver/tsuperlink.py | 228 +++++ 4 files changed, 1709 insertions(+), 6 deletions(-) create mode 100644 pipedream_solver/_tsuperlink.py create mode 100644 pipedream_solver/tsuperlink.py diff --git a/pipedream_solver/_tsuperlink.py b/pipedream_solver/_tsuperlink.py new file mode 100644 index 0000000..57a1d33 --- /dev/null +++ b/pipedream_solver/_tsuperlink.py @@ -0,0 +1,1407 @@ +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_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 + +@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/diagnostics.py b/pipedream_solver/diagnostics.py index d5d468a..0850c9a 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -78,13 +78,17 @@ def _compute_error(self, *args, **kwargs): 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 could cause problems, need to make sure this stays updated at each step + # 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) @@ -139,12 +143,13 @@ def __on_step_end__(self, *args, **kwargs): class ConvergenceTracker(BaseCallback): - def __init__(self, model, rtol=1e-5, atol=1e-8): + def __init__(self, model, rtol=1e-5, atol=1e-8, learning_rate=0.5): self.model = model self.prior_guess = 0. self.next_guess = 0. self.rtol = rtol self.atol = atol + self.learning_rate = learning_rate def _convergence_met(self, prior_guess, next_guess, rtol=None, atol=None): if rtol is None: @@ -158,11 +163,9 @@ def _convergence_met(self, prior_guess, next_guess, rtol=None, atol=None): return condition def _compute_prior_guess(self): - #return np.copy(self.model.H_j) return self.model.state_vector def _compute_next_guess(self): - #return np.copy(self.model.H_j) return self.model.state_vector 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, @@ -201,11 +204,13 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, break self.model.iter_elapsed = iter_elapsed + class LegacyConvergenceTracker(ConvergenceTracker): def __init__(self, model): self.model = model self.prior_guess = 0. self.next_guess = 0. + self.learning_rate = 0.5 def _convergence_met(self, prior_guess, next_guess, head_tol=0.0015): e = np.abs(next_guess - prior_guess) @@ -213,13 +218,20 @@ def _convergence_met(self, prior_guess, next_guess, head_tol=0.0015): return condition def _compute_prior_guess(self): - #return np.copy(self.model.H_j) return self.model.H_j def _compute_next_guess(self): - #return np.copy(self.model.H_j) 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): @@ -234,6 +246,7 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, self.model.iter_count -= 1 self.model.t -= dt self.prior_guess = self._compute_prior_guess() + self.prior_states = self.model.return_state() 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, @@ -245,8 +258,10 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, 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 @@ -328,6 +343,21 @@ def momentum_error_uk_2(model, dt): 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'] @@ -365,6 +395,21 @@ def momentum_error_dk_2(model, dt): 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'] @@ -435,6 +480,17 @@ def momentum_error_ik(model, dt): 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'] diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index a5572e0..0da54c2 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -4003,6 +4003,18 @@ def state_space_system(self, _dt=None): 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): """ Save current model state to dict stored in self.states. diff --git a/pipedream_solver/tsuperlink.py b/pipedream_solver/tsuperlink.py new file mode 100644 index 0000000..0d8db75 --- /dev/null +++ b/pipedream_solver/tsuperlink.py @@ -0,0 +1,228 @@ +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 + +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 \ No newline at end of file From 4e78c8448e2af373e2794e2ada9e4ebded55222c Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 15 Aug 2025 01:00:25 -0500 Subject: [PATCH 22/27] Add velocity limiter --- pipedream_solver/_nsuperlink.py | 23 ++++++-- pipedream_solver/nsuperlink.py | 100 +++++++++++++++++++++++++++++++- pipedream_solver/superlink.py | 15 +++++ 3 files changed, 129 insertions(+), 9 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 05d13d9..0e0da69 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -669,19 +669,30 @@ def numba_solve_internals_ls(_h_Ik, NK, nk, _k_1k, _i_1k, _I_1k, _U, _X, _b): _h_Ik[jstart+1:jstart+nlinks] = _h_inner return _h_Ik -@njit(float64[:](float64[:], float64[:], float64[:]), +@njit(float64[:](float64[:], float64[:], float64[:], float64[:]), cache=True) -def numba_u_ik(_Q_ik, _A_ik, _u_ik): +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] - if _A_i > SMALLEST_NORMAL: - _u_ik[i] = _Q_i / _A_i - else: - _u_ik[i] = 0. + _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): diff --git a/pipedream_solver/nsuperlink.py b/pipedream_solver/nsuperlink.py index bf46fb0..6afdaf0 100644 --- a/pipedream_solver/nsuperlink.py +++ b/pipedream_solver/nsuperlink.py @@ -482,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. @@ -722,12 +813,15 @@ def node_velocities(self): _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 - _u_ik = 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, _u_uk) - _u_dk = numba_u_ik(_Q_dk, _A_dk, _u_dk) + _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, _dx_uk, _u_uk, _link_start, _ki, _u_Ik) # Compute velocities for end nodes (2 -> Nk+1) diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 0da54c2..d4e2ff9 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -773,6 +773,21 @@ def __init__(self, superlinks, superjunctions, 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 From d999c2852a3bc7249a974e69ebcf8f77cd9b1328 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 15 Aug 2025 01:02:07 -0500 Subject: [PATCH 23/27] Change direction of recurrence based on Froude condition --- pipedream_solver/_tsuperlink.py | 97 ++++++++++++++++++++++----------- pipedream_solver/tsuperlink.py | 72 +++++++++++++++++++++++- 2 files changed, 136 insertions(+), 33 deletions(-) diff --git a/pipedream_solver/_tsuperlink.py b/pipedream_solver/_tsuperlink.py index 57a1d33..eb6d5ae 100644 --- a/pipedream_solver/_tsuperlink.py +++ b/pipedream_solver/_tsuperlink.py @@ -517,6 +517,14 @@ def Q_i_b(h_Ik, h_Np1k, X_Ik, Y_Ik, Z_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): @@ -526,44 +534,60 @@ def h_i_b(Q_ik, h_Np1k, X_Ik, Y_Ik, Z_Ik): return result @njit(int64(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], - float64[:], float64[:], float64[:], int64[:], int64[:], int64[:], int64, - float64, float64[:], boolean), + float64[:], float64[:], float64[:], int64[:], int64[:], int64[:], int64[:], int64), 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): +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 - 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]) + # 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: - # 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]) + raise ValueError('Invalid Froude Number.') return 1 @@ -640,19 +664,30 @@ def numba_solve_internals_ls(_h_Ik, NK, nk, _k_1k, _i_1k, _I_1k, _U, _X, _b): _h_Ik[jstart+1:jstart+nlinks] = _h_inner return _h_Ik -@njit(float64[:](float64[:], float64[:], float64[:]), +@njit(float64[:](float64[:], float64[:], float64[:], float64[:]), cache=True) -def numba_u_ik(_Q_ik, _A_ik, _u_ik): +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] - if _A_i > SMALLEST_NORMAL: - _u_ik[i] = _Q_i / _A_i - else: - _u_ik[i] = 0. + _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): diff --git a/pipedream_solver/tsuperlink.py b/pipedream_solver/tsuperlink.py index 0d8db75..a64537f 100644 --- a/pipedream_solver/tsuperlink.py +++ b/pipedream_solver/tsuperlink.py @@ -1,7 +1,7 @@ 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 +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, @@ -225,4 +225,72 @@ def superlink_boundary_flow_coefficients(self, _dt=None): self._chi_uk = chi_uk self._alpha_dk = alpha_dk self._beta_dk = beta_dk - self._chi_dk = chi_dk \ No newline at end of file + 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 From 645d2a9a673dadb9962e4be32148b6ead9f12ce8 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 15 Aug 2025 01:02:40 -0500 Subject: [PATCH 24/27] Add functions to volume tracker; enforce minimum depth --- pipedream_solver/diagnostics.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index 0850c9a..f2ae346 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -124,6 +124,33 @@ def __init__(self, model): 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.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) + + @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 __on_step_end__(self, *args, **kwargs): model = self.model @@ -211,6 +238,7 @@ def __init__(self, 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) @@ -247,6 +275,8 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, 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, From f3f94306f4c54493240156264dc63b7e6adb3e11 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 15 Aug 2025 03:04:29 -0500 Subject: [PATCH 25/27] Implement flux limiter in Newton-Raphson iterations --- pipedream_solver/diagnostics.py | 128 +++++++++++++++++++++++++------- 1 file changed, 101 insertions(+), 27 deletions(-) diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index f2ae346..91a773f 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -170,35 +170,95 @@ def __on_step_end__(self, *args, **kwargs): class ConvergenceTracker(BaseCallback): - def __init__(self, model, rtol=1e-5, atol=1e-8, learning_rate=0.5): + def __init__(self, model, xtol=1e-6, max_learning_rate=0.5, min_learning_rate=0.01, beta=2.): self.model = model - self.prior_guess = 0. - self.next_guess = 0. - self.rtol = rtol - self.atol = atol - self.learning_rate = learning_rate - - def _convergence_met(self, prior_guess, next_guess, rtol=None, atol=None): - if rtol is None: - rtol = self.rtol - if atol is None: - atol = self.atol - e = np.abs(next_guess - prior_guess) - ewt = np.maximum(rtol * np.maximum(np.abs(prior_guess), np.abs(next_guess)), atol) - valid = (e > 0.) & (ewt > 0.) - condition = (np.log(e[valid]) - np.log(ewt[valid])).max() <= 0. + 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.convergence_metric = np.inf + + def _compute_step_ratio(self, dx): + 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 + # Compute binding step ratio + ratio_ub = max(max((dx['Q_ik'] / Q_ik_ub).max(), 0.), + max((dx['Q_uk'] / Q_uk_ub).max(), 0.), + max((dx['Q_dk'] / Q_dk_ub).max(), 0.)) + ratio_lb = max(max(-(dx['Q_ik'] / Q_ik_lb).min(), 0.), + max(-(dx['Q_uk'] / Q_uk_lb).max(), 0.), + max(-(dx['Q_dk'] / Q_dk_lb).max(), 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): - return self.model.state_vector + 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): - return self.model.state_vector + 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.prior_guess = self._compute_prior_guess() + 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, @@ -206,14 +266,17 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, # 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, rtol=rtol, atol=atol) - if not convergence_met: + self.convergence_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.prior_guess = self._compute_prior_guess() + self.x_old = self._compute_prior_guess() 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, @@ -224,15 +287,21 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, except: self.model.load_state() raise - self.next_guess = self._compute_next_guess() + 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) iter_elapsed += 1 - convergence_met = self._convergence_met(self.prior_guess, self.next_guess, rtol=rtol, atol=atol) - if convergence_met: + if self.success: break self.model.iter_elapsed = iter_elapsed -class LegacyConvergenceTracker(ConvergenceTracker): +class LegacyConvergenceTracker(BaseCallback): def __init__(self, model): self.model = model self.prior_guess = 0. @@ -260,6 +329,11 @@ def _set_states(self, prior_states, next_states): 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): From 9fd58e99f04ce87594e7f9fb595cfba6d2355710 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Thu, 18 Dec 2025 02:59:53 -0600 Subject: [PATCH 26/27] Add Newton convergence flux limiter code --- pipedream_solver/_nsuperlink.py | 3 ++ pipedream_solver/diagnostics.py | 71 +++++++++++++++++++++++++++------ pipedream_solver/nquality.py | 26 ++---------- pipedream_solver/superlink.py | 8 ++-- 4 files changed, 70 insertions(+), 38 deletions(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 0e0da69..88baa93 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -669,6 +669,7 @@ def numba_solve_internals_ls(_h_Ik, NK, nk, _k_1k, _i_1k, _I_1k, _U, _X, _b): _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): @@ -693,6 +694,7 @@ def numba_u_ik(_Q_ik, _A_ik, _A_ik_min, _u_ik): # _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): @@ -716,6 +718,7 @@ def numba_u_Ik(_dx_ik, _u_ik, _dx_uk, _u_uk, _link_start, _ki, _u_Ik): _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): diff --git a/pipedream_solver/diagnostics.py b/pipedream_solver/diagnostics.py index 91a773f..c54d88c 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -124,11 +124,7 @@ def __init__(self, model): 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.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) + self.set_init_volume() @property def init_volume(self): @@ -152,6 +148,14 @@ 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 @@ -182,9 +186,12 @@ def __init__(self, model, xtol=1e-6, max_learning_rate=0.5, min_learning_rate=0. 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] @@ -199,13 +206,32 @@ def _compute_step_ratio(self, dx): + (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((dx['Q_ik'] / Q_ik_ub).max(), 0.), - max((dx['Q_uk'] / Q_uk_ub).max(), 0.), - max((dx['Q_dk'] / Q_dk_ub).max(), 0.)) - ratio_lb = max(max(-(dx['Q_ik'] / Q_ik_lb).min(), 0.), - max(-(dx['Q_uk'] / Q_uk_lb).max(), 0.), - max(-(dx['Q_dk'] / Q_dk_lb).max(), 0.)) + 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 @@ -267,6 +293,7 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, 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) @@ -277,6 +304,9 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, 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, @@ -285,7 +315,8 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, first_time=first_time, implicit=implicit, banded=banded, first_iter=False) except: - self.model.load_state() + 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) @@ -295,10 +326,14 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, 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 LegacyConvergenceTracker(BaseCallback): @@ -673,3 +708,15 @@ def total_volume(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/nquality.py b/pipedream_solver/nquality.py index 3a4ad90..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) @@ -223,7 +223,7 @@ def __init__(self, hydraulics, superjunction_params, superlink_params, self.N_process_sigma = superlink_params['KF_process_sigma'].values[0].astype(np.float64) self.N_measure_sigma = superlink_params['KF_measure_sigma'].values[0].astype(np.float64) self.step(dt=1e-6) - + # TODO: It might be safer to have these as @properties def import_hydraulic_states(self, _dt): self._H_j_next = self.hydraulics.H_j @@ -277,27 +277,7 @@ def import_hydraulic_states(self, _dt): self.hydraulics._geom_codes, self.hydraulics._g1_ik, self._H_j_next, self._z_inv_j, self._H_j_prev, self._Q_dk_next, self._B_dk, self._dx_dk, self._Q_uk_up_next, self._Q_uk_dn_next, self._Q_dk_up_next, self._Q_dk_dn_next, self._A_uk_next, self._A_dk_next, self._A_uk_prev, self._A_dk_prev) - - @property - def c(self): - # NOTE: more performant to write in-place - result = np.concatenate([self._c_j, self._c_Ik, self._c_ik, - self._c_uk, self._c_dk]) - return result - - @c.setter - def c(self, value): - n1 = self._c_j.size - n2 = n1 + self._c_Ik.size - n3 = n2 + self._c_ik.size - n4 = n3 + self._c_uk.size - n5 = n4 + self._c_dk.size - self._c_j[:] = np.asarray(value[:n1]) - self._c_Ik[:] = np.asarray(value[n1:n2]) - self._c_ik[:] = np.asarray(value[n2:n3]) - self._c_uk[:] = np.asarray(value[n3:n4]) - self._c_dk[:] = np.asarray(value[n4:n5]) - + @property def c_j(self): return self._c_j diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index d4e2ff9..3926dee 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -9,7 +9,7 @@ import pipedream_solver.storage import pipedream_solver.visualization from pipedream_solver.callbacks import BaseCallback -from pipedream_solver.diagnostics import LegacyConvergenceTracker +from pipedream_solver.diagnostics import ConvergenceTracker, VolumeTracker class SuperLink(): """ @@ -796,8 +796,10 @@ def __init__(self, superlinks, superjunctions, # Compute bandwidth self._compute_bandwidth() # Initialize to stable state - self.convergence_tracker = LegacyConvergenceTracker(self) + self.convergence_tracker = ConvergenceTracker(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 @@ -4280,7 +4282,7 @@ def step(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, u_p=None, d first_time=first_time, implicit=implicit, banded=banded, first_iter=first_iter) except: - self.load_state() + #self.load_state() raise first_iter = False From 798ca504c1bf56de75740752669781620a532412 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Thu, 18 Dec 2025 13:40:03 -0600 Subject: [PATCH 27/27] Modify convergence tracker formulation --- pipedream_solver/_nsuperlink.py | 77 +++++++++++++++++++++++++++++++++ pipedream_solver/diagnostics.py | 62 ++++++++++++++++++++++++++ pipedream_solver/superlink.py | 3 +- 3 files changed, 141 insertions(+), 1 deletion(-) diff --git a/pipedream_solver/_nsuperlink.py b/pipedream_solver/_nsuperlink.py index 88baa93..d9d8175 100644 --- a/pipedream_solver/_nsuperlink.py +++ b/pipedream_solver/_nsuperlink.py @@ -1599,3 +1599,80 @@ def numba_reposition_junctions(_x_Ik, _z_inv_Ik, _h_Ik, _dx_ik, _Q_ik, _H_dk, _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/diagnostics.py b/pipedream_solver/diagnostics.py index c54d88c..c4dee67 100644 --- a/pipedream_solver/diagnostics.py +++ b/pipedream_solver/diagnostics.py @@ -1,7 +1,11 @@ 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 @@ -336,6 +340,64 @@ def __on_step_end__(self, H_bc=None, Q_in=None, Q_0Ik=None, u_o=None, u_w=None, 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 diff --git a/pipedream_solver/superlink.py b/pipedream_solver/superlink.py index 3926dee..f9ab035 100644 --- a/pipedream_solver/superlink.py +++ b/pipedream_solver/superlink.py @@ -10,6 +10,7 @@ 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(): """ @@ -796,7 +797,7 @@ def __init__(self, superlinks, superjunctions, # Compute bandwidth self._compute_bandwidth() # Initialize to stable state - self.convergence_tracker = ConvergenceTracker(self) + 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')