{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 423
        },
        "id": "DaR3P2HHTYG9",
        "outputId": "4eed57d7-3e23-4ba0-d9e1-382b821613d1"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 37704 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from sklearn.datasets import make_swiss_roll\n",
        "from mpl_toolkits.mplot3d import Axes3D\n",
        "\n",
        "# seed_value = 50\n",
        "# np.random.seed(seed_value)\n",
        "plt.style.use('dark_background')  # Dark background style\n",
        "\n",
        "mean = np.array([0, 0, 0])\n",
        "covariance_matrix = np.array([[0.08, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])\n",
        "\n",
        "data = np.random.multivariate_normal(mean, covariance_matrix, 2000).transpose()\n",
        "\n",
        "fig = plt.figure()\n",
        "fig.patch.set_facecolor('white')  # Set figure background to white\n",
        "ax = fig.add_subplot(111, projection='3d')\n",
        "\n",
        "ax.set_xlim(-3, 3)\n",
        "ax.set_ylim(-3, 3)\n",
        "ax.set_zlim(-3, 3)\n",
        "\n",
        "ax.tick_params(axis='x', colors='black')\n",
        "ax.tick_params(axis='y', colors='black')\n",
        "ax.tick_params(axis='z', colors='black')\n",
        "\n",
        "ax.set_facecolor('white')  # Set 3D axis facecolor to white\n",
        "\n",
        "# Change the grid background color to gray\n",
        "ax.xaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "ax.yaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "ax.zaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "\n",
        "# Set grid line color for better visibility\n",
        "ax.grid(color='white')\n",
        "\n",
        "# Set axis labels\n",
        "ax.set_xlabel('X', labelpad=10, color='blue')\n",
        "ax.set_ylabel('Y', labelpad=10, color='blue')\n",
        "ax.set_zlabel('Z', labelpad=10, color='blue')\n",
        "\n",
        "# Adjust view angle\n",
        "ax.view_init(elev=5, azim=-90)\n",
        "\n",
        "# Scatter plot\n",
        "ax.scatter(*data, color='pink', edgecolor='pink', zorder=1, marker='o', facecolors='none')\n",
        "\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q8ehaMrjT3AY"
      },
      "outputs": [],
      "source": [
        "Tuning=1e-3\n",
        "tuning_k=1e-6\n",
        "tunp=1e-6\n",
        "sigma=0.01\n",
        "def denoising_score_matching_reg(scorenet,k, samples,Tuning,tuning_k,sigma=0.01):\n",
        "    perturbed_samples = samples + torch.randn_like(samples) * sigma\n",
        "    target = - 1 / (sigma ** 2) * (perturbed_samples - samples)\n",
        "    scores = k*scorenet(perturbed_samples)\n",
        "    target = target.view(target.shape[0], -1)\n",
        "    scores = scores.view(scores.shape[0], -1)\n",
        "    l1_norm = sum(p.abs().sum() for p in scorenet.parameters())\n",
        "\n",
        "    loss = 1 / 2. * ((scores - target) ** 2).sum(dim=-1).mean(dim=0)+Tuning*(scores.norm(p=1).sum(dim=-1))+(tuning_k*torch.abs(k))+(tunp*l1_norm)\n",
        "    return loss\n",
        "\n",
        "def denoising_score_matching_Or(scorenet,k, samples,sigma=0.01):\n",
        "    perturbed_samples = samples + torch.randn_like(samples) * sigma\n",
        "    target = - 1 / (sigma ** 2) * (perturbed_samples - samples)\n",
        "    scores = scorenet(perturbed_samples)\n",
        "    target = target.view(target.shape[0], -1)\n",
        "    scores = scores.view(scores.shape[0], -1)\n",
        "    loss = 1 / 2. * ((scores - target) ** 2).sum(dim=-1).mean(dim=0)\n",
        "    return loss\n",
        "\n",
        "\n",
        "# def denoising_score_matching_reg_m(scorenet, samples, noise_levels, tuning_parameter):\n",
        "#     total_loss = 0.0\n",
        "\n",
        "#     for sigma in noise_levels:\n",
        "#         perturbed_samples = samples + torch.randn_like(samples) * sigma\n",
        "#         target = -1 / (sigma ** 2) * (perturbed_samples - samples)\n",
        "#         scores = scorenet(perturbed_samples)\n",
        "#         target = target.view(target.shape[0], -1)\n",
        "#         scores = scores.view(scores.shape[0], -1)\n",
        "#         loss = 1 / 2. * ((scores - target) ** 2).sum(dim=-1).mean(dim=0) + tuning_parameter * (scores.norm(p=1).sum(dim=-1))\n",
        "#         total_loss += loss\n",
        "\n",
        "#     # Average the loss over all noise levels\n",
        "#     avg_loss = total_loss / len(noise_levels)\n",
        "\n",
        "#     return avg_loss\n",
        "\n",
        "# noise_levels = [0.2, 0.1, 0.01, 0.005]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {
        "id": "bW0P5FYJT8d_"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "# Our approximation model\n",
        "# modelS = nn.Sequential(\n",
        "#     nn.Linear(3, 128), nn.Softplus(),\n",
        "#     nn.Linear(128, 128), nn.Softplus(),\n",
        "#     nn.Linear(128, 3)\n",
        "# )\n",
        "\n",
        "k = torch.tensor(0.0, requires_grad=True)  ## Create ADAM optimizer over our model\n",
        "def training(modelS,Tuning,tuning_k,flag):\n",
        "  parameters_to_optimize = list(modelS.parameters()) + [k]\n",
        "  optimizer = optim.Adam(parameters_to_optimize, lr=1e-4)\n",
        "  dataset = torch.tensor(data.T).float()\n",
        "  if flag==True:\n",
        "     loss = denoising_score_matching_reg(modelS,k, dataset,Tuning,tuning_k)\n",
        "  else:\n",
        "     loss = denoising_score_matching_Or(modelS,k, dataset)\n",
        "     k.requires_grad = False\n",
        " \n",
        "     \n",
        "  for t in range(10000):\n",
        "    #   # Compute the loss.\n",
        "    #   loss = denoising_score_matching_reg(modelS,k, dataset,Tuning,tuning_k)\n",
        "      # Before the backward pass, zero all of the network gradients\n",
        "      if flag==True:\n",
        "        loss = denoising_score_matching_reg(modelS,k, dataset,Tuning,tuning_k)\n",
        "      else:\n",
        "        loss = denoising_score_matching_Or(modelS,k, dataset)\n",
        "\n",
        "      # Backward pass: compute gradient of the loss with respect to parameters\n",
        "      optimizer.zero_grad()\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "      # Print loss\n",
        "      if ((t % 1000) == 0):\n",
        "          print(loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 64,
      "metadata": {
        "id": "I57fKl8fZQek"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def plot_gradients3D(model, k_t,data, ax,plot_scatter=True):\n",
        "    # Generate meshgrid for 3D data\n",
        "    x = np.linspace(-3.0, 3.0, 15)\n",
        "    y = np.linspace(-3.0, 3.0, 15)\n",
        "    z = np.linspace(-3.0, 3.0, 15)\n",
        "    xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')\n",
        "    meshgrid_points = np.stack([xx, yy, zz], axis=-1).reshape(-1, 3)\n",
        "\n",
        "    # Convert meshgrid to tensor and compute gradients\n",
        "    meshgrid_points_tensor = torch.tensor(meshgrid_points, dtype=torch.float32, requires_grad=True)\n",
        "    scores = k_t*model(meshgrid_points_tensor).detach().numpy()\n",
        "    # gradients = torch.autograd.grad(scores, meshgrid_points_tensor)[0].detach().numpy()\n",
        "\n",
        "    # Normalize gradients\n",
        "    gradients_norm = np.linalg.norm(scores, axis=-1, ord=2, keepdims=True)\n",
        "    gradients_normalized = scores / (gradients_norm + 1e-9) * np.log1p(gradients_norm)\n",
        "\n",
        "    # Perform the plots\n",
        "    # fig = plt.figure()\n",
        "    # ax = fig.add_subplot(111, projection='3d')\n",
        "    # plt.figure(figsize=(16, 12))\n",
        "    ax.scatter(*data[:,1:400],color='pink', edgecolor='pink',zorder=1,marker='o', facecolors='none')\n",
        "\n",
        "    # ax.quiver(meshgrid_points[:, 0], meshgrid_points[:, 1], meshgrid_points[:, 2],\n",
        "    #           gradients_normalized[:, 0], gradients_normalized[:, 1], gradients_normalized[:, 2],length=0.1, linewidths=0.3,normalize=False, color='red',zorder=1)\n",
        "\n",
        "    # ax.set_xlim(-2.0, 2.0)\n",
        "    # ax.set_ylim(-2.0, 2.0)\n",
        "    # ax.set_zlim(-2.0, 2.0)\n",
        "\n",
        "    # plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {
        "id": "CryHw-3kQ5rv"
      },
      "outputs": [],
      "source": [
        "def sample_langevin3D(model, k,x, n_steps=60, eps=0.3, decay=.99, temperature=0.05):\n",
        "    eps=58*eps/n_steps\n",
        "    x_sequence = [x.unsqueeze(0)]\n",
        "    for s in range(n_steps):\n",
        "        z_t = torch.rand(x.size())\n",
        "        x = x + (eps / 2) * k*model(x) + (np.sqrt(eps) * temperature * z_t)\n",
        "        x_sequence.append(x.unsqueeze(0))\n",
        "        eps *= decay\n",
        "    return torch.cat(x_sequence)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {
        "id": "5MDiBij3RS2-"
      },
      "outputs": [],
      "source": [
        "def sampling(ax,model,k):\n",
        "  start=torch.Tensor([1.2, 2.2,2.2])\n",
        "  # Sample x ∈ [-1.5, 1.5]\n",
        "  # x = torch.empty(1).uniform_(-2, 2)\n",
        "\n",
        "  # # Sample y, z ∈ [-3, 3]\n",
        "  # yz = torch.empty(2).uniform_(-3, 3)\n",
        "\n",
        "  # # Combine into a single start tensor\n",
        "  # start_random = torch.cat([x, yz])  # shape: (3,)\n",
        "  # start=start_random\n",
        "\n",
        "  x = start\n",
        "  samples = sample_langevin3D(model,k, x).detach()\n",
        "  # plot_gradients3D(model, data)\n",
        "\n",
        "  # fig = plt.figure(figsize=(16, 12))\n",
        "  # ax = fig.add_subplot(111, projection='3d')\n",
        "  ax.scatter(samples[:, 0], samples[:, 1], samples[:, 2], color='green', edgecolor='green', s=50,zorder=3)\n",
        "  # ax.scatter(samples[0, 0], samples[0, 1], samples[0, 2], color='blue', edgecolor='blue', s=50,zorder=3)\n",
        "  ax.scatter(start[ 0], start[1], start[2], color='black', edgecolor='black', s=70,zorder=3)\n",
        "\n",
        "\n",
        "  # Draw arrows for each MCMC step\n",
        "  deltas = (samples[1:] - samples[:-1])\n",
        "\n",
        "  for i in range(len(deltas)):\n",
        "      arrow_start = samples[i]\n",
        "      arrow_direction = deltas[i]\n",
        "      arrow_color = 'red'\n",
        "\n",
        "      ax.quiver(arrow_start[0], arrow_start[1], arrow_start[2],\n",
        "                arrow_direction[0], arrow_direction[1], arrow_direction[2],\n",
        "                color=arrow_color, arrow_length_ratio=0.2,zorder=2)\n",
        "\n",
        "  plot_gradients3D(model, k, data, ax)\n",
        "\n",
        "  # Set axis limits\n",
        "  ax.set_xlim(-3, 3)\n",
        "  ax.set_ylim(-3, 3)\n",
        "  ax.set_zlim(-3, 3)\n",
        "\n",
        "  # # Set axis labels\n",
        "  ax.set_xlabel('X')\n",
        "  ax.set_ylabel('Y')\n",
        "  ax.set_zlabel('Z')\n",
        "  # ax.view_init(elev=30, azim=-90)\n",
        "  # plot_gradients3D(model, data,ax)\n",
        "  return samples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AltfnHPAgRjQ",
        "outputId": "1f64781b-5a0e-4239-86f8-c38121210037"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(14759.0781, grad_fn=<AddBackward0>)\n",
            "tensor(14823.3555, grad_fn=<AddBackward0>)\n",
            "tensor(15000.2627, grad_fn=<AddBackward0>)\n",
            "tensor(15074.1035, grad_fn=<AddBackward0>)\n",
            "tensor(15209.1680, grad_fn=<AddBackward0>)\n",
            "tensor(15043.2510, grad_fn=<AddBackward0>)\n",
            "tensor(14459.2139, grad_fn=<AddBackward0>)\n",
            "tensor(15367.4990, grad_fn=<AddBackward0>)\n",
            "tensor(15099.1426, grad_fn=<AddBackward0>)\n",
            "tensor(15350.5498, grad_fn=<AddBackward0>)\n"
          ]
        }
      ],
      "source": [
        "#traing with regularization\n",
        "model_reg = nn.Sequential(\n",
        "    nn.Linear(3, 128), nn.Softplus(),\n",
        "    nn.Linear(128, 128), nn.Softplus(),\n",
        "    nn.Linear(128, 3)\n",
        ")\n",
        "training(model_reg,Tuning=5e-4,tuning_k=5e-5,flag=True)\n",
        "# training(model_reg,Tuning=4e-4,tuning_k=0,flag=True)\n",
        "\n",
        "k_reg=k.detach()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 68,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 60876 characters]",
            "text/plain": [
              "<Figure size 1200x600 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from mpl_toolkits.mplot3d import Axes3D\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "np.random.seed() \n",
        "\n",
        "# Set the size of the figure with a white background\n",
        "fig = plt.figure(figsize=(12, 6), facecolor='white')\n",
        "\n",
        "# Create subplots\n",
        "ax2 = fig.add_subplot(122, projection='3d')\n",
        "\n",
        "# Set titles for each subplot with adjusted padding\n",
        "ax2.set_title('Regularized score matching', pad=10)\n",
        "\n",
        "# Set axis labels for the first subplot with label colors\n",
        "\n",
        "\n",
        "# Set axis labels for the second subplot with label colors\n",
        "ax2.set_xlabel('X-axis Label', labelpad=10, color='blue')\n",
        "ax2.set_ylabel('Y-axis Label', labelpad=10, color='blue')\n",
        "ax2.set_zlabel('Z-axis Label', labelpad=10, color='blue')\n",
        "\n",
        "# Set the axis (chart area) background color to white\n",
        "ax2.set_facecolor('white')\n",
        "\n",
        "# Change the grid background color to gray\n",
        "ax2.xaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "ax2.yaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "ax2.zaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "\n",
        "# Set grid line color for better visibility\n",
        "ax2.grid(color='white')\n",
        "\n",
        "# Set tick labels' color for both subplots\n",
        "\n",
        "ax2.tick_params(axis='x', colors='black')\n",
        "ax2.tick_params(axis='y', colors='black')\n",
        "ax2.tick_params(axis='z', colors='black')\n",
        "ax2.view_init(elev=5, azim=-90)\n",
        "\n",
        "\n",
        "# Assuming `k_org` and `k_reg` are defined\n",
        "# k_org = k_org.detach()\n",
        "k_reg = k_reg.detach()\n",
        "\n",
        "\n",
        "# Plot on the second subplot\n",
        "\n",
        "samples = sampling(ax2, model_reg, k_reg)\n",
        "    \n",
        "  \n",
        "# Adjust layout to prevent overlapping labels\n",
        "plt.tight_layout()\n",
        "\n",
        "# Save the figure with a white background\n",
        "plt.savefig('Toy', facecolor='white')\n",
        "\n",
        "# Show the figure with both subplots\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 69,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KykMxEurSnrv",
        "outputId": "ad60b42e-cfdd-44fd-f486-73b51219fd2a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(15009.1289, grad_fn=<MulBackward0>)\n",
            "tensor(14740.4248, grad_fn=<MulBackward0>)\n",
            "tensor(15053.3027, grad_fn=<MulBackward0>)\n",
            "tensor(14931.0059, grad_fn=<MulBackward0>)\n",
            "tensor(15194.0156, grad_fn=<MulBackward0>)\n",
            "tensor(14542.6436, grad_fn=<MulBackward0>)\n",
            "tensor(14830.3760, grad_fn=<MulBackward0>)\n",
            "tensor(14551.7920, grad_fn=<MulBackward0>)\n",
            "tensor(15239.9141, grad_fn=<MulBackward0>)\n",
            "tensor(15011.3633, grad_fn=<MulBackward0>)\n"
          ]
        }
      ],
      "source": [
        "#traing withought regularization\n",
        "model_org = nn.Sequential(\n",
        "    nn.Linear(3, 128), nn.Softplus(),\n",
        "    nn.Linear(128, 128), nn.Softplus(),\n",
        "    nn.Linear(128, 3)\n",
        ")\n",
        "\n",
        "training(model_org,Tuning=0,tuning_k=0,flag=False)\n",
        "k_org=1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {},
      "outputs": [],
      "source": [
        "def sample_langevin3D(model, k,x, n_steps=60, eps=0.3, decay=.9, temperature=0.05):\n",
        "    x_sequence = [x.unsqueeze(0)]\n",
        "    for s in range(n_steps):\n",
        "        z_t = torch.rand(x.size())\n",
        "        x = x + (eps / 2) * k*model(x) + (np.sqrt(eps) * temperature * z_t)\n",
        "        x_sequence.append(x.unsqueeze(0))\n",
        "        eps *= decay\n",
        "    return torch.cat(x_sequence)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 61744 characters]",
            "text/plain": [
              "<Figure size 1200x600 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from mpl_toolkits.mplot3d import Axes3D\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "# random.seed(100) \n",
        "\n",
        "# Set the size of the figure with a white background\n",
        "fig = plt.figure(figsize=(12, 6), facecolor='white')\n",
        "\n",
        "# Create subplots\n",
        "ax1 = fig.add_subplot(121, projection='3d')\n",
        "\n",
        "# Set titles for each subplot with adjusted padding\n",
        "ax1.set_title('Original score matching', pad=10)\n",
        "\n",
        "# Set axis labels for the first subplot with label colors\n",
        "ax1.set_xlabel('X-axis Label', labelpad=10, color='blue')\n",
        "ax1.set_ylabel('Y-axis Label', labelpad=10, color='blue')\n",
        "ax1.set_zlabel('Z-axis Label', labelpad=10, color='blue')\n",
        "ax1.view_init(elev=5, azim=-90)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# Set the axis (chart area) background color to white\n",
        "ax1.set_facecolor('white')  # Ensures subplot background is white\n",
        "\n",
        "# Change the grid background color to gray\n",
        "ax1.xaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))  # RGB + Alpha for gray\n",
        "ax1.yaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "ax1.zaxis.pane.set_facecolor((0.8, 0.8, 0.8, 1.0))\n",
        "\n",
        "\n",
        "# Set grid line color for better visibility\n",
        "ax1.grid(color='white')\n",
        "\n",
        "# Set tick labels' color for both subplots\n",
        "ax1.tick_params(axis='x', colors='black')\n",
        "ax1.tick_params(axis='y', colors='black')\n",
        "ax1.tick_params(axis='z', colors='black')\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "samples = sampling(ax1, model_org, k_org)\n",
        "    \n",
        "   \n",
        "\n",
        "# Adjust layout to prevent overlapping labels\n",
        "plt.tight_layout()\n",
        "\n",
        "# Save the figure with a white background\n",
        "plt.savefig('Toy', facecolor='white')\n",
        "\n",
        "# Show the figure with both subplots\n",
        "plt.show()\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "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.9.6"
    },
    "latex_envs": {
      "LaTeX_envs_menu_present": true,
      "autoclose": false,
      "autocomplete": true,
      "bibliofile": "biblio.bib",
      "cite_by": "apalike",
      "current_citInitial": 1,
      "eqLabelWithNumbers": true,
      "eqNumInitial": 1,
      "hotkeys": {
        "equation": "Ctrl-E",
        "itemize": "Ctrl-I"
      },
      "labels_anchors": false,
      "latex_user_defs": false,
      "report_style_numbering": false,
      "user_envs_cfg": false
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
