Skip to content

Commit

Permalink
Refactoring of geometry and state space representation
Browse files Browse the repository at this point in the history
  • Loading branch information
@mdbartos
mdbartos committed Aug 9, 2024
1 parent dacd658 commit 73006f4
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 40 deletions.
82 changes: 48 additions & 34 deletions pipedream_solver/_nsuperlink.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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[:]),
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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[:],
Expand Down
10 changes: 6 additions & 4 deletions pipedream_solver/ngeometry.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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.
Expand All @@ -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
Expand Down
13 changes: 11 additions & 2 deletions pipedream_solver/superlink.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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):
"""
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 73006f4

Please sign in to comment.