Generate all plots

flood.py

@@ -9,7 +9,7 @@ Plot a sympy expression f(x_sym) with respect to variable x_sym.
 x_sym in range lim = [low, high].
 N is the number of sample points.
 """
-def plot_expr(f, x_sym, lim, N):
+def plot_expr(ax, f, x_sym, lim, N):
     F = lambdify(x_sym, f)
     x = []
     y = []
@@ -17,8 +17,7 @@ def plot_expr(f, x_sym, lim, N):
         xi = i * (lim[1]-lim[0]) / (N-1) + lim[0]
         x.append(xi)
         y.append(F(xi))
-    plt.plot(x, y)
-    plt.show()
+    ax.plot(x, y)
 
 
 """
@@ -205,6 +204,10 @@ def solve_characteristics(Q, c_sym, c_init, x_sym, delta_x, L, delta_t, T):
     # Get expression for char. velocity
     vc = diff(Q, c_sym)
 
+    #vc_diff = diff(vc, c_sym)
+    #plot_expr(vc_diff, c_sym, [0, 1000], 1000)
+    #exit(0)
+
     # Derivative of vc(c_init(x)) wrt x
     vc_diff = diff(vc, c_sym).subs(c_sym, c_init) * diff(c_init, x_sym)
 
@@ -212,6 +215,8 @@ def solve_characteristics(Q, c_sym, c_init, x_sym, delta_x, L, delta_t, T):
     x0_shock_arr = get_all_argmin(vc_diff, x_sym, (0, L))
     shocks = []
 
+    print("Solving shocks... ", end="", flush=True)
+    
     # Solve each shock
     for x0_shock in x0_shock_arr:
         vc_shock = vc.subs(c_sym, c_init.subs(x_sym, x0_shock).evalf()).evalf()
@@ -221,15 +226,20 @@ def solve_characteristics(Q, c_sym, c_init, x_sym, delta_x, L, delta_t, T):
         shock_sln = solve_shock(Q, c_sym, c_init, x_sym, t_shock, x_shock, x0_shock,
                                 delta_x, L, delta_t, T)
         shocks.append({"t": t_shock, "x": x_shock, "x0": x0_shock, "sln": shock_sln})
-        
 
+    print("done")
+
+    print(f"Shock starts at ({shocks[0]['x']}, {shocks[0]['t']})")
+    
     c_of_x0 = lambdify(x_sym, c_init)
     vc_of_c = lambdify(c_sym, vc)
 
     chars = []
+
+    print("Generating characteristics... [", end="", flush=True)
     
     # Generate characteristics starting from uniform x-ordinates
-    ch_spacing = 4
+    ch_spacing = int(L/25)
     for j in range(int(L/ch_spacing)): # For each characteristic
         x0 = x = ch_spacing*j
 
@@ -267,59 +277,72 @@ def solve_characteristics(Q, c_sym, c_init, x_sym, delta_x, L, delta_t, T):
 
         # Store
         chars.append({"x": x_ch, "t": t_ch})
+        if (int(20*len(chars)/int(L/ch_spacing)) >
+            int(20*(len(chars)-1)/int(L/ch_spacing))):
+            print("=", end="", flush=True)
 
+    print("] done")
+    
     return chars, shocks
 
 def plot_shocks(shocks, delta_t, T, ax):
+    print("Plotting shocks... ", end="", flush=True)
+
     # Plot the shock solutions
     for s in shocks:
-        t_arr = [i*delta_t+s["t"] for i in range(int((T-s["t"])/delta_t))]
-        y = [s["sln"](t)[0] for t in t_arr]
-        ax.plot(y, t_arr, color="blue")
+        if s["t"] < T:
+            t_arr = [i*delta_t+s["t"] for i in range(int((T-s["t"])/delta_t))]
+            y = s["sln"](np.array(t_arr))[0,:] #[s["sln"](t)[0] for t in t_arr]
+            ax.plot(y, t_arr, color="blue")
+
+    print("done")
 
 def plot_chars(chars, ax):
+    print("Plotting characteristics... ", end="", flush=True)
     for ch in chars:
         ax.plot(ch["x"], ch["t"], color="red")
+    print("done")
     
-    
-def main():
+def gen_fig(prefix, g_val, alpha_val, f_val, w_expr, y_sym, a_sym, a_val,
+            delta_x, L, delta_t, T,
+            A_init, x_sym):
     g, alpha, f = symbols("g,alpha,f")
     
     # Width of channel wrt vertical coordinate (y)
-    y = Symbol("y", positive=True, real=True)
+    y = y_sym #Symbol("y", positive=True, real=True)
     a = Symbol("a", positive=True, real=True)
-    w = a #sqrt(y/a)
-
+    w = w_expr #a #sqrt(y/a)
+
+    plt.figure()
+    ax = plt.gca()
+    plot_expr(ax, A_init, x_sym, [0, L], int(L/delta_x))
+    plt.xlabel("Position (m)")
+    plt.ylabel("Initial Area (m^2)")
+    plt.savefig(prefix+"-init.pdf", bbox_inches="tight")
+    
     A_sym = Symbol("A", positive=True, real=True)
 
     l_of_A = derive_wetted_perim(w, y, A_sym)
-    #print(l_of_A)
 
     u = simplify(sqrt(A_sym*g*sin(alpha) / (f * l_of_A)))
-    #print("u(A) = ")
-    #pprint(u)
     
-    Q = (A_sym*u).subs([("g", 9.81), ("alpha", 0.001), ("f", 0.05), (a, 500)])
-
-    delta_x = 0.5
-    L = 300
+    Q = (A_sym*u).subs([("g", g_val), ("alpha", alpha_val), ("f", f_val), (a, a_val)])
 
-    delta_t = 0.5
-    T = 200
-
-    x_sym = Symbol("x", real=True)
-    A_init = 1 + 10000/(10+(x_sym-100)**2)
-    
     chars, shocks = solve_characteristics(Q, A_sym, A_init,
                                           x_sym, delta_x, L, delta_t, T)
 
-    ax0 = plt.subplot(2, 1, 1)
+    plt.figure()
+    ax0 = plt.gca()
 
     plot_chars(chars, ax0)
     plot_shocks(shocks, delta_t, T, ax0)
     ax0.axis([0, L, 0, T])
-    
-    ax1 = plt.subplot(2, 1, 2)
+    plt.xlabel("Position (m)")
+    plt.ylabel("Time (s)")
+    plt.savefig(prefix+"-shock.pdf", bbox_inches="tight")
+
+    plt.figure()
+    ax1 = plt.gca()
     
     #plot_expr(diff(Q, A_sym), A_sym, [0,30], 1000)
     #return
@@ -327,27 +350,68 @@ def main():
 
     x = [i*delta_x for i in range(int(L/delta_x))]
 
+    A_inflow = float(A_init.evalf(subs={x_sym: 0}))
     A_init_lamb = lambdify(x_sym, A_init)
     A_init_data = [A_init_lamb(xi) for xi in x]
-    A_inflow = 1
 
     t = [i*delta_t for i in range(int(T/delta_t))]
-    
-    c_sln = godunov_solve(Q_lamb, A_init_data, A_inflow, delta_x, t)
 
-    n_xticks = 21
-    n_tticks = 11
+    print("Solving Godunov... ", end="", flush=True)
+    c_sln = godunov_solve(Q_lamb, A_init_data, A_inflow, delta_x, t)
+    print("done")
+    
+    n_xticks = 6
+    n_tticks = 6
     x_ticks = [int(i*L/delta_x/(n_xticks-1)) for i in range(n_xticks)]
     x_tick_labels = [delta_x*i for i in x_ticks]
     t_ticks = [int(i*T/delta_t/(n_tticks-1)) for i in range(n_tticks)]
     t_tick_labels = [delta_t*i for i in t_ticks]
+
+    print("Plotting Godunov... ", end="", flush=True)
     
     plt.colorbar(plt.pcolor(c_sln))
     ax1.imshow(c_sln, origin="lower", aspect="auto")
     ax1.set_xticks(x_ticks, x_tick_labels)
     ax1.set_yticks(t_ticks, t_tick_labels)
+    plt.xlabel("Position (m)")
+    plt.ylabel("Time (s)")
+    plt.savefig(prefix+"-godunov.pdf", bbox_inches="tight")
+
+    print("done")
 
-    plt.show()
 
+def main():
+    y_sym = Symbol("y", positive=True, real=True)
+    x_sym = Symbol("x", positive=True, real=True)
+    a_sym = Symbol("a", positive=True, real=True)
+
+    gen_A_init = lambda A0,A1,s_slope,s_grad: (A0+A1)/2 - (A0-A1)/2*tanh((x_sym-s_slope)*s_grad)
+    
+    plot_defaults = {
+        "g_val": 9.81, "alpha_val": 0.01, "f_val": 0.05, "w_expr": a_sym,
+        "y_sym": y_sym, "a_sym": a_sym, "a_val": 50,
+        "delta_x": 1, "L": 300, "delta_t": 1, "T": 200,
+        "A_init": gen_A_init(20,1,100,1/50),
+        "x_sym": x_sym
+    }
+
+    plots = [
+        #plot_defaults | { "prefix": "fig/baseline" },
+        #plot_defaults | { "prefix": "alpha-0.02", "alpha_val": 0.02 },
+        #plot_defaults | { "prefix": "f-0.1", "f_val": 0.1 },
+        #plot_defaults | { "prefix": "dx_0.25,dt_0.25", "delta_x": 0.25, "delta_t": 0.25 },
+        #plot_defaults | { "prefix": "fig/a-5", "a_val": 5 },
+        #plot_defaults | { "prefix": "fig/a-3", "a_val": 3, "L": 700, "T": 400 },
+        #plot_defaults | { "prefix": "fig/a-10", "a_val": 10 },
+        #plot_defaults | { "prefix": "fig/a-1", "a_val": 1 },
+        #plot_defaults | { "prefix": "fig/slope-0.1", "A_init": gen_A_init(20,1,100,0.1) },
+        #plot_defaults | { "prefix": "fig/A-pulse",  "A_init": 1 + 19*1000/(1*1000+(x_sym-100)**2) },
+        plot_defaults | { "prefix": "fig/A-pulse-short",  "A_init": 1 + 19*10/(1*10+(x_sym-100)**2) }
+    ]
+
+    plt.rcParams.update({'font.size': 22})
+
+    for p in plots:
+        gen_fig(**p)
 
 main()