app.py

import time
import torch
import warnings
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt

# Import Burgers' equation components
from data_burgers import exact_solution as exact_solution_burgers
from model_io_burgers import load_model
from model_v2 import Encoder, Decoder, Propagator_concat as Propagator, Model

# Import Advection-Diffusion components
from data_adv_dif import exact_solution as exact_solution_adv_dif
from model_io_adv_dif import load_model as load_model_adv_dif
from model_adv_dif import Encoder as Encoder2D, Decoder as Decoder2D, Propagator_concat as Propagator2D, Model as Model2D

warnings.filterwarnings("ignore")

# ========== Burgers' Equation Setup ==========
def get_burgers_model(input_dim, latent_dim):
    encoder = Encoder(input_dim, latent_dim)
    decoder = Decoder(latent_dim, input_dim)
    propagator = Propagator(latent_dim)
    return Model(encoder, decoder, propagator)

flexi_prop_model = get_burgers_model(128, 2)
checkpoint = torch.load("dcp_best.pt", map_location='cpu')
flexi_prop_model.load_state_dict(checkpoint['model_state_dict'])
flexi_prop_model.eval()


# Precompute contour plots
z1_vals = np.linspace(-10, 0.5, 250)
z2_vals = np.linspace(5, 32, 250)
Z1, Z2 = np.meshgrid(z1_vals, z2_vals)
latent_grid = np.stack([Z1.ravel(), Z2.ravel()], axis=1)

# Convert to tensor for decoding
latent_tensors = torch.tensor(latent_grid, dtype=torch.float32)

# Decode latent vectors and compute properties
with torch.no_grad():
    decoded_signals = flexi_prop_model.decoder(latent_tensors)

sharpness = []
peak_positions = []
x_vals = np.linspace(0, 2, decoded_signals.shape[1])
dx = x_vals[1] - x_vals[0]

for signal in decoded_signals.numpy():
    grad_u = np.gradient(signal, dx)
    sharpness.append(np.max(np.abs(grad_u)))
    peak_positions.append(x_vals[np.argmax(signal)])

sharpness = np.array(sharpness).reshape(Z1.shape)
peak_positions = np.array(peak_positions).reshape(Z1.shape)

def plot_burgers_comparison(Re, tau, t_0):
    dt = 2.0 / 500.0
    t_final = t_0 + tau * dt
    x_exact = exact_solution_burgers(Re, t_final)
    
    tau_tensor, Re_tensor, xt = torch.tensor([tau]).float()[:, None], torch.tensor([Re]).float()[:, None], torch.tensor([exact_solution_burgers(Re, t_0)]).float()[:, None]

    with torch.no_grad():
        _, x_hat_tau, *_ = flexi_prop_model(xt, tau_tensor, Re_tensor)

    fig, ax = plt.subplots(figsize=(9, 5))
    ax.plot(xt.squeeze(), '--', linewidth=3, alpha=0.5, color="C0")
    ax.plot(x_hat_tau.squeeze(), 'D', markersize=5, color="C2")
    ax.plot(x_exact.squeeze(), linewidth=2, alpha=0.5, color="Black")
    ax.set_title(f"Comparison ($t_0$={t_0:.2f} → $t_f$={t_final:.2f}), τ={tau}", fontsize=14)
    ax.legend(["Initial", "Flexi-Prop", "True"])
    return fig

def burgers_update(Re, tau, t0):
    fig1 = plot_burgers_comparison(Re, tau, t0)

    _ = flexi_prop_model(torch.randn(1, 1, 128), torch.tensor([[tau]]), torch.tensor([[Re]]))

    # Latent space visualization
    latent_fig = plot_latent_interpretation(Re, tau, t0)

    return fig1, latent_fig

def plot_latent_interpretation(Re, tau, t_0):
    tau_tensor = torch.tensor([tau]).float()[:, None]
    Re_tensor = torch.tensor([Re]).float()[:, None]
    x_t = exact_solution_burgers(Re, t_0)
    xt = torch.tensor([x_t]).float()[:, None]

    with torch.no_grad():
        _, _, _, _, z_tau = flexi_prop_model(xt, tau_tensor, Re_tensor)
    
    z_tau = z_tau.squeeze().numpy()

    fig, axes = plt.subplots(1, 2, figsize=(10, 3))

    # Sharpness Plot
    c1 = axes[0].pcolormesh(Z1, Z2, sharpness, cmap='plasma', shading='gouraud')
    axes[0].scatter(z_tau[0], z_tau[1], color='red', marker='o', s=50, label="Current State")
    axes[0].set_ylabel("$Z_2$", fontsize=14)
    axes[0].set_title("Sharpness Encoding", fontsize=14)
    fig.colorbar(c1, ax=axes[0])
    axes[0].legend()

    # Peak Position Plot
    c2 = axes[1].pcolormesh(Z1, Z2, peak_positions, cmap='viridis', shading='gouraud')
    axes[1].scatter(z_tau[0], z_tau[1], color='red', marker='o', s=50, label="Current State")
    axes[1].set_title("Peak position Encoding", fontsize=14)
    fig.colorbar(c2, ax=axes[1], label="Peak Position")
    
    # Remove redundant y-axis labels on the second plot for better aesthetics
    axes[1].set_yticklabels([])

    # Set a single x-axis label centered below both plots
    fig.supxlabel("$Z_1$", fontsize=14)

    return fig

# ========== Advection-Diffusion Setup ==========
def get_adv_dif_model(latent_dim, output_dim):
    encoder = Encoder2D(latent_dim)
    decoder = Decoder2D(latent_dim)
    propagator = Propagator2D(latent_dim)
    return Model2D(encoder, decoder, propagator)

adv_dif_model = get_adv_dif_model(3, 128)
adv_dif_model, _, _, _ = load_model_adv_dif(
    "adv_diff_2d_best.pt", 
    adv_dif_model,
    map_location='cpu'
)


def generate_3d_visualization(Re, t_0, tau):
    dt = 2 / 500
    t = t_0 + tau * dt

    U_initial = exact_solution_adv_dif(Re, t_0)
    U_evolved = exact_solution_adv_dif(Re, t)

    if np.isnan(U_initial).any() or np.isnan(U_evolved).any():
        return None

    fig3d = plt.figure(figsize=(12, 5))
    ax3d = fig3d.add_subplot(111, projection='3d')

    x_vals = np.linspace(-2, 2, U_initial.shape[1])
    y_vals = np.linspace(-2, 2, U_initial.shape[0])
    X, Y = np.meshgrid(x_vals, y_vals)

    surf1 = ax3d.plot_surface(X, Y, U_initial, cmap="viridis", alpha=0.6, label="Initial")
    surf2 = ax3d.plot_surface(X, Y, U_evolved, cmap="plasma", alpha=0.8, label="Evolved")

    ax3d.set_xlim(-3, 3)
    ax3d.set_xlabel("x")
    ax3d.set_ylabel("y")
    ax3d.set_zlabel("u(x,y,t)")
    ax3d.view_init(elev=25, azim=-45)
    ax3d.set_box_aspect((2,1,1))

    fig3d.colorbar(surf1, ax=ax3d, shrink=0.5, label="Initial")
    fig3d.colorbar(surf2, ax=ax3d, shrink=0.5, label="Evolved")
    ax3d.set_title(f"Solution Evolution\nInitial ($t_0$={t_0:.2f}) vs Evolved ($t_f$={t:.2f})")

    plt.tight_layout()
    plt.close(fig3d)
    return fig3d

def adv_dif_comparison(Re, t_0, tau):
    dt = 2 / 500
    exact_initial = exact_solution_adv_dif(Re, t_0)
    exact_final = exact_solution_adv_dif(Re, t_0 + tau * dt)

    if np.isnan(exact_initial).any() or np.isnan(exact_final).any():
        return None

    x_in = torch.tensor(exact_initial, dtype=torch.float32)[None, None, :, :]
    Re_in = torch.tensor([[Re]], dtype=torch.float32)
    tau_in = torch.tensor([[tau]], dtype=torch.float32)

    with torch.no_grad():
        x_hat, x_hat_tau, *_ = adv_dif_model(x_in, tau_in, Re_in)

    pred = x_hat_tau.squeeze().numpy()
    if pred.shape != exact_final.shape:
        return None

    mse = np.square(pred - exact_final)

    fig, axs = plt.subplots(1, 3, figsize=(15, 4))

    for ax, (data, title) in zip(axs, [(pred, "Model Prediction"),
                                       (exact_final, "Exact Solution"),
                                       (mse, "MSE Error")]):
        if title == "MSE Error":
            im = ax.imshow(data, cmap="viridis", vmin=0, vmax=1e-2)
            plt.colorbar(im, ax=ax, fraction=0.075)
        else:
            im = ax.imshow(data, cmap="jet")

        ax.set_title(title)
        ax.axis("off")

    plt.tight_layout()
    plt.close(fig)
    return fig

def update_initial_plot(Re, t_0):
    exact_initial = exact_solution_adv_dif(Re, t_0)
    fig, ax = plt.subplots(figsize=(5, 5))
    im = ax.imshow(exact_initial, cmap='jet')
    plt.colorbar(im, ax=ax)
    ax.set_title('Initial State')
    return fig

# Gradio Interface
with gr.Blocks(title="Latent Evolution Operator (LEO): PDE Prediction Suite") as app:
    gr.Markdown("# Latent Evolution Operator (LEO): Unified PDE Prediction Interface")

    with gr.Tabs():
        # 1D Burgers' Equation Tab
        with gr.Tab("1D Burgers' Equation"):
            gr.Markdown(r"""
                    ## LEO: Single-Shot Prediction for Nonlinear PDEs.
                    **Governing Equation (1D Burgers' Equation):**
                    $$
                    \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} = \nu \frac{\partial^2 u}{\partial x^2}
                    $$
                    **Key Advantages:**  
                    ✔️ ✔️ **Learns a state-to-state operator**: \(\mathfrak{F}:\boldsymbol{u}(t,\zeta)\mapsto \boldsymbol{u}(t+\tau,\zeta)\), \(t\in(0,T]\), \(\tau\ge0\)
                    ✔️ **Parametric control**: Embeds system parameters in latent space  
                    ✔️ **50-100× faster** than AE-LSTM baselines  
                    
                    **Physically Interpretable Latent Space - Disentanglement:**  
                    <div align="left">
                    $$
                    Z_1 \text{ Encodes Peak Location, } Z_2 \text{ Predominantly Encodes Re (Sharpness)}
                    $$
                    </div>

                """)
            
            with gr.Row():
                with gr.Column():
                    re_burgers = gr.Slider(425, 2350, 1040, label="Reynolds Number (Re)")
                    gr.Markdown("""
                    <span style="font-size: 14px; color: #444;">
                    • <b>Interpolation Zone</b>: [822, 1022]  
                    • <b>Left Extrapolation</b>: [400, 600]  
                    • <b>Right Extrapolation</b>: [2200, 2400]  
                    </span>
                    """)
                    
                    tau_burgers = gr.Slider(150, 450, 315, label="Time Steps (τ)")
                    gr.Markdown("""
                    <span style="font-size: 14px; color: #444;">
                    • <b>Interpolation Zone</b>: [340, 368]  
                    • <b>Left Extrapolation</b>: [150, 178]  
                    • <b>Right Extrapolation</b>: [398, 425]  
                    </span>
                    """)
                
                    t0_burgers = gr.Number(0.4, label="Initial Time")
                    latent_plot = gr.Plot(label="Latent Space Dynamics")

                with gr.Column():
                    burgers_plot = gr.Plot()
            # with gr.Row():
            #     latent_plot = gr.Plot(label="Latent Space Dynamics")

            re_burgers.change(burgers_update, [re_burgers, tau_burgers, t0_burgers], 
                            [burgers_plot, latent_plot])
            tau_burgers.change(burgers_update, [re_burgers, tau_burgers, t0_burgers], 
                            [burgers_plot, latent_plot])
            t0_burgers.change(burgers_update, [re_burgers, tau_burgers, t0_burgers], 
                            [burgers_plot, latent_plot])

        # 2D Advection-Diffusion Tab
        with gr.Tab("2D Advection-Diffusion"):
            gr.Markdown(r"""
                ## 🌪️ 2D Advection-Diffusion Visualization
                **Governing Equation:**
                $$
                \frac{\partial u}{\partial t} + c \frac{\partial u}{\partial x} = \nu \left( \frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2} \right)
                $$
                """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    re_adv = gr.Slider(1, 10, 9, label="Advection Coefficient (α)")
                    gr.Markdown("""
                    <span style="font-size: 14px; color: #444;">
                    • <b>Interpolation Zone</b>: [2.09, 2.99]  
                    • <b>Left Extrapolation</b>: [1.0, 1.9]  
                    • <b>Right Extrapolation</b>: [9.1, 10.0]  
                    </span>
                    """)
                
                    t0_adv = gr.Number(0.45, label="Initial Time")
                
                    tau_adv = gr.Slider(150, 425, 225, label="Time Steps (τ)")
                    gr.Markdown("""
                    <span style="font-size: 14px; color: #444;">
                    • <b>Interpolation Zone</b>: [364.7, 392.2]  
                    • <b>Left Extrapolation</b>: [150, 177.5]  
                    • <b>Right Extrapolation</b>: [397.5, 425]  
                    </span>
                    """)
                
                    initial_plot_adv = gr.Plot(label="Initial State")

                
                with gr.Column(scale=3):
                    with gr.Row():
                        three_d_plot_adv = gr.Plot(label="3D Evolution")
                    with gr.Row():
                        comparison_plots_adv = gr.Plot(label="Model Comparison")

            def adv_update(Re, t0, tau):
                return (
                    generate_3d_visualization(Re, t0, tau),
                    adv_dif_comparison(Re, t0, tau),
                    update_initial_plot(Re, t0)
                )

            for component in [re_adv, t0_adv, tau_adv]:
                component.change(adv_update, [re_adv, t0_adv, tau_adv], 
                               [three_d_plot_adv, comparison_plots_adv, initial_plot_adv])

            app.load(lambda: adv_update(8, 0.35, 225), 
                   outputs=[three_d_plot_adv, comparison_plots_adv, initial_plot_adv])

app.launch()