app.py

# app.py

import gradio as gr
from utils import initialize_gmm, generate_grid, generate_contours, generate_intermediate_points, plot_samples_and_contours, create_animation
import matplotlib.pyplot as plt

def visualize_gmm(mu_list, Sigma_list, pi_list, dx, dtheta, T, N):
    gmm = initialize_gmm(mu_list, Sigma_list, pi_list)
    grid_points = generate_grid(dx)
    std_normal_contours = generate_contours(dtheta)
    gmm_samples = gmm.sample(500)
    intermediate_points = generate_intermediate_points(gmm, grid_points, std_normal_contours, gmm_samples, T, N)
    
    fig1, ax1 = plot_samples_and_contours(gmm_samples, std_normal_contours, grid_points, "GMM Samples and Contours")
    fig2, ax2 = plot_samples_and_contours(gmm_samples, std_normal_contours, grid_points, "Standard Normal Samples and Contours")
    
    anim1 = create_animation(fig1, ax1, N, *intermediate_points[:3])
    anim2 = create_animation(fig2, ax2, N, *intermediate_points[3:])
    
    return fig1, fig2, anim1.to_jshtml(), anim2.to_jshtml()

demo = gr.Interface(
    fn=visualize_gmm,
    inputs=[
        gr.Textbox(label="Mu List", placeholder="Enter means as a list of lists, e.g., [[0,0], [1,1]]"),
        gr.Textbox(label="Sigma List", placeholder="Enter covariances as a list of lists, e.g., [[[0.2, 0.1], [0.1, 0.3]], [[1.0, -0.1], [-0.1, 0.1]]]"),
        gr.Textbox(label="Pi List", placeholder="Enter weights as a list, e.g., [0.5, 0.5]"),
        gr.Slider(minimum=0.01, maximum=1.0, label="dx", default=0.1),
        gr.Slider(minimum=0.01, maximum=0.1, label="dtheta", default=0.01),
        gr.Slider(minimum=1, maximum=100, label="T", default=10),
        gr.Slider(minimum=1, maximum=500, label="N", default=100)
    ],
    outputs=[
        gr.Plot(label="GMM to Normal Flow"),
        gr.Plot(label="Normal to GMM Flow"),
        gr.HTML(label="GMM to Normal Animation"),
        gr.HTML(label="Normal to GMM Animation")
    ],
    live=True
)

demo.launch()