Update app.py

app.py

@@ -6,69 +6,71 @@ import gradio as gr
 import plotly.graph_objects as go
 
 # --- Simulation Core ---
-def solve_2d_heat_equation(Lx: float,
-                           Ly: float,
-                           t_max: float,
-                           Gamma: float = 0.1,
-                           Nx: int = 50,
-                           Ny: int = 50,
-                           initial: str = "gaussian",
-                           bc: str = "dirichlet"):
+def solve_3d_heat_equation(Lx: float, Ly: float, Lz: float,
+                           t_max: float, Gamma: float = 0.1,
+                           Nx: int = 50, Ny: int = 50, Nz: int = 50,
+                           initial: str = "gaussian", bc: str = "dirichlet"):
     """
-    Solve the 2D heat equation u_t = Gamma*(u_xx + u_yy) and return the solution array U and time step dt.
+    Solve the 3D heat equation u_t = Gamma*(u_xx + u_yy + u_zz) and return the solution array U and time step dt.
     """
     # Spatial grid
     x = np.linspace(0, Lx, Nx)
     y = np.linspace(0, Ly, Ny)
-    dx, dy = x[1] - x[0], y[1] - y[0]
-    if dx == 0 or dy == 0:
-        raise ValueError("Nx and Ny must be > 1.")
+    z = np.linspace(0, Lz, Nz)
+    dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
+    if dx <= 0 or dy <= 0 or dz <= 0:
+        raise ValueError("Nx, Ny, and Nz must be > 1.")
 
     # Time stepping for stability
-    dt = 1.0 / (2 * Gamma * (1/dx**2 + 1/dy**2))
+    dt = 1.0 / (2 * Gamma * (1/dx**2 + 1/dy**2 + 1/dz**2))
     Nt = int(np.ceil(t_max / dt)) + 1
-    rx, ry = Gamma * dt / dx**2, Gamma * dt / dy**2
+    rx, ry, rz = Gamma * dt / dx**2, Gamma * dt / dy**2, Gamma * dt / dz**2
 
     # Initial condition
-    X, Y = np.meshgrid(x, y, indexing='ij')
+    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
     if initial == "gaussian":
-        u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2) / (2*(Lx/10)**2)))
+        u = np.exp(-((X - Lx/2)**2 / (2*(Lx/10)**2) + (Y - Ly/2)**2 / (2*(Ly/10)**2) + (Z - Lz/2)**2 / (2*(Lz/10)**2)))
     elif initial == "random":
-        u = np.random.rand(Nx, Ny)
+        u = np.random.rand(Nx, Ny, Nz)
     elif initial == "sinusoidal":
-        kx, ky = 2 * np.pi / Lx, 2 * np.pi / Ly
-        u = np.sin(kx * X) * np.sin(ky * Y)
+        kx, ky, kz = 2 * np.pi / Lx, 2 * np.pi / Ly, 2 * np.pi / Lz
+        u = np.sin(kx * X) * np.sin(ky * Y) * np.sin(kz * Z)
     elif initial == "step":
-        u = np.where((X < Lx/2) & (Y < Ly/2), 1.0, 0.0)
+        u = np.where((X < Lx/2) & (Y < Ly/2) & (Z < Lz/2), 1.0, 0.0)
     else:
         raise ValueError(f"Unknown initial condition: {initial}")
 
     # Storage for solution
-    U = np.zeros((Nt, Nx, Ny))
+    U = np.zeros((Nt, Nx, Ny, Nz))
     U[0] = u.copy()
 
     # Time-stepping loop
     for n in range(1, Nt):
         un = u.copy()
         # Interior update
-        u[1:-1, 1:-1] = (
-            un[1:-1, 1:-1]
-            + rx * (un[2:, 1:-1] - 2 * un[1:-1, 1:-1] + un[:-2, 1:-1])
-            + ry * (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, :-2])
+        u[1:-1, 1:-1, 1:-1] = (
+            un[1:-1, 1:-1, 1:-1]
+            + rx * (un[2:, 1:-1, 1:-1] - 2 * un[1:-1, 1:-1, 1:-1] + un[:-2, 1:-1, 1:-1])
+            + ry * (un[1:-1, 2:, 1:-1] - 2 * un[1:-1, 1:-1, 1:-1] + un[1:-1, :-2, 1:-1])
+            + rz * (un[1:-1, 1:-1, 2:] - 2 * un[1:-1, 1:-1, 1:-1] + un[1:-1, 1:-1, :-2])
         )
         # Boundary conditions
         if bc == "dirichlet":
-            u[0, :] = u[-1, :] = u[:, 0] = u[:, -1] = 0.0
+            u[0, :, :] = u[-1, :, :] = u[:, 0, :] = u[:, -1, :] = u[:, :, 0] = u[:, :, -1] = 0.0
         elif bc == "neumann":
-            u[0, :] = u[1, :]
-            u[-1, :] = u[-2, :]
-            u[:, 0] = u[:, 1]
-            u[:, -1] = u[:, -2]
+            u[0, :, :] = u[1, :, :]
+            u[-1, :, :] = u[-2, :, :]
+            u[:, 0, :] = u[:, 1, :]
+            u[:, -1, :] = u[:, -2, :]
+            u[:, :, 0] = u[:, :, 1]
+            u[:, :, -1] = u[:, :, -2]
         elif bc == "periodic":
-            u[0, :] = un[-2, :]
-            u[-1, :] = un[1, :]
-            u[:, 0] = un[:, -2]
-            u[:, -1] = un[:, 1]
+            u[0, :, :] = un[-2, :, :]
+            u[-1, :, :] = un[1, :, :]
+            u[:, 0, :] = un[:, -2, :]
+            u[:, -1, :] = un[:, 1, :]
+            u[:, :, 0] = un[:, :, -2]
+            u[:, :, -1] = un[:, :, 1]
         else:
             raise ValueError(f"Unknown bc: {bc}")
         U[n] = u.copy()
@@ -76,22 +78,22 @@ def solve_2d_heat_equation(Lx: float,
     return U, dt
 
 # --- Animation Generator ---
-def create_animation_gif(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
+def create_animation_gif(U, Lx, Ly, Lz, initial, bc, Gamma, frame_skip, dt, z_slice):
     """
-    Create and save a GIF animation from the solution array U.
+    Create and save a GIF animation from the 3D solution array U at a fixed z-slice.
     """
-    Nt, Nx, Ny = U.shape
+    Nt, Nx, Ny, Nz = U.shape
     fig, ax = plt.subplots()
     x = np.linspace(0, Lx, Nx)
     y = np.linspace(0, Ly, Ny)
-    im = ax.imshow(U[0].T, cmap='viridis', origin='lower', extent=[0, Lx, 0, Ly], vmin=U.min(), vmax=U.max())
-    ax.set_title(f"2D Heat Eq — init={initial}, bc={bc}, Gamma={Gamma:.2f}")
+    im = ax.imshow(U[0, :, :, z_slice].T, cmap='viridis', origin='lower', extent=[0, Lx, 0, Ly], vmin=U.min(), vmax=U.max())
+    ax.set_title(f"2D Slice at z={z_slice} — init={initial}, bc={bc}, Gamma={Gamma:.2f}")
     ax.set_xlabel("x")
     ax.set_ylabel("y")
     plt.colorbar(im, ax=ax, label="u")
 
     def update(frame):
-        im.set_data(U[frame].T)
+        im.set_data(U[frame, :, :, z_slice].T)
         return [im]
 
     idx = list(range(0, Nt, frame_skip))
@@ -108,33 +110,34 @@ def create_animation_gif(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
     return gif_path
 
 # --- Plotly Figure Generator ---
-def create_plotly_figure(u, Lx, Ly, time_label):
+def create_plotly_figure(u, Lx, Ly, Lz, time_label, z_slice):
     """
-    Create an interactive Plotly heatmap for a given time slice u.
+    Create an interactive Plotly heatmap for a given time slice u at a fixed z-slice.
     """
     x = np.linspace(0, Lx, u.shape[0])
     y = np.linspace(0, Ly, u.shape[1])
     fig = go.Figure(data=go.Heatmap(
-        z=u.T,
+        z=u[:, :, z_slice].T,
         x=x,
         y=y,
         colorscale='viridis'
     ))
     fig.update_layout(
-        title=f"Heat Distribution at t={time_label}",
+        title=f"Heat Distribution at t={time_label}, z={z_slice}",
         xaxis_title='x',
         yaxis_title='y'
     )
     return fig
 
-# --- Gradio Interface Logic ---
-def gradio_interface(lx, ly, t_max, gamma, nx, ny, initial, bc, frame_skip):
-    nx, ny, frame_skip = int(nx), int(ny), int(frame_skip)
-    U, dt = solve_2d_heat_equation(
-        Lx=lx, Ly=ly, t_max=t_max, Gamma=gamma, Nx=nx, Ny=ny,
+# --- Simulation Runner ---
+def run_simulation(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip, z_slice_percent):
+    U, dt = solve_3d_heat_equation(
+        Lx=lx, Ly=ly, Lz=lz, t_max=t_max, Gamma=gamma, Nx=nx, Ny=ny, Nz=nz,
         initial=initial, bc=bc
     )
     Nt = U.shape[0]
+    # Compute z_slice from percentage
+    z_slice = int((z_slice_percent / 100) * (nz - 1))
     # Compute indices for t=0, t/4, 3t/4, t
     idx0 = 0
     idx1 = round((Nt - 1) / 4)
@@ -145,25 +148,32 @@ def gradio_interface(lx, ly, t_max, gamma, nx, ny, initial, bc, frame_skip):
     u1 = U[idx1]
     u2 = U[idx2]
     u3 = U[idx3]
-    # Create Plotly figures
-    fig0 = create_plotly_figure(u0, lx, ly, "0")
-    fig1 = create_plotly_figure(u1, lx, ly, f"{idx1*dt:.2f}")
-    fig2 = create_plotly_figure(u2, lx, ly, f"{idx2*dt:.2f}")
-    fig3 = create_plotly_figure(u3, lx, ly, f"{idx3*dt:.2f}")
-    # Create GIF
-    gif_path = create_animation_gif(U, lx, ly, initial, bc, gamma, frame_skip, dt)
+    # Create Plotly figures for the z_slice
+    fig0 = create_plotly_figure(u0, lx, ly, lz, "0", z_slice)
+    fig1 = create_plotly_figure(u1, lx, ly, lz, f"{idx1*dt:.2f}", z_slice)
+    fig2 = create_plotly_figure(u2, lx, ly, lz, f"{idx2*dt:.2f}", z_slice)
+    fig3 = create_plotly_figure(u3, lx, ly, lz, f"{idx3*dt:.2f}", z_slice)
+    # Create GIF for the z_slice
+    gif_path = create_animation_gif(U, lx, ly, lz, initial, bc, gamma, frame_skip, dt, z_slice)
     return gif_path, fig0, fig1, fig2, fig3
 
+# --- Gradio Interface Logic ---
+def gradio_interface(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip, z_slice_percent):
+    nx, ny, nz, frame_skip = int(nx), int(ny), int(nz), int(frame_skip)
+    return run_simulation(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip, z_slice_percent)
+
 # --- Gradio UI Layout ---
-with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
-    gr.Markdown("# ♨️ 2D Heat Equation Simulator\nAdjust parameters and run the simulation.")
+with gr.Blocks(theme=gr.themes.Soft(), title="3D Heat Simulator") as demo:
+    gr.Markdown("# ♨️ 3D Heat Equation Simulator\nAdjust parameters and run the simulation.")
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("## Domain & Grid")
             lx_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Lx")
             ly_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Ly")
+            lz_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Lz")
             nx_slider = gr.Slider(3, 200, 50, 1, label="Nx")
             ny_slider = gr.Slider(3, 200, 50, 1, label="Ny")
+            nz_slider = gr.Slider(3, 200, 50, 1, label="Nz")
 
             gr.Markdown("## Simulation")
             t_slider = gr.Slider(0.01, 5.0, 0.5, 0.01, label="t_max")
@@ -179,6 +189,8 @@ with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
 
             gr.Markdown("## Animation")
             frame_skip_slider = gr.Slider(1, 50, 5, 1, label="Frame Skip")
+            z_slice_percent_slider = gr.Slider(0, 100, 50, 1, label="Z-Slice Position (%)")
+
             run_btn = gr.Button("Run Simulation", variant="primary")
         with gr.Column(scale=3):
             gif_output = gr.Image(label="Animation")
@@ -188,15 +200,15 @@ with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
             with gr.Row():
                 plot3 = gr.Plot(label="t=3t/4")
                 plot4 = gr.Plot(label="t=t")
-    inputs_list = [lx_slider, ly_slider, t_slider, gamma_slider,
-                   nx_slider, ny_slider, initial_dropdown, bc_dropdown, frame_skip_slider]
+    inputs_list = [lx_slider, ly_slider, lz_slider, t_slider, gamma_slider,
+                   nx_slider, ny_slider, nz_slider, initial_dropdown, bc_dropdown, frame_skip_slider, z_slice_percent_slider]
     outputs_list = [gif_output, plot1, plot2, plot3, plot4]
     run_btn.click(fn=gradio_interface, inputs=inputs_list, outputs=outputs_list)
     gr.Examples(
         examples=[
-            [1.0, 1.0, 0.5, 0.1, 50, 50, "gaussian", "dirichlet", 5],
-            [2.0, 1.0, 1.0, 0.05, 60, 30, "sinusoidal", "periodic", 10],
-            [1.0, 1.0, 0.2, 0.2, 80, 80, "step", "neumann", 2],
+            [1.0, 1.0, 1.0, 0.5, 0.1, 50, 50, 50, "gaussian", "dirichlet", 5, 50],
+            [2.0, 1.0, 1.5, 1.0, 0.05, 60, 30, 40, "sinusoidal", "periodic", 10, 50],
+            [1.0, 1.0, 1.0, 0.2, 0.2, 80, 80, 80, "step", "neumann", 2, 50],
         ],
         inputs=inputs_list,
         outputs=outputs_list,