@@ -99,19 +99,31 @@ def info(s):
99
99
100
100
# coordinate arrays
101
101
x_t_2d = fromfunction (
102
- lambda i , j : xmin + i * dx + dx / 2 , (nx , ny ), dtype = dtype , device = device
102
+ lambda i , j : xmin + i * dx + dx / 2 ,
103
+ (nx , ny ),
104
+ dtype = dtype ,
105
+ device = device ,
103
106
)
104
107
y_t_2d = fromfunction (
105
- lambda i , j : ymin + j * dy + dy / 2 , (nx , ny ), dtype = dtype , device = device
108
+ lambda i , j : ymin + j * dy + dy / 2 ,
109
+ (nx , ny ),
110
+ dtype = dtype ,
111
+ device = device ,
106
112
)
107
113
x_u_2d = fromfunction (
108
114
lambda i , j : xmin + i * dx , (nx + 1 , ny ), dtype = dtype , device = device
109
115
)
110
116
y_u_2d = fromfunction (
111
- lambda i , j : ymin + j * dy + dy / 2 , (nx + 1 , ny ), dtype = dtype , device = device
117
+ lambda i , j : ymin + j * dy + dy / 2 ,
118
+ (nx + 1 , ny ),
119
+ dtype = dtype ,
120
+ device = device ,
112
121
)
113
122
x_v_2d = fromfunction (
114
- lambda i , j : xmin + i * dx + dx / 2 , (nx , ny + 1 ), dtype = dtype , device = device
123
+ lambda i , j : xmin + i * dx + dx / 2 ,
124
+ (nx , ny + 1 ),
125
+ dtype = dtype ,
126
+ device = device ,
115
127
)
116
128
y_v_2d = fromfunction (
117
129
lambda i , j : ymin + j * dy , (nx , ny + 1 ), dtype = dtype , device = device
@@ -189,7 +201,9 @@ def bathymetry(x_t_2d, y_t_2d, lx, ly):
189
201
return bath * np .full (T_shape , 1.0 , dtype , device = device )
190
202
191
203
# inital elevation
192
- u0 , v0 , e0 = exact_solution (0 , x_t_2d , y_t_2d , x_u_2d , y_u_2d , x_v_2d , y_v_2d )
204
+ u0 , v0 , e0 = exact_solution (
205
+ 0 , x_t_2d , y_t_2d , x_u_2d , y_u_2d , x_v_2d , y_v_2d
206
+ )
193
207
e [:, :] = e0
194
208
u [:, :] = u0
195
209
v [:, :] = v0
@@ -229,10 +243,14 @@ def rhs(u, v, e):
229
243
hu [1 :- 1 , :] = 0.5 * (H [:- 1 , :] + H [1 :, :]) * u [1 :- 1 , :]
230
244
hv [:, 1 :- 1 ] = 0.5 * (H [:, :- 1 ] + H [:, 1 :]) * v [:, 1 :- 1 ]
231
245
232
- dedt = - 1.0 * ((hu [1 :, :] - hu [:- 1 , :]) / dx + (hv [:, 1 :] - hv [:, :- 1 ]) / dy )
246
+ dedt = - 1.0 * (
247
+ (hu [1 :, :] - hu [:- 1 , :]) / dx + (hv [:, 1 :] - hv [:, :- 1 ]) / dy
248
+ )
233
249
234
250
# total depth at F points
235
- H_at_f [1 :- 1 , 1 :- 1 ] = 0.25 * (H [1 :, 1 :] + H [:- 1 , 1 :] + H [1 :, :- 1 ] + H [:- 1 , :- 1 ])
251
+ H_at_f [1 :- 1 , 1 :- 1 ] = 0.25 * (
252
+ H [1 :, 1 :] + H [:- 1 , 1 :] + H [1 :, :- 1 ] + H [:- 1 , :- 1 ]
253
+ )
236
254
H_at_f [0 , 1 :- 1 ] = 0.5 * (H [0 , 1 :] + H [0 , :- 1 ])
237
255
H_at_f [- 1 , 1 :- 1 ] = 0.5 * (H [- 1 , 1 :] + H [- 1 , :- 1 ])
238
256
H_at_f [1 :- 1 , 0 ] = 0.5 * (H [1 :, 0 ] + H [:- 1 , 0 ])
@@ -312,6 +330,7 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
312
330
initial_v = None
313
331
initial_e = None
314
332
tic = time_mod .perf_counter ()
333
+ block_tic = 0
315
334
for i in range (nt + 1 ):
316
335
sync ()
317
336
t = i * dt
@@ -356,9 +375,12 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
356
375
info (
357
376
f"{ i_export :2d} { i :4d} { t :.3f} elev={ elev_max :7.5f} "
358
377
f"u={ u_max :7.5f} q={ q_max :8.5f} dV={ diff_v : 6.3e} "
359
- f"PE={ total_pe :7.2e} KE={ total_ke :7.2e} dE={ diff_e : 6.3e} " + tcpu_str
378
+ f"PE={ total_pe :7.2e} KE={ total_ke :7.2e} dE={ diff_e : 6.3e} "
379
+ + tcpu_str
360
380
)
361
- if not benchmark_mode and (elev_max > 1e3 or not math .isfinite (elev_max )):
381
+ if not benchmark_mode and (
382
+ elev_max > 1e3 or not math .isfinite (elev_max )
383
+ ):
362
384
raise ValueError (f"Invalid elevation value: { elev_max } " )
363
385
i_export += 1
364
386
next_t_export = i_export * t_export
@@ -372,7 +394,9 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
372
394
duration = time_mod .perf_counter () - tic
373
395
info (f"Duration: { duration :.2f} s" )
374
396
375
- e_exact = exact_solution (t , x_t_2d , y_t_2d , x_u_2d , y_u_2d , x_v_2d , y_v_2d )[2 ]
397
+ e_exact = exact_solution (t , x_t_2d , y_t_2d , x_u_2d , y_u_2d , x_v_2d , y_v_2d )[
398
+ 2
399
+ ]
376
400
err2 = (e_exact - e ) * (e_exact - e ) * dx * dy / lx / ly
377
401
err_L2 = math .sqrt (float (np .sum (err2 , all_axes )))
378
402
info (f"L2 error: { err_L2 :7.15e} " )
0 commit comments