From 73006f47c7d4293a2e376da1a8d53301413ac503 Mon Sep 17 00:00:00 2001 From: Matt Bartos Date: Fri, 9 Aug 2024 01:54:12 -0500 Subject: [PATCH] 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: