Update app.py

app.py

@@ -1,85 +1,60 @@
 import numpy as np
-import plotly.graph_objects as go # Changed from matplotlib
 import gradio as gr
-# Removed: matplotlib.pyplot, matplotlib.animation.FuncAnimation, tempfile
-
-def solve_and_create_plotly_animation(Lx: float,
-                                     Ly: float,
-                                     t_max: float,
-                                     Gamma: float = 0.1,
-                                     Nx: int = 50,
-                                     Ny: int = 50,
-                                     initial: str = "gaussian",
-                                     bc: str = "dirichlet",
-                                     frame_skip: int = 1):
+import plotly.graph_objects as go
+
+# --- 1. Simulation Core (Modified to return raw data) ---
+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"):
     """
-    Solve the 2D heat equation u_t = Gamma*(u_xx + u_yy) and return an interactive Plotly animation.
-    Initial conditions: {"gaussian", "random", "sinusoidal", "step"}
-    Boundary conditions: {"dirichlet", "neumann", "periodic"}
+    Solves the 2D heat equation and returns the entire time-series data.
     """
     # Spatial grid
-    x_coords = np.linspace(0, Lx, Nx)
-    y_coords = np.linspace(0, Ly, Ny)
-    dx, dy = x_coords[1] - x_coords, y_coords[1] - y_coords
-
-    if dx == 0 or dy == 0: # Should not happen with Nx, Ny >= 3 from Gradio sliders
-        # Create an empty figure with an error message for Gradio
-        fig = go.Figure()
-        fig.update_layout(title_text="Error: Nx and Ny must be > 1 for dx, dy calculation.",
-                          xaxis_showticklabels=False, yaxis_showticklabels=False)
-        return fig
-
+    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
-    # Ensure dt is positive, handle potential division by zero if Gamma is zero or dx/dy are problematic
-    denominator = (2 * Gamma * (1/dx**2 + 1/dy**2))
-    if denominator <= 0: # Avoid division by zero or negative dt
-        fig = go.Figure()
-        fig.update_layout(title_text=f"Error: Unstable dt parameters (Gamma={Gamma}, dx={dx}, dy={dy}). Check Gamma.",
-                          xaxis_showticklabels=False, yaxis_showticklabels=False)
-        return fig
-    dt = 1.0 / denominator
-    
+    # 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
-    if Nt <=1: # Ensure there's at least an initial and one computed frame for animation
-        Nt = 2
-
     rx, ry = Gamma * dt / dx**2, Gamma * dt / dy**2
 
     # Initial condition
-    X, Y = np.meshgrid(x_coords, y_coords, indexing='ij') # Use x_coords, y_coords
-    u = np.zeros((Nx, Ny)) # Initialize u
+    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)))
     elif initial == "random":
         u = np.random.rand(Nx, Ny)
     elif initial == "sinusoidal":
-        # Ensure Lx and Ly are not zero to avoid division by zero
-        kx = 2 * np.pi / Lx if Lx > 0 else 0
-        ky = 2 * np.pi / Ly if Ly > 0 else 0
+        kx, ky = 2 * np.pi / Lx, 2 * np.pi / Ly
         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: # Should not happen due to dropdown
-        fig = go.Figure()
-        fig.update_layout(title_text=f"Error: Unknown initial condition: {initial}",
-                          xaxis_showticklabels=False, yaxis_showticklabels=False)
-        return fig
+    else:
+        raise ValueError(f"Unknown initial condition: {initial}")
 
-    # Storage for solution frames to be animated
-    # We will store only the frames selected by frame_skip
-    frames_to_animate_data =
-    frames_to_animate_data.append(u.copy().T) # Transpose for heatmap like imshow(origin='lower')
+    # 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
+        # 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
@@ -89,168 +64,169 @@ def solve_and_create_plotly_animation(Lx: float,
             u[:, 0] = u[:, 1]
             u[:, -1] = u[:, -2]
         elif bc == "periodic":
-            # Ensure indices are valid for periodic BCs (Nx, Ny >= 3)
-            if Nx >= 3 and Ny >=3:
-                u[0, :] = un[-2, :]
-                u[-1, :] = un[1, :]
-                u[:, 0] = un[:, -2]
-                u[:, -1] = un[:, 1]
-            else: # Fallback for very small grids where periodic BCs are ill-defined
-                u[0, :] = u[1, :] # Effectively Neumann
-                u[-1, :] = u[-2, :]
-                u[:, 0] = u[:, 1]
-                u[:, -1] = u[:, -2]
-        else: # Should not happen
-            fig = go.Figure()
-            fig.update_layout(title_text=f"Error: Unknown bc: {bc}",
-                              xaxis_showticklabels=False, yaxis_showticklabels=False)
-            return fig
-        
-        if n % frame_skip == 0 or n == Nt -1: # Store frame based on frame_skip
-            frames_to_animate_data.append(u.copy().T) # Transpose
-
-    if not frames_to_animate_data:
-        fig = go.Figure()
-        fig.update_layout(title_text="Error: No frames generated for animation.",
-                          xaxis_showticklabels=False, yaxis_showticklabels=False)
-        return fig
-
-    # Determine global min/max for consistent colorscale
-    global_min = np.min([np.min(frame) for frame in frames_to_animate_data])
-    global_max = np.max([np.max(frame) for frame in frames_to_animate_data])
-    if global_min == global_max: # Avoid issues if all values are the same
-        global_max += 1e-6 if global_max == 0 else abs(global_max * 0.01) # Add a tiny epsilon
-
-    # Create Plotly animation
+            # Implement periodic boundary conditions correctly
+            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 the full data history and the time step size
+    return U, dt
+
+# --- 2. Plotly Animation Generator (New Function) ---
+def create_plotly_animation(U, Lx, Ly, initial, bc, Gamma, frame_skip, dt):
+    """Creates an interactive Plotly animation from the raw simulation data."""
+    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
+        idx.append(Nt - 1)
+
+    # Subsample the data for animation frames
+    frames_to_animate_data = U[idx]
+
+    # Create the figure with frames
     fig = go.Figure(
-        data=[go.Heatmap(
-            z=frames_to_animate_data,
-            x=x_coords, # Use actual coordinates for axes
-            y=y_coords, # Use actual coordinates for axes
-            colorscale='Viridis',
-            zmin=global_min,
-            zmax=global_max,
-            showscale=True,
-            colorbar_title_text="u"
-        )],
-        layout=go.Layout(
-            title_text=f"2D Heat Eq: init={initial}, bc={bc}, Gamma={Gamma:.3f}",
-            xaxis_title_text="x",
-            yaxis_title_text="y",
-            # yaxis_scaleanchor="x", # Makes pixels square if dx=dy and Lx=Ly
-            font=dict(size=10),
-            # Ensure plot updates don't change axis ranges during animation
-            xaxis_range=[0, Lx],
-            yaxis_range=[0, Ly],
-        ),
-        frames=[
-            go.Frame(
-                data=[go.Heatmap(
-                    z=frame_data,
-                    x=x_coords,
-                    y=y_coords,
-                    colorscale='Viridis',
-                    zmin=global_min,
-                    zmax=global_max
-                )],
-                name=f"frame{k}"
-            ) for k, frame_data in enumerate(frames_to_animate_data)
-        ]
+        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 animation control buttons
+    # 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',
+        zmin=vmin,
+        zmax=vmax,
+        colorbar=dict(title="u")
+    ))
+
+    # Create and configure the play/pause button
     fig.update_layout(
-        updatemenus=), # No transition smoothing
-                dict(label="Pause",
-                     method="animate",
-                     args=[[None], {"frame": {"duration": 0, "redraw": False},
-                                    "mode": "immediate",
-                                    "transition": {"duration": 0}}])
-            ],
+        updatemenus=[dict(
+            type="buttons",
             direction="left",
-            pad={"r": 10, "t": 70}, # Adjust padding
-            x=0.1, xanchor="left",
-            y=0, yanchor="top"
+            x=0.1,
+            xanchor="left",
+            y=1.15,
+            yanchor="top",
+            buttons=[dict(label="Play",
+                          method="animate",
+                          args=[None, {"frame": {"duration": 50, "redraw": True},
+                                       "fromcurrent": True, "transition": {"duration": 0}}]),
+                     dict(label="Pause",
+                          method="animate",
+                          args=[[None], {"frame": {"duration": 0, "redraw": False},
+                                         "mode": "immediate", "transition": {"duration": 0}}])]
         )]
     )
 
-    # Add frame slider
+    # Create and configure the frame slider
+    sliders = [dict(
+        active=0,
+        yanchor="top",
+        xanchor="left",
+        currentvalue=dict(
+            font=dict(size=16),
+            prefix="Time: ",
+            visible=True,
+            xanchor="right"
+        ),
+        pad=dict(b=10, t=50),
+        len=0.9,
+        x=0.1,
+        y=0,
+        steps=[dict(
+            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
+            method="animate")
+            for i, f in enumerate(fig.frames)]
+    )]
+
     fig.update_layout(
-        sliders=[dict(
-            active=0,
-            steps=[
-                dict(label=str(k),
-                     method="animate",
-                     args=[[f"frame{k}"],
-                           {"mode": "immediate",
-                            "frame": {"duration": 100, "redraw": True}, # Duration if played
-                            "transition": {"duration": 0}}]) # No transition for slider drag
-                for k in range(len(frames_to_animate_data))
-            ],
-            currentvalue={"prefix": "Frame: ", "visible": True, "xanchor": "right"},
-            pad={"t": 60, "b": 10} # Adjust padding
-        )]
+        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) ---
 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 new function that returns a Plotly figure
-    return solve_and_create_plotly_animation(
+    
+    # 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, frame_skip=frame_skip
+        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
 
+# --- 4. Gradio UI Layout (Modified to use gr.Plot) ---
 with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
-    gr.Markdown("# ♨️ 2D Heat Equation Simulator (Interactive with Plotly)\nAdjust parameters and run the simulation.")
+    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")
             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 (min 3)") # Min 3 for some BCs
-            ny_slider = gr.Slider(3, 200, 50, 1, label="Ny (min 3)") # Min 3 for some BCs
+            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 (Diffusion Coeff.)")
+            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"], value="gaussian", label="Initial Condition"
+                ["gaussian", "random", "sinusoidal", "step"], "gaussian", label="Initial"
             )
             bc_dropdown = gr.Dropdown(
-                ["dirichlet", "neumann", "periodic"], value="dirichlet", label="Boundary Condition"
+                ["dirichlet", "neumann", "periodic"], "dirichlet", label="Boundary"
             )
+            
             gr.Markdown("## Animation")
-            frame_skip_slider = gr.Slider(1, 50, 5, 1, label="Frame Skip (Higher = Fewer Frames)")
+            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):
-            # Changed from gr.Image to gr.Plot
-            plot_output = gr.Plot(label="Interactive Heatmap Animation")
+            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
 
     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],
+        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],
-            [0.5, 0.5, 0.1, 0.5, 30, 30, "random", "dirichlet", 1],
+        ],
         inputs=inputs_list,
-        outputs=[plot_output], # Output is now a single Plot component
-        fn=gradio_interface,
-        # cache_examples=True # Consider enabling if simulations are slow and inputs are identical
+        outputs=[plot_output],
+        fn=gradio_interface
     )
-    gr.Markdown("### How to Interact:"
-                "\n- **Play/Pause Buttons:** Control the animation playback (located below the plot title)."
-                "\n- **Slider:** Drag to scrub through frames (located below the plot)."
-                "\n- **Plotly Modebar (top right of plot on hover):**"
-                "\n  - **Zoom:** Use zoom tools (box zoom, zoom in/out icons, scroll wheel)."
-                "\n  - **Pan:** Use the pan tool to move the view when zoomed."
-                "\n  - **Autoscale/Reset:** Return to the default view or reset axes."
-               )
 
 if __name__ == "__main__":
+    # To run this, you will need to install plotly: pip install plotly
     demo.launch()