working version

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -1,42 +1,84 @@
 # app.py
 
 import gradio as gr
-from utils import initialize_gmm, generate_grid, generate_contours, generate_intermediate_points, plot_samples_and_contours, create_animation
+from utils import initialize_gmm, generate_grid, generate_contours, generate_intermediate_points, plot_samples_and_contours
 import matplotlib.pyplot as plt
+import torch
+import numpy as np
+
+def validate_inputs(mu_list, Sigma_list, pi_list):
+    try:
+        mu = eval(mu_list)
+        Sigma = eval(Sigma_list)
+        pi = eval(pi_list)
+        
+        if not (isinstance(mu, list) and all(isinstance(i, list) for i in mu)):
+            return False, "Mu list is invalid."
+        if not (isinstance(Sigma, list) and all(isinstance(i, list) for i in Sigma)):
+            return False, "Sigma list is invalid."
+        if not isinstance(pi, list):
+            return False, "Pi list is invalid."
+        
+        if not torch.isclose(torch.tensor(pi).sum(), torch.tensor(1.0)):
+            return False, "Mixture weights must sum to 1."
+        
+        return True, ""
+    except Exception as e:
+        return False, str(e)
 
 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()
+    is_valid, error_message = validate_inputs(mu_list, Sigma_list, pi_list)
+    if not is_valid:
+        fig, ax = plt.subplots()
+        ax.text(0.5, 0.5, f'Invalid input: {error_message}', horizontalalignment='center', verticalalignment='center')
+        ax.set_xlim(-5, 5)
+        ax.set_ylim(-5, 5)
+        ax.set_aspect('equal', adjustable='box')
+        plt.close(fig)
+        return fig, fig
+
+    try:
+        gmm = initialize_gmm(eval(mu_list), eval(Sigma_list), eval(pi_list))
+        grid_points = generate_grid(dx)
+        std_normal_contours = generate_contours(dtheta)
+        gmm_samples = gmm.sample(500)
+        normal_samples = torch.distributions.MultivariateNormal(torch.zeros(2), torch.eye(2)).sample((500,))
+        (intermediate_points_gmm_to_normal, contour_intermediate_points_gmm_to_normal, grid_intermediate_points_gmm_to_normal,
+            intermediate_points_normal_to_gmm, contour_intermediate_points_normal_to_gmm, grid_intermediate_points_normal_to_gmm) = \
+                    generate_intermediate_points(gmm, grid_points, std_normal_contours, gmm_samples, normal_samples, T, N)
+        
+        final_frame_gmm_to_normal = intermediate_points_gmm_to_normal.cpu().detach().numpy()
+        final_frame_normal_to_gmm = intermediate_points_normal_to_gmm.cpu().detach().numpy()
+
+        fig1, ax1 = plot_samples_and_contours(final_frame_gmm_to_normal, contour_intermediate_points_gmm_to_normal.cpu().detach().numpy(), grid_intermediate_points_gmm_to_normal.cpu().detach().numpy(), "GMM to Normal Final Frame")
+        fig2, ax2 = plot_samples_and_contours(final_frame_normal_to_gmm, contour_intermediate_points_normal_to_gmm.cpu().detach().numpy(), grid_intermediate_points_normal_to_gmm.cpu().detach().numpy(), "Normal to GMM Final Frame")
+        
+        return fig1, fig2
+    except Exception as e:
+        fig, ax = plt.subplots()
+        ax.text(0.5, 0.5, f'Error during visualization: {str(e)}', horizontalalignment='center', verticalalignment='center')
+        ax.set_xlim(-5, 5)
+        ax.set_ylim(-5, 5)
+        ax.set_aspect('equal', adjustable='box')
+        plt.close(fig)
+        return fig, fig
 
 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)
+        gr.Textbox(label="Mu List", value="[[2, 1], [-1, -2], [3, -2]]", placeholder="Enter means as a list of lists, e.g., [[0,0], [1,1]]"),
+        gr.Textbox(label="Sigma List", value="[[[0.2, 0.1], [0.1, 0.3]], [[1.0, -0.1], [-0.1, 0.1]], [[0.05, 0.0], [0.0, 0.05]]]", 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", value="[0.05, 0.8, 0.15]", placeholder="Enter weights as a list, e.g., [0.5, 0.5]"),
+        gr.Slider(minimum=0.01, maximum=1.0, label="dx", value=0.1),
+        gr.Slider(minimum=2*np.pi/3600, maximum=2*np.pi/36, label="dtheta", value=2*np.pi/360),
+        gr.Slider(minimum=1, maximum=100, label="T", value=10),
+        gr.Slider(minimum=1, maximum=500, label="N", value=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")
+        gr.Plot(label="GMM to Normal Flow Final Frame"),
+        gr.Plot(label="Normal to GMM Flow Final Frame")
     ],
     live=True
 )
 
-demo.launch()
+demo.launch()

gmm.py

@@ -42,7 +42,7 @@ class GaussianMixtureModel:
         samples = []
         for _ in range(n_samples):
             # Choose a component based on mixture weights
-            k = torch.multinomial(self.pi.squeeze(), 1).item()
+            k = torch.multinomial(self.pi.reshape(self.pi.shape[0]), 1).item()
             sample = torch.distributions.MultivariateNormal(self.mu[k].squeeze(), self.Sigma[k]).sample()
             samples.append(sample)
         return torch.stack(samples)

utils.py

@@ -3,7 +3,6 @@
 import torch
 import numpy as np
 import matplotlib.pyplot as plt
-from matplotlib.animation import FuncAnimation
 from gmm import GaussianMixtureModel
 
 def initialize_gmm(mu_list, Sigma_list, pi_list):
@@ -15,8 +14,10 @@ def initialize_gmm(mu_list, Sigma_list, pi_list):
 def generate_grid(dx):
     x_positions = np.arange(-10, 10.5, 0.5)
     y_positions = np.arange(-10, 10.5, 0.5)
-    vertical_lines = [np.stack([np.full(int((10 - (-10))/ dx + 1), x), np.arange(-10, 10 + dx, dx)], axis=1) for x in x_positions]
-    horizontal_lines = [np.stack([np.arange(-10, 10 + dx, dx), np.full(int((10 - (-10)) / dx + 1), y)], axis=1) for y in y_positions]
+    fine_points = np.arange(-10, 10 + dx, dx)
+    ones_same_size = np.ones_like(fine_points)
+    vertical_lines = [np.stack([x*ones_same_size, fine_points], axis=1) for x in x_positions]
+    horizontal_lines = [np.stack([fine_points, y*ones_same_size], axis=1) for y in y_positions]
     grid_points = np.concatenate(vertical_lines + horizontal_lines, axis=0)
     return torch.tensor(grid_points, dtype=torch.float32)
 
@@ -25,14 +26,22 @@ def generate_contours(dtheta):
     std_normal_contours = np.concatenate([np.stack([r * np.cos(angles), r * np.sin(angles)], axis=1) for r in range(1, 4)], axis=0)
     return torch.tensor(std_normal_contours, dtype=torch.float32)
 
-def generate_intermediate_points(gmm, grid_points, std_normal_contours, gmm_samples, T, N):
-    intermediate_points_gmm_to_normal = gmm.flow_gmm_to_normal(grid_points, T, N)
-    contour_intermediate_points_gmm_to_normal = gmm.flow_gmm_to_normal(std_normal_contours, T, N)
-    grid_intermediate_points_gmm_to_normal = gmm.flow_gmm_to_normal(grid_points, T, N)
-
-    intermediate_points_normal_to_gmm = gmm.flow_normal_to_gmm(gmm_samples, T, N)
-    contour_intermediate_points_normal_to_gmm = gmm.flow_normal_to_gmm(std_normal_contours, T, N)
-    grid_intermediate_points_normal_to_gmm = gmm.flow_normal_to_gmm(grid_points, T, N)
+def transform_std_to_gmm_contours(std_contours, mu, Sigma):
+    gmm_contours = []
+    for k in range(mu.shape[0]):
+        L = torch.linalg.cholesky(Sigma[k])
+        gmm_contours.append(mu[k] + torch.matmul(std_contours, L.T))
+    return torch.cat(gmm_contours, dim=0)
+
+def generate_intermediate_points(gmm, grid_points, std_normal_contours, gmm_samples, normal_samples, T, N):
+    gmm_contours = transform_std_to_gmm_contours(std_normal_contours, gmm.mu.squeeze(), gmm.Sigma)
+    intermediate_points_gmm_to_normal = gmm.flow_gmm_to_normal(gmm_samples.clone(), T, N)
+    contour_intermediate_points_gmm_to_normal = gmm.flow_gmm_to_normal(gmm_contours.clone(), T, N)
+    grid_intermediate_points_gmm_to_normal = gmm.flow_gmm_to_normal(grid_points.clone(), T, N)
+
+    intermediate_points_normal_to_gmm = gmm.flow_normal_to_gmm(normal_samples.clone(), T, N)
+    contour_intermediate_points_normal_to_gmm = gmm.flow_normal_to_gmm(std_normal_contours.clone(), T, N)
+    grid_intermediate_points_normal_to_gmm = gmm.flow_normal_to_gmm(grid_points.clone(), T, N)
     
     return (intermediate_points_gmm_to_normal, contour_intermediate_points_gmm_to_normal, grid_intermediate_points_gmm_to_normal,
             intermediate_points_normal_to_gmm, contour_intermediate_points_normal_to_gmm, grid_intermediate_points_normal_to_gmm)
@@ -52,17 +61,3 @@ def plot_samples_and_contours(samples, contours, grid_points, title):
     ax.set_aspect('equal', adjustable='box')
     plt.close(fig)
     return fig, ax
-
-def create_animation(fig, ax, frames, intermediate_points, intermediate_samples, intermediate_contours, intermediate_grid):
-    scatter_grid = ax.scatter([], [], c='black', alpha=0.5, s=1, label='Grid Points')
-    contour_scatter = ax.scatter([], [], c='blue', alpha=0.5, s=3, label='Contours')
-    scatter_samples = ax.scatter([], [], c='red', alpha=0.5, label='Samples')
-
-    def update(frame):
-        scatter_grid.set_offsets(intermediate_points[frame].numpy())
-        scatter_samples.set_offsets(intermediate_samples[frame].numpy())
-        contour_scatter.set_offsets(intermediate_contours[frame].numpy())
-        return scatter_grid, scatter_samples, contour_scatter
-
-    anim = FuncAnimation(fig, update, frames=frames, blit=True)
-    return anim