1717def _log_base_q (x , q ):
1818 return torch .log (x ) / torch .log (torch .tensor (q , dtype = x .dtype ))
1919
20- def _get_velocity (q_t , _n , _p_spatial , width , _q_spatial , _s0 , velocity_lb , depth_lb ) -> torch .Tensor :
20+ def _get_trapezoid_velocity (
21+ q_t ,
22+ _n : torch .Tensor ,
23+ top_width : torch .Tensor ,
24+ side_slope : torch .Tensor ,
25+ _s0 : torch .Tensor ,
26+ p_spatial : torch .Tensor ,
27+ _q_spatial : torch .Tensor ,
28+ velocity_lb : torch .Tensor ,
29+ depth_lb : torch .Tensor ,
30+ _btm_width_lb : torch .Tensor ,
31+ ) -> torch .Tensor :
2132 """Calculate flow velocity using Manning's equation.
22-
23- Parameters
24- ----------
25- q_t : torch.Tensor
26- Discharge at time t.
27- _n : torch.Tensor
28- Manning's roughness coefficient.
29- _q_spatial : torch.Tensor
30- Spatial discharge parameter.
31- _s0 : torch.Tensor
32- Channel slope.
33- p_spatial : torch.Tensor
34- Spatial parameter for width calculation.
35-
36- Returns
37- -------
38- torch.Tensor
39- Celerity (wave speed) of the flow.
40-
41- Notes
42- -----
43- The function first calculates flow depth using Manning's equation, then
44- computes velocity and finally celerity. The celerity is clamped between
45- 0.3 and 15 m/s and scaled by 5/3 according to kinematic wave theory.
4633 """
47- depth = _log_base_q (width / _p_spatial , _q_spatial )
48- v = torch .div (1 , _n ) * torch .pow (depth , (2 / 3 )) * torch .pow (_s0 , (1 / 2 ))
49- c_ = torch .clamp (v , velocity_lb , 15 )
34+ numerator = q_t * _n * (_q_spatial + 1 )
35+ denominator = p_spatial * torch .pow (_s0 , 0.5 )
36+ depth = torch .clamp (
37+ torch .pow (
38+ torch .div (numerator , denominator + 1e-8 ),
39+ torch .div (3.0 , 5.0 + 3.0 * _q_spatial ),
40+ ),
41+ min = depth_lb ,
42+ )
43+
44+ # For z:1 side slopes (z horizontal : 1 vertical)
45+ _bottom_width = top_width - (2 * side_slope * depth )
46+ bottom_width = torch .clamp (_bottom_width , min = _btm_width_lb )
47+
48+ # Area = (top_width + bottom_width)*depth/2
49+ area = (top_width + bottom_width ) * depth / 2
50+
51+ # Side length = sqrt(1 + z^2) * depth
52+ # Since for every 1 unit vertical, we go z units horizontal
53+ wetted_p = bottom_width + 2 * depth * torch .sqrt (1 + side_slope ** 2 )
54+
55+ # Calculate hydraulic radius
56+ R = area / wetted_p
57+
58+ v = torch .div (1 , _n ) * torch .pow (R , (2 / 3 )) * torch .pow (_s0 , (1 / 2 ))
59+ c_ = torch .clamp (v , min = velocity_lb , max = 15 )
5060 c = c_ * 5 / 3
5161 return c
5262
@@ -76,19 +86,17 @@ def __init__(
7686 # Base routing parameters
7787 self .n = None
7888 self .q_spatial = None
79- self .p_spatial = None
8089
8190 # Routing state
82- self .length = None
83- self .slope = None
84- self .velocity = None
8591 self ._discharge_t = None
86- self .adjacency_matrix = None
92+ self .network = None
8793
8894 self .parameter_bounds = self .cfg .params .parameter_ranges .range
95+ self .p_spatial = torch .tensor (self .cfg .params .defaults .p , device = self .device_num )
8996 self .velocity_lb = torch .tensor (self .cfg .params .attribute_minimums .velocity , device = self .device_num )
9097 self .depth_lb = torch .tensor (self .cfg .params .attribute_minimums .depth , device = self .device_num )
9198 self .discharge_lb = torch .tensor (self .cfg .params .attribute_minimums .discharge , device = self .device_num )
99+ self .bottom_width_lb = torch .tensor (self .cfg .params .attribute_minimums .bottom_width , device = self .device_num )
92100
93101 def forward (self , ** kwargs ) -> dict [str , torch .Tensor ]:
94102 """The forward pass for the dMC model
@@ -106,18 +114,14 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]:
106114 # gage_information = hydrofabric.network.gage_information
107115 # TODO: create a dynamic gauge look up
108116 gage_indices = torch .tensor ([- 1 ])
109- self .adjacency_matrix = hydrofabric .adjacency_matrix
117+ self .network = hydrofabric .adjacency_matrix
110118
111119 # Set up base parameters
112120 self .n = denormalize (value = kwargs ["spatial_parameters" ]["n" ], bounds = self .parameter_bounds ["n" ])
113121 self .q_spatial = denormalize (
114122 value = kwargs ["spatial_parameters" ]["q_spatial" ],
115123 bounds = self .parameter_bounds ["q_spatial" ],
116124 )
117- self .p_spatial = denormalize (
118- value = kwargs ["spatial_parameters" ]["p_spatial" ],
119- bounds = self .parameter_bounds ["p_spatial" ],
120- )
121125
122126 # Initialize discharge
123127 self ._discharge_t = q_prime [0 ].to (self .device_num )
@@ -147,8 +151,9 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]:
147151 hydrofabric .slope .to (self .device_num ).to (torch .float32 ),
148152 min = self .cfg .params .attribute_minimums .slope ,
149153 )
150- width = hydrofabric .length .to (self .device_num ).to (torch .float32 )
151- x_storage = hydrofabric .length .to (self .device_num ).to (torch .float32 )
154+ top_width = hydrofabric .top_width .to (self .device_num ).to (torch .float32 )
155+ side_slope = hydrofabric .side_slope .to (self .device_num ).to (torch .float32 )
156+ x_storage = hydrofabric .x .to (self .device_num ).to (torch .float32 )
152157
153158 desc = "Running dMC Routing"
154159 for timestep in tqdm (
@@ -161,22 +166,24 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]:
161166 ):
162167 q_prime_sub = q_prime [timestep - 1 ].clone ()
163168 q_prime_clamp = torch .clamp (q_prime_sub , min = self .cfg .params .attribute_minimums .discharge )
164- velocity = _get_velocity (
169+ velocity = _get_trapezoid_velocity (
165170 q_t = self ._discharge_t ,
166171 _n = self .n ,
167- _q_spatial = self .q_spatial ,
172+ top_width = top_width ,
173+ side_slope = side_slope ,
168174 _s0 = slope ,
169- _p_spatial = self .p_spatial ,
170- width = width ,
175+ p_spatial = self .p_spatial ,
176+ _q_spatial = self . q_spatial ,
171177 velocity_lb = self .velocity_lb ,
172178 depth_lb = self .depth_lb ,
179+ _btm_width_lb = self .bottom_width_lb ,
173180 )
174181 k = torch .div (length , velocity )
175182 denom = (2.0 * k * (1.0 - x_storage )) + self .t
176183 c_2 = (self .t + (2.0 * k * x_storage )) / denom
177184 c_3 = ((2.0 * k * (1.0 - x_storage )) - self .t ) / denom
178185 c_4 = (2.0 * self .t ) / denom
179- i_t = torch .matmul (self .adjacency_matrix , self ._discharge_t )
186+ i_t = torch .matmul (self .network , self ._discharge_t )
180187 q_l = q_prime_clamp
181188
182189 b_array = (c_2 * i_t ) + (c_3 * self ._discharge_t ) + (c_4 * q_l )
@@ -204,3 +211,44 @@ def forward(self, **kwargs) -> dict[str, torch.Tensor]:
204211 }
205212
206213 return output_dict
214+
215+ def fill_op (self , data_vector : torch .Tensor ):
216+ """A fill operation function for the sparse matrix
217+
218+ The equation we want to solve
219+ (I - C_1*N) * Q_t+1 = c_2*(N*Q_t_1) + c_3*Q_t + c_4*Q`
220+ (I - C_1*N) * Q_t+1 = b(t)
221+
222+ Parameters
223+ ----------
224+ data_vector: torch.Tensor
225+ The data vector to fill the sparse matrix with
226+ """
227+ identity_matrix = self ._sparse_eye (self .network .shape [0 ])
228+ vec_diag = self ._sparse_diag (data_vector )
229+ # vec_filled = bnb.matmul(vec_diag, self.network, threshold=6.0)
230+ vec_filled = torch .matmul (vec_diag .cpu (), self .network .cpu ()).to (self .device_num )
231+ A = identity_matrix + vec_filled
232+ return A
233+
234+ def _sparse_eye (self , n ):
235+ indices = torch .arange (n , dtype = torch .int32 , device = self .device_num )
236+ values = torch .ones (n , device = self .device_num )
237+ identity_coo = torch .sparse_coo_tensor (
238+ indices = torch .vstack ([indices , indices ]),
239+ values = values ,
240+ size = (n , n ),
241+ device = self .device_num ,
242+ )
243+ return identity_coo .to_sparse_csr ()
244+
245+ def _sparse_diag (self , data ):
246+ n = len (data )
247+ indices = torch .arange (n , dtype = torch .int32 , device = self .device_num )
248+ diagonal_coo = torch .sparse_coo_tensor (
249+ indices = torch .vstack ([indices , indices ]),
250+ values = data ,
251+ size = (n , n ),
252+ device = self .device_num ,
253+ )
254+ return diagonal_coo .to_sparse_csr ()
0 commit comments