Update app.py

app.py

@@ -1,302 +1,224 @@
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib.animation import FuncAnimation
-import tempfile
+import pyvista as pv
 import gradio as gr
-import plotly.graph_objects as go
-import base64
-from fastapi import FastAPI
-from pydantic import BaseModel
 
-# --- Simulation Core (3D) ---
-def solve_3d_heat_equation(Lx: float, Ly: float, Lz: float,
-                           t_max: float,
-                           Gamma: float = 0.1,
-                           Nx: int = 30, Ny: int = 30, Nz: int = 30,
-                           initial: str = "gaussian",
-                           bc: str = "dirichlet"):
+# --- 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
+                               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.
+    """
+    # --- 1. Simulation Setup ---
+    # Spatial grid
     x = np.linspace(0, Lx, Nx)
     y = np.linspace(0, Ly, Ny)
-    z = np.linspace(0, Lz, Nz)
-    
-    # Corrected dx, dy, dz calculation
-    dx = x[1] - x[0] if Nx > 1 else Lx
-    dy = y[1] - y[0] if Ny > 1 else Ly
-    dz = z[1] - z[0] if Nz > 1 else Lz
-
-    if dx == 0 or dy == 0 or dz == 0:
-        raise ValueError("Grid spacing (dx, dy, dz) cannot be zero. Ensure Nx, Ny, Nz > 1.")
+    dx, dy = x[1] - x[0], y[1] - y[0]
+    if dx == 0 or dy == 0:
+        raise ValueError("Nx and Ny must be > 1.")
 
-    # Stability condition for 3D FTCS scheme
-    dt = 0.5 / (Gamma * (1/dx**2 + 1/dy**2 + 1/dz**2))
+    # 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, rz = Gamma * dt / dx**2, Gamma * dt / dy**2, Gamma * dt / dz**2
-    
-    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
+    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 + (Z - Lz/2)**2) / (2*(max(Lx,Ly,Lz)/10)**2)))
+        u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2) / (2*(Lx/10)**2)))
     elif initial == "random":
-        u = np.random.rand(Nx, Ny, Nz)
+        u = np.random.rand(Nx, Ny)
     elif initial == "sinusoidal":
-        kx = 2 * np.pi / Lx if Lx > 0 else 0
-        ky = 2 * np.pi / Ly if Ly > 0 else 0
-        kz = 2 * np.pi / Lz if Lz > 0 else 0
-        u = np.sin(kx * X) * np.sin(ky * Y) * np.sin(kz * Z)
+        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) & (Z < Lz/2), 1.0, 0.0)
+        u = np.where((X < Lx/2) & (Y < Ly/2), 1.0, 0.0)
     else:
         raise ValueError(f"Unknown initial condition: {initial}")
 
-    U = np.zeros((Nt, Nx, Ny, Nz))
-    U[0] = u.copy()  # Store initial condition
+    # --- 2. Solve the Heat Equation ---
+    # Select M equidistant time indices to store for visualization
+    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()
-        if Nx > 2 and Ny > 2 and Nz > 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])
-            )
-
+        # 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":
-            if Nx > 0: u[0, :, :] = u[-1, :, :] = 0.0
-            if Ny > 0: u[:, 0, :] = u[:, -1, :] = 0.0
-            if Nz > 0: u[:, :, 0] = u[:, :, -1] = 0.0
+            u[0, :] = u[-1, :] = u[:, 0] = u[:, -1] = 0.0
         elif bc == "neumann":
-            if Nx > 1: u[0, :, :] = u[1, :, :]; u[-1, :, :] = u[-2, :, :]
-            if Ny > 1: u[:, 0, :] = u[:, 1, :]; u[:, -1, :] = u[:, -2, :]
-            if Nz > 1: 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":
-            if Nx > 1: u[0, :, :] = un[-2, :, :]; u[-1, :, :] = un[1, :, :]
-            if Ny > 1: u[:, 0, :] = un[:, -2, :]; u[-1, :] = un[:, 1, :]
-            if Nz > 1: u[:, :, 0] = un[:, :, -2]; u[:, :, -1] = un[:, :, 1]
-        else:
-            raise ValueError(f"Unknown bc: {bc}")
-        U[n] = u.copy()
-    return U, dt
-
-# --- Animation Generator (3D Slice) ---
-def create_animation_gif_3d_slice(U, Lx, Ly, Lz, initial, bc, Gamma, frame_skip, dt):
-    Nt, Nx, Ny, Nz = U.shape
-    fig, ax = plt.subplots()
-    
-    slice_z_idx = Nz // 2 if Nz > 0 else 0
-    z_coord_slice = np.linspace(0, Lz, Nz)[slice_z_idx] if Nz > 0 else 0
-    
-    data_slice = U[0, :, :, slice_z_idx].T if Nt > 0 and Nz > 0 else np.zeros((Ny, Nx))
+            # 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]
+        
+        # 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'
+    )
     
-    vmin_val = U.min() if Nt > 0 and U.size > 0 and np.all(np.isfinite(U)) else 0
-    vmax_val = U.max() if Nt > 0 and U.size > 0 and np.all(np.isfinite(U)) else 1
-    if vmin_val == vmax_val: vmax_val = vmin_val + 1
-
-    im = ax.imshow(data_slice, cmap='viridis', origin='lower', 
-                   extent=[0, Lx, 0, Ly], vmin=vmin_val, vmax=vmax_val)
-    ax.set_title(f"3D Heat Eq (xy-slice at z={z_coord_slice:.2f})\ninit={initial}, bc={bc}, Gamma={Gamma:.2f}")
-    ax.set_xlabel("x")
-    ax.set_ylabel("y")
-    plt.colorbar(im, ax=ax, label="u")
-
-    def update(frame):
-        if Nz > 0:
-            im.set_data(U[frame, :, :, slice_z_idx].T)
-        return [im]
-
-    # Corrected idx generation
-    if Nt <= 1:
-        idx = [0]  # Only initial frame if Nt is 0 or 1
-    else:
-        current_frame_skip = max(1, frame_skip)
-        idx = list(range(0, Nt, current_frame_skip))
-        if (Nt - 1) not in idx:
-            idx.append(Nt - 1)
-        idx = sorted(list(set(idx)))
-
-    ani = FuncAnimation(fig, update, frames=idx, blit=True)
-
-    with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as tmpfile:
-        ani.save(tmpfile.name, writer='pillow', fps=max(1, 30 // max(1, frame_skip)))
-        gif_path = tmpfile.name
+    # Display the plotter window. This is a blocking call.
+    # The script will pause here until you close the PyVista window.
+    plotter.show()
 
-    plt.close(fig)
-    return gif_path
 
-# --- Plotly Figure Generator (3D Volume) ---
-def create_plotly_figure_3d(u_3d, Lx, Ly, Lz, time_label):
-    Nx, Ny, Nz = u_3d.shape
-    if Nx == 0 or Ny == 0 or Nz == 0:
-        return go.Figure(layout_title_text=f"3D Heat Distribution (No Data) at t={time_label}")
-
-    x_coords = np.linspace(0, Lx, Nx)
-    y_coords = np.linspace(0, Ly, Ny)
-    z_coords = np.linspace(0, Lz, Nz)
+# --- 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)
     
-    X, Y, Z = np.meshgrid(x_coords, y_coords, z_coords, indexing='ij')
-
-    vmin = np.min(u_3d[np.isfinite(u_3d)]) if np.any(np.isfinite(u_3d)) else 0
-    vmax = np.max(u_3d[np.isfinite(u_3d)]) if np.any(np.isfinite(u_3d)) else 1
-    if vmin == vmax: vmax = vmin + 0.1
-
-    fig = go.Figure(data=go.Volume(
-        x=X.flatten(),
-        y=Y.flatten(),
-        z=Z.flatten(),
-        value=u_3d.flatten(),
-        isomin=vmin, 
-        isomax=vmax,
-        opacity=0.1, 
-        surface_count=17, 
-        colorscale='viridis'
-    ))
-    fig.update_layout(
-        title=f"3D Heat Distribution at t={time_label}",
-        scene=dict(
-            xaxis_title='x',
-            yaxis_title='y',
-            zaxis_title='z',
-            aspectmode='cube' 
-        )
-    )
-    return fig
-
-# --- Simulation Runner (Extracted Logic for 3D) ---
-def run_simulation_3d(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip):
-    nx = max(3, int(nx))
-    ny = max(3, int(ny))
-    nz = max(3, int(nz))
-
-    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
+    # Run the simulation and launch the interactive plot
+    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
     )
-    Nt = U.shape[0]
-    
-    idx0 = 0
-    idx1 = round((Nt - 1) / 4) if Nt > 1 else 0
-    idx2 = round(3 * (Nt - 1) / 4) if Nt > 1 else 0
-    idx3 = Nt - 1 if Nt > 1 else 0
-    
-    u0 = U[idx0]
-    u1 = U[idx1]
-    u2 = U[idx2]
-    u3 = U[idx3]
-    
-    fig0 = create_plotly_figure_3d(u0, lx, ly, lz, "0")
-    fig1 = create_plotly_figure_3d(u1, lx, ly, lz, f"{idx1*dt:.2f}")
-    fig2 = create_plotly_figure_3d(u2, lx, ly, lz, f"{idx2*dt:.2f}")
-    fig3 = create_plotly_figure_3d(u3, lx, ly, lz, f"{idx3*dt:.2f}")
-    
-    gif_path = create_animation_gif_3d_slice(U, lx, ly, lz, initial, bc, gamma, frame_skip, dt)
-    return gif_path, fig0, fig1, fig2, fig3
+    # This function no longer needs to return anything to Gradio
+    return "Simulation window launched. Please check your desktop."
 
-# --- Gradio Interface Logic (3D) ---
-def gradio_interface_3d(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip):
-    nx_int, ny_int, nz_int = int(nx), int(ny), int(nz)
-    frame_skip_int = max(1, int(frame_skip))
 
-    nx_int = max(3, nx_int)
-    ny_int = max(3, ny_int)
-    nz_int = max(3, nz_int)
+# --- 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.")
 
-    gif_path, fig0, fig1, fig2, fig3 = run_simulation_3d(
-        lx, ly, lz, t_max, gamma, nx_int, ny_int, nz_int, initial, bc, frame_skip_int
-    )
-    return gif_path, fig0, fig1, fig2, fig3
-
-# --- Gradio UI Layout (3D) ---
-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. Animation shows a central xy-slice. Grid (Nx,Ny,Nz) min 3.")
     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, 60, 20, 1, label="Nx (min 3, e.g., 20-40)")
-            ny_slider = gr.Slider(3, 60, 20, 1, label="Ny (min 3, e.g., 20-40)")
-            nz_slider = gr.Slider(3, 60, 20, 1, label="Nz (min 3, e.g., 20-40)")
+            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, 1.0, 0.1, 0.01, label="t_max (e.g., 0.1-0.5)")
-            gamma_slider = gr.Slider(0.001, 1.0, 0.1, 0.001, label="Gamma")
+            t_slider = gr.Slider(0.01, 5.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"], "gaussian", label="Initial"
+                ["gaussian", "random", "sinusoidal", "step"], value="gaussian", label="Initial Condition"
             )
             bc_dropdown = gr.Dropdown(
-                ["dirichlet", "neumann", "periodic"], "dirichlet", label="Boundary"
+                ["dirichlet", "neumann", "periodic"], value="dirichlet", label="Boundary Condition"
             )
 
-            gr.Markdown("## Animation")
-            frame_skip_slider = gr.Slider(1, 50, 5, 1, label="Frame Skip (e.g., 5-10)")
-            run_btn = gr.Button("Run 3D Simulation", variant="primary")
+            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)")
             
-        with gr.Column(scale=3):
-            gif_output = gr.Image(label="Animation (Central XY Slice)")
-            with gr.Row():
-                plot1 = gr.Plot(label="3D Volume at t=0")
-                plot2 = gr.Plot(label="3D Volume at t=t/4")
-            with gr.Row():
-                plot3 = gr.Plot(label="3D Volume at t=3t/4")
-                plot4 = gr.Plot(label="3D Volume at t=t_max")
-                
-    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]
-    outputs_list = [gif_output, plot1, plot2, plot3, plot4]
+            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.
+                """
+            )
+
+
+    # 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_3d, inputs=inputs_list, outputs=outputs_list)
+    run_btn.click(fn=gradio_interface, inputs=inputs_list, outputs=status_output)
     
+    # Define some example configurations
     gr.Examples(
         examples=[
-            [1.0, 1.0, 1.0, 0.1, 0.1, 20, 20, 20, "gaussian", "dirichlet", 5],
-            [1.5, 1.0, 0.5, 0.2, 0.05, 25, 20, 15, "sinusoidal", "periodic", 10],
-            [1.0, 1.0, 1.0, 0.05, 0.2, 15, 15, 15, "step", "neumann", 2]
+            [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"],
         ],
         inputs=inputs_list,
-        outputs=outputs_list,
-        fn=gradio_interface_3d,
-        cache_examples=False 
+        outputs=[status_output],
+        fn=gradio_interface,
+        cache_examples=False # It's better to rerun live simulations
     )
 
-# --- FastAPI Setup for API Endpoint (3D) ---
-app = FastAPI()
-
-app = gr.mount_gradio_app(app, demo, path="/")
-
-class SimulationParams3D(BaseModel):
-    lx: float
-    ly: float
-    lz: float 
-    t_max: float
-    gamma: float
-    nx: int
-    ny: int
-    nz: int 
-    initial: str
-    bc: str
-    frame_skip: int
-
-@app.post("/simulate_3d") 
-def simulate_3d_api(params: SimulationParams3D):
-    params.frame_skip = max(1, params.frame_skip)
-    params.nx = max(3, params.nx)
-    params.ny = max(3, params.ny)
-    params.nz = max(3, params.nz)
-
-    gif_path, fig0, fig1, fig2, fig3 = run_simulation_3d(**params.dict())
-    with open(gif_path, "rb") as f:
-        gif_data = base64.b64encode(f.read()).decode('utf-8')
-    return {
-        "gif_base64": gif_data,
-        "plot0_3d_volume": fig0.to_json(), 
-        "plot1_3d_volume": fig1.to_json(),
-        "plot2_3d_volume": fig2.to_json(),
-        "plot3_3d_volume": fig3.to_json()
-    }
-
 if __name__ == "__main__":
     demo.launch()