simulation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SUytbAd92HaR"
   },
   "source": [
    "# Numerical Simulation for Generlized Exponential Function (GEF)\n",
    "from the GES paper https://abdullahamdi.com/ges/\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 12515,
     "status": "ok",
     "timestamp": 1695394339532,
     "user": {
      "displayName": "Abdullah Hamdi",
      "userId": "15226589023402472176"
     },
     "user_tz": -60
    },
    "id": "maK_YB6AwFjM",
    "outputId": "6444a685-766f-4347-bede-e43a9c5be9f3",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# !pip install open_clip_torch celluloid matplotlib tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## plotting the GEF functions family\n",
    "\n",
    "### $f_{\\beta}(x) = Ae^{-\\left(\\frac{|x - \\mu|}{\\alpha}\\right)^\\beta}$\n",
    "\n",
    "* when $\\beta=2$, GEF reduces to Gaussians"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.cm as cm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Define the parameters\n",
    "mu = 0  # location\n",
    "alpha = 1  # scale\n",
    "beta_values = [0.5, 1, 1.5, 2, 3,10]  # shape parameters\n",
    "\n",
    "# Define the x range\n",
    "x = np.linspace(-5, 5, 1000)\n",
    "\n",
    "# Plotting\n",
    "colors = cm.viridis(np.linspace(0, 1, len(beta_values)))\n",
    "\n",
    "plt.figure(figsize=(12, 8))\n",
    "\n",
    "for beta, color in zip(beta_values, colors):\n",
    "    # Compute the PDF\n",
    "    A = 1.0 # (beta / (2 * alpha * gamma(1 / beta)))\n",
    "    pdf = A * np.exp(-np.abs((x - mu) / alpha) ** beta)\n",
    "    plt.plot(x, pdf, label=f'$\\\\beta$ = {beta}', linewidth=2)\n",
    "\n",
    "# Customize the plot\n",
    "# plt.title('Generalized Exponential Function for Different $\\\\beta$ Values', fontsize=20)\n",
    "plt.xlabel('x', fontsize=22)\n",
    "plt.ylabel(f\"f(x | $\\\\beta$)\", fontsize=22)\n",
    "plt.legend(fontsize=20)\n",
    "plt.grid(True)\n",
    "plt.xticks(fontsize=18)\n",
    "plt.yticks(fontsize=18)\n",
    "\n",
    "# Show plot\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 1D Simulation "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### models "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import math\n",
    "class GaussianMixture(nn.Module):\n",
    "    def __init__(self, N, positive_weights=False):\n",
    "        super(GaussianMixture, self).__init__()\n",
    "        \n",
    "        self.means = nn.Parameter(torch.randn(N))\n",
    "        self.variances = nn.Parameter(torch.abs(0.1 * torch.randn(N) + 1))\n",
    "        self.weights = nn.Parameter(torch.randn(N))\n",
    "        \n",
    "        self.positive_weights = positive_weights\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.unsqueeze(-1)\n",
    "        epsilon = 1e-8\n",
    "        \n",
    "        # Apply softplus if positive_weights flag is True\n",
    "        weights = F.softplus(self.weights) if self.positive_weights else self.weights\n",
    "        gaussians = weights * torch.exp(-(x - self.means)**2 / (2 * self.variances + epsilon))\n",
    "        \n",
    "        return gaussians.sum(dim=-1)\n",
    "class DoGMixture(nn.Module):\n",
    "    def __init__(self, N, positive_weights=False, ratio=4.0):\n",
    "        super(DoGMixture, self).__init__()\n",
    "        \n",
    "        self.means = nn.Parameter(torch.randn(N))\n",
    "        self.scales = nn.Parameter(torch.abs(0.1 * torch.randn(N) + 1))  # Single scale parameter for both variances\n",
    "        self.weights = nn.Parameter(torch.randn(N))\n",
    "        \n",
    "        self.ratio = ratio\n",
    "        self.positive_weights = positive_weights\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.unsqueeze(-1)\n",
    "        epsilon = 1e-8\n",
    "        \n",
    "        var1 = self.scales**2\n",
    "        var2 = var1 / self.ratio  # Keeping a fixed ratio between the variances\n",
    "        \n",
    "        weights = F.softplus(self.weights) if self.positive_weights else self.weights\n",
    "        dog_components = weights * (torch.exp(-(x - self.means)**2 / (2 * var1 + epsilon)) -\n",
    "                                    torch.exp(-(x - self.means)**2 / (2 * var2 + epsilon)))\n",
    "        \n",
    "        return dog_components.sum(dim=-1)\n",
    "    \n",
    "class LoGMixture(nn.Module):\n",
    "    def __init__(self, N, positive_weights=False):\n",
    "        super(LoGMixture, self).__init__()\n",
    "        \n",
    "        self.means = nn.Parameter(torch.randn(N))\n",
    "        self.scales = nn.Parameter(torch.abs(0.1 * torch.randn(N) + 1))  # Scale parameter, equivalent to sigma\n",
    "        self.weights = nn.Parameter(torch.randn(N))\n",
    "        \n",
    "        self.positive_weights = positive_weights\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.unsqueeze(-1)\n",
    "        epsilon = 1e-8\n",
    "        \n",
    "        weights = F.softplus(self.weights) if self.positive_weights else self.weights\n",
    "        log_components = weights * (- (x - self.means)**2 / (self.scales**2) + 1) * torch.exp(-(x - self.means)**2 / (2 * self.scales**2 + epsilon))\n",
    "        \n",
    "        return log_components.sum(dim=-1)        \n",
    "        return dog_components.sum(dim=-1)\n",
    "class GeneralMixture(nn.Module):\n",
    "    def __init__(self, N, positive_weights=False, learn_beta=False, fixed_beta=2.0):\n",
    "        super(GeneralMixture, self).__init__()\n",
    "\n",
    "        # Parameters for the means, variances, and weights of the Gaussian components.\n",
    "        self.means = nn.Parameter(torch.randn(N))\n",
    "        self.variances = nn.Parameter(torch.abs(0.1 * torch.randn(N) + 1))\n",
    "        self.weights = nn.Parameter(torch.randn(N))\n",
    "\n",
    "        # Additional parameter beta for the power of the exponent in each Gaussian component.\n",
    "        if learn_beta:\n",
    "            # If beta is learnable, we initialize it as a parameter.\n",
    "            # We will ensure that beta is always greater than 0.5 during optimization.\n",
    "            self.beta = nn.Parameter(torch.rand(N) * 0.5 + 0.5)  # Initialize in the range [0.5, 1.0]\n",
    "        else:\n",
    "            # If beta is fixed, we use a tensor that does not require gradients.\n",
    "            self.register_buffer('beta', torch.full((N,), fixed_beta))\n",
    "\n",
    "        self.positive_weights = positive_weights\n",
    "        self.learn_beta = learn_beta  # Keep track of whether beta is learnable\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.unsqueeze(-1)  # Make sure that x is treated as a column vector\n",
    "        epsilon = 1e-8  # Small constant to prevent division by zero\n",
    "\n",
    "        # Apply softplus if positive_weights flag is True\n",
    "        weights = F.softplus(self.weights) if self.positive_weights else self.weights\n",
    "\n",
    "        # Apply transformation to beta if it is learnable to ensure it's greater than 0.5\n",
    "        beta = (F.softplus(self.beta) + 0.5) if self.learn_beta else self.beta\n",
    "\n",
    "        # Compute the Gaussian components with the given beta exponent power\n",
    "        gaussians = weights * torch.exp(-((x - self.means).abs() ** beta) / (2 * self.variances + epsilon))\n",
    "\n",
    "        # Sum over the Gaussian components to form the final output\n",
    "        return gaussians.sum(dim=-1)\n",
    "def train_model(model, optimizer, loss_func, x, y, epochs):\n",
    "    for epoch in range(epochs):\n",
    "        optimizer.zero_grad()\n",
    "        y_pred = model(x)\n",
    "        loss = loss_func(y_pred, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return model , y_pred , loss.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_results_gaussian(model, x, y, y_pred, N, loss_rec, positive_weights, show_fig=True, signal_type=\"square\"):\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    \n",
    "    plt.plot(x.numpy(), y.numpy(), 'r-', label='True ' + signal_type, linewidth=2)\n",
    "    plt.plot(x.numpy(), y_pred.detach().numpy(), 'b--', label='Gaussian Mixture', linewidth=2)\n",
    "\n",
    "    for i in range(N):\n",
    "        weights = F.softplus(model.weights) if model.positive_weights else model.weights\n",
    "        component = weights[i].item() * torch.exp(-(x - model.means[i].item())**2 / (2 * model.variances[i].item()))\n",
    "        plt.plot(x.numpy(), component.numpy(), 'g-.', alpha=0.5, linewidth=2)\n",
    "\n",
    "    plt.legend(fontsize=20,loc='upper right')\n",
    "    plt.title(f\"Overfitting Gaussian Mixture to a {signal_type} Function, N={N}, loss={100*loss_rec:.2f}\", fontsize=22)\n",
    "    plt.xlabel(\"x\", fontsize=22)\n",
    "    plt.ylabel(\"y\", fontsize=22)\n",
    "    plt.xticks(fontsize=18)\n",
    "    plt.yticks(fontsize=18)\n",
    "    plt.grid(True)\n",
    "    \n",
    "    if not os.path.exists(os.path.join(\"comparisons\", \"gaussians\")):\n",
    "        os.makedirs(os.path.join(\"comparisons\", \"gaussians\"))\n",
    "    \n",
    "    suffix = \"P\" if positive_weights else \"N\"\n",
    "    plt.savefig(os.path.join(\"comparisons\", \"gaussians\", signal_type + \"_\" + suffix + str(N) + \".pdf\"), bbox_inches=\"tight\")\n",
    "    if show_fig:\n",
    "        plt.show()\n",
    "def plot_results_dog(model, x, y, y_pred, N, loss_rec, positive_weights, show_fig=True, signal_type=\"square\"):\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    \n",
    "    plt.plot(x.numpy(), y.numpy(), 'r-', label='True ' + signal_type, linewidth=2)\n",
    "    plt.plot(x.numpy(), y_pred.detach().numpy(), 'b--', label='DoG Mixture', linewidth=2)\n",
    "\n",
    "    for i in range(N):\n",
    "        weights = F.softplus(model.weights) if model.positive_weights else model.weights\n",
    "        var1 = model.scales[i].item()**2\n",
    "        var2 = var1 / model.ratio\n",
    "        \n",
    "        component = weights[i].item() * (torch.exp(-(x - model.means[i].item())**2 / (2 * var1)) -\n",
    "                                         torch.exp(-(x - model.means[i].item())**2 / (2 * var2)))\n",
    "        plt.plot(x.numpy(), component.numpy(), 'g-.', alpha=0.5, linewidth=2)\n",
    "\n",
    "    plt.legend(fontsize=20)\n",
    "    plt.title(f\"Overfitting DoG Mixture to a {signal_type} Function, N={N}, loss={100*loss_rec:.2f}\", fontsize=22)\n",
    "    plt.xlabel(\"x\", fontsize=22)\n",
    "    plt.ylabel(\"y\", fontsize=22)\n",
    "    plt.xticks(fontsize=18)\n",
    "    plt.yticks(fontsize=18)\n",
    "    plt.grid(True)\n",
    "    \n",
    "    if not os.path.exists(os.path.join(\"comparisons\", \"dogs\")):\n",
    "        os.makedirs(os.path.join(\"comparisons\", \"dogs\"))\n",
    "    \n",
    "    suffix = \"DP\" if positive_weights else \"DN\"\n",
    "    plt.savefig(os.path.join(\"comparisons\", \"dogs\", signal_type + \"_\" + suffix + str(N) + \".pdf\"), bbox_inches=\"tight\")\n",
    "    if show_fig:\n",
    "        plt.show()\n",
    "def plot_results_log(model, x, y, y_pred, N, loss_rec, positive_weights, show_fig=True, signal_type=\"square\"):\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    \n",
    "    plt.plot(x.numpy(), y.numpy(), 'r-', label='True ' + signal_type, linewidth=2)\n",
    "    plt.plot(x.numpy(), y_pred.detach().numpy(), 'b--', label='LoG Mixture', linewidth=2)\n",
    "\n",
    "    for i in range(N):\n",
    "        weights = F.softplus(model.weights) if model.positive_weights else model.weights\n",
    "        component = weights[i].item() * (- (x - model.means[i].item())**2 / (model.scales[i].item()**2) + 1) * \\\n",
    "                    torch.exp(-(x - model.means[i].item())**2 / (2 * model.scales[i].item()**2))\n",
    "        plt.plot(x.numpy(), component.numpy(), 'g-.', alpha=0.5, linewidth=2)\n",
    "\n",
    "    plt.legend(fontsize=20)\n",
    "    plt.title(f\"Overfitting LoG Mixture to a {signal_type} Function, N={N}, loss={100*loss_rec:.2f}\", fontsize=22)\n",
    "    plt.xlabel(\"x\", fontsize=22)\n",
    "    plt.ylabel(\"y\", fontsize=22)\n",
    "    plt.xticks(fontsize=18)\n",
    "    plt.yticks(fontsize=18)\n",
    "    plt.grid(True)\n",
    "    \n",
    "    if not os.path.exists(os.path.join(\"comparisons\", \"logs\")):\n",
    "        os.makedirs(os.path.join(\"comparisons\", \"logs\"))\n",
    "    \n",
    "    suffix = \"LP\" if positive_weights else \"LN\"\n",
    "    plt.savefig(os.path.join(\"comparisons\", \"logs\", signal_type + \"_\" + suffix + str(N) + \".pdf\"), bbox_inches=\"tight\")\n",
    "    if show_fig:\n",
    "        plt.show()\n",
    "def plot_results_general(model, x, y, y_pred, N, loss_rec, positive_weights, show_fig=True, signal_type=\"square\"):\n",
    "    plt.figure(figsize=(12, 8))\n",
    "\n",
    "    plt.plot(x.numpy(), y.numpy(), 'r-', label='True ' + signal_type, linewidth=2)\n",
    "    plt.plot(x.numpy(), y_pred.detach().numpy(), 'b--', label='GEF Mixture', linewidth=2)\n",
    "\n",
    "    for i in range(N):\n",
    "        weights = F.softplus(model.weights) if model.positive_weights else model.weights\n",
    "        component = weights[i].item() * torch.exp(-((x - model.means[i].item()).abs() ** model.beta[i].item()) / (2 * model.variances[i].item()))\n",
    "        plt.plot(x.numpy(), component.detach().numpy(), 'g-.', alpha=0.5, linewidth=2)\n",
    "\n",
    "    plt.legend(fontsize=20)\n",
    "    plt.title(f\"Overfitting General Mixture to a {signal_type} Function, N={N}, loss={100*loss_rec:.2f}\", fontsize=22)\n",
    "    plt.xlabel(\"x\", fontsize=22)\n",
    "    plt.ylabel(\"y\", fontsize=22)\n",
    "    plt.xticks(fontsize=18)\n",
    "    plt.yticks(fontsize=18)\n",
    "    plt.grid(True)\n",
    "\n",
    "    if not os.path.exists(os.path.join(\"comparisons\", \"general_mixtures\")):\n",
    "        os.makedirs(os.path.join(\"comparisons\", \"general_mixtures\"))\n",
    "    \n",
    "    suffix = \"P\" if positive_weights else \"N\"\n",
    "    plt.savefig(os.path.join(\"comparisons\", \"general_mixtures\", signal_type + \"_\" + suffix + str(N) + \".pdf\"), bbox_inches=\"tight\")\n",
    "    if show_fig:\n",
    "        plt.show()\n",
    "def triangle_wave(x, width=2):\n",
    "    period = width / 2.0\n",
    "    phase = (x + period) % width - period\n",
    "    return torch.where(phase < 0, 1 + phase / period, 1 - phase / period)\n",
    "def get_signal(x, signal_type, width=2):\n",
    "    if signal_type == \"square\":\n",
    "        return torch.where((x > -width/2) & (x < width/2), torch.ones_like(x), torch.zeros_like(x))\n",
    "    elif signal_type == \"triangle\":\n",
    "        return torch.where((x > -width/2) & (x < width/2), width/2 - torch.abs(x), torch.zeros_like(x))\n",
    "    elif signal_type == \"parabolic\":\n",
    "        return torch.where((x > -width/2) & (x < width/2), (width/2)**2 - x**2, torch.zeros_like(x))\n",
    "    elif signal_type == \"half_sinusoid\":\n",
    "        return torch.where((x > -width/2) & (x < width/2), torch.sin((x + width/2) * (math.pi / width)), torch.zeros_like(x))\n",
    "    elif signal_type == \"exponential\":\n",
    "        return torch.where((x > -width/2) & (x < width/2), torch.exp(-torch.abs(x)), torch.zeros_like(x))\n",
    "    elif signal_type == \"gaussian\":\n",
    "        return torch.exp(-x.pow(2) / (2 * (width/3)**2))\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown signal_type: {signal_type}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### single fit "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "N = 5\n",
    "positive_weights = True\n",
    "epochs = 10000\n",
    "data_size  = 1000\n",
    "data_extent = 10\n",
    "signal_width = 6\n",
    "learn_beta = True # only for general mixture  \n",
    "fixed_beta = 2.0 # only for general mixture  \n",
    "signal_types = [\"square\", \"triangle\", \"parabolic\", \"half_sinusoid\", \"gaussian\",\"exponential\"]\n",
    "signal_type = \"square\" \n",
    "x = torch.linspace(-data_extent, data_extent, data_size).unsqueeze(-1).cuda()\n",
    "y = get_signal(x,signal_type,signal_width).cuda()\n",
    "model = GaussianMixture(N,positive_weights=positive_weights).cuda()\n",
    "# model = DoGMixture(N,positive_weights=positive_weights).cuda()\n",
    "# model = LoGMixture(N,positive_weights=positive_weights).cuda()\n",
    "# model = GeneralMixture(N,positive_weights=positive_weights,learn_beta=learn_beta,fixed_beta=fixed_beta ).cuda()\n",
    "\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
    "loss_func = nn.MSELoss()\n",
    "# Training\n",
    "model,y_pred,loss_rec = train_model(model, optimizer, loss_func, x, y, epochs=epochs)\n",
    "\n",
    "plot_results_gaussian(model, x.cpu(), y.cpu(), y_pred.cpu(), N=N, loss_rec=loss_rec, positive_weights=positive_weights,show_fig=False,signal_type=signal_type)\n",
    "# plot_results_dog(model, x.cpu(), y.cpu(), y_pred.cpu(), N=N, loss_rec=loss_rec, positive_weights=positive_weights,show_fig=False,signal_type=signal_type)\n",
    "# plot_results_log(model, x.cpu(), y.cpu(), y_pred.cpu(), N=N, loss_rec=loss_rec, positive_weights=positive_weights,show_fig=False,signal_type=signal_type)\n",
    "# plot_results_general(model, x.cpu(), y.cpu(), y_pred.cpu(), N=N, loss_rec=loss_rec, positive_weights=positive_weights,show_fig=False,signal_type=signal_type)\n",
    "print(\"loss:   \",loss_rec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### multiple fits "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# Configurations\n",
    "signal_types = [\"square\", \"triangle\", \"parabolic\", \"half_sinusoid\", \"gaussian\",\"exponential\"]\n",
    "Ns =  [2,5,8, 10, 15,20,50,100]\n",
    "positive_weights_list = [True, False]\n",
    "model_types = ['gaussian', 'dog', 'log','general']\n",
    "epochs = 10000\n",
    "runs_per_config = 20  # run each configuration multiple times to account for variance\n",
    "data_size  = 1000\n",
    "data_extent = 10\n",
    "signal_width = 6\n",
    "learn_beta = True # if gernal mixture used \n",
    "# Create a directory to store the results\n",
    "if not os.path.exists(\"comparisons\"):\n",
    "    os.makedirs(\"comparisons\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results_per_signal = {}  # Dictionary to hold results for each signal type\n",
    "\n",
    "for signal_type in signal_types:\n",
    "    # Data\n",
    "    x = torch.linspace(-data_extent, data_extent, data_size).unsqueeze(-1).cuda()\n",
    "    y = get_signal(x,signal_type,signal_width).cuda()\n",
    "    results = []  # To store the average loss and nan_counts per configuration\n",
    "    \n",
    "    for N in Ns:\n",
    "        for positive_weights in positive_weights_list:\n",
    "            for model_type in model_types:\n",
    "\n",
    "                # Keeping track of how many times the training encounters 'nan' loss\n",
    "                nan_counts = 0\n",
    "                total_loss = 0  # Sum of losses for the current configuration\n",
    "\n",
    "                for run in range(runs_per_config):\n",
    "                    # Add a condition to handle the 'general' model type\n",
    "                    if model_type == 'gaussian':\n",
    "                        model = GaussianMixture(N, positive_weights=positive_weights).cuda()\n",
    "                    elif model_type == 'dog':\n",
    "                        model = DoGMixture(N, positive_weights=positive_weights).cuda()\n",
    "                    elif model_type == 'log':\n",
    "                        model = LoGMixture(N, positive_weights=positive_weights).cuda()\n",
    "                    elif model_type == 'general':\n",
    "                        # Adjust parameters as needed for GeneralMixture\n",
    "                        model = GeneralMixture(N, positive_weights=positive_weights, learn_beta=learn_beta).cuda()  # or other parameters\n",
    "\n",
    "\n",
    "                    optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
    "                    loss_func = nn.MSELoss()\n",
    "\n",
    "                    model, y_pred, loss_rec = train_model(model, optimizer, loss_func, x, y, epochs=epochs)\n",
    "\n",
    "                    if np.isnan(loss_rec):\n",
    "                        nan_counts += 1\n",
    "                    else:\n",
    "                        total_loss += loss_rec\n",
    "\n",
    "                    # Plot and save\n",
    "                    plot_func = {\n",
    "                        'gaussian': plot_results_gaussian,\n",
    "                        'dog': plot_results_dog,\n",
    "                        'log': plot_results_log,\n",
    "                        'general': plot_results_general  # make sure you have a plot function for GeneralMixture\n",
    "                    }[model_type]\n",
    "\n",
    "                    plot_func(model, x.cpu(), y.cpu(), y_pred.cpu(), N=N, loss_rec=loss_rec,\n",
    "                              positive_weights=positive_weights, show_fig=False,signal_type=signal_type)\n",
    "\n",
    "                average_loss = total_loss / (runs_per_config - nan_counts) if runs_per_config - nan_counts != 0 else np.nan\n",
    "                results.append({\n",
    "                    'Model': model_type,\n",
    "                    'N': N,\n",
    "                    'Positive Weights': positive_weights,\n",
    "                    'NaN Loss Counts': nan_counts,\n",
    "                    'Average Loss': average_loss\n",
    "                })\n",
    "\n",
    "    for r in results:\n",
    "        print(f\"Signal: {signal_type}, Model: {r['Model']}, N: {r['N']}, Positive Weights: {r['Positive Weights']}, NaN Loss Counts: {r['NaN Loss Counts']}/{runs_per_config}, Average Loss: {r['Average Loss']:.2f}\")\n",
    "    \n",
    "    results_per_signal[signal_type] = results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for signal_type, results in results_per_signal.items():\n",
    "\n",
    "    # Initializing a dictionary for holding results\n",
    "    results_dict = {\n",
    "        'gaussian': {\n",
    "            'positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns},\n",
    "            'non-positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns}\n",
    "        },\n",
    "        'dog': {\n",
    "            'positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns},\n",
    "            'non-positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns}\n",
    "        },\n",
    "        'log': {\n",
    "            'positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns},\n",
    "            'non-positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns}\n",
    "        },\n",
    "         'general': {  # Add the 'general' model type\n",
    "            'positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns},\n",
    "            'non-positive': {N_val: {'stability': 0, 'loss': 0} for N_val in Ns}\n",
    "        }\n",
    "        \n",
    "    }\n",
    "    # Assuming the results list is populated correctly as before:\n",
    "    for result in results:\n",
    "        model_type = result['Model'].lower()  # Making sure we match keys in results_dict\n",
    "        weight_type = 'positive' if result['Positive Weights'] else 'non-positive'\n",
    "        N_val = result['N']\n",
    "        stability = 100 * (runs_per_config - result['NaN Loss Counts']) / runs_per_config\n",
    "        results_dict[model_type][weight_type][N_val]['stability'] = stability\n",
    "        results_dict[model_type][weight_type][N_val]['loss'] = result['Average Loss']\n",
    "    # Step 3: Extract data for plotting\n",
    "    # Extracting stability and loss values\n",
    "    def extract_values(results_dict, model_type, weight_type, Ns):\n",
    "        return [results_dict[model_type][weight_type][N_val]['stability'] for N_val in Ns], \\\n",
    "               [results_dict[model_type][weight_type][N_val]['loss'] for N_val in Ns]\n",
    "\n",
    "    stability_gaussian_positive, loss_gaussian_positive = extract_values(results_dict, 'gaussian', 'positive', Ns)\n",
    "    stability_gaussian_non_positive, loss_gaussian_non_positive = extract_values(results_dict, 'gaussian', 'non-positive', Ns)\n",
    "\n",
    "    stability_dog_positive, loss_dog_positive = extract_values(results_dict, 'dog', 'positive', Ns)\n",
    "    stability_dog_non_positive, loss_dog_non_positive = extract_values(results_dict, 'dog', 'non-positive', Ns)\n",
    "\n",
    "    stability_log_positive, loss_log_positive = extract_values(results_dict, 'log', 'positive', Ns)\n",
    "    stability_log_non_positive, loss_log_non_positive = extract_values(results_dict, 'log', 'non-positive', Ns)\n",
    "\n",
    "    stability_general_positive, loss_general_positive = extract_values(results_dict, 'general', 'positive', Ns)\n",
    "    stability_general_non_positive, loss_general_non_positive = extract_values(results_dict, 'general', 'non-positive', Ns)\n",
    "\n",
    "    # Now, let's plot the results:\n",
    "\n",
    "    color_scheme = {\n",
    "        'gaussian': 'b',\n",
    "        'dog': 'g',\n",
    "        'log': 'r',\n",
    "        'general': 'm'  # A color for the 'general' model, e.g., magenta\n",
    "\n",
    "    }\n",
    "\n",
    "    # Define line styles for weight types\n",
    "    line_styles = {\n",
    "        'positive': '-',\n",
    "        'non-positive': '--'\n",
    "    }\n",
    "    if not os.path.exists(os.path.join(\"comparisons\", \"stats\")):\n",
    "        os.makedirs(os.path.join(\"comparisons\", \"stats\"))\n",
    "\n",
    "    # Stability vs Number of Components\n",
    "    plt.figure(figsize=(12, 8))\n",
    "\n",
    "    # Use linewidth=2 for all plots\n",
    "    plt.plot(Ns, stability_gaussian_positive, color_scheme['gaussian'] + line_styles['positive'] + 'o', label='Gaussian (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, stability_gaussian_non_positive, color_scheme['gaussian'] + line_styles['non-positive'] + 'x', label='Gaussian (Real)', linewidth=2)\n",
    "    \n",
    "    plt.plot(Ns, stability_dog_positive, color_scheme['dog'] + line_styles['positive'] + 'o', label='DoG (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, stability_dog_non_positive, color_scheme['dog'] + line_styles['non-positive'] + 'x', label='DoG (Real)', linewidth=2)\n",
    "\n",
    "    plt.plot(Ns, stability_log_positive, color_scheme['log'] + line_styles['positive'] + 'o', label='LoG (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, stability_log_non_positive, color_scheme['log'] + line_styles['non-positive'] + 'x', label='LoG (Real)', linewidth=2)\n",
    "\n",
    "    plt.plot(Ns, stability_general_positive, color_scheme['general'] + line_styles['positive'] + 'o', label='GEF (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, stability_general_non_positive, color_scheme['general'] + line_styles['non-positive'] + 'x', label='GEF (Real)', linewidth=2)\n",
    "    \n",
    "    \n",
    "    plt.title(\"Stability vs Number of Components on {} signal\".format(signal_type), fontsize=22)\n",
    "    plt.xlabel(\"Number of Components (N)\", fontsize=22)\n",
    "    plt.ylabel(\"Stability (%)\", fontsize=22)\n",
    "    plt.xticks(fontsize=18)\n",
    "    plt.yticks(fontsize=18)\n",
    "    plt.legend(fontsize=20)\n",
    "    plt.grid(True)\n",
    "    plt.savefig(os.path.join(\"comparisons\", \"stats\", \"stability_{}.pdf\".format(signal_type)), bbox_inches=\"tight\")\n",
    "    plt.show()\n",
    "\n",
    "    # Loss Values vs Number of Components\n",
    "    plt.figure(figsize=(12, 8))\n",
    "\n",
    "    # Use linewidth=2 for all plots\n",
    "    plt.plot(Ns, loss_gaussian_positive, color_scheme['gaussian'] + line_styles['positive'] + 'o', label='Gaussian (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, loss_gaussian_non_positive, color_scheme['gaussian'] + line_styles['non-positive'] + 'x', label='Gaussian (Real)', linewidth=2)\n",
    "    plt.plot(Ns, loss_dog_positive, color_scheme['dog'] + line_styles['positive'] + 'o', label='DoG (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, loss_dog_non_positive, color_scheme['dog'] + line_styles['non-positive'] + 'x', label='DoG (Real)', linewidth=2)\n",
    "\n",
    "    plt.plot(Ns, loss_log_positive, color_scheme['log'] + line_styles['positive'] + 'o', label='LoG (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, loss_log_non_positive, color_scheme['log'] + line_styles['non-positive'] + 'x', label='LoG (Real)', linewidth=2)\n",
    "\n",
    "    plt.plot(Ns, loss_general_positive, color_scheme['general'] + line_styles['positive'] + 'o', label='GEF (Positive)', linewidth=2)\n",
    "    plt.plot(Ns, loss_general_non_positive, color_scheme['general'] + line_styles['non-positive'] + 'x', label='GEF (Real)', linewidth=2)\n",
    "    \n",
    "    \n",
    "    plt.title(\"Loss Value vs Number of Components on {} signal\".format(signal_type), fontsize=22)\n",
    "    plt.xlabel(\"Number of Components (N)\", fontsize=22)\n",
    "    plt.ylabel(\"Average Loss (log scale)\", fontsize=22)\n",
    "    plt.yscale('log')\n",
    "    plt.xticks(fontsize=18)\n",
    "    plt.yticks(fontsize=18)\n",
    "    plt.legend(fontsize=20)\n",
    "    plt.grid(True)\n",
    "    plt.savefig(os.path.join(\"comparisons\", \"stats\", \"loss_{}.pdf\".format(signal_type)), bbox_inches=\"tight\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Miscellaneous visualizations "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Frequency-based image filter $\\mathbf{M}_{\\omega}$ for target freq $\\omega$\n",
    "\n",
    "from the GES paper https://abdullahamdi.com/ges/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import torch\n",
    "import os\n",
    "import torch.nn.functional as F\n",
    "def tensor_info(tensor):\n",
    "    \"\"\"\n",
    "    Prints information about a PyTorch tensor including min, max, mean, std, and shape.\n",
    "    \n",
    "    Args:\n",
    "        tensor (torch.Tensor): Input tensor\n",
    "    \"\"\"\n",
    "    # print(\"\\nname:\", f\"{tensor=}\")\n",
    "    print(\"\\nShape:\", tensor.shape)\n",
    "    print(\"Datatype:\", tensor.dtype)\n",
    "    print(\"Device:\", tensor.device)\n",
    "    print(\"Requires grad:\", tensor.requires_grad)\n",
    "    print(\"Min value:\", tensor.min().item())\n",
    "    print(\"Max value:\", tensor.max().item())\n",
    "    print(\"Mean value:\", tensor.mean().item())\n",
    "    print(\"Standard deviation:\", tensor.std().item())\n",
    "\n",
    "def show(img,c_path=None):\n",
    "    img = img - img.min()  # Normalize to [0, max]\n",
    "    img = img / img.max()  # Normalize to [0, 1]\n",
    "    npimg = img.detach().cpu().numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')\n",
    "    plt.show()\n",
    "def show_mask(img,c_path=None):\n",
    "    img = img - img.min()  # Normalize to [0, max]\n",
    "    img = img / img.max()  # Normalize to [0, 1]\n",
    "    \n",
    "    # Detach tensor from computation graph and move to CPU\n",
    "    npimg = img.detach().cpu().squeeze().numpy()  # Squeeze is used to remove any singleton dimensions\n",
    "    \n",
    "    plt.imshow(npimg, interpolation='nearest', cmap='inferno')\n",
    "    plt.colorbar()\n",
    "    plt.show()\n",
    "def show_with_masks(image, masks, mask_names=None, alpha=0.5,c_path=None):\n",
    "    \"\"\"\n",
    "    Overlay masks on an image.\n",
    "    \n",
    "    Args:\n",
    "    - image: PyTorch tensor of shape (C, H, W)\n",
    "    - masks: List of PyTorch tensors, each of shape (H, W)\n",
    "    - mask_names: List of names for each mask for the legend\n",
    "    - alpha: Transparency level for masks\n",
    "    \"\"\"\n",
    "    # Normalize image\n",
    "    image = image - image.min()\n",
    "    image = image / image.max()\n",
    "    npimg = image.detach().cpu().numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')\n",
    "    \n",
    "    colors = ['red', 'green', 'blue', 'yellow', 'purple', 'cyan']  # You can add more colors if needed\n",
    "    \n",
    "    for idx, mask in enumerate(masks):\n",
    "        mask = mask - mask.min()\n",
    "        mask = mask / mask.max()\n",
    "        npmask = mask.detach().cpu().numpy()\n",
    "\n",
    "        # Overlay the binary mask with a color\n",
    "        # First, create a RGB version of the mask where it's colored\n",
    "        mask_colored = np.zeros((npmask.shape[0], npmask.shape[1], 3))\n",
    "        for i in range(3):  # for R, G, B channels\n",
    "            mask_colored[..., i] = npmask * matplotlib.colors.to_rgb(colors[idx % len(colors)])[i]\n",
    "        \n",
    "        plt.imshow(mask_colored, interpolation='nearest', alpha=alpha)\n",
    "    \n",
    "    if mask_names:\n",
    "        patches = [plt.Rectangle((0,0),1,1, color=colors[i % len(colors)]) for i in range(len(masks))]\n",
    "        plt.legend(patches, mask_names, loc='upper left')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "class DoGFilter(nn.Module):\n",
    "    def __init__(self, channels, sigma1):\n",
    "        super(DoGFilter, self).__init__()\n",
    "        self.channels = channels\n",
    "        self.sigma1 = sigma1\n",
    "        self.sigma2 = 2 * sigma1  # Ensure the 1:2 ratio\n",
    "        self.kernel_size1 = int(2 * round(3 * self.sigma1) + 1)\n",
    "        self.kernel_size2 = int(2 * round(3 * self.sigma2) + 1)\n",
    "        self.padding1 = (self.kernel_size1 - 1) // 2\n",
    "        self.padding2 = (self.kernel_size2 - 1) // 2\n",
    "        self.weight1 = self.get_gaussian_kernel(self.kernel_size1, self.sigma1)\n",
    "        self.weight2 = self.get_gaussian_kernel(self.kernel_size2, self.sigma2)\n",
    "\n",
    "\n",
    "    def get_gaussian_kernel(self, kernel_size, sigma):\n",
    "        x_cord = torch.arange(kernel_size)\n",
    "        x_grid = x_cord.repeat(kernel_size).view(kernel_size, kernel_size)\n",
    "        y_grid = x_grid.t()\n",
    "        xy_grid = torch.stack([x_grid, y_grid], dim=-1).float()\n",
    "\n",
    "        mean = (kernel_size - 1) / 2.\n",
    "        variance = sigma**2.\n",
    "        \n",
    "        kernel = torch.exp(-(xy_grid - mean).pow(2).sum(dim=-1) / (2 * variance))\n",
    "        kernel = kernel / kernel.sum()  # Normalize the kernel\n",
    "        kernel = kernel.repeat(self.channels, 1, 1, 1)\n",
    "        \n",
    "        return kernel\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def forward(self, x):\n",
    "        gaussian1 = F.conv2d(x, self.weight1.to(x.device), bias=None, stride=1, padding=self.padding1, groups=self.channels)\n",
    "        gaussian2 = F.conv2d(x, self.weight2.to(x.device), bias=None, stride=1, padding=self.padding2, groups=self.channels)\n",
    "        return gaussian1 - gaussian2\n",
    "def apply_dog_filter(batch, freq=50, scale_factor=0.5):\n",
    "    \"\"\"\n",
    "    Apply a Difference of Gaussian filter to a batch of images.\n",
    "    \n",
    "    Args:\n",
    "        batch: torch.Tensor, shape (B, C, H, W)\n",
    "        freq: Control variable ranging from 0 to 100.\n",
    "              - 0 means original image\n",
    "              - 1.0 means smoother difference\n",
    "              - 100 means sharpest difference\n",
    "        scale_factor: Factor by which the image is downscaled before applying DoG.\n",
    "    \n",
    "    Returns:\n",
    "        torch.Tensor: Processed image using DoG.\n",
    "    \"\"\"\n",
    "    # Convert to grayscale if it's a color image\n",
    "    if batch.size(1) == 3:\n",
    "        batch = torch.mean(batch, dim=1, keepdim=True)\n",
    "\n",
    "    # Downscale the image\n",
    "    downscaled = F.interpolate(batch, scale_factor=scale_factor, mode='bilinear', align_corners=False)\n",
    "\n",
    "    channels = downscaled.size(1)\n",
    "\n",
    "    # Set sigma1 value based on freq parameter. sigma2 will be 2*sigma1.\n",
    "    sigma1 = 0.1 + (100 - freq) * 0.1 if freq >=50 else 0.1 + freq * 0.1\n",
    "\n",
    "    dog_filter = DoGFilter(channels, sigma1)\n",
    "    mask = dog_filter(downscaled)\n",
    "\n",
    "    # Upscale the mask back to original size\n",
    "    upscaled_mask = F.interpolate(mask, size=batch.shape[-2:], mode='bilinear', align_corners=False)\n",
    "\n",
    "    upscaled_mask = upscaled_mask - upscaled_mask.min()\n",
    "    upscaled_mask = upscaled_mask / upscaled_mask.max() if freq >=50 else  1.0 - upscaled_mask / upscaled_mask.max()\n",
    "    \n",
    "    upscaled_mask = (upscaled_mask >=0.5).to(torch.float)\n",
    "    return upscaled_mask[:,0,...]\n",
    "    \n",
    "class LoGFilter(nn.Module):\n",
    "    def __init__(self, channels, sigma, max_kernel_size=None):\n",
    "        super(LoGFilter, self).__init__()\n",
    "        self.channels = channels\n",
    "        self.sigma = sigma\n",
    "        self.kernel_size = int(2 * round(3 * sigma) + 1)\n",
    "        if max_kernel_size:\n",
    "            self.kernel_size = min(self.kernel_size, max_kernel_size)\n",
    "        self.padding = (kernel_size - 1) // 2\n",
    "        self.weight = self.get_log_kernel()\n",
    "\n",
    "    def get_log_kernel(self):\n",
    "        x_cord = torch.arange(self.kernel_size)\n",
    "        x_grid = x_cord.repeat(self.kernel_size).view(self.kernel_size, self.kernel_size)\n",
    "        y_grid = x_grid.t()\n",
    "        xy_grid = torch.stack([x_grid, y_grid], dim=-1).float()\n",
    "\n",
    "        mean = (self.kernel_size - 1) / 2.\n",
    "        variance = self.sigma**2.\n",
    "        \n",
    "        # Laplacian of Gaussian formula\n",
    "        kernel = (-1 / (3.14 * variance**4)) * (1 - (xy_grid - mean).pow(2).sum(dim=-1) / (2 * variance))\n",
    "        kernel *= torch.exp(-(xy_grid - mean).pow(2).sum(dim=-1) / (2 * variance))\n",
    "        \n",
    "        kernel = kernel - kernel.mean()\n",
    "        kernel = kernel.repeat(self.channels, 1, 1, 1)\n",
    "        \n",
    "        return kernel\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def forward(self, x):\n",
    "        return F.conv2d(x, self.weight, bias=None, stride=1, padding=self.padding, groups=self.channels)\n",
    "\n",
    "# Global dictionary to store kernels\n",
    "kernels = {}\n",
    "\n",
    "def apply_log_filter(batch, sigma=1.0):\n",
    "    \"\"\"\n",
    "    Apply a Laplacian of Gaussian filter to a batch of images to highlight high-frequency areas.\n",
    "    Args:\n",
    "        batch: torch.Tensor, shape (B, C, H, W)\n",
    "        sigma: control variable that determines the frequency band of the highlighted region.\n",
    "    Returns:\n",
    "        torch.Tensor: a grayscale mask highlighting high-frequency areas.\n",
    "    \"\"\"\n",
    "    # Convert to grayscale if it's a color image\n",
    "    if batch.size(1) == 3:\n",
    "        batch = torch.mean(batch, dim=1, keepdim=True)\n",
    "    \n",
    "    channels = batch.size(1)\n",
    "    kernel_size = int(2 * round(3 * sigma) + 1)  # Ensure kernel size is odd\n",
    "\n",
    "    # Use existing kernel if it's in the dictionary\n",
    "    if (channels, kernel_size, sigma) in kernels:\n",
    "        weight = kernels[(channels, kernel_size, sigma)]\n",
    "    else:\n",
    "        log_filter = LoGFilter(channels, kernel_size, sigma)\n",
    "        weight = log_filter.weight\n",
    "        kernels[(channels, kernel_size, sigma)] = weight\n",
    "\n",
    "    padding = (kernel_size - 1) // 2\n",
    "    mask = F.conv2d(batch, weight, bias=None, stride=1, padding=padding, groups=channels)\n",
    "    \n",
    "    # Normalize the mask to [0, 1]\n",
    "    mask = mask - mask.min()\n",
    "    mask = mask / mask.max()\n",
    "    \n",
    "    return mask[:,0,...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Example usage:\n",
    "import torchvision\n",
    "freq = 20\n",
    "scale_factor = 0.2\n",
    "# batch_images = torch.randn(16, 3, 224, 224)  # Example batch of images\n",
    "batch_images = torchvision.io.read_image(os.path.join(\".\",\"assets\",\"example.png\"))[None,...].to(torch.float).to(\"cuda\")\n",
    "tensor_info(batch_images)\n",
    "\n",
    "# mask = apply_log_filter(batch_images, sigma = sigma)\n",
    "for freq in range(5,101,5):\n",
    "    mask = apply_dog_filter(batch_images, freq = freq,scale_factor=scale_factor)\n",
    "    tensor_info(mask)\n",
    "    show(batch_images[0])\n",
    "    show_mask(mask[0])\n",
    "    show_with_masks(batch_images[0], [mask[0]], mask_names=[\"freq:{}%\".format(int(freq))], alpha=0.5,c_path=os.path.join(\"output\",f\"mask_{freq}.png\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## studying fourier domain of different signals "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Time-domain signal definitions\n",
    "def get_signal(x, signal_type, width=2):\n",
    "    if signal_type == \"square\":\n",
    "        return np.where((x > -width/2) & (x < width/2), np.ones_like(x), np.zeros_like(x))\n",
    "    elif signal_type == \"triangle\":\n",
    "        return np.where((x > -width/2) & (x < width/2), width/2 - np.abs(x), np.zeros_like(x))\n",
    "    elif signal_type == \"parabolic\":\n",
    "        return np.where((x > -width/2) & (x < width/2), (width/2)**2 - x**2, np.zeros_like(x))\n",
    "    elif signal_type == \"half_sinusoid\":\n",
    "        return np.where((x > -width/2) & (x < width/2), np.sin((x + width/2) * (np.pi / width)), np.zeros_like(x))\n",
    "    elif signal_type == \"exponential\":\n",
    "        return np.where((x > -width/2) & (x < width/2), np.exp(-np.abs(x)), np.zeros_like(x))\n",
    "    elif signal_type == \"gaussian\":\n",
    "        return np.exp(-x**2 / (2 * (width/3)**2))\n",
    "\n",
    "# Fourier transform definitions\n",
    "def fourier_square_wave(x, width):\n",
    "    return np.sinc(x * width / np.pi)\n",
    "\n",
    "def fourier_triangle_wave(x, width):\n",
    "    return (np.sinc(x * width / (2 * np.pi)))**2\n",
    "\n",
    "def fourier_parabolic_wave(x, width):\n",
    "    return (3 * (np.sinc(x * width / (2 * np.pi)))**2) / (np.pi**2 * x**2)\n",
    "\n",
    "def fourier_half_sinusoid(x, width):\n",
    "    return np.where(x == 0, width / 2, width * np.sin(np.pi * x * width) / (np.pi**2 * x**2))\n",
    "\n",
    "def fourier_exponential(x, width):\n",
    "    return width / (x**2 + (width/2)**2)\n",
    "\n",
    "def fourier_gaussian(x, width):\n",
    "    sigma = width / 3\n",
    "    return np.sqrt(2 * np.pi) * sigma * np.exp(-2 * (np.pi**2) * sigma**2 * x**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "###### Signal types and their Fourier transform functions\n",
    "signal_types = [\"square\", \"triangle\", \"parabolic\", \"half_sinusoid\", \"exponential\", \"gaussian\"]\n",
    "# signal_types = [\"square\", \"triangle\",\"gaussian\"]\n",
    "\n",
    "fourier_functions = {\n",
    "  \"square\": fourier_square_wave, \n",
    "  \"triangle\": fourier_triangle_wave, \n",
    "  \"parabolic\": fourier_parabolic_wave, \n",
    "  \"half_sinusoid\": fourier_half_sinusoid, \n",
    "  \"exponential\": fourier_exponential, \n",
    "  \"gaussian\": fourier_gaussian\n",
    "}\n",
    "\n",
    "# Time and frequency domain range\n",
    "x_time = np.linspace(-10, 10, 1000)\n",
    "x_freq = np.linspace(-10, 10, 1000)\n",
    "width = 2\n",
    "\n",
    "# Create the plots\n",
    "plt.figure(figsize=(20, 12))\n",
    "# plt.figure(figsize=(20, 6))\n",
    "\n",
    "\n",
    "for i, signal_type in enumerate(signal_types):\n",
    "    # Time-domain signal\n",
    "    signal = get_signal(x_time, signal_type, width)\n",
    "\n",
    "    # Fourier transform\n",
    "    fourier_transform = fourier_functions[signal_type](x_freq, width)\n",
    "\n",
    "    # Plotting time-domain signal\n",
    "    plt.subplot(len(signal_types), 2, 2*i + 1)\n",
    "    plt.plot(x_time, signal, label=f\"{signal_type} Signal\", linewidth=2)\n",
    "    plt.title(f\"{signal_type.capitalize()} Signal\", fontsize=18)\n",
    "    plt.xlabel(\"Time\", fontsize=16)\n",
    "    plt.ylabel(\"Amplitude\", fontsize=16)\n",
    "    plt.grid(True, which=\"both\", ls=\"--\")\n",
    "    # plt.legend(fontsize=16)\n",
    "    plt.tick_params(axis='both', which='major', labelsize=15)\n",
    "    plt.tick_params(axis='both', which='minor', labelsize=15)\n",
    "\n",
    "    # Plotting Fourier transform\n",
    "    plt.subplot(len(signal_types), 2, 2*i + 2)\n",
    "    plt.plot(x_freq, fourier_transform, label=f\"Fourier of {signal_type}\", linewidth=2)\n",
    "    plt.title(f\"Fourier Transform of {signal_type.capitalize()}\", fontsize=18)\n",
    "    plt.xlabel(\"Frequency\", fontsize=16)\n",
    "    plt.ylabel(\"Magnitude\", fontsize=16)\n",
    "    plt.grid(True, which=\"both\", ls=\"--\")\n",
    "    # plt.legend(fontsize=16)\n",
    "    plt.tick_params(axis='both', which='major', labelsize=15)\n",
    "    plt.tick_params(axis='both', which='minor', labelsize=15)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"output/part_signals_and_fourier_transforms.pdf\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "",
   "provenance": [
    {
     "file_id": "1YniEH63VfZPuRGTddviUvNH48cDaLqtg",
     "timestamp": 1695128610581
    },
    {
     "file_id": "https://github.com/camenduru/gaussian-splatting-colab/blob/main/gaussian_splatting_colab.ipynb",
     "timestamp": 1694175856974
    }
   ],
   "version": ""
  },
  "kernelspec": {
   "display_name": "Python [conda env:magicxl]",
   "language": "python",
   "name": "conda-env-magicxl-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0850e8e8f23c49b7b55bac8a0ec5324d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "115dd7c9db75492eadfcb930b08cf3bb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_d1170afaac5b4948bfd87e9d93c9d097",
       "IPY_MODEL_2ff878a986fd40c8aba2462d38fe1c9c",
       "IPY_MODEL_1cbdc3cffac649d287eb4c709d9ad2e0"
      ],
      "layout": "IPY_MODEL_7797ef7414b44bb5bd2224e31dc30af6"
     }
    },
    "1372fb5f72934f659b2c03cb0eb24252": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "1cbdc3cffac649d287eb4c709d9ad2e0": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_fb0d9d71286040fd9568c0f6d85a12c6",
      "placeholder": "​",
      "style": "IPY_MODEL_f891466eb89248f7b25d68fb2ba744f2",
      "value": " 605M/605M [00:06&lt;00:00, 61.8MB/s]"
     }
    },
    "2ff878a986fd40c8aba2462d38fe1c9c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_c052561baf38426389a540b250b47d2b",
      "max": 605219813,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_1372fb5f72934f659b2c03cb0eb24252",
      "value": 605219813
     }
    },
    "7797ef7414b44bb5bd2224e31dc30af6": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a1c5ba59b9c5459787485d734ec93467": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "c052561baf38426389a540b250b47d2b": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d1170afaac5b4948bfd87e9d93c9d097": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a1c5ba59b9c5459787485d734ec93467",
      "placeholder": "​",
      "style": "IPY_MODEL_0850e8e8f23c49b7b55bac8a0ec5324d",
      "value": "Downloading (…)ip_pytorch_model.bin: 100%"
     }
    },
    "f891466eb89248f7b25d68fb2ba744f2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "fb0d9d71286040fd9568c0f6d85a12c6": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}