2024-04-23 DDPM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from inspect import isfunction\n",
    "from functools import partial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.auto import tqdm\n",
    "from einops import rearrange, reduce\n",
    "from einops.layers.torch import Rearrange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, einsum\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import torchvision.models as models\n",
    "# from skimage import util\n",
    "# #import skimage\n",
    "# # skimage import util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "from torchvision.transforms import Compose, ToTensor, Lambda, ToPILImage, CenterCrop, Resize\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "device\n",
    "# use seed for reproducability\n",
    "torch.manual_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#   !pip install -q -U einops datasets matplotlib tqdm\n",
    "import math\n",
    "from inspect import isfunction\n",
    "from functools import partial\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.auto import tqdm\n",
    "from einops import rearrange, reduce\n",
    "from einops.layers.torch import Rearrange\n",
    "\n",
    "import torch\n",
    "from torch import nn, einsum\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torchvision.models as models\n",
    "#from skimage.util import random_noise\n",
    "\n",
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "from torchvision.transforms import Compose, ToTensor, Lambda, ToPILImage, CenterCrop, Resize\n",
    "import numpy as np\n",
    "\n",
    "#import ddpm\n",
    "\n",
    "def exists(x):\n",
    "    return x is not None\n",
    "\n",
    "def default(val, d):\n",
    "    if exists(val):\n",
    "        return val\n",
    "    return d() if isfunction(d) else d\n",
    "\n",
    "\n",
    "def num_to_groups(num, divisor):\n",
    "    groups = num // divisor\n",
    "    remainder = num % divisor\n",
    "    arr = [divisor] * groups\n",
    "    if remainder > 0:\n",
    "        arr.append(remainder)\n",
    "    return arr\n",
    "\n",
    "\n",
    "class Residual(nn.Module):\n",
    "    def __init__(self, fn):\n",
    "        super().__init__()\n",
    "        self.fn = fn\n",
    "\n",
    "    def forward(self, x, *args, **kwargs):\n",
    "        return self.fn(x, *args, **kwargs) + x\n",
    "\n",
    "\n",
    "def Upsample(dim, dim_out=None):\n",
    "    return nn.Sequential(\n",
    "        nn.Upsample(scale_factor=2, mode=\"nearest\"),\n",
    "        nn.Conv2d(dim, default(dim_out, dim), 3, padding=1),\n",
    "    )\n",
    "\n",
    "\n",
    "def Downsample(dim, dim_out=None):\n",
    "    # No More Strided Convolutions or Pooling\n",
    "    return nn.Sequential(\n",
    "        Rearrange(\"b c (h p1) (w p2) -> b (c p1 p2) h w\", p1=2, p2=2),\n",
    "        nn.Conv2d(dim * 4, default(dim_out, dim), 1),\n",
    "    )\n",
    "\n",
    "class SinusoidalPositionEmbeddings(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "    def forward(self, time):\n",
    "        device = time.device\n",
    "        half_dim = self.dim // 2\n",
    "        embeddings = math.log(10000) / (half_dim - 1)\n",
    "        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)\n",
    "        embeddings = time[:, None] * embeddings[None, :]\n",
    "        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)\n",
    "        return embeddings\n",
    "\n",
    "class WeightStandardizedConv2d(nn.Conv2d):\n",
    "    \"\"\"\n",
    "    https://arxiv.org/abs/1903.10520\n",
    "    weight standardization purportedly works synergistically with group normalization\n",
    "    \"\"\"\n",
    "\n",
    "    def forward(self, x):\n",
    "        eps = 1e-5 if x.dtype == torch.float32 else 1e-3\n",
    "\n",
    "        weight = self.weight\n",
    "        mean = reduce(weight, \"o ... -> o 1 1 1\", \"mean\")\n",
    "        var = reduce(weight, \"o ... -> o 1 1 1\", partial(torch.var, unbiased=False))\n",
    "        normalized_weight = (weight - mean) / (var + eps).rsqrt()\n",
    "\n",
    "        return F.conv2d(\n",
    "            x,\n",
    "            normalized_weight,\n",
    "            self.bias,\n",
    "            self.stride,\n",
    "            self.padding,\n",
    "            self.dilation,\n",
    "            self.groups,\n",
    "        )\n",
    "\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, dim, dim_out, groups=8):\n",
    "        super().__init__()\n",
    "        self.proj = WeightStandardizedConv2d(dim, dim_out, 3, padding=1)\n",
    "        self.norm = nn.GroupNorm(groups, dim_out)\n",
    "        self.act = nn.SiLU()\n",
    "\n",
    "    def forward(self, x, scale_shift=None):\n",
    "        x = self.proj(x)\n",
    "        x = self.norm(x)\n",
    "\n",
    "        if exists(scale_shift):\n",
    "            scale, shift = scale_shift\n",
    "            x = x * (scale + 1) + shift\n",
    "\n",
    "        x = self.act(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class ResnetBlock(nn.Module):\n",
    "    \"\"\"https://arxiv.org/abs/1512.03385\"\"\"\n",
    "\n",
    "    def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):\n",
    "        super().__init__()\n",
    "        self.mlp = (\n",
    "            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out * 2))\n",
    "            if exists(time_emb_dim)\n",
    "            else None\n",
    "        )\n",
    "\n",
    "        self.block1 = Block(dim, dim_out, groups=groups)\n",
    "        self.block2 = Block(dim_out, dim_out, groups=groups)\n",
    "        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()\n",
    "\n",
    "    def forward(self, x, time_emb=None):\n",
    "        scale_shift = None\n",
    "        if exists(self.mlp) and exists(time_emb):\n",
    "            time_emb = self.mlp(time_emb)\n",
    "            time_emb = rearrange(time_emb, \"b c -> b c 1 1\")\n",
    "            scale_shift = time_emb.chunk(2, dim=1)\n",
    "\n",
    "        h = self.block1(x, scale_shift=scale_shift)\n",
    "        h = self.block2(h)\n",
    "        return h + self.res_conv(x)\n",
    "    \n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim, heads=4, dim_head=32):\n",
    "        super().__init__()\n",
    "        self.scale = dim_head**-0.5\n",
    "        self.heads = heads\n",
    "        hidden_dim = dim_head * heads\n",
    "        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)\n",
    "        self.to_out = nn.Conv2d(hidden_dim, dim, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, c, h, w = x.shape\n",
    "        qkv = self.to_qkv(x).chunk(3, dim=1)\n",
    "        q, k, v = map(\n",
    "            lambda t: rearrange(t, \"b (h c) x y -> b h c (x y)\", h=self.heads), qkv\n",
    "        )\n",
    "        q = q * self.scale\n",
    "\n",
    "        sim = einsum(\"b h d i, b h d j -> b h i j\", q, k)\n",
    "        sim = sim - sim.amax(dim=-1, keepdim=True).detach()\n",
    "        attn = sim.softmax(dim=-1)\n",
    "\n",
    "        out = einsum(\"b h i j, b h d j -> b h i d\", attn, v)\n",
    "        out = rearrange(out, \"b h (x y) d -> b (h d) x y\", x=h, y=w)\n",
    "        return self.to_out(out)\n",
    "\n",
    "class LinearAttention(nn.Module):\n",
    "    def __init__(self, dim, heads=4, dim_head=32):\n",
    "        super().__init__()\n",
    "        self.scale = dim_head**-0.5\n",
    "        self.heads = heads\n",
    "        hidden_dim = dim_head * heads\n",
    "        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)\n",
    "\n",
    "        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), \n",
    "                                    nn.GroupNorm(1, dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, c, h, w = x.shape\n",
    "        qkv = self.to_qkv(x).chunk(3, dim=1)\n",
    "        q, k, v = map(\n",
    "            lambda t: rearrange(t, \"b (h c) x y -> b h c (x y)\", h=self.heads), qkv\n",
    "        )\n",
    "\n",
    "        q = q.softmax(dim=-2)\n",
    "        k = k.softmax(dim=-1)\n",
    "\n",
    "        q = q * self.scale\n",
    "        context = torch.einsum(\"b h d n, b h e n -> b h d e\", k, v)\n",
    "\n",
    "        out = torch.einsum(\"b h d e, b h d n -> b h e n\", context, q)\n",
    "        out = rearrange(out, \"b h c (x y) -> b (h c) x y\", h=self.heads, x=h, y=w)\n",
    "        return self.to_out(out)\n",
    "    \n",
    "class PreNorm(nn.Module):\n",
    "    def __init__(self, dim, fn):\n",
    "        super().__init__()\n",
    "        self.fn = fn\n",
    "        self.norm = nn.GroupNorm(1, dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.norm(x)\n",
    "        return self.fn(x)\n",
    "    \n",
    "class Unet(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim,\n",
    "        init_dim=None,\n",
    "        out_dim=None,\n",
    "        dim_mults=(1, 2, 4, 8),\n",
    "        channels=1,\n",
    "        self_condition=False,\n",
    "        resnet_block_groups=1,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        # determine dimensions\n",
    "        self.channels = channels\n",
    "        self.self_condition = self_condition\n",
    "        input_channels = channels * (2 if self_condition else 1)\n",
    "\n",
    "        init_dim = default(init_dim, dim)\n",
    "        self.init_conv = nn.Conv2d(input_channels, init_dim, 3, padding=1)  # 修改卷积核的大小和填充\n",
    "\n",
    "        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]\n",
    "        in_out = list(zip(dims[:-1], dims[1:]))\n",
    "\n",
    "        block_klass = partial(ResnetBlock, groups=resnet_block_groups)\n",
    "\n",
    "        # time embeddings\n",
    "        time_dim = dim * 4\n",
    "\n",
    "        self.time_mlp = nn.Sequential(\n",
    "            SinusoidalPositionEmbeddings(dim),\n",
    "            nn.Linear(dim, time_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(time_dim, time_dim),\n",
    "        )\n",
    "\n",
    "        # layers\n",
    "        self.downs = nn.ModuleList([])\n",
    "        self.ups = nn.ModuleList([])\n",
    "        num_resolutions = len(in_out)\n",
    "\n",
    "        for ind, (dim_in, dim_out) in enumerate(in_out):\n",
    "            is_last = ind >= (num_resolutions - 1)\n",
    "\n",
    "            self.downs.append(\n",
    "                nn.ModuleList(\n",
    "                    [\n",
    "                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),\n",
    "                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),\n",
    "                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),\n",
    "                        Downsample(dim_in, dim_out)\n",
    "                        if not is_last\n",
    "                        else nn.Conv2d(dim_in, dim_out, 3, padding=1),\n",
    "                    ]\n",
    "                )\n",
    "            )\n",
    "\n",
    "        mid_dim = dims[-1]\n",
    "        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)\n",
    "        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))\n",
    "        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)\n",
    "\n",
    "        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):\n",
    "            is_last = ind == (len(in_out) - 1)\n",
    "\n",
    "            self.ups.append(\n",
    "                nn.ModuleList(\n",
    "                    [\n",
    "                        block_klass(dim_out + dim_in, dim_out, time_emb_dim=time_dim),\n",
    "                        block_klass(dim_out + dim_in, dim_out, time_emb_dim=time_dim),\n",
    "                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),\n",
    "                        Upsample(dim_out, dim_in)\n",
    "                        if not is_last\n",
    "                        else nn.Conv2d(dim_out, dim_in, 3, padding=1),\n",
    "                    ]\n",
    "                )\n",
    "            )\n",
    "\n",
    "        self.out_dim = default(out_dim, channels)\n",
    "\n",
    "        self.final_res_block = block_klass(dim * 2, dim, time_emb_dim=time_dim)\n",
    "        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)\n",
    "\n",
    "    def forward(self, x, time, x_self_cond=None):\n",
    "        if self.self_condition:\n",
    "            x_self_cond = default(x_self_cond, lambda: torch.zeros_like(x))\n",
    "            x = torch.cat((x_self_cond, x), dim=1)\n",
    "\n",
    "        x = self.init_conv(x)\n",
    "        r = x.clone()\n",
    "\n",
    "        t = self.time_mlp(time)\n",
    "\n",
    "        h = []\n",
    "\n",
    "        for block1, block2, attn, downsample in self.downs:\n",
    "            x = block1(x, t)\n",
    "            h.append(x)\n",
    "\n",
    "            x = block2(x, t)\n",
    "            x = attn(x)\n",
    "            h.append(x)\n",
    "\n",
    "            x = downsample(x)\n",
    "\n",
    "        x = self.mid_block1(x, t)\n",
    "        x = self.mid_attn(x)\n",
    "        x = self.mid_block2(x, t)\n",
    "\n",
    "        for block1, block2, attn, upsample in self.ups:\n",
    "            x = torch.cat((x, h.pop()), dim=1)\n",
    "            x = block1(x, t)\n",
    "\n",
    "            x = torch.cat((x, h.pop()), dim=1)\n",
    "            x = block2(x, t)\n",
    "            x = attn(x)\n",
    "\n",
    "            x = upsample(x)\n",
    "\n",
    "        x = torch.cat((x, r), dim=1)\n",
    "\n",
    "        x = self.final_res_block(x, t)\n",
    "        return self.final_conv(x)\n",
    "\n",
    "    \n",
    "def cosine_beta_schedule(timesteps, s=0.008):\n",
    "    \"\"\"\n",
    "    cosine schedule as proposed in https://arxiv.org/abs/2102.09672\n",
    "    \"\"\"\n",
    "    steps = timesteps + 1\n",
    "    x = torch.linspace(0, timesteps, steps)\n",
    "    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2\n",
    "    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]\n",
    "    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])\n",
    "    return torch.clip(betas, 0.0001, 0.9999)\n",
    "\n",
    "def linear_beta_schedule(timesteps):\n",
    "    beta_start = 0.0001\n",
    "    beta_end = 0.02\n",
    "    return torch.linspace(beta_start, beta_end, timesteps)\n",
    "\n",
    "def quadratic_beta_schedule(timesteps):\n",
    "    beta_start = 0.0001\n",
    "    beta_end = 0.02\n",
    "    return torch.linspace(beta_start**0.5, beta_end**0.5, timesteps) ** 2\n",
    "\n",
    "def sigmoid_beta_schedule(timesteps):\n",
    "    beta_start = 0.0001\n",
    "    beta_end = 0.02\n",
    "    betas = torch.linspace(-6, 6, timesteps)\n",
    "    return torch.sigmoid(betas) * (beta_end - beta_start) + beta_start\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timesteps = 300\n",
    "\n",
    "# define beta schedule\n",
    "betas = linear_beta_schedule(timesteps=timesteps)\n",
    "\n",
    "# define alphas \n",
    "alphas = 1. - betas\n",
    "alphas_cumprod = torch.cumprod(alphas, axis=0)\n",
    "alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)\n",
    "sqrt_recip_alphas = torch.sqrt(1.0 / alphas)\n",
    "\n",
    "# calculations for diffusion q(x_t | x_{t-1}) and others\n",
    "sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)\n",
    "sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)\n",
    "\n",
    "# calculations for posterior q(x_{t-1} | x_t, x_0)\n",
    "posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)\n",
    "\n",
    "def extract(a, t, x_shape):\n",
    "    batch_size = t.shape[0]\n",
    "    out = a.gather(-1, t.cpu())\n",
    "    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward diffusion (using the nice property)\n",
    "def q_sample(x_start, t, noise=None):\n",
    "    if noise is None:\n",
    "        noise = torch.randn_like(x_start)\n",
    "\n",
    "    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)\n",
    "    sqrt_one_minus_alphas_cumprod_t = extract(\n",
    "        sqrt_one_minus_alphas_cumprod, t, x_start.shape\n",
    "    )\n",
    "\n",
    "    return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_noisy(x_start, t):\n",
    "  # add noise\n",
    "  x_noisy = q_sample(x_start, t=t)\n",
    "\n",
    "  return x_noisy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#calculate the loss of noise\n",
    "def p_losses(denoise_model, x_start, t, noise=None, loss_type=\"l1\"):\n",
    "    if noise is None:\n",
    "        noise = torch.randn_like(x_start)\n",
    "\n",
    "    x_noisy = q_sample(x_start=x_start, t=t, noise=noise)\n",
    "    predicted_noise = denoise_model(x_noisy, t)\n",
    "\n",
    "    if loss_type == 'l1':\n",
    "        loss = F.l1_loss(noise, predicted_noise)\n",
    "    elif loss_type == 'l2':\n",
    "        loss = F.mse_loss(noise, predicted_noise)\n",
    "    elif loss_type == \"huber\":\n",
    "        loss = F.smooth_l1_loss(noise, predicted_noise)\n",
    "    else:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    return loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df1 = pd.read_csv('data.csv')\n",
    "\n",
    "df1.info()\n",
    "use_df = df1[:]\n",
    "use_df.head(10)\n",
    "feature = df1['feature'].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def protein_to_onehot(protein_sequence):\n",
    "    # Define 21 amino acid types\n",
    "    amino_acids = 'ACDEFGHIKLMNPQRSTVWYX'\n",
    "    num_amino_acids = len(amino_acids)\n",
    "    \n",
    "    # Extended the protein sequence length to 60 and made up the deficiency with the 21st amino acid\n",
    "    padded_sequence = protein_sequence.ljust(60, amino_acids[-1])\n",
    "    \n",
    "    # Initialize a onehot encoding matrix with all zeros\n",
    "    onehot = np.zeros((len(padded_sequence), num_amino_acids))\n",
    "    \n",
    "    # Mapping amino acids in protein sequences to corresponding onehot codes\n",
    "    for i, amino_acid in enumerate(padded_sequence):\n",
    "        if amino_acid in amino_acids:\n",
    "            onehot[i, amino_acids.index(amino_acid)] = 1\n",
    "        else:\n",
    "            # If an unknown amino acid is encountered, it is represented by a vector of all zeros\n",
    "            pass\n",
    "    \n",
    "    # Change onehot coding to plus or minus one\n",
    "    onehot[onehot == 0] = -1\n",
    "    new_array = -1 * np.ones((60, 60))\n",
    "    new_array[:60, :21] = onehot\n",
    "    reshaped_tensor = np.expand_dims(new_array, axis=0)\n",
    "    return reshaped_tensor\n",
    "\n",
    "# example sequence\n",
    "protein_sequence = \"AC\"\n",
    "\n",
    "# Converted to a plus or minus 1 code of length 60\n",
    "onehot_encoding = protein_to_onehot(protein_sequence)\n",
    "\n",
    "print(onehot_encoding.shape)\n",
    "onehot_encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom Dataset Classes\n",
    "class ProteinDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        # Get feature sequence and convert to onehot encoding\n",
    "        feature_sequence = self.data[idx]\n",
    "        onehot = protein_to_onehot(feature_sequence)\n",
    "        # Convert to PyTorch's Tensor and return\n",
    "        return torch.tensor(onehot, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a dataset instance\n",
    "protein_dataset = ProteinDataset(df2['feature'].tolist())\n",
    "\n",
    "# Creating a Data Loader\n",
    "batch_size = 32\n",
    "data_loader = DataLoader(protein_dataset, batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size2 = 1\n",
    "data_loader2 = DataLoader(protein_dataset, batch_size=batch_size2, shuffle=True)\n",
    "data = next(iter(data_loader2))\n",
    "print(data)\n",
    "print(data.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def p_sample(model, x, t, t_index):\n",
    "    betas_t = extract(betas, t, x.shape)\n",
    "    sqrt_one_minus_alphas_cumprod_t = extract(\n",
    "        sqrt_one_minus_alphas_cumprod, t, x.shape\n",
    "    )\n",
    "    sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)\n",
    "    \n",
    "    # Equation 11 in the paper\n",
    "    # Use our model (noise predictor) to predict the mean\n",
    "    model_mean = sqrt_recip_alphas_t * (\n",
    "        x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t\n",
    "    )\n",
    "\n",
    "    if t_index == 0:\n",
    "        return model_mean\n",
    "    else:\n",
    "        posterior_variance_t = extract(posterior_variance, t, x.shape)\n",
    "        noise = torch.randn_like(x)\n",
    "        # Algorithm 2 line 4:\n",
    "        return model_mean + torch.sqrt(posterior_variance_t) * noise \n",
    "\n",
    "# Algorithm 2 (including returning all images)\n",
    "@torch.no_grad()\n",
    "def p_sample_loop(model, shape):\n",
    "    device = next(model.parameters()).device\n",
    "\n",
    "    b = shape[0]\n",
    "    # start from pure noise (for each example in the batch)\n",
    "    img = torch.randn(shape, device=device)\n",
    "    imgs = []\n",
    "\n",
    "    for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):\n",
    "        img = p_sample(model, img, torch.full((b,), i, device=device, dtype=torch.long), i)\n",
    "        imgs.append(img.cpu().numpy())\n",
    "    return imgs\n",
    "\n",
    "@torch.no_grad()\n",
    "def sample(model, image_size, batch_size=16, channels=3):\n",
    "    return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_size = 60\n",
    "image_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import Adam\n",
    "\n",
    "\n",
    "model = Unet(\n",
    "    dim=image_size,\n",
    "    channels=1,\n",
    "    dim_mults=(1, 2, 4,)\n",
    ")\n",
    "model.to(device)\n",
    "\n",
    "optimizer = Adam(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.utils import save_image\n",
    "\n",
    "epochs = 5\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    for step, batch in enumerate(data_loader):\n",
    "      optimizer.zero_grad()\n",
    "\n",
    "      batch_size = batch.shape[0]\n",
    "      batch = batch.to(device)\n",
    "\n",
    "      # Algorithm 1 line 3: sample t uniformally for every example in the batch\n",
    "      t = torch.randint(0, timesteps, (batch_size,), device=device).long()\n",
    "\n",
    "      loss = p_losses(model, batch, t, loss_type=\"huber\")\n",
    "\n",
    "      if step % 100 == 0:\n",
    "        print(\"epoch:\", epoch, \"Loss:\", loss.item())\n",
    "\n",
    "      loss.backward()\n",
    "      optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model, 'DDPM.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.load('DDPM.pth').to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Algorithm 2 (including returning all images)\n",
    "@torch.no_grad()\n",
    "def single_p_sample_loop(model, x_sample, shape, t = 300):\n",
    "    timesteps = t\n",
    "    device = next(model.parameters()).device\n",
    "\n",
    "    b = shape[0]\n",
    "    # start from pure noise (for each example in the batch)\n",
    "    #img = torch.randn(shape, device=device)\n",
    "    img = x_sample\n",
    "    #print(x_sample.shape)\n",
    "    #print(img.shape)\n",
    "    imgs = []\n",
    "    \n",
    "    for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):\n",
    "        img = p_sample(model, img, torch.full((b,), i, device=device, dtype=torch.long), i)\n",
    "        imgs.append(img.cpu().numpy())\n",
    "    \n",
    "    return imgs\n",
    "\n",
    "@torch.no_grad()\n",
    "def single_sample(model, x_sample, image_size, batch_size=16, channels=3, t = 300):\n",
    "    return single_p_sample_loop(model, x_sample, shape=(batch_size, channels, image_size, image_size), t = t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample 64 images\n",
    "samples = sample(model, image_size=image_size, batch_size=1, channels=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# demo protein sequence\n",
    "protein_sequence = \"VDNKFNKEQQNAFYEILHLPNLTEEQRNAFIQDLKDDPSQSANLLAEAKKLNDAQAPK\"\n",
    "\n",
    "# Convert protein sequence to positive negative 1 encoding with a length of 60\n",
    "proteinData = protein_to_onehot(protein_sequence)\n",
    "\n",
    "proteinData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def showData(data):\n",
    "    rotated_array = np.rot90(data)\n",
    "    result_array = rotated_array[-21:, :]\n",
    "    return result_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display grayscale image\n",
    "plt.imshow(showData(proteinData[0]), cmap='gray', vmin=-1, vmax=1)  # Set the range of the grayscale image to [-1, 1]\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(imgs, with_orig=False, row_title=None, **imshow_kwargs):\n",
    "    if not isinstance(imgs[0], list):\n",
    "        # Make a 2d grid even if there's just 1 row\n",
    "        imgs = [[img] for img in imgs]\n",
    "        \n",
    "    num_rows = len(imgs)\n",
    "    num_cols = 1 + with_orig\n",
    "    fig, axs = plt.subplots(figsize=(200,200), nrows=num_rows, ncols=num_cols, squeeze=False)\n",
    "    for row_idx, row in enumerate(imgs):\n",
    "        if with_orig:\n",
    "            ax = axs[row_idx, 0]\n",
    "            ax.imshow(showData(np.asarray(row[0]).squeeze()), **imshow_kwargs, cmap='gray', vmin=-1, vmax=1)\n",
    "            ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])\n",
    "            ax.set(title='Original image')\n",
    "            ax.title.set_size(8)\n",
    "            row = row[1:]\n",
    "        for col_idx, img in enumerate(row):\n",
    "            ax = axs[row_idx, col_idx + with_orig]\n",
    "            ax.imshow(showData(np.asarray(img).squeeze()), **imshow_kwargs, cmap='gray', vmin=-1, vmax=1)\n",
    "            ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])\n",
    "\n",
    "    if row_title is not None:\n",
    "        for row_idx in range(num_rows):\n",
    "            axs[row_idx, 0].set(ylabel=row_title[row_idx])\n",
    "\n",
    "    plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proteinSample = torch.tensor(proteinData[0], device='cpu')\n",
    "proteinNoise = get_noisy(proteinSample, torch.tensor([50]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot([get_noisy(proteinSample, torch.tensor([t])) for t in [0, 10, 20, 30, 40, 50]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# Display grayscale image\n",
    "plt.imshow(showData(proteinNoise.numpy()), cmap='gray', vmin=-1, vmax=1)  # Set the range of the grayscale image to [-1, 1]\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proteinNoise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_sample = proteinNoise.to(device).reshape(1, 1, 60, 60).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proteinDenoise = single_sample(model, x_sample, image_size=image_size, batch_size=1, channels=1, t = 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(showData(proteinDenoise[-1].squeeze()), cmap='gray', vmin=-1, vmax=1)  # Set the range of the grayscale image to [-1, 1]\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.set_printoptions(threshold=np.inf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def normalize_onehot(matrix):\n",
    "    # Find the index with the maximum value for each row\n",
    "    max_values = np.max(matrix[:,:,:,:21], axis=-1, keepdims=True)\n",
    "    \n",
    "    # Initialize an array of the same size as the input matrix to store the converted encoding\n",
    "    onehot_matrix = np.where(matrix == max_values, 1, 0)\n",
    "    \n",
    "    return onehot_matrix\n",
    "\n",
    "# Convert to positive negative 1 encoding\n",
    "onehot_encoding = normalize_onehot(proteinDenoise[-1])\n",
    "print(onehot_encoding.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "denoiseList = []\n",
    "for i in range(50):\n",
    "    denoiseList.append(proteinDenoise[i])\n",
    "denoiseList[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(showData(onehot_encoding[-1].squeeze()), cmap='gray', vmin=0, vmax=1)  # Set the range of the grayscale image to [-1, 1]\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot([proteinDenoise[t].squeeze() for t in [0, 10, 20, 30, 40, 49]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "onehot_encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def decode_onehot(matrix):\n",
    "    amino_acids = 'ACDEFGHIKLMNPQRSTVWY'\n",
    "    decoded_sequences = []\n",
    "    for sample in matrix:\n",
    "        sequence = ''\n",
    "        for column in sample[0]:\n",
    "            index = np.argmax(column)\n",
    "            if index < len(amino_acids):\n",
    "                sequence += amino_acids[index]\n",
    "            else:\n",
    "                sequence += ''\n",
    "        decoded_sequences.append(sequence)\n",
    "    return decoded_sequences\n",
    "\n",
    "decoded_sequences = decode_onehot(onehot_encoding)\n",
    "print(decoded_sequences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "e42634819b8c191a5d07eaf23810ff32516dd8d3875f28ec3e488928fbd3c187"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}