Update app.py

app.py

@@ -2,7 +2,7 @@ import numpy as np
 import gradio as gr
 import plotly.graph_objects as go
 
-# --- 1. Simulation Core (Modified to return raw data) ---
+# --- 1. Simulation Core ---
 def solve_2d_heat_equation(Lx: float,
                            Ly: float,
                            t_max: float,
@@ -11,23 +11,15 @@ def solve_2d_heat_equation(Lx: float,
                            Ny: int = 50,
                            initial: str = "gaussian",
                            bc: str = "dirichlet"):
-    """
-    Solves the 2D heat equation and returns the entire time-series data.
-    """
-    # Spatial grid
+    # [Unchanged code from original]
     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
-    # A small factor (0.9) is added for extra stability margin
     dt = 0.9 / (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')
     if initial == "gaussian":
         u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2) / (2*(Lx/10)**2)))
@@ -40,22 +32,15 @@ def solve_2d_heat_equation(Lx: float,
         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))
     U[0] = u.copy()
-
-    # Time-stepping loop
     for n in range(1, Nt):
         un = u.copy()
-        # Interior update using a vectorized operation
         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
         if bc == "dirichlet":
             u[0, :] = u[-1, :] = u[:, 0] = u[:, -1] = 0.0
         elif bc == "neumann":
@@ -64,7 +49,6 @@ def solve_2d_heat_equation(Lx: float,
             u[:, 0] = u[:, 1]
             u[:, -1] = u[:, -2]
         elif bc == "periodic":
-            # Implement periodic boundary conditions correctly
             u[0, :] = un[-2, :]
             u[-1, :] = un[1, :]
             u[:, 0] = un[:, -2]
@@ -72,31 +56,21 @@ def solve_2d_heat_equation(Lx: float,
         else:
             raise ValueError(f"Unknown bc: {bc}")
         U[n] = u.copy()
-
-    # Return the full data history and the time step size
     return U, dt
 
-# --- 2. Plotly Animation Generator (New Function) ---
+# --- 2. Plotly Animation Generator ---
 def create_plotly_animation(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
-    """Creates an interactive Plotly animation from the raw simulation data."""
+    # [Unchanged code from original]
     Nt, Nx, Ny = U.shape
     vmin, vmax = U.min(), U.max()
-
-    # Create a list of frame indices to animate based on frame_skip
     idx = list(range(0, Nt, frame_skip))
-    if not idx or idx[-1] != Nt - 1: # Ensure the last frame is always included
+    if not idx or idx[-1] != Nt - 1:
         idx.append(Nt - 1)
-
-    # Subsample the data for animation frames
     frames_to_animate_data = U[idx]
-
-    # Create the figure with frames
     fig = go.Figure(
         frames=[go.Frame(data=go.Heatmap(z=frame_data.T, zmin=vmin, zmax=vmax), name=str(i))
                 for i, frame_data in enumerate(frames_to_animate_data)]
     )
-
-    # Add the initial heatmap trace (frame 0) that will be displayed first
     fig.add_trace(go.Heatmap(
         z=frames_to_animate_data[0].T,
         colorscale='viridis',
@@ -104,8 +78,6 @@ def create_plotly_animation(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
         zmax=vmax,
         colorbar=dict(title="u")
     ))
-
-    # Create and configure the play/pause button
     fig.update_layout(
         updatemenus=[dict(
             type="buttons",
@@ -124,8 +96,6 @@ def create_plotly_animation(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
                                          "mode": "immediate", "transition": {"duration": 0}}])]
         )]
     )
-
-    # Create and configure the frame slider
     sliders = [dict(
         active=0,
         yanchor="top",
@@ -144,45 +114,40 @@ def create_plotly_animation(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
             args=[[f.name], {"frame": {"duration": 0, "redraw": True},
                               "mode": "immediate",
                               "transition": {"duration": 0}}],
-            label=f"{idx[i]*dt:.2f}s", # Display time in seconds on slider
+            label=f"{idx[i]*dt:.2f}s",
             method="animate")
             for i, f in enumerate(fig.frames)]
     )]
-
     fig.update_layout(
         title=f"2D Heat Eq — init={initial}, bc={bc}, Gamma={Gamma:.2f}",
         xaxis_title="x",
         yaxis_title="y",
         sliders=sliders,
-        # Set aspect ratio to match the domain
         yaxis=dict(scaleanchor="x", scaleratio=Ly/Lx if Lx > 0 else 1)
     )
-
     return fig
 
-
-# --- 3. Gradio Interface Logic (Modified to connect new functions) ---
+# --- 3. Gradio Interface Logic ---
 def gradio_interface(lx, ly, t_max, gamma, nx, ny, initial, bc, frame_skip):
     """Main function for the Gradio interface."""
     nx, ny, frame_skip = int(nx), int(ny), int(frame_skip)
-    
-    # Call the solver to get the raw simulation data and time step
-    U, dt = solve_2d_heat_equation(
-        Lx=lx, Ly=ly, t_max=t_max, Gamma=gamma, Nx=nx, Ny=ny,
-        initial=initial, bc=bc
-    ) # <<< THIS IS THE CORRECTED LINE with the closing parenthesis
-
-    # Create the Plotly figure from the data
-    fig = create_plotly_animation(
-        U=U, Lx=lx, Ly=ly, initial=initial, bc=bc, Gamma=gamma,
-        frame_skip=frame_skip, dt=dt
-    )
-    return fig
+    try:
+        U, dt = solve_2d_heat_equation(
+            Lx=lx, Ly=ly, t_max=t_max, Gamma=gamma, Nx=nx, Ny=ny,
+            initial=initial, bc=bc
+        )
+        fig = create_plotly_animation(
+            U=U, Lx=lx, Ly=ly, initial=initial, bc=bc, Gamma=gamma,
+            frame_skip=frame_skip, dt=dt
+        )
+        return fig
+    except Exception as e:
+        print(f"Error in simulation: {e}")
+        return None
 
-# --- 4. Gradio UI Layout (Modified to use gr.Plot) ---
+# --- 4. Gradio UI Layout ---
 with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
     gr.Markdown("# ♨️ Interactive 2D Heat Equation Simulator\nAdjust parameters and run the simulation.")
-    
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("## Domain & Grid")
@@ -190,11 +155,9 @@ with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
             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"
@@ -202,20 +165,15 @@ with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
             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):
             gr.Markdown("### Interactive Heatmap Animation\nUse the play/pause buttons, drag the slider, or use your mouse/trackpad to zoom and pan the plot.")
-            plot_output = gr.Plot(label="Interactive Heatmap") # Changed from gr.Image to gr.Plot
-
+            plot_output = gr.Plot(label="Interactive Heatmap")
     inputs_list = [lx_slider, ly_slider, t_slider, gamma_slider,
                    nx_slider, ny_slider, initial_dropdown, bc_dropdown, frame_skip_slider]
-    
     run_btn.click(fn=gradio_interface, inputs=inputs_list, outputs=plot_output)
-    
     gr.Examples(
         examples=[
             [1.0, 1.0, 0.5, 0.1, 50, 50, "gaussian", "dirichlet", 5],
@@ -228,5 +186,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
     )
 
 if __name__ == "__main__":
-    # To run this, you will need to install plotly: pip install plotly
     demo.launch()