initial commit of app

app.py

@@ -0,0 +1,42 @@
+# 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()

gmm.py

@@ -0,0 +1,151 @@
+import torch
+import numpy as np
+
+class GaussianMixtureModel:
+    def __init__(self, mu, Sigma, pi):
+        # Ensure all inputs are torch tensors
+        assert isinstance(mu, torch.Tensor), "mu must be a torch.Tensor."
+        assert isinstance(Sigma, torch.Tensor), "Sigma must be a torch.Tensor."
+        assert isinstance(pi, torch.Tensor), "pi must be a torch.Tensor."
+        
+        # Ensure the dtype is torch.float32
+        assert mu.dtype == torch.float32, "mu must have dtype torch.float32."
+        assert Sigma.dtype == torch.float32, "Sigma must have dtype torch.float32."
+        assert pi.dtype == torch.float32, "pi must have dtype torch.float32."
+        
+        self.K, self.d = mu.shape[:2]
+        
+        # Check the shape of mu
+        if mu.shape == (self.K, self.d):
+            mu = mu.unsqueeze(-1)  # Shape: (K, d, 1)
+        assert mu.shape == (self.K, self.d, 1), "mu must have shape (K, d, 1)."
+        
+        # Check the shape of Sigma
+        assert Sigma.shape == (self.K, self.d, self.d), "Sigma must have shape (K, d, d)."
+        
+        # Check the shape of pi and fix it if necessary
+        if pi.shape == (self.K,):
+            pi = pi.view(self.K, 1, 1)
+        elif pi.shape == (self.K, 1):
+            pi = pi.unsqueeze(-1)
+        assert pi.shape == (self.K, 1, 1), "pi must have shape (K, 1, 1)."
+        
+        # Ensure pi sums to 1
+        assert torch.isclose(torch.sum(pi), torch.tensor(1.0)), "Mixture weights must sum to 1."
+        
+        self.mu = mu
+        self.Sigma = Sigma
+        self.pi = pi
+        
+    def sample(self, n_samples):
+        # Sample from the mixture model
+        samples = []
+        for _ in range(n_samples):
+            # Choose a component based on mixture weights
+            k = torch.multinomial(self.pi.squeeze(), 1).item()
+            sample = torch.distributions.MultivariateNormal(self.mu[k].squeeze(), self.Sigma[k]).sample()
+            samples.append(sample)
+        return torch.stack(samples)
+    
+    def log_prob(self, x):
+        # Compute the log probability of a given sample x
+        x = x.view(1, self.d, 1)  # Shape: (1, d, 1)
+        diff = x - self.mu  # Shape: (K, d, 1)
+        inv_Sigma = torch.inverse(self.Sigma)  # Shape: (K, d, d)
+        
+        exponent = -0.5 * torch.bmm(torch.bmm(diff.transpose(1, 2), inv_Sigma), diff).squeeze()  # Shape: (K,)
+        normalization = torch.log(torch.det(2 * torch.pi * self.Sigma)) / 2  # Shape: (K,)
+        log_probs = torch.log(self.pi.squeeze()) + exponent - normalization  # Shape: (K,)
+        return torch.logsumexp(log_probs, dim=0)
+    
+    def score(self, x):
+        # Compute the score function (gradient of log probability) for a batch of samples x
+        B = x.shape[0]
+        x = x.view(B, 1, self.d, 1)  # Shape: (B, 1, d, 1)
+        diff = x - self.mu.unsqueeze(0)  # Shape: (B, K, d, 1)
+        inv_Sigma = torch.inverse(self.Sigma).unsqueeze(0)  # Shape: (1, K, d, d)
+        
+        diff_t = diff.transpose(-2, -1).contiguous().view(B * self.K, 1, self.d)  # Shape: (B*K, 1, d)
+        inv_Sigma_flat = inv_Sigma.view(self.K, self.d, self.d).expand(B, self.K, self.d, self.d).contiguous().view(B * self.K, self.d, self.d)  # Shape: (B*K, d, d)
+        exponent = -0.5 * torch.bmm(torch.bmm(diff_t, inv_Sigma_flat), diff.view(B * self.K, self.d, 1)).view(B, self.K)  # Shape: (B, K)
+        
+        normalization = torch.log(torch.det(2 * torch.pi * self.Sigma)).unsqueeze(0) / 2  # Shape: (1, K)
+        probs = torch.exp(torch.log(self.pi.squeeze()).unsqueeze(0) + exponent - normalization)  # Shape: (B, K)
+        norm_probs = probs / torch.sum(probs, dim=1, keepdim=True)  # Shape: (B, K)
+        
+        gradients = []
+        for k in range(self.K):
+            gradient = -torch.bmm(inv_Sigma[:, k].expand(B, self.d, self.d), diff[:, k])  # Shape: (B, d, 1)
+            gradients.append(norm_probs[:, k].unsqueeze(1) * gradient.squeeze(-1))  # Shape: (B, d)
+        return torch.sum(torch.stack(gradients, dim=1), dim=1)  # Shape: (B, d)
+    
+    def forward_diffusion(self, t):
+        # Compute the evolved mean and covariance
+        mu_t = self.mu * torch.tensor(np.exp(-0.5 * t), dtype=torch.float32)  # Shape: (K, d, 1)
+        exp_neg_beta_t = torch.tensor(np.exp(-t), dtype=torch.float32).view(1, 1, 1)  # Shape: (1, 1, 1)
+        Sigma_t = self.Sigma * exp_neg_beta_t + torch.eye(self.d, dtype=torch.float32) * (1 - exp_neg_beta_t)  # Shape: (K, d, d)
+        return GaussianMixtureModel(mu_t, Sigma_t, self.pi)
+    
+    def flow(self, x_t, t, dt, num_steps):
+        for _ in range(num_steps):
+            x_t = self.probability_flow_ode(x_t, t, dt)
+            t += dt
+        return x_t
+    
+    def flow_gmm_to_normal(self, x_0, T=5, N=32):
+        dt = T / N
+        return self.flow(x_0, 0, dt, N)
+    
+    def flow_normal_to_gmm(self, x_T, T=5, N=32):
+        dt = -T / N
+        return self.flow(x_T, T, dt, N)
+    
+    def probability_flow_ode(self, x_t, t, dt):
+        # Compute the evolved x_t based on the probability flow ODE using the Euler method
+        # Forward diffusion to get the evolved GMM parameters
+        gmm_t = self.forward_diffusion(t)
+        
+        # Compute the score function at x_t
+        score = gmm_t.score(x_t)
+        
+        # Compute the drift term
+        drift = -0.5 * x_t - 0.5 * score
+        
+        # Euler update
+        x_t_plus_dt = x_t + drift * dt
+        
+        return x_t_plus_dt
+
+# Example usage
+if __name__ == "__main__":
+    mu = torch.tensor([[0, 0], [1, 1]], dtype=torch.float32)  # Shape: (K, d)
+    Sigma = torch.stack([torch.eye(2), torch.eye(2)], dim=0).float()  # Shape: (K, d, d)
+    pi = torch.tensor([0.5, 0.5], dtype=torch.float32)  # Shape: (K,)
+
+    gmm = GaussianMixtureModel(mu, Sigma, pi)
+
+    # Perform forward diffusion
+    t = 1.0
+    gmm_t = gmm.forward_diffusion(t)
+
+    # Sampling from the diffused GMM
+    samples = gmm_t.sample(10)
+    print("Samples after forward diffusion:\n", samples)
+
+    # Log probability of a sample after forward diffusion
+    log_prob = gmm_t.log_prob(torch.tensor([0.5, 0.5], dtype=torch.float32))
+    print("Log Probability after forward diffusion:", log_prob)
+
+    # Score of a sample after forward diffusion
+    score = gmm_t.score(torch.tensor([[0.5, 0.5]], dtype=torch.float32))
+    print("Score after forward diffusion:", score)
+
+    # Using the flow_gmm_to_normal
+    x_0 = torch.tensor([[0.5, 0.5]], dtype=torch.float32)
+    x_T_normal = gmm.flow_gmm_to_normal(x_0)
+    print("x_T_normal after flowing from GMM to normal:", x_T_normal)
+
+    # Using the flow_normal_to_gmm
+    x_T = torch.tensor([[0.0, 0.0]], dtype=torch.float32)
+    x_0_gmm = gmm.flow_normal_to_gmm(x_T)
+    print("x_0_gmm after flowing from normal to GMM:", x_0_gmm)

utils.py

@@ -0,0 +1,68 @@
+# utils.py
+
+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):
+    mu = torch.tensor(mu_list, dtype=torch.float32)
+    Sigma = torch.tensor(Sigma_list, dtype=torch.float32)
+    pi = torch.tensor(pi_list, dtype=torch.float32)
+    return GaussianMixtureModel(mu, Sigma, pi)
+
+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]
+    grid_points = np.concatenate(vertical_lines + horizontal_lines, axis=0)
+    return torch.tensor(grid_points, dtype=torch.float32)
+
+def generate_contours(dtheta):
+    angles = np.linspace(0, 2 * np.pi, int(2 * np.pi / 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)
+    
+    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)
+
+def plot_samples_and_contours(samples, contours, grid_points, title):
+    fig, ax = plt.subplots(figsize=(8, 6))
+    ax.scatter(grid_points[:, 0], grid_points[:, 1], alpha=0.5, c='black', s=1, label='Grid Points')
+    ax.scatter(contours[:, 0], contours[:, 1], alpha=0.5, s=3, c='blue', label='Contours')
+    ax.scatter(samples[:, 0], samples[:, 1], alpha=0.5, c='red', label='Samples')
+    ax.set_title(title)
+    ax.set_xlabel("x1")
+    ax.set_ylabel("x2")
+    ax.grid(True)
+    ax.legend(loc='upper right')
+    ax.set_xlim(-5, 5)
+    ax.set_ylim(-5, 5)
+    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