Init commit

app.py

@@ -0,0 +1,479 @@
+import gradio as gr
+import numpy as np
+import os
+import time
+import plotly.graph_objs as go
+import matplotlib.pyplot as plt
+import shutil
+from colorama import Fore
+
+# Path to the npz folder
+npz_folder = "npz"
+
+glob_a = -2
+glob_b = -2
+glob_c = -4
+glob_d = 7
+
+
+def clear_folder(folder_path=npz_folder):
+    for filename in os.listdir(folder_path):
+        file_path = os.path.join(folder_path, filename)
+        try:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+                os.unlink(file_path)  # Remove the file or symbolic link
+            elif os.path.isdir(file_path):
+                shutil.rmtree(file_path)  # Remove the directory and its contents
+        except Exception as e:
+            print(f"Failed to delete {file_path}. Reason: {e}")
+
+
+def complex_heat_eq_solution(x, t, a=glob_a, b=glob_b, c=glob_c, d=glob_d, k=0.5):
+    global glob_a, glob_b, glob_c, glob_d
+    return (
+        np.exp(-k * t) * np.sin(np.pi * x)
+        + 0.5 * np.exp(glob_a * k * t) * np.sin(glob_b * np.pi * x)
+        + 0.25 * np.exp(glob_c * k * t) * np.sin(glob_d * np.pi * x)
+    )
+
+
+def plot_heat_equation(m, approx_type):
+    # Define grid dimensions
+    n_x = 32  # Fixed spatial grid resolution
+    n_t = 50
+
+    try:
+        loaded_values = np.load(f"npz/{approx_type}_m{m}.npz")
+    except:
+        raise gr.Error(f"First train the coefficients for {approx_type} and m = {m}")
+    alpha = loaded_values["alpha"]
+    Phi = loaded_values["Phi"]
+
+    # Create grids for x and t
+    x = np.linspace(0, 1, n_x)  # Spatial grid
+    t = np.linspace(0, 5, n_t)  # Temporal grid
+    X, T = np.meshgrid(x, t)
+
+    # Compute the real solution over the grid
+    U_real = complex_heat_eq_solution(X, T)
+
+    # Compute the selected approximation
+    U_approx = np.zeros_like(U_real)
+    for i, t_val in enumerate(t):
+        Phi_gff_at_t = Phi[i * n_x : (i + 1) * n_x]
+        U_approx[i, :] = np.dot(Phi_gff_at_t, alpha)
+
+    # Create the 3D plot with Plotly
+    traces = []
+
+    # Real solution surface with a distinct color (e.g., 'Viridis')
+    traces.append(
+        go.Surface(
+            z=U_real,
+            x=X,
+            y=T,
+            colorscale="Blues",
+            showscale=False,
+            name="Real Solution",
+            showlegend=True,
+        )
+    )
+
+    # Approximation surface with a distinct color (e.g., 'Plasma')
+    traces.append(
+        go.Surface(
+            z=U_approx,
+            x=X,
+            y=T,
+            colorscale="Reds",
+            reversescale=True,
+            showscale=False,
+            name=f"{approx_type} Approximation",
+            showlegend=True,
+        )
+    )
+
+    # Layout for the Plotly plot without controls
+    layout = go.Layout(
+        title=f"Heat Equation Approximation | Kernel = {approx_type} | m = {m}",
+        scene=dict(
+            camera=dict(
+                eye=dict(x=0, y=-2, z=0),  # Front view
+            ),
+            xaxis_title="x",
+            yaxis_title="t",
+            zaxis_title="u",
+        ),
+    )
+
+    # Config to remove modebar buttons except the save image button
+    config = {
+        "modeBarButtonsToRemove": [
+            "pan",
+            "resetCameraLastSave",
+            "hoverClosest3d",
+            "hoverCompareCartesian",
+            "zoomIn",
+            "zoomOut",
+            "select2d",
+            "lasso2d",
+            "zoomIn2d",
+            "zoomOut2d",
+            "sendDataToCloud",
+            "zoom3d",
+            "orbitRotation",
+            "tableRotation",
+        ],
+        "displayModeBar": True,  # Keep the modebar visible
+        "displaylogo": False,  # Hide the Plotly logo
+    }
+
+    # Create the figure
+    fig = go.Figure(data=traces, layout=layout)
+
+    fig.show(config=config)
+
+
+def plot_errors(m, approx_type):
+    # Define grid dimensions
+    n_x = 32  # Fixed spatial grid resolution
+    n_t = 50
+
+    try:
+        loaded_values = np.load(f"npz/{approx_type}_m{m}.npz")
+    except:
+        raise gr.Error(f"First train the coefficients for {approx_type} and m = {m}")
+    alpha = loaded_values["alpha"]
+    Phi = loaded_values["Phi"]
+
+    # Create grids for x and t
+    x = np.linspace(0, 1, n_x)  # Spatial grid
+    t = np.linspace(0, 5, n_t)  # Temporal grid
+    X, T = np.meshgrid(x, t)
+
+    # Compute the real solution over the grid
+    U_real = complex_heat_eq_solution(X, T)
+
+    # Compute the selected approximation
+    U_approx = np.zeros_like(U_real)
+    for i, t_val in enumerate(t):
+        Phi_gff_at_t = Phi[i * n_x : (i + 1) * n_x]
+        U_approx[i, :] = np.dot(Phi_gff_at_t, alpha)
+
+    U_err = abs(U_approx - U_real)
+
+    # Create the 3D plot with Plotly
+    traces = []
+
+    # Real solution surface with a distinct color (e.g., 'Viridis')
+    traces.append(
+        go.Surface(
+            z=U_err,
+            x=X,
+            y=T,
+            colorscale="Viridis",
+            showscale=False,
+            name=f"Absolute Error",
+            showlegend=True,
+        )
+    )
+
+    # Layout for the Plotly plot without controls
+    layout = go.Layout(
+        title=f"Heat Equation Approximation Error | Kernel = {approx_type} | m = {m}",
+        scene=dict(
+            camera=dict(
+                eye=dict(x=0, y=-2, z=0),  # Front view
+            ),
+            xaxis_title="x",
+            yaxis_title="t",
+            zaxis_title="u",
+        ),
+    )
+
+    # Config to remove modebar buttons except the save image button
+    config = {
+        "modeBarButtonsToRemove": [
+            "pan",
+            "resetCameraLastSave",
+            "hoverClosest3d",
+            "hoverCompareCartesian",
+            "zoomIn",
+            "zoomOut",
+            "select2d",
+            "lasso2d",
+            "zoomIn2d",
+            "zoomOut2d",
+            "sendDataToCloud",
+            "zoom3d",
+            "orbitRotation",
+            "tableRotation",
+        ],
+        "displayModeBar": True,  # Keep the modebar visible
+        "displaylogo": False,  # Hide the Plotly logo
+    }
+
+    # Create the figure
+    fig = go.Figure(data=traces, layout=layout)
+
+    fig.show(config=config)
+
+
+# Function to get the available .npz files in the npz folder
+def get_available_approx_files():
+    files = os.listdir(npz_folder)
+    npz_files = [f for f in files if f.endswith(".npz")]
+    return npz_files
+
+
+def generate_data(n_x=32, n_t=50):
+    """Generate training data."""
+    x = np.linspace(0, 1, n_x)  # spatial points
+    t = np.linspace(0, 5, n_t)  # temporal points
+    X, T = np.meshgrid(x, t)
+    a_train = np.c_[X.ravel(), T.ravel()]  # shape (n_x * n_t, 2)
+    u_train = complex_heat_eq_solution(
+        a_train[:, 0], a_train[:, 1]
+    )  # shape (n_x * n_t,)
+    return a_train, u_train, x, t
+
+
+def random_features(a, theta_j, kernel="SINE", k=0.5, t=1.0):
+    """Compute random features with adjustable kernel width."""
+    if kernel == "SINE":
+        return np.sin(t * np.linalg.norm(a - theta_j, axis=-1))
+    elif kernel == "GFF":
+        return np.log(np.linalg.norm(a - theta_j, axis=-1)) / (2 * np.pi)
+    else:
+        raise ValueError("Unsupported kernel type!")
+
+
+def design_matrix(a, theta, kernel):
+    """Construct design matrix."""
+    return np.array([random_features(a, theta_j, kernel=kernel) for theta_j in theta]).T
+
+
+def learn_coefficients(Phi, u):
+    """Learn coefficients alpha via least squares."""
+    return np.linalg.lstsq(Phi, u, rcond=None)[0]
+
+
+def approximate_solution(a, alpha, theta, kernel):
+    """Compute the approximation."""
+    Phi = design_matrix(a, theta, kernel)
+    return Phi @ alpha
+
+
+def polyfit2d(x, y, z, kx=3, ky=3, order=None):
+    # grid coords
+    x, y = np.meshgrid(x, y)
+    # coefficient array, up to x^kx, y^ky
+    coeffs = np.ones((kx + 1, ky + 1))
+
+    # solve array
+    a = np.zeros((coeffs.size, x.size))
+
+    # for each coefficient produce array x^i, y^j
+    for index, (j, i) in enumerate(np.ndindex(coeffs.shape)):
+        # do not include powers greater than order
+        if order is not None and i + j > order:
+            arr = np.zeros_like(x)
+        else:
+            arr = coeffs[i, j] * x**i * y**j
+        a[index] = arr.ravel()
+
+    # do leastsq fitting and return leastsq result
+    return np.linalg.lstsq(a.T, np.ravel(z), rcond=None)
+
+
+def train_coefficients(m, kernel):
+    # Start time for training
+    start_time = time.time()
+
+    # Generate data
+    n_x, n_t = 32, 50
+    a_train, u_train, x, t = generate_data(n_x, n_t)
+
+    # Define random features
+    theta = np.column_stack(
+        (
+            np.random.uniform(-1, 1, size=m),  # First dimension: [-1, 1]
+            np.random.uniform(-5, 5, size=m),  # Second dimension: [-5, 5]
+        )
+    )
+
+    # Construct design matrix and learn coefficients
+    Phi = design_matrix(a_train, theta, kernel)
+    alpha = learn_coefficients(Phi, u_train)
+    # Validate and animate results
+    u_real = np.array([complex_heat_eq_solution(x, t_i) for t_i in t])
+    a_test = np.c_[np.meshgrid(x, t)[0].ravel(), np.meshgrid(x, t)[1].ravel()]
+    u_approx = approximate_solution(a_test, alpha, theta, kernel).reshape(n_t, n_x)
+
+    # Save values to the npz folder
+    np.savez(
+        f"{npz_folder}/{kernel}_m{m}.npz",
+        alpha=alpha,
+        kernel=kernel,
+        Phi=Phi,
+        theta=theta,
+    )
+
+    # Compute average error
+    avg_err = np.mean(np.abs(u_real - u_approx))
+
+    return f"Training completed in {time.time() - start_time:.2f} seconds. The average error is {avg_err}."
+
+
+def plot_function(a, b, c, d, k=0.5):
+    global glob_a, glob_b, glob_c, glob_d
+
+    glob_a, glob_b, glob_c, glob_d = a, b, c, d
+    
+    x = np.linspace(0, 1, 100)
+    t = np.linspace(0, 5, 500)
+    X, T = np.meshgrid(x, t)  # Create the mesh grid
+    Z = complex_heat_eq_solution(X, T, a, b, c, d)
+    
+    traces = []
+    traces.append(
+        go.Surface(
+            z=Z,
+            x=X,
+            y=T,
+            colorscale="Viridis",
+            showscale=False,
+            showlegend=False,
+        )
+    )
+
+    # Layout for the Plotly plot without controls
+    layout = go.Layout(
+        scene=dict(
+            camera=dict(
+                eye=dict(x=1.25, y=-1.75, z=0.3),  # Front view
+            ),
+            xaxis_title="x",
+            yaxis_title="t",
+            zaxis_title="u",
+        ),
+        margin=dict(l=0, r=0, t=0, b=0),  # Reduce margins
+    )
+    
+    # Create the figure
+    fig = go.Figure(data=traces, layout=layout)
+
+    # fig.show(config=config)
+    fig.update_layout(
+        modebar_remove=[
+            "pan",
+            "resetCameraLastSave",
+            "hoverClosest3d",
+            "hoverCompareCartesian",
+            "zoomIn",
+            "zoomOut",
+            "select2d",
+            "lasso2d",
+            "zoomIn2d",
+            "zoomOut2d",
+            "sendDataToCloud",
+            "zoom3d",
+            "orbitRotation",
+            "tableRotation",
+            "toImage",
+            "resetCameraDefault3d"
+        ]
+    )
+
+    return fig
+
+
+# Gradio interface
+def create_gradio_ui():
+    # Get the initial available files
+    with gr.Blocks() as demo:
+        gr.Markdown("# Learn the Coefficients for the Heat Equation using the RFM")
+
+        # Function parameter inputs
+        gr.Markdown(
+            """
+        ## Function: $$u_k(x, t)\\coloneqq\\exp(-kt)\\cdot\\sin(\\pi x)+0.5\\cdot\\exp(\\textcolor{red}{a}kt)\\cdot\\sin(\\textcolor{red}{b}\\pi x)+0.25\\cdot\\exp(\\textcolor{red}{c}kt)\\cdot\\sin(\\textcolor{red}{d}\\pi x)$$
+        
+        Adjust the values for <span style='color: red;'>a</span>, <span style='color: red;'>b</span>, <span style='color: red;'>c</span> and <span style='color: red;'>d</span> with the sliders below.
+        """
+        )
+
+        with gr.Row():
+            with gr.Column():
+                a_slider = gr.Slider(minimum=-10, maximum=-1, step=1, value=-2, label="a")
+                b_slider = gr.Slider(minimum=-10, maximum=10, step=1, value=-2, label="b")
+                c_slider = gr.Slider(minimum=-10, maximum=-1, step=1, value=-4, label="c")
+                d_slider = gr.Slider(minimum=-10, maximum=10, step=1, value=7, label="d")
+
+            plot_output = gr.Plot()
+
+        a_slider.change(
+            fn=plot_function,
+            inputs=[a_slider, b_slider, c_slider, d_slider],
+            outputs=[plot_output],
+        )
+        b_slider.change(
+            fn=plot_function,
+            inputs=[a_slider, b_slider, c_slider, d_slider],
+            outputs=[plot_output],
+        )
+        c_slider.change(
+            fn=plot_function,
+            inputs=[a_slider, b_slider, c_slider, d_slider],
+            outputs=[plot_output],
+        )
+        d_slider.change(
+            fn=plot_function,
+            inputs=[a_slider, b_slider, c_slider, d_slider],
+            outputs=[plot_output],
+        )
+
+        with gr.Column():
+            with gr.Row():
+                # Kernel selection and slider for m
+                kernel_dropdown = gr.Dropdown(
+                    label="Choose Kernel", choices=["SINE", "GFF"], value="SINE"
+                )
+                m_slider = gr.Dropdown(
+                    label="Number of Random Features (m)",
+                    choices=[50, 250, 1000, 5000, 10000, 25000],
+                    value=1000,
+                )
+
+                # Output to show status
+                output = gr.Textbox(label="Status", interactive=False)
+
+            with gr.Row():
+                # Button to train coefficients
+                train_button = gr.Button("Train Coefficients")
+                # Function to trigger training and update dropdown
+                train_button.click(
+                    fn=train_coefficients,
+                    inputs=[m_slider, kernel_dropdown],
+                    outputs=output,
+                )
+            with gr.Row():
+                approx_button = gr.Button("Plot Approximation")
+                approx_button.click(
+                    fn=plot_heat_equation, inputs=[m_slider, kernel_dropdown], outputs=None
+                )
+
+                error_button = gr.Button("Plot Errors")
+                error_button.click(
+                    fn=plot_errors, inputs=[m_slider, kernel_dropdown], outputs=None
+                )
+        demo.load(fn=clear_folder, inputs=None, outputs=None)
+        demo.load(fn=plot_function, inputs=[a_slider, b_slider, c_slider, d_slider], outputs=[plot_output])
+
+    return demo
+
+
+# Launch Gradio app
+if __name__ == "__main__":
+    interface = create_gradio_ui()
+    interface.launch(share=False)