@@ -99,19 +99,31 @@ def info(s):
9999
100100 # coordinate arrays
101101 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 ,
103106 )
104107 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 ,
106112 )
107113 x_u_2d = fromfunction (
108114 lambda i , j : xmin + i * dx , (nx + 1 , ny ), dtype = dtype , device = device
109115 )
110116 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 ,
112121 )
113122 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 ,
115127 )
116128 y_v_2d = fromfunction (
117129 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):
189201 return bath * np .full (T_shape , 1.0 , dtype , device = device )
190202
191203 # 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+ )
193207 e [:, :] = e0
194208 u [:, :] = u0
195209 v [:, :] = v0
@@ -229,10 +243,14 @@ def rhs(u, v, e):
229243 hu [1 :- 1 , :] = 0.5 * (H [:- 1 , :] + H [1 :, :]) * u [1 :- 1 , :]
230244 hv [:, 1 :- 1 ] = 0.5 * (H [:, :- 1 ] + H [:, 1 :]) * v [:, 1 :- 1 ]
231245
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+ )
233249
234250 # 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+ )
236254 H_at_f [0 , 1 :- 1 ] = 0.5 * (H [0 , 1 :] + H [0 , :- 1 ])
237255 H_at_f [- 1 , 1 :- 1 ] = 0.5 * (H [- 1 , 1 :] + H [- 1 , :- 1 ])
238256 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):
312330 initial_v = None
313331 initial_e = None
314332 tic = time_mod .perf_counter ()
333+ block_tic = 0
315334 for i in range (nt + 1 ):
316335 sync ()
317336 t = i * dt
@@ -356,9 +375,12 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
356375 info (
357376 f"{ i_export :2d} { i :4d} { t :.3f} elev={ elev_max :7.5f} "
358377 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
360380 )
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+ ):
362384 raise ValueError (f"Invalid elevation value: { elev_max } " )
363385 i_export += 1
364386 next_t_export = i_export * t_export
@@ -372,7 +394,9 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
372394 duration = time_mod .perf_counter () - tic
373395 info (f"Duration: { duration :.2f} s" )
374396
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+ ]
376400 err2 = (e_exact - e ) * (e_exact - e ) * dx * dy / lx / ly
377401 err_L2 = math .sqrt (float (np .sum (err2 , all_axes )))
378402 info (f"L2 error: { err_L2 :7.15e} " )
0 commit comments