Update app.py

app.py

@@ -4,32 +4,35 @@ from matplotlib.animation import FuncAnimation
 import tempfile
 import gradio as gr
 import plotly.graph_objects as go
+import base64
+from fastapi import FastAPI
+from pydantic import BaseModel
 
-# --- Simulation Core ---
+# --- Simulation Core (3D) ---
 def solve_3d_heat_equation(Lx: float, Ly: float, Lz: float,
-                           t_max: float, Gamma: float = 0.1,
-                           Nx: int = 50, Ny: int = 50, Nz: int = 50,
-                           initial: str = "gaussian", bc: str = "dirichlet"):
-    """
-    Solve the 3D heat equation u_t = Gamma*(u_xx + u_yy + u_zz) and return the solution array U and time step dt.
-    """
-    # Spatial grid
+                           t_max: float,
+                           Gamma: float = 0.1,
+                           Nx: int = 30, Ny: int = 30, Nz: int = 30, # Reduced default for 3D
+                           initial: str = "gaussian",
+                           bc: str = "dirichlet"):
     x = np.linspace(0, Lx, Nx)
     y = np.linspace(0, Ly, Ny)
     z = np.linspace(0, Lz, Nz)
-    dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
-    if dx <= 0 or dy <= 0 or dz <= 0:
+    dx, dy, dz = x[1] - x, y[1] - y, z[1] - z
+
+    if dx == 0 or dy == 0 or dz == 0:
         raise ValueError("Nx, Ny, and Nz must be > 1.")
 
-    # Time stepping for stability
-    dt = 1.0 / (2 * Gamma * (1/dx**2 + 1/dy**2 + 1/dz**2))
+    # Stability condition for 3D FTCS scheme
+    dt = 0.5 / (Gamma * (1/dx**2 + 1/dy**2 + 1/dz**2)) # Adjusted for 3D
     Nt = int(np.ceil(t_max / dt)) + 1
+    
     rx, ry, rz = Gamma * dt / dx**2, Gamma * dt / dy**2, Gamma * dt / dz**2
-
-    # Initial condition
+    
     X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
+
     if initial == "gaussian":
-        u = np.exp(-((X - Lx/2)**2 / (2*(Lx/10)**2) + (Y - Ly/2)**2 / (2*(Ly/10)**2) + (Z - Lz/2)**2 / (2*(Lz/10)**2)))
+        u = np.exp(-(((X - Lx/2)**2 + (Y - Ly/2)**2 + (Z - Lz/2)**2) / (2*(Lx/10)**2)))
     elif initial == "random":
         u = np.random.rand(Nx, Ny, Nz)
     elif initial == "sinusoidal":
@@ -40,24 +43,23 @@ def solve_3d_heat_equation(Lx: float, Ly: float, Lz: float,
     else:
         raise ValueError(f"Unknown initial condition: {initial}")
 
-    # Storage for solution
     U = np.zeros((Nt, Nx, Ny, Nz))
-    U[0] = u.copy()
+    U = u.copy()
 
-    # Time-stepping loop
     for n in range(1, Nt):
         un = u.copy()
-        # Interior update
         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])
         )
-        # Boundary conditions
+
         if bc == "dirichlet":
-            u[0, :, :] = u[-1, :, :] = u[:, 0, :] = u[:, -1, :] = u[:, :, 0] = u[:, :, -1] = 0.0
-        elif bc == "neumann":
+            u[0, :, :] = u[-1, :, :] = 0.0
+            u[:, 0, :] = u[:, -1, :] = 0.0
+            u[:, :, 0] = u[:, :, -1] = 0.0
+        elif bc == "neumann": # Zero flux
             u[0, :, :] = u[1, :, :]
             u[-1, :, :] = u[-2, :, :]
             u[:, 0, :] = u[:, 1, :]
@@ -74,109 +76,125 @@ def solve_3d_heat_equation(Lx: float, Ly: float, Lz: float,
         else:
             raise ValueError(f"Unknown bc: {bc}")
         U[n] = u.copy()
-
     return U, dt
 
-# --- Animation Generator ---
-def create_animation_gif(U, Lx, Ly, Lz, initial, bc, Gamma, frame_skip, dt, z_slice):
-    """
-    Create and save a GIF animation from the 3D solution array U at a fixed z-slice.
-    """
+# --- 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()
-    x = np.linspace(0, Lx, Nx)
-    y = np.linspace(0, Ly, Ny)
-    im = ax.imshow(U[0, :, :, z_slice].T, cmap='viridis', origin='lower', extent=[0, Lx, 0, Ly], vmin=U.min(), vmax=U.max())
-    ax.set_title(f"2D Slice at z={z_slice} — init={initial}, bc={bc}, Gamma={Gamma:.2f}")
+    
+    # We'll animate the central xy-slice
+    slice_z_idx = Nz // 2
+    z_coord_slice = np.linspace(0, Lz, Nz)[slice_z_idx]
+    
+    data_slice = U[0, :, :, slice_z_idx].T
+    
+    im = ax.imshow(data_slice, cmap='viridis', origin='lower', 
+                   extent=[0, Lx, 0, Ly], vmin=U.min(), vmax=U.max())
+    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):
-        im.set_data(U[frame, :, :, z_slice].T)
+        im.set_data(U[frame, :, :, slice_z_idx].T)
         return [im]
 
     idx = list(range(0, Nt, frame_skip))
-    if idx[-1] != Nt - 1:
-        idx.append(Nt - 1)
+    if not idx or (idx[-1]!= Nt - 1 and Nt > 1) : # Ensure last frame is included if Nt > 1
+         if Nt-1 not in idx: idx.append(Nt - 1)
+    if not idx and Nt ==1: # Handle case with only one time step
+        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=30)
+        ani.save(tmpfile.name, writer='pillow', fps=max(1, 30 // frame_skip)) # Adjust fps based on skip
         gif_path = tmpfile.name
 
     plt.close(fig)
     return gif_path
 
-# --- Plotly Figure Generator ---
-def create_plotly_figure(u, Lx, Ly, Lz, time_label, z_slice):
-    """
-    Create an interactive Plotly heatmap for a given time slice u at a fixed z-slice.
-    """
-    x = np.linspace(0, Lx, u.shape[0])
-    y = np.linspace(0, Ly, u.shape[1])
-    fig = go.Figure(data=go.Heatmap(
-        z=u[:, :, z_slice].T,
-        x=x,
-        y=y,
+# --- Plotly Figure Generator (3D Volume) ---
+def create_plotly_figure_3d(u_3d, Lx, Ly, Lz, time_label):
+    Nx, Ny, Nz = u_3d.shape
+    x_coords = np.linspace(0, Lx, Nx)
+    y_coords = np.linspace(0, Ly, Ny)
+    z_coords = np.linspace(0, Lz, Nz)
+    
+    X, Y, Z = np.meshgrid(x_coords, y_coords, z_coords, indexing='ij')
+
+    fig = go.Figure(data=go.Volume(
+        x=X.flatten(),
+        y=Y.flatten(),
+        z=Z.flatten(),
+        value=u_3d.flatten(),
+        isomin=u_3d.min(),
+        isomax=u_3d.max(),
+        opacity=0.1, # needs to be small to see through all surfaces
+        surface_count=20, # needs to be a large number for good volume rendering
         colorscale='viridis'
     ))
     fig.update_layout(
-        title=f"Heat Distribution at t={time_label}, z={z_slice}",
-        xaxis_title='x',
-        yaxis_title='y'
+        title=f"3D Heat Distribution at t={time_label}",
+        scene=dict(
+            xaxis_title='x',
+            yaxis_title='y',
+            zaxis_title='z'
+        )
     )
     return fig
 
-# --- Simulation Runner ---
-def run_simulation(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip, z_slice_percent):
+# --- Simulation Runner (Extracted Logic for 3D) ---
+def run_simulation_3d(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip):
     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
+        Lx=lx, Ly=ly, Lz=lz, t_max=t_max, Gamma=gamma, 
+        Nx=nx, Ny=ny, Nz=nz, initial=initial, bc=bc
     )
-    Nt = U.shape[0]
-    # Compute z_slice from percentage
-    z_slice = int((z_slice_percent / 100) * (nz - 1))
-    # Compute indices for t=0, t/4, 3t/4, t
+    Nt = U.shape
+    
     idx0 = 0
-    idx1 = round((Nt - 1) / 4)
-    idx2 = round(3 * (Nt - 1) / 4)
-    idx3 = Nt - 1
-    # Extract slices
+    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]
-    # Create Plotly figures for the z_slice
-    fig0 = create_plotly_figure(u0, lx, ly, lz, "0", z_slice)
-    fig1 = create_plotly_figure(u1, lx, ly, lz, f"{idx1*dt:.2f}", z_slice)
-    fig2 = create_plotly_figure(u2, lx, ly, lz, f"{idx2*dt:.2f}", z_slice)
-    fig3 = create_plotly_figure(u3, lx, ly, lz, f"{idx3*dt:.2f}", z_slice)
-    # Create GIF for the z_slice
-    gif_path = create_animation_gif(U, lx, ly, lz, initial, bc, gamma, frame_skip, dt, z_slice)
+    
+    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
 
-# --- Gradio Interface Logic ---
-def gradio_interface(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip, z_slice_percent):
+# --- Gradio Interface Logic (3D) ---
+def gradio_interface_3d(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip):
     nx, ny, nz, frame_skip = int(nx), int(ny), int(nz), int(frame_skip)
-    return run_simulation(lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip, z_slice_percent)
+    gif_path, fig0, fig1, fig2, fig3 = run_simulation_3d(
+        lx, ly, lz, t_max, gamma, nx, ny, nz, initial, bc, frame_skip
+    )
+    return gif_path, fig0, fig1, fig2, fig3
 
-# --- Gradio UI Layout ---
+# --- 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.")
+    gr.Markdown("# 🔥 3D Heat Equation Simulator\nAdjust parameters and run the simulation. Animation shows a central xy-slice.")
     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, 200, 50, 1, label="Nx")
-            ny_slider = gr.Slider(3, 200, 50, 1, label="Ny")
-            nz_slider = gr.Slider(3, 200, 50, 1, label="Nz")
+            lz_slider = gr.Slider(0.1, 5.0, 1.0, 0.1, label="Lz (New)") # New
+            nx_slider = gr.Slider(3, 100, 30, 1, label="Nx") # Max reduced for 3D
+            ny_slider = gr.Slider(3, 100, 30, 1, label="Ny") # Max reduced for 3D
+            nz_slider = gr.Slider(3, 100, 30, 1, label="Nz (New)") # New, max reduced
 
             gr.Markdown("## Simulation")
-            t_slider = gr.Slider(0.01, 5.0, 0.5, 0.01, label="t_max")
+            t_slider = gr.Slider(0.01, 2.0, 0.2, 0.01, label="t_max") # Max t_max reduced
             gamma_slider = gr.Slider(0.001, 1.0, 0.1, 0.001, label="Gamma")
 
             gr.Markdown("## Conditions")
@@ -188,32 +206,66 @@ with gr.Blocks(theme=gr.themes.Soft(), title="3D Heat Simulator") as demo:
             )
 
             gr.Markdown("## Animation")
-            frame_skip_slider = gr.Slider(1, 50, 5, 1, label="Frame Skip")
-            z_slice_percent_slider = gr.Slider(0, 100, 50, 1, label="Z-Slice Position (%)")
-
-            run_btn = gr.Button("Run Simulation", variant="primary")
+            frame_skip_slider = gr.Slider(1, 50, 10, 1, label="Frame Skip") # Default skip increased
+            run_btn = gr.Button("Run 3D Simulation", variant="primary")
+            
         with gr.Column(scale=3):
-            gif_output = gr.Image(label="Animation")
+            gif_output = gr.Image(label="Animation (Central XY Slice)")
             with gr.Row():
-                plot1 = gr.Plot(label="t=0")
-                plot2 = gr.Plot(label="t=t/4")
+                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="t=3t/4")
-                plot4 = gr.Plot(label="t=t")
+                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, z_slice_percent_slider]
+                   nx_slider, ny_slider, nz_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)
+    
+    run_btn.click(fn=gradio_interface_3d, inputs=inputs_list, outputs=outputs_list)
+    
     gr.Examples(
-        examples=[
-            [1.0, 1.0, 1.0, 0.5, 0.1, 50, 50, 50, "gaussian", "dirichlet", 5, 50],
-            [2.0, 1.0, 1.5, 1.0, 0.05, 60, 30, 40, "sinusoidal", "periodic", 10, 50],
-            [1.0, 1.0, 1.0, 0.2, 0.2, 80, 80, 80, "step", "neumann", 2, 50],
-        ],
+        examples=[1.0, 1.0, 1.0, 0.2, 0.1, 20, 20, 20, "gaussian", "dirichlet", 10], # Reduced Nx,Ny,Nz for example
+            [1.5, 1.0, 0.5, 0.3, 0.05, 25, 20, 15, "sinusoidal", "periodic", 15],
+            [1.0, 1.0, 1.0, 0.1, 0.2, 30, 30, 30, "step", "neumann", 5],
         inputs=inputs_list,
         outputs=outputs_list,
-        fn=gradio_interface
+        fn=gradio_interface_3d
     )
 
+# --- FastAPI Setup for API Endpoint (3D) ---
+app = FastAPI()
+
+# Mount Gradio app to FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
+
+# Define the simulation parameters model for 3D
+class SimulationParams3D(BaseModel):
+    lx: float
+    ly: float
+    lz: float # New
+    t_max: float
+    gamma: float
+    nx: int
+    ny: int
+    nz: int # New
+    initial: str
+    bc: str
+    frame_skip: int
+
+# Custom API endpoint to run 3D simulation and return results
+@app.post("/simulate_3d") # Renamed endpoint
+def simulate_3d_api(params: SimulationParams3D):
+    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(), # Indicate 3D plot
+        "plot1_3d_volume": fig1.to_json(),
+        "plot2_3d_volume": fig2.to_json(),
+        "plot3_3d_volume": fig3.to_json()
+    }
+
 if __name__ == "__main__":
     demo.launch()