Update app.py

app.py

@@ -1,48 +1,31 @@
 import numpy as np
 import pyvista as pv
+import plotly.graph_objects as go
 import gradio as gr
+from scipy.spatial import Delaunay
 
-# --- Core Simulation and Plotting Function ---
 def solve_and_plot_interactive(Lx: float,
                                Ly: float,
                                t_max: float,
-                               M: int, # New: Number of time steps for the slider
+                               M: int,
                                Gamma: float = 0.1,
                                Nx: int = 50,
                                Ny: int = 50,
                                initial: str = "gaussian",
                                bc: str = "dirichlet"):
     """
-    Solves the 2D heat equation and displays the result in an interactive
-    PyVista window with a time slider.
-
-    Args:
-        Lx (float): Length of the domain in the x-direction.
-        Ly (float): Length of the domain in the y-direction.
-        t_max (float): Maximum simulation time.
-        M (int): Number of equidistant time steps to be available on the slider.
-        Gamma (float): Thermal diffusivity.
-        Nx (int): Number of grid points in the x-direction.
-        Ny (int): Number of grid points in the y-direction.
-        initial (str): Initial condition type.
-        bc (str): Boundary condition type.
+    Solves the 2D heat equation and returns an interactive Plotly figure
+    that can be rendered in a web browser.
     """
-    # --- 1. Simulation Setup ---
-    # Spatial grid
+    # --- 1. Simulation Setup (Same as before) ---
     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 more robust stability
     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')
+
     u = np.zeros((Nx, Ny))
     if initial == "gaussian":
         u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2) / (2*(Lx/10)**2)))
@@ -53,171 +36,140 @@ def solve_and_plot_interactive(Lx: float,
         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}")
 
-    # --- 2. Solve the Heat Equation ---
-    # Select M equidistant time indices to store for visualization
+    # --- 2. Solve the Heat Equation (Same as before) ---
     time_indices = np.linspace(0, Nt - 1, M, dtype=int)
     U_slider = np.zeros((M, Nx, Ny))
     store_idx = 0
-
     if 0 in time_indices:
         U_slider[store_idx] = u.copy()
         store_idx += 1
 
-    # Time-stepping loop
     for n in range(1, Nt):
         un = u.copy()
-        # Interior update using finite differences
         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":
-            u[0, :] = u[1, :]
-            u[-1, :] = u[-2, :]
-            u[:, 0] = u[:, 1]
-            u[:, -1] = u[:, -2]
+            u[0, :], u[-1, :], u[:, 0], u[:, -1] = u[1, :], u[-2, :], u[:, 1], u[:, -2]
         elif bc == "periodic":
-            # Note: A true periodic BC often uses ghost cells. This is a simpler implementation.
-            u[0, :] = un[-2, :]
-            u[-1, :] = un[1, :]
-            u[:, 0] = un[:, -2]
-            u[:, -1] = un[:, 1]
+            u[0, :], u[-1, :], u[:, 0], u[:, -1] = un[-2, :], un[1, :], un[:, -2], un[:, 1]
         
-        # Store the frame if it's one of the selected time steps
-        if n in time_indices:
-            if store_idx < M:
-                U_slider[store_idx] = u.copy()
-                store_idx += 1
-
-    # --- 3. Interactive Visualization with PyVista ---
-    # Create a 2D mesh (a structured grid) in 3D space (with Z=0)
-    grid = pv.StructuredGrid()
-    grid.points = np.c_[X.flatten('F'), Y.flatten('F'), np.zeros(Nx * Ny)]
-    grid.dimensions = [Nx, Ny, 1]
-
-    # Add the initial temperature data as scalars to the grid
-    # The data needs to be flattened in 'C' (row-major) order for PyVista
-    grid['temperature'] = U_slider[0, :, :].flatten('C')
-
-    # Set up the plotter
-    plotter = pv.Plotter()
-    plotter.add_mesh(grid, scalars='temperature', cmap='viridis',
-                     scalar_bar_args={'title': 'Temperature'})
-    plotter.view_xy() # Set camera to look down the Z-axis
-
-    # Define the callback function that updates the plot when the slider moves
-    def update_plot(time_step_index):
-        # Get the integer index from the slider
-        idx = int(time_step_index)
-        # Update the scalars on the grid
-        grid['temperature'] = U_slider[idx, :, :].flatten('C')
-        # Optional: Add a text annotation for the current time
-        time_value = (time_indices[idx] / (Nt-1)) * t_max if Nt > 1 else 0
-        plotter.add_text(f"Time: {time_value:.2f}s", name='time_label')
-
-    # Add the slider widget to the plotter
-    plotter.add_slider_widget(
-        callback=update_plot,
-        rng=[0, M - 1], # The slider range corresponds to the indices of U_slider
-        value=0,
-        title="Time Step",
-        style='modern'
+        if n in time_indices and store_idx < M:
+            U_slider[store_idx] = u.copy()
+            store_idx += 1
+            
+    # --- 3. Create a Plotly Figure for Web ---
+    # We use Delaunay triangulation to create the mesh faces for Plotly
+    points_2d = np.vstack([X.ravel(), Y.ravel()]).T
+    tri = Delaunay(points_2d)
+    
+    # Create the figure
+    fig = go.Figure()
+
+    # Add one mesh trace for each time step. We'll make only the first one visible.
+    for i in range(M):
+        time_value = (time_indices[i] / (Nt-1)) * t_max if Nt > 1 else 0
+        z_data = U_slider[i, :, :].flatten()
+        fig.add_trace(
+            go.Mesh3d(
+                x=X.flatten(), y=Y.flatten(), z=z_data,
+                i=tri.simplices[:, 0], j=tri.simplices[:, 1], k=tri.simplices[:, 2],
+                intensity=z_data,
+                colorscale='Viridis',
+                name=f'Time: {time_value:.2f}s',
+                showscale=True if i == 0 else False, # Show colorbar only once
+                visible=(i == 0) # Make only the first trace visible
+            )
+        )
+
+    # Create the slider
+    steps = []
+    for i in range(len(fig.data)):
+        time_value = (time_indices[i] / (Nt-1)) * t_max if Nt > 1 else 0
+        step = dict(
+            method="update",
+            args=[{"visible": [False] * len(fig.data)}],  # hide all traces
+            label=f"{time_value:.2f}s"
+        )
+        step["args"][0]["visible"][i] = True  # show the i-th trace
+        steps.append(step)
+
+    sliders = [dict(
+        active=0,
+        currentvalue={"prefix": "Time: "},
+        pad={"t": 50},
+        steps=steps
+    )]
+
+    # Update the layout of the figure
+    fig.update_layout(
+        title=f'2D Heat Eq — init={initial}, bc={bc}',
+        scene=dict(
+            xaxis_title='X',
+            yaxis_title='Y',
+            zaxis_title='Temperature'
+        ),
+        sliders=sliders
     )
     
-    # Display the plotter window. This is a blocking call.
-    # The script will pause here until you close the PyVista window.
-    plotter.show()
+    return fig
 
 
 # --- Gradio Interface Function ---
 def gradio_interface(lx, ly, t_max, m_steps, gamma, nx, ny, initial, bc):
-    """Wrapper function to connect Gradio inputs to the simulation."""
-    # Ensure integer types for grid dimensions
     nx, ny, m_steps = int(nx), int(ny), int(m_steps)
-    
-    # Run the simulation and launch the interactive plot
-    solve_and_plot_interactive(
+    # This function now returns a Plotly figure object
+    return solve_and_plot_interactive(
         Lx=lx, Ly=ly, t_max=t_max, M=m_steps, Gamma=gamma, Nx=nx, Ny=ny,
         initial=initial, bc=bc
     )
-    # This function no longer needs to return anything to Gradio
-    return "Simulation window launched. Please check your desktop."
 
 
 # --- Gradio UI Definition ---
 with gr.Blocks(theme=gr.themes.Soft(), title="2D Heat Simulator") as demo:
     gr.Markdown("# ♨️ 2D Heat Equation Simulator")
-    gr.Markdown("Adjust parameters and click 'Run' to launch an interactive window where you can pan, zoom, and scrub through time.")
+    gr.Markdown("Adjust parameters and click 'Run' to generate an interactive plot directly in your browser.")
 
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("## Domain & Grid")
+            gr.Markdown("## Simulation Parameters")
             lx_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Lx (Domain Length X)")
             ly_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Ly (Domain Length Y)")
-            nx_slider = gr.Slider(10, 200, 50, 1, label="Nx (Grid Points X)")
-            ny_slider = gr.Slider(10, 200, 50, 1, label="Ny (Grid Points Y)")
-
-            gr.Markdown("## Simulation")
-            t_slider = gr.Slider(0.01, 5.0, 0.5, 0.01, label="t_max (Total Time)")
+            nx_slider = gr.Slider(10, 100, 40, 1, label="Nx (Grid Points X)") # Reduced max for performance
+            ny_slider = gr.Slider(10, 100, 40, 1, label="Ny (Grid Points Y)")
+            t_slider = gr.Slider(0.01, 2.0, 0.5, 0.01, label="t_max (Total Time)")
             gamma_slider = gr.Slider(0.001, 1.0, 0.1, 0.001, label="Gamma (Diffusivity)")
-
-            gr.Markdown("## Conditions")
-            initial_dropdown = gr.Dropdown(
-                ["gaussian", "random", "sinusoidal", "step"], value="gaussian", label="Initial Condition"
-            )
-            bc_dropdown = gr.Dropdown(
-                ["dirichlet", "neumann", "periodic"], value="dirichlet", label="Boundary Condition"
-            )
-
-            gr.Markdown("## Interactive Plot")
-            # New slider to control the number of time steps in the interactive window
-            m_slider = gr.Slider(2, 200, 40, 1, label="M (Time Steps on Slider)")
+            m_slider = gr.Slider(10, 100, 30, 1, label="M (Time Steps on Slider)") # Reduced max
+            
+            with gr.Row():
+                initial_dropdown = gr.Dropdown(["gaussian", "random", "sinusoidal", "step"], value="gaussian", label="Initial")
+                bc_dropdown = gr.Dropdown(["dirichlet", "neumann", "periodic"], value="dirichlet", label="Boundary")
             
             run_btn = gr.Button("Run Simulation", variant="primary")
 
-        with gr.Column(scale=2):
-            # The output is now a simple text confirmation
-            status_output = gr.Textbox(label="Status")
-            gr.Markdown(
-                """
-                ### How to Use:
-                1.  Set your desired simulation parameters on the left.
-                2.  `M (Time Steps on Slider)` controls how many time points will be available in the interactive view. Higher values give smoother time control but use more memory.
-                3.  Click **Run Simulation**.
-                4.  A new window will open.
-                    - **Left-Click + Drag**: Rotate the view.
-                    - **Right-Click + Drag**: Pan the view.
-                    - **Scroll Wheel**: Zoom in and out.
-                    - **Use the Slider**: To move through simulation time.
-                5. Close the interactive window to run another simulation.
-                """
-            )
-
+        with gr.Column(scale=3):
+            # The output is now a gr.Plot component that will render the Plotly figure
+            plot_output = gr.Plot(label="Interactive Heatmap")
 
-    # Connect the button to the interface function
     inputs_list = [lx_slider, ly_slider, t_slider, m_slider, gamma_slider,
                    nx_slider, ny_slider, initial_dropdown, bc_dropdown]
     
-    run_btn.click(fn=gradio_interface, inputs=inputs_list, outputs=status_output)
+    run_btn.click(fn=gradio_interface, inputs=inputs_list, outputs=plot_output)
     
-    # Define some example configurations
     gr.Examples(
         examples=[
-            [1.0, 1.0, 0.5, 50, 0.1, 50, 50, "gaussian", "dirichlet"],
-            [2.0, 1.0, 1.0, 80, 0.05, 60, 30, "sinusoidal", "periodic"],
-            [1.0, 1.0, 0.2, 40, 0.2, 80, 80, "step", "neumann"],
+            [1.0, 1.0, 0.5, 30, 0.1, 40, 40, "gaussian", "dirichlet"],
+            [2.0, 1.0, 1.0, 50, 0.05, 50, 25, "sinusoidal", "periodic"],
         ],
         inputs=inputs_list,
-        outputs=[status_output],
+        outputs=plot_output,
         fn=gradio_interface,
-        cache_examples=False # It's better to rerun live simulations
+        cache_examples=True # Caching is fine for Plotly objects
     )
 
 if __name__ == "__main__":