Update app.py

app.py

@@ -1,26 +1,36 @@
-# app.py
 import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+import tempfile
 import gradio as gr
+import plotly.graph_objects as go
 
-# --- Simulation Core (No changes needed here) ---
-def solve_2d_heat_equation(Lx, Ly, t_max, Gamma, Nx, Ny, initial, bc):
+# --- 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"):
     """
-    Solves the 2D heat equation and returns the data and time step.
-    (This function is the same as in your original code)
+    Solve the 2D heat equation u_t = Gamma*(u_xx + u_yy) 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.")
 
     # Time stepping for stability
-    dt = 0.9 / (2 * Gamma * (1/dx**2 + 1/dy**2))
+    dt = 1.0 / (2 * Gamma * (1/dx**2 + 1/dy**2))
     Nt = int(np.ceil(t_max / dt)) + 1
     rx, ry = Gamma * dt / dx**2, Gamma * dt / dy**2
 
     # Initial condition
     X, Y = np.meshgrid(x, y, indexing='ij')
-    u = np.zeros((Nx, Ny))
     if initial == "gaussian":
         u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2) / (2*(Lx/10)**2)))
     elif initial == "random":
@@ -30,6 +40,8 @@ def solve_2d_heat_equation(Lx, Ly, t_max, Gamma, Nx, Ny, initial, bc):
         u = np.sin(kx * X) * np.sin(ky * Y)
     elif initial == "step":
         u = np.where((X < Lx/2) & (Y < Ly/2), 1.0, 0.0)
+    else:
+        raise ValueError(f"Unknown initial condition: {initial}")
 
     # Storage for solution
     U = np.zeros((Nt, Nx, Ny))
@@ -38,12 +50,13 @@ def solve_2d_heat_equation(Lx, Ly, t_max, Gamma, Nx, Ny, initial, bc):
     # 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])
         )
-        # Boundary conditions (simplified for brevity)
+        # Boundary conditions
         if bc == "dirichlet":
             u[0, :] = u[-1, :] = u[:, 0] = u[:, -1] = 0.0
         elif bc == "neumann":
@@ -51,57 +64,144 @@ def solve_2d_heat_equation(Lx, Ly, t_max, Gamma, Nx, Ny, initial, bc):
             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]
+        else:
+            raise ValueError(f"Unknown bc: {bc}")
         U[n] = u.copy()
-    
-    # Return raw data - convert numpy arrays to lists for JSON compatibility
-    return U.tolist(), dt
 
-# --- Gradio API Function ---
-def run_simulation_api(lx, ly, t_max, gamma, nx, ny, initial, bc):
+    return U, dt
+
+# --- Animation Generator ---
+def create_animation_gif(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
     """
-    API function that runs the simulation and returns data as a dictionary.
+    Create and save a GIF animation from the solution array U.
     """
+    Nt, Nx, Ny = 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}")
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+    plt.colorbar(im, ax=ax, label="u")
+
+    def update(frame):
+        im.set_data(U[frame].T)
+        return [im]
+
+    idx = list(range(0, Nt, frame_skip))
+    if idx[-1] != Nt - 1:
+        idx.append(Nt - 1)
+
+    ani = FuncAnimation(fig, update, frames=idx, blit=True)
+
+    with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as tmpfile:
+        ani.save(tmpfile.name, writer='pillow', fps=30)
+        gif_path = tmpfile.name
+
+    plt.close(fig)
+    return gif_path
+
+# --- Plotly Figure Generator ---
+def create_plotly_figure(u, Lx, Ly, time_label):
+    """
+    Create an interactive Plotly heatmap for a given time slice u.
+    """
+    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,
+        x=x,
+        y=y,
+        colorscale='viridis'
+    ))
+    fig.update_layout(
+        title=f"Heat Distribution at t={time_label}",
+        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=int(nx), Ny=int(ny),
+        Lx=lx, Ly=ly, t_max=t_max, Gamma=gamma, Nx=nx, Ny=ny,
         initial=initial, bc=bc
     )
-    
-    # Return data in a format that's easy to use in JavaScript
-    return {
-        "simulation_data": U,
-        "time_step": dt,
-        "domain": {"Lx": lx, "Ly": ly},
-        "params": {"initial": initial, "bc": bc, "gamma": gamma}
-    }
-
-# --- Gradio Interface (API only) ---
-with gr.Blocks(title="2D Heat Simulator API") as demo:
-    # We define the inputs for the API
-    lx_slider = gr.Slider(0.1, 5.0, 1.0, label="Lx")
-    ly_slider = gr.Slider(0.1, 5.0, 1.0, label="Ly")
-    nx_slider = gr.Slider(3, 200, 50, label="Nx")
-    ny_slider = gr.Slider(3, 200, 50, label="Ny")
-    t_slider = gr.Slider(0.01, 5.0, 0.5, label="t_max")
-    gamma_slider = gr.Slider(0.001, 1.0, 0.1, label="Gamma")
-    initial_dropdown = gr.Dropdown(["gaussian", "random", "sinusoidal", "step"], value="gaussian", label="Initial")
-    bc_dropdown = gr.Dropdown(["dirichlet", "neumann"], value="dirichlet", label="Boundary")
-    
-    # We only need a JSON output component for the API
-    output_json = gr.JSON(label="Simulation Output")
-    
-    # This is the crucial part: we create a hidden button to host our API endpoint
-    api_btn = gr.Button("Run Simulation")
-    
-    api_btn.click(
-        fn=run_simulation_api,
-        inputs=[
-            lx_slider, ly_slider, t_slider, gamma_slider,
-            nx_slider, ny_slider, initial_dropdown, bc_dropdown
+    Nt = U.shape[0]
+    # Compute indices for t=0, t/4, 3t/4, t
+    idx0 = 0
+    idx1 = round((Nt - 1) / 4)
+    idx2 = round(3 * (Nt - 1) / 4)
+    idx3 = Nt - 1
+    # Extract slices
+    u0 = U[idx0]
+    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)
+    return gif_path, fig0, fig1, fig2, fig3
+
+# --- 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.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")
+            nx_slider = gr.Slider(3, 200, 50, 1, label="Nx")
+            ny_slider = gr.Slider(3, 200, 50, 1, label="Ny")
+
+            gr.Markdown("## Simulation")
+            t_slider = gr.Slider(0.01, 5.0, 0.5, 0.01, label="t_max")
+            gamma_slider = gr.Slider(0.001, 1.0, 0.1, 0.001, label="Gamma")
+
+            gr.Markdown("## Conditions")
+            initial_dropdown = gr.Dropdown(
+                ["gaussian", "random", "sinusoidal", "step"], "gaussian", label="Initial"
+            )
+            bc_dropdown = gr.Dropdown(
+                ["dirichlet", "neumann", "periodic"], "dirichlet", label="Boundary"
+            )
+
+            gr.Markdown("## Animation")
+            frame_skip_slider = gr.Slider(1, 50, 5, 1, label="Frame Skip")
+            run_btn = gr.Button("Run Simulation", variant="primary")
+        with gr.Column(scale=3):
+            gif_output = gr.Image(label="Animation")
+            with gr.Row():
+                plot1 = gr.Plot(label="t=0")
+                plot2 = gr.Plot(label="t=t/4")
+            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]
+    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],
         ],
-        outputs=output_json,
-        api_name="run_simulation"  # This makes it available at /run/run_simulation
+        inputs=inputs_list,
+        outputs=outputs_list,
+        fn=gradio_interface
     )
 
 if __name__ == "__main__":
-    # Launch the Gradio app
     demo.launch()