glm_pca.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "9832442e-1dec-4eed-bbc8-94d75f628c1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import copy\n",
    "import numpy as np\n",
    "import scipy\n",
    "import pandas as pd\n",
    "import anndata\n",
    "import scanpy as sc\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import sklearn.decomposition\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "317075b2-65f9-405a-8c9d-dc60237edc92",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/anders/pyenv_torch/lib/python3.11/site-packages/anndata/_core/anndata.py:121: ImplicitModificationWarning: Transforming to str index.\n",
      "  warnings.warn(\"Transforming to str index.\", ImplicitModificationWarning)\n",
      "/home/anders/pyenv_torch/lib/python3.11/site-packages/anndata/_core/anndata.py:121: ImplicitModificationWarning: Transforming to str index.\n",
      "  warnings.warn(\"Transforming to str index.\", ImplicitModificationWarning)\n"
     ]
    }
   ],
   "source": [
    "ad = anndata.AnnData(\n",
    "    X = scipy.io.mmread('/home/anders/sds/sd17l002/u/anders/ifnagrko/ifnagrko_raw_counts.mtx.gz').tocsr(),\n",
    "    obs = pd.read_csv('/home/anders/sds/sd17l002/u/anders/ifnagrko/ifnagrko_obs.csv.gz'),\n",
    "    var = pd.read_csv('/home/anders/sds/sd17l002/u/anders/ifnagrko/ifnagrko_var.csv') )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7f4087d4-e0d1-41be-8e2e-753ae961f7ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 18302 × 2000\n",
       "    obs: 'barcode', 'celltype', 'celltype1', 'celltype2', 'dpt_pseudotime', 'svz_frac', 'hto_sum', 'leiden', 'dbscan', 'frac_mito', 'counts', 'genes', 'scrub_score', 'age', 'genotype', 'UMAP1', 'UMAP2'\n",
       "    var: 'Unnamed: 0', 'gene_name', 'gene_id', 'highly_variable', 'highly_variable_rank', 'means', 'variances', 'variances_norm'\n",
       "    uns: 'hvg'\n",
       "    layers: 'counts'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ad.layers[\"counts\"] = ad.X.copy()\n",
    "sc.pp.normalize_total( ad )\n",
    "sc.pp.highly_variable_genes(\n",
    "    ad,\n",
    "    flavor=\"seurat_v3\",\n",
    "    n_top_genes=2000,\n",
    "    layer=\"counts\",\n",
    "    subset=True,\n",
    ")\n",
    "ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "702d0196-c461-4d81-ad79-83bf7f19110c",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = torch.tensor( ad.layers[\"counts\"].todense(), dtype=torch.float32 )\n",
    "totals = counts.sum(1)\n",
    "lncounts = torch.log( counts/totals[:,None] + 1e-4 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6fb4361e-a00c-44f5-a35c-63c608e13419",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(20, 2000)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pca = sklearn.decomposition.PCA(20).fit( lncounts.numpy() )\n",
    "pca.components_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "a0a909a0-cc8f-4781-8eae-1b94355845a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_centers = lncounts.mean(0).clone().requires_grad_(True)\n",
    "encoder_matrix = torch.tensor( pca.components_.T, requires_grad=True, dtype=torch.float32 )\n",
    "decoder_matrix = torch.tensor( pca.components_, requires_grad=True, dtype=torch.float32 )\n",
    "decoder_centers = lncounts.mean(0).clone().requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0fea10b8-2d3a-4cf0-96e9-ad6cbba95370",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.1630e+01,  2.7517e+01,  2.3102e+00,  ...,  9.2211e-01,\n",
       "         -1.2445e+00, -5.8396e-01],\n",
       "        [-2.1500e+01, -4.7937e+00, -1.7415e+00,  ...,  1.4030e+00,\n",
       "         -1.4130e+00, -2.9863e+00],\n",
       "        [ 2.3597e+01, -8.0399e+00, -1.2087e-01,  ...,  2.4568e+00,\n",
       "          8.8939e-01,  4.5007e+00],\n",
       "        ...,\n",
       "        [-1.3697e+01, -1.7856e+00,  1.7581e+01,  ...,  2.4613e+00,\n",
       "          3.8042e-01, -9.2547e-01],\n",
       "        [-1.5142e+00,  9.7162e+00, -5.6973e+00,  ..., -2.0212e+00,\n",
       "          2.5672e-02, -6.3906e-01],\n",
       "        [-1.4272e+01, -3.9927e+00,  2.3455e+01,  ...,  1.0099e+00,\n",
       "          2.9713e-01, -2.1273e+00]], grad_fn=<MmBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "latent = ( lncounts - encoder_centers ) @ encoder_matrix\n",
    "latent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f792c0ec-b912-4000-b1d6-d86b2711cc8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_hat = torch.exp( latent @ decoder_matrix + decoder_centers ) * totals[:,None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "74d8e9d2-bcf3-4dfe-b8ad-8d220cc6d0cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2.5049e-01, 2.0214e-01, 4.7346e-01,  ..., 5.2245e+01, 2.4171e-01,\n",
       "         2.3430e-01],\n",
       "        [9.0665e-02, 2.3785e-01, 1.3015e-01,  ..., 5.7705e+01, 9.3463e-02,\n",
       "         1.4912e+00],\n",
       "        [1.5239e-01, 1.5394e-01, 2.1735e+00,  ..., 8.6692e-01, 1.5844e-01,\n",
       "         3.3152e-01],\n",
       "        ...,\n",
       "        [1.5659e-01, 2.0947e-01, 1.1862e-01,  ..., 1.6465e+02, 1.5789e-01,\n",
       "         7.7177e+00],\n",
       "        [3.8133e-02, 3.7210e-02, 5.0348e-02,  ..., 4.0214e+00, 9.5131e-02,\n",
       "         5.5177e-02],\n",
       "        [1.5728e-01, 2.2616e-01, 1.6303e-01,  ..., 9.0928e+01, 1.7404e-01,\n",
       "         6.9949e+00]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lambda_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "82059156-b90c-4ac3-8a39-77694a359c45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  0.,   0.,   0.,  ...,  46.,   0.,   0.],\n",
       "        [  0.,   0.,   0.,  ...,  90.,   0.,   0.],\n",
       "        [  0.,   0.,   0.,  ...,   0.,   0.,   0.],\n",
       "        ...,\n",
       "        [  0.,   0.,   0.,  ..., 109.,   0.,  12.],\n",
       "        [  0.,   0.,   0.,  ...,  21.,   0.,   0.],\n",
       "        [  0.,   0.,   0.,  ..., 151.,   0.,   9.]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "00f27d23-6d91-4e25-a621-e14b4c101689",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = .2**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "4cc08ff0-b6ab-4a09-90d7-d1d2d13b4522",
   "metadata": {},
   "outputs": [],
   "source": [
    "nbprob = 1. / ( lambda_hat * alpha + 1. )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cffcb8eb-8300-4c02-9b24-56db7f9b66f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-16732323., grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_likelihood = torch.lgamma( counts + 1./alpha ) - math.lgamma(1./alpha) - torch.lgamma( counts + 1 ) + counts * torch.log( 1 - nbprob ) + (1./alpha) * torch.log( nbprob )\n",
    "log_likelihood.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "de25cdcf-df05-4efd-aa2e-58ad3cc1414d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-16732323., grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_likelihood = torch.lgamma( counts + 1./alpha ) - math.lgamma(1./alpha) - torch.lgamma( counts + 1 ) + \\\n",
    "    torch.xlogy( counts, 1 - nbprob ) + torch.xlogy( 1./alpha, nbprob )\n",
    "log_likelihood.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "656cf36d-49e9-441c-ace3-f2757918d348",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5815576314926147"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nbprob0 = 1. / ( counts * alpha + 1 )\n",
    "log_likelihood_saturated = torch.lgamma( counts + 1./alpha ) - math.lgamma(1./alpha) - torch.lgamma( counts + 1 ) + \\\n",
    "    torch.xlogy( counts, 1 - nbprob0 ) + torch.xlogy( 1./alpha, nbprob0 )\n",
    "( -2 * (log_likelihood - log_likelihood_saturated ) ).mean().item()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3b6fbd31-f42f-4eb6-8cd0-3cf585d5769d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(54075896., grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_likelihood_part = counts * torch.log( 1 - nbprob ) + (1./alpha) * torch.log( nbprob )\n",
    "loss = -log_likelihood_part.sum()\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "711b9453-dbe6-480a-a402-d93061273809",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a2588954-ef05-49ae-a810-867c18a5c8f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-939.5082, 1316.0903, 3787.9824,  ..., 7095.9502, -346.8978,\n",
       "        6605.7959])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder_centers.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "82b9637b-8dbb-4e14-be16-bbd0ad8cb507",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-4.3974e-06, -1.1479e-02,  5.7341e-02,  ..., -8.9799e-02,\n",
       "          9.9221e-04, -6.2656e-02],\n",
       "        [ 1.4553e-03, -4.5472e-03, -4.3717e-02,  ...,  1.2532e-02,\n",
       "          1.1846e-03, -3.7184e-02],\n",
       "        [-3.7185e-04,  1.1363e-02, -3.1827e-03,  ...,  3.8812e-02,\n",
       "         -6.9147e-04,  6.2706e-02],\n",
       "        ...,\n",
       "        [-4.8306e-03,  1.8585e-02,  1.8708e-02,  ..., -5.7901e-02,\n",
       "         -7.0025e-04,  1.8238e-02],\n",
       "        [ 4.0338e-04,  3.7371e-02,  1.3643e-02,  ...,  3.1685e-03,\n",
       "         -5.1020e-04,  1.1436e-02],\n",
       "        [-4.9728e-04,  1.7212e-03, -3.5354e-02,  ...,  2.1039e-02,\n",
       "         -7.3830e-04,  1.2716e-01]], requires_grad=True)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoder_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "65b09be6-2706-4df1-965f-f984f18f10ca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0705e+04,  4.9684e+04, -1.3327e+05,  ...,  1.1444e+04,\n",
       "          8.1080e+03,  2.1411e+05],\n",
       "        [ 3.1664e+03,  9.8356e+03,  3.3381e+04,  ..., -5.4091e+03,\n",
       "          3.1639e+03,  4.1589e+04],\n",
       "        [ 6.3209e+03, -4.2251e+03,  6.2486e+03,  ..., -4.0955e+04,\n",
       "          6.6417e+03, -6.2264e+04],\n",
       "        ...,\n",
       "        [ 1.7448e+02, -1.0185e+03, -5.9606e+02,  ..., -5.5355e+03,\n",
       "          3.9751e+01, -2.3006e+03],\n",
       "        [-9.8762e+01, -1.7503e+03, -1.7647e+03,  ..., -5.6964e+03,\n",
       "         -6.9122e+01, -2.0445e+02],\n",
       "        [ 1.9211e+01,  2.7887e+02,  1.7133e+03,  ...,  7.1400e+03,\n",
       "          6.0949e+01, -4.7210e+03]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoder_matrix.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d2464a0f-0434-414a-8959-116c71705cec",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    encoder_centers -= 1e-8 * encoder_centers.grad\n",
    "    encoder_matrix  -= 1e-8 * encoder_matrix.grad\n",
    "    decoder_matrix  -= 1e-8 * decoder_matrix.grad    \n",
    "    decoder_centers -= 1e-8 * decoder_centers.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "9495bbc0-f41e-48c3-bbfa-b203cb51266c",
   "metadata": {},
   "outputs": [],
   "source": [
    "latent = ( lncounts - encoder_centers ) @ encoder_matrix\n",
    "lambda_hat = torch.exp( latent @ decoder_matrix + decoder_centers ) * totals[:,None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "04dfa533-0cfb-45ff-a255-d4d77cf6cf02",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(53987828., grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nbprob = 1. / ( lambda_hat * alpha + 1 )\n",
    "log_likelihood_part = counts * torch.log( 1 - nbprob ) + (1./alpha) * torch.log( nbprob )\n",
    "loss = -log_likelihood_part.sum()\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "017e407f-9d20-44c2-90fd-21d127aecc1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam( ( encoder_centers, encoder_matrix, decoder_matrix, decoder_centers ), lr=1e-8 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "3e35e33a-0a0f-4481-8e17-7b8e6461ef57",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([-9.1620, -9.0162, -8.1960,  ..., -4.8811, -9.1269, -7.6952],\n",
       "        device='cuda:0'),\n",
       " tensor([[-0.0229, -0.0136, -0.0276,  ...,  0.0174,  0.0304, -0.0231],\n",
       "         [-0.0061,  0.0132,  0.0337,  ...,  0.0030,  0.0129,  0.0197],\n",
       "         [ 0.0642, -0.0506, -0.0197,  ...,  0.0146,  0.0099, -0.0345],\n",
       "         ...,\n",
       "         [-0.1027,  0.0014,  0.0586,  ..., -0.0757,  0.0235,  0.0279],\n",
       "         [-0.0119, -0.0004, -0.0293,  ...,  0.0207,  0.0335, -0.0244],\n",
       "         [-0.0655, -0.0440,  0.0801,  ...,  0.0246,  0.0319,  0.1345]],\n",
       "        device='cuda:0'),\n",
       " tensor([[-0.0149, -0.0324,  0.0714,  ..., -0.0739, -0.0046, -0.0775],\n",
       "         [-0.0159, -0.0234, -0.0434,  ...,  0.0194, -0.0099, -0.0539],\n",
       "         [-0.0206,  0.0141, -0.0156,  ...,  0.0416, -0.0178,  0.0615],\n",
       "         ...,\n",
       "         [ 0.0117,  0.0413,  0.0080,  ..., -0.0390,  0.0284,  0.0251],\n",
       "         [ 0.0287,  0.0625,  0.0155,  ...,  0.0161,  0.0276,  0.0330],\n",
       "         [-0.0288, -0.0220, -0.0378,  ...,  0.0116, -0.0286,  0.1216]],\n",
       "        device='cuda:0'),\n",
       " tensor([-9.2202, -9.0133, -8.1998,  ..., -4.8273, -9.1867, -7.6519],\n",
       "        device='cuda:0')]"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "73dce710-f801-4a3c-bff2-4d2dfcb17f76",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_centers = lncounts.mean(0).clone().requires_grad_(True)\n",
    "encoder_matrix = torch.tensor( pca.components_.T, requires_grad=True, dtype=torch.float32 )\n",
    "decoder_matrix = torch.tensor( pca.components_, requires_grad=True, dtype=torch.float32 )\n",
    "decoder_centers = lncounts.mean(0).clone().requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "d3839136-4824-4703-ba88-4d0734d8b9f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0.5815567374229431, 0.2480899691581726, -1.52587890625e-05, 55982.703125)\n",
      "(0.3472716510295868, 0.01691269688308239, -3.0517578125e-05, 967.3753662109375)\n",
      "(0.3406349718570709, 0.013564876280725002, -4.57763671875e-05, 495.1505126953125)\n",
      "(0.33918288350105286, 0.012897128239274025, -1.52587890625e-05, 420.5322570800781)\n",
      "(0.3385719060897827, 0.012658648192882538, -1.52587890625e-05, 354.87255859375)\n",
      "(0.33810749650001526, 0.012473827227950096, -4.57763671875e-05, 337.72418212890625)\n",
      "(0.3377564549446106, 0.012348627671599388, -4.57763671875e-05, 334.24237060546875)\n",
      "(0.33765503764152527, 0.012171745300292969, -3.0517578125e-05, 333.337158203125)\n",
      "(0.3375290334224701, 0.012098225764930248, -1.52587890625e-05, 327.39349365234375)\n",
      "(0.33724069595336914, 0.01214592158794403, -2.288818359375e-05, 325.6202392578125)\n",
      "(0.3371332287788391, 0.012116111814975739, -3.0517578125e-05, 323.88092041015625)\n",
      "(0.3370274007320404, 0.012092264369130135, -1.52587890625e-05, 322.64410400390625)\n",
      "(0.3369239568710327, 0.012062454596161842, -3.0517578125e-05, 321.851318359375)\n",
      "(0.3368293046951294, 0.012008797377347946, -1.52587890625e-05, 322.12860107421875)\n",
      "(0.3367539644241333, 0.011996873654425144, -1.52587890625e-05, 322.1785888671875)\n",
      "(0.33663567900657654, 0.01199091225862503, -1.52587890625e-05, 321.08489990234375)\n",
      "(0.3365321457386017, 0.01199091225862503, -3.0517578125e-05, 319.23504638671875)\n",
      "(0.33644694089889526, 0.011961102485656738, -3.0517578125e-05, 316.89520263671875)\n",
      "(0.33644938468933105, 0.012104188092052937, -3.0517578125e-05, 316.4998779296875)\n",
      "(0.3362773358821869, 0.012032645754516125, -3.0517578125e-05, 314.3897705078125)\n",
      "(0.3361998498439789, 0.011967064812779427, -1.52587890625e-05, 315.83551025390625)\n",
      "(0.33612731099128723, 0.01195514015853405, -1.52587890625e-05, 315.2091064453125)\n",
      "(0.3360527753829956, 0.011937255039811134, -6.103515625e-05, 314.0787353515625)\n",
      "(0.3361693024635315, 0.011931292712688446, -1.52587890625e-05, 311.969970703125)\n",
      "(0.33593621850013733, 0.011943216435611248, -4.57763671875e-05, 312.09722900390625)\n",
      "(0.3359326720237732, 0.011925331316888332, -3.0517578125e-05, 315.652099609375)\n",
      "(0.33590105175971985, 0.011899948120117188, -1.52587890625e-05, 313.19268798828125)\n",
      "(0.3358522355556488, 0.011895522475242615, -3.0517578125e-05, 312.794677734375)\n",
      "(0.3358103334903717, 0.011877636425197124, -3.0517578125e-05, 312.1876220703125)\n",
      "(0.33580291271209717, 0.011823979206383228, -3.0517578125e-05, 312.38104248046875)\n",
      "(0.33572331070899963, 0.011806094087660313, -1.52587890625e-05, 310.75384521484375)\n",
      "(0.33587875962257385, 0.011740513145923615, -1.52587890625e-05, 313.41900634765625)\n",
      "(0.3357557952404022, 0.01192474365234375, -1.52587890625e-05, 310.56475830078125)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[155], line 27\u001b[0m\n\u001b[1;32m     25\u001b[0m loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39mlog_likelihood_part\u001b[38;5;241m.\u001b[39msum()\n\u001b[1;32m     26\u001b[0m \u001b[38;5;66;03m#loss = deviance.sum()\u001b[39;00m\n\u001b[0;32m---> 27\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43misnan\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitem\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m:\n\u001b[1;32m     28\u001b[0m     \u001b[38;5;28mprint\u001b[39m( \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnan loss\u001b[39m\u001b[38;5;124m\"\u001b[39m )\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;28;01mbreak\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "counts = counts.detach().cuda()\n",
    "lncounts = lncounts.detach().cuda()\n",
    "totals = totals.detach().cuda()\n",
    "encoder_centers = encoder_centers.detach().cuda().requires_grad_(True)\n",
    "encoder_matrix = encoder_matrix.detach().cuda().requires_grad_(True)\n",
    "decoder_matrix = ( decoder_matrix.detach() + torch.randn_like(decoder_matrix)/1e5 ).cuda().requires_grad_(True) \n",
    "decoder_centers = decoder_centers.detach().cuda().requires_grad_(True)\n",
    "\n",
    "optimizer = torch.optim.Adam( \n",
    "    ( encoder_centers, encoder_matrix, decoder_matrix, decoder_centers ), \n",
    "    #( encoder_matrix, decoder_matrix ), \n",
    "    lr=1e-3 )\n",
    "\n",
    "nbprob0 = 1. / ( counts * alpha + 1 )\n",
    "log_likelihood_saturated_part = torch.xlogy( counts, 1 - nbprob0 ) + torch.xlogy( 1./alpha, nbprob0 )\n",
    "\n",
    "for step in range(100000):\n",
    "    latent = ( lncounts - encoder_centers ) @ encoder_matrix\n",
    "    lambda_hat = torch.exp( latent @ decoder_matrix + decoder_centers ) * totals[:,None]\n",
    "    nbprob = 1. / ( lambda_hat * alpha + 1 )\n",
    "    nbprob = torch.clamp( nbprob, 1e-5, 1-1e-5 )\n",
    "    log_likelihood_part = torch.xlogy( counts, 1 - nbprob ) + torch.xlogy( 1./alpha, nbprob )\n",
    "    #deviance = -2* ( log_likelihood_part - log_likelihood_saturated_part )\n",
    "    #deviance = torch.clamp( deviance, -1e-5, 100 )\n",
    "    loss = -log_likelihood_part.sum()\n",
    "    #loss = deviance.sum()\n",
    "    if torch.isnan(loss).item():\n",
    "        print( \"nan loss\" )\n",
    "        break\n",
    "    if step % 100 == 0:\n",
    "        with torch.no_grad():\n",
    "            #nbprob0 = 1. / ( counts * alpha + 1 )\n",
    "            #log_likelihood_saturated_part = torch.xlogy( counts, 1 - nbprob0 ) + torch.xlogy( 1./alpha, nbprob0 )\n",
    "            deviance = -2* (log_likelihood_part - log_likelihood_saturated_part )\n",
    "            print( ( deviance.mean().item(), deviance.median().item(), deviance.min().item(), deviance.max().item() ) )\n",
    "        \n",
    "    prev = [ m.clone().detach() for m in (encoder_centers, encoder_matrix, decoder_matrix, decoder_centers) ]\n",
    "    optimizer.zero_grad()    \n",
    "    loss.backward()\n",
    "    nn.utils.clip_grad_norm_((encoder_centers, encoder_matrix, decoder_matrix, decoder_centers), max_norm=1.0)\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "fc5dd2ad-0521-4b83-8d1d-c1fc69078796",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0167, 0.0332, 1.2245,  ..., 0.1054, 0.0273, 0.0978],\n",
       "        [0.0211, 1.3585, 0.1299,  ..., 1.5236, 0.0200, 4.5722],\n",
       "        [0.0169, 0.0289, 4.9223,  ..., 3.2385, 0.0414, 0.7346],\n",
       "        ...,\n",
       "        [0.0203, 0.2231, 0.1014,  ..., 3.9115, 0.0329, 0.0523],\n",
       "        [0.0132, 0.0195, 0.0203,  ..., 3.5928, 0.2352, 0.0882],\n",
       "        [0.0171, 0.3543, 0.1363,  ..., 1.6039, 0.0327, 0.0080]],\n",
       "       device='cuda:0', grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "deviance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "e89e119a-531a-4e9e-90d6-c729c5e52d9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cur = ( encoder_centers, encoder_matrix, decoder_matrix, decoder_centers )\n",
    "encoder_centers, encoder_matrix, decoder_matrix, decoder_centers = prev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "5b8aa86c-aa19-4a10-b30e-6bd88b8f7527",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in prev:\n",
    "    m.requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "6d063f66-f5b0-488f-93f2-7f665e5092f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "latent = ( lncounts - encoder_centers ) @ encoder_matrix\n",
    "lambda_hat = torch.exp( latent @ decoder_matrix + decoder_centers ) * totals[:,None]\n",
    "nbprob = 1. / ( lambda_hat * alpha + 1 )\n",
    "nbprob = n\n",
    "log_likelihood_part = torch.xlogy( counts, 1 - nbprob ) + torch.xlogy( 1./alpha, nbprob )\n",
    "deviance = -2* ( log_likelihood_part - log_likelihood_saturated_part )\n",
    "deviance = torch.clamp( deviance, -50, 0 )\n",
    "loss = -log_likelihood_part.sum()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "61c7a4cd-e757-4661-81f6-6f2438dfbf6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "a30fa2d5-b890-4f4c-934e-4c66ce7a073b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.9622, 0.9812, 0.8026,  ..., 0.9776, 1.0000, 0.9652], device='cuda:0',\n",
       "       grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nbprob[:,630]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "39293c12-94c4-48db-b849-55459fa4aa36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x7f6024a96d90>"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiIAAAGdCAYAAAAvwBgXAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAmMElEQVR4nO3dfXSU5Z3/8c8MSSYBMiGEBIpEwoMCivIQlJLFHw+C/qpIaQu2uwsIRQm7rv5KlRoWK2CL0JOq3W37W1OLbFNsWYlbtehiqYBiedAEAUUQg0ETCCmOkGR4uEnIvX/YRMJDTGDuuWbueb/OmcPJnZnc38zJYT7nur7XdXls27YFAABggNd0AQAAIHYRRAAAgDEEEQAAYAxBBAAAGEMQAQAAxhBEAACAMQQRAABgDEEEAAAYE2e6gJY0NDTo0KFDSk5OlsfjMV0OAABoBdu2VVtbq+7du8vrbXnMI6KDyKFDh5SZmWm6DAAAcAnKy8vVo0ePFp8T0UEkOTlZ0ue/iN/vN1wNAABojZqaGmVmZjZ9jrckooNI43SM3+8niAAAEGVa01ZBsyoAADCGIAIAAIwhiAAAAGMIIgAAwBiCCAAAMIYgAgAAjCGIAAAAYwgiAADAGIIIAAAwhiACAACMIYgAAABjCCIAAMAYgggAADEqELRU8Pp+BYKWsRoIIgAAxKiikgpt3h9QUUmFsRrijN0ZAAAYNTm7R7N/TSCIAAAQo9I6+pQ7qo/RGpiaAQAAxhBEAACAMQQRAABgDEEEAAAYQxABAADGEEQAAIAxBBEAAGAMQQQAABhDEAEAAMYQRAAAgDEEEQAAYAxBBAAAGEMQAQAAxhBEAACAMQQRAABgDEEEAAAYQxABAADGEEQAAIAxBBEAAGAMQQQAABhDEAEAAMYQRAAAgDEEEQAAYAxBBACAMAkELRW8vl+BoGW6lIhBEAEAIEyKSiq0eX9ARSUVpkuJGHGmCwAAIFZMzu7R7F8QRAAACJu0jj7ljupjuoyIwtQMAAAwhiACAACMIYgAAABjCCIAAMAYgggAADCGIAIAAIwhiAAAAGMIIgAAtAHbtIcWQQQAgDZgm/bQYmdVAADagG3aQ4sgAgBAG7BNe2gxNQMAAIwhiAAAAGMIIgAAwBhHg0hWVpY8Hk+zx7Jly5y8JQAAiCKON6s++uijuueee5q+Tk5OdvqWAAAgSjgeRJKTk9WtWzenbwMAAKKQ4z0iy5YtU1pamoYMGaL8/HzV19c7fUsAABAlHB0Ruf/++zV06FB17txZmzdv1vz581VZWaknnnjigs+3LEuW9cWWuTU1NU6WBwAADPPYtm235QV5eXn6yU9+0uJz9uzZo/79+593/ZlnnlFubq6CwaB8Pt9531+0aJEWL1583vXq6mr5/f62lAkAAAypqalRSkpKqz6/2xxEjhw5okAg0OJzevfurYSEhPOu7969WwMHDtTevXvVr1+/875/oRGRzMxMgggAAFGkLUGkzVMz6enpSk9Pv6TCduzYIa/Xq4yMjAt+3+fzXXCkBAAAuJNjPSJbtmzRtm3bNGbMGCUnJ2vLli2aO3eupk6dqtTUVKduCwAAoohjQcTn82nVqlVatGiRLMtSr169NHfuXH3/+9936pYAACDKOBZEhg4dqq1btzr14wEAgAtw1gwAADCGIAIAAIwhiAAAAGMIIgAAwBiCCAAAMIYgAgAAjCGIAAAAYwgiAICoEQhaKnh9vwJB68ufjKhAEAEARI2ikgpt3h9QUUmF6VIQIo7trAoAQKhNzu7R7F9EP4IIACBqpHX0KXdUH9NlIISYmgEAAMYQRAAAgDEEEQBAxGBVTOwhiAAAIgarYmIPzaoAgIjBqpjYQxABAEQMVsXEHqZmAACAMQQRAABgDEEEAAAYQxABAADGEEQAAIAxBBEAAGAMQQQAABhDEAEAAMYQRAAAIcNZMWgrgggAIGQ4KwZtxRbvAICQ4awYtBVBBAAQMpwVg7ZiagYAABhDEAEAAMYQRAAArcKKGDiBIAIAaBVWxMAJNKsCAFqFFTFwAkEEANAqrIiBE5iaAQAAxhBEAACAMQQRAABgDEEEAAAYQxABAADGEEQAAIAxBBEAAGAMQQQAABhDEAEAAMYQRAAAgDEEEQBwMU7MRaQjiACAi3FiLiIdh94BgItxYi4iHSMiAOAi507FNJ6Ym9bRZ7gy4MIIIgDgIkzFINowNQMALsJUDKINQQQAXKRxKgaIFkzNAAAAYwgiAADAGIIIAAAwhiACAFGEnVLhNjSrAkAUCAQtFZVU6MTpeu0or5YkmlLhCoyIAEAUaNwfRLaU0yeN5blwDUeDyMsvv6zhw4crKSlJqampmjRpkpO3AwBXKS4LaHT+Bv3wD+/q5v4ZyumTpuk5WeyUCldxbGrm+eef1z333KPHHntMY8eOVX19vd577z2nbgcArvPQf7+rTz47oaLtFerRuT1TMXAlj23bdqh/aH19vbKysrR48WLNmjXrkn9OTU2NUlJSVF1dLb/fH8IKASDyFZcF9GDRLt3Ut4u+N/5qRkEQNdry+e3IiMj27dt18OBBeb1eDRkyRIcPH9bgwYOVn5+vgQMHXvR1lmXJsr7oBK+pqXGiPACICsN6pWnjvDGmywAc5UiPyEcffSRJWrRokR5++GGtWbNGqampGj16tD777LOLvm7p0qVKSUlpemRmZjpRHgAAiBBtCiJ5eXnyeDwtPvbu3auGhgZJ0oIFC/Stb31L2dnZWrFihTwej1avXn3Rnz9//nxVV1c3PcrLyy/vtwMAABGtTVMzDzzwgGbMmNHic3r37q3KykpJ0jXXXNN03efzqXfv3vrkk08u+lqfzyefjzlQALGhcW+Qydk96P9AzGpTEElPT1d6evqXPi87O1s+n08ffPCBRo4cKUmqq6vTgQMH1LNnz0urFABcpmlvELE5GWKXI82qfr9fc+bM0cKFC5WZmamePXsqPz9fkjRlyhQnbgkAUeHsUZDGTcnYnAyxzLF9RPLz8xUXF6dp06bp5MmTGj58uNavX6/U1FSnbgkAEe/cURBGQhDrHNlHJFTYRwSAmwSClgo3H5A80vQRWfSFwLXa8vnNWTMAECZFJRXaUVGt9glxhBDgbzh9FwDChJ4Q4HwEEQAIk7SOPnpCgHMwNQMAIRYIWip4fb8CQevLnwzEOIIIAIRY48qYopIK06UAEY8gAgCX6dwRkMnZPZTTJ41eEKAVCCIAcJnOHQFp7AVhZQzw5WhWBYDLxGoY4NIRRADgMrEaBrh0TM0AAABjCCIAAMAYgggAADCGIAIAX4INygDnEEQA4EuwQRngHFbNAMAFBIKWikoqNDm7B8tzAQcRRADgLIGgpcItB1R84Kga7M+v5Y7qw/JcwCFMzQDAWYpKKvTHnZWqOHpSifFeRkEAhzEiAgD6Yirm5v4ZOnG6XrKl6TlZbNMOOIwgAgD6oiFVkuaO72e4GiB2EEQAQJwXA5hCEAEAcV4MYArNqgAAwBiCCAAAMIYgAiAmsE07EJnoEQHgao3Lck+crteO8mpJohcEiCAEEQCuFQhaenD1Tp2qa9CNWanK6ZPGqhggwhBEALhWUUmFrPoGJcZ72ZwMiFAEEQCudfbeIIQQIDLRrArANc5tSG3cG4QQAkQugggA12jcpr2opMJ0KQBaiakZAFGtcVXM5OwebNMORCGCCICodvZhdbmj+rA0F4gyBBEAUalxJOTm/hmSGAUBohVBBEBUOnckBEB0IogAiBr0gwDuQxABEDXoBwHchyACIGowCgK4D/uIAIhYbFAGuB8jIgAiTmlVrZa8skd9MzpqX1VQEg2pgFsRRABElEDQUu7KEgWCp3W6vkH/5+p0pmIAFyOIAIgYgaClB1fvVEff5/81LZ54rfp2TTZcFQAnEUQARIyikgpZ9Q3q3CFBz8y4gV4QIAYQRABEjLNXxRBCgNhAEAEQMRpXxQCIHSzfBRB25y7LBRC7CCIAwq5xh9SikgrTpQAwjKkZAGHHDqkAGhFEAIQdvSAAGjE1AwAAjCGIAAAAYwgiAADAGIIIAAAwhmZVACERCFoq3HxA8kjTR2SxMyqAViGIAAiJws0H9NutHys5MV7tE+JYFQOgVZiaARAaHik5MV6ZnZPYHwRAqzEiAuCSBYKWikoqNDm7h6aPyFL7hDgOrAPQJgQRAG3WGEBOWPXaUVEtScod1YfpGABt5tjUzMaNG+XxeC74ePvtt526LQCHBYKWHly9Uxs/OCJ5pJw+aUzFALhkjo2I5OTkqLKystm1H/7wh3rttdc0bNgwp24LwGFFJRWy6huUGO9ldQyAy+ZYEElISFC3bt2avq6rq9OLL76o++67Tx6Px6nbAnDA2b0gZx9YRwgBcLnC1iPy0ksvKRAIaObMmRd9jmVZsiyr6euamppwlAagBY1TMVZ9gyR6QQCEVtiCyPLly3XrrbeqR4+LzyUvXbpUixcvDldJAFpwdkPqqbrPp2LoBQEQam1uVs3Ly7toE2rjY+/evc1eU1FRoVdffVWzZs1q8WfPnz9f1dXVTY/y8vK2lgcgBM5tSB3dL10/nTKIqRgAIdfmEZEHHnhAM2bMaPE5vXv3bvb1ihUrlJaWpokTJ7b4Op/PJ5+P/+gA02hIBRAubQ4i6enpSk9Pb/XzbdvWihUrNH36dMXHx7f1dgDCpLSqVkte2aMFtw2gIRVA2Di+xfv69etVVlamu+++2+lbAbgMS17Zox3lx7TklT1K6+hT7qg+hBAAjnM8iCxfvlw5OTnq37+/07cCcAkCQUsFr+/XvaP7aHBmJy24bYDpkgDEEMdXzfzud79z+hYALkNRSYU27w9IklbMvNFwNQBiDWfNADHu7H4QAAg3x6dmAESW4rKAbn58o4rLPh8FoR8EgEkEESCGFJcF9J2nt6nsyHE99N/vmi4HAJiaAWLBa+8f1rzndyne65Vt2/J6PfrJN68zXRYAEEQAtwsELd37u3d0qr5BHX3tlNWlg37yzes0rFea6dIAgCACuNXZZ8X4k+J1Omjp/17bTT+9c7Dp0gCgCUEEcKnGZbmDM1P09zdmSrY0PSfLdFkA0AxBBHAptmkHEA1YNQO4SOMuqYGgxbJcAFGBIAK4SON0TFFJhelSAKBVmJoBolwgaKlwywHJlu4Y1F0Su6QCiB4EESDKFZVU6I87KyVJ7X1xyh3Vx3BFANB6BBEgSjUuz725f4ZOnK6XbEZCAEQfgggQpc4+NXfu+H6GqwGAS0MQAaJEIGjpqdf36/1DNVo88VpOzQXgCqyaAaJE4ZYDWrn1Y5V8fFRLXtnD8lwArkAQASJccVlANz++UQePnlRaB5+6pSRqwW0DTJcFACHB1AwQoRqbUX+/7RNVHDupM2dsTRvRk51SAbgKQQSIUIWbD+iPuyo1LCtV7co9nJgLwJUIIkCEaTo1t+6MJOmK1CROzAXgWgQRIEKUVtVq0R93q66+Qe3aeXVDz1R9+4ZMVsUAcDWCCGBYIGipcPMBvbjzkA5Xn1JifDsNubKTpudk0QsCwPUIIoBhhVsO6LdbP5Yvrp26pSRq/ICumjOaZbkAYgNBBDDNlpIT45WZmqR/+/shBBAAMYUgAhjS2JR6x6Duau+LY1kugJhEEAHCqDF8TM7u0eysGE7MBRCrCCJAmASClh5cvVNWfYMkcVYMAIggAoRFaVWtcleWqKMvTp07JDRNwzASAiDWEUQAhwWClnJXluhIrSVJembGDfSCAMDfcOgd4KDSqlrdWbBFSfHtlJ7sU8HUbEIIAJyFERHAQUte2aNPg5a6dPTpudwRhBAAOAcjIkAIlVbV6s6nNuvbBZtVWlWrBbcN0NArUxkJAYCLIIgAIbTklT3a/skxlXx8TEte2aO+XZO1YuaN6ts12XRpABCRmJoBQqC0qlZLXtmjqcOv1HGrXh6PtOC2AabLAoCIRxABQmDJK3u0o/yYJOm5OTlmiwGAKMLUDHCJissCuvnxjSouC2jBbQM0OLMToyAA0EYe27Zt00VcTE1NjVJSUlRdXS2/32+6HKBJcVlA3/7VVjXYUs+09to4b4zpkgAgYrTl85upGaANissCmvPbEn16oq7p2k19uxisCACiG0EEaKXisoC+8/Q21Td8MYg4bfiV+t74qw1WBQDRjSACtEIgaGlWYbHqG2x5JHm9Hi2eMEBTc3qZLg0AohpBBGhBaVWtFrzwrg4dO6W6Mw3ySBqcmaI/3DvSdGkA4AoEEeAiAkFLd/+mWB9/dkJer9Q9JUlXZrbX4onXmi4NAFyDIAKcIxC0VFRSoROn61Vv20qI8yoj2acVM25gh1QACDGCCHCOpzbu1/PbK3TbdV/Rt4ZeIdnS9JwszooBAAcQRIC/aRwJ2XnwmE7WndGBwHH9+BvXmS4LAFyNIIKYV1pVq4Uv7VZdQ4PivF4NuqKTOiTEsUsqAIQBQQQxrbSqVnf+aouOn6pXkq+dhmSmas7oPkzDAECYEEQQ0/71D++q5mSd2ifEaUp2puaMIoQAQDgRRBCTGvtB6hoaJHk04CvJWnD7NabLAoCYQxBBTAkELT21cb/W7alSl44+DevZWZ2SEugHAQBDCCKIGa+9f1j/8vt35PV6dOaMrXZeD1MxAGCY13QBQLjMe36XTtY16MyZBg3LSlXB1GxCCAAYxogIXK20qlaLXtqtAd39euT2AXr05T3K/9b1uvmabqZLAwCIIAIXK62q1bf+Y7NqTtVr18Fq3Tumr7b/8BbTZQEAzsLUDFwpELSUu7JEQatekpSSFK/J2T0MVwUAOBcjInClopIKde6QoNP1DereKVFLJl1HPwgARCCCCFyjcW+Qydk9mkY/Jmf3IIAAQAQjiCDqBYKWnnp9v/60u0oZ/s9DR+6oPsod1cdwZQCAL+NYj8i+ffv09a9/XV26dJHf79fIkSO1YcMGp26HGBUIWnpw9U6tLq5QVc0pfXb8NL0gABBFHAsiEyZMUH19vdavX6+SkhINGjRIEyZM0OHDh526JWJQUUmFTtU1qF/XjhrWk71BACDaeGzbtkP9Qz/99FOlp6frjTfe0E033SRJqq2tld/v17p16zRu3LhW/ZyamhqlpKSourpafr8/1GUiipVW1WrJK3t07+g+KvnkGL0gABBB2vL57UiPSFpamvr166fCwkINHTpUPp9PBQUFysjIUHZ29kVfZ1mWLMtq+rqmpsaJ8hDFGhtSN+07ot2VNfrlxv1aMfNG02UBAC6RI0HE4/Hoz3/+syZNmqTk5GR5vV5lZGRo7dq1Sk1Nvejrli5dqsWLFztRElyguCyg2StLdGVqe93Qq7Pi47wcVgcAUa5NPSJ5eXnyeDwtPvbu3SvbtnXvvfcqIyNDmzZt0ltvvaVJkybpjjvuUGVl5UV//vz581VdXd30KC8vv+xfENEvELQ077kdmlKwVcdO1OmToyc0Z1QfrZh5o/p2TTZdHgDgMrSpR+TIkSMKBAItPqd3797atGmTbrnlFh09erTZ3NBVV12lWbNmKS8vr1X3o0cEklTw+n7lv/qB6htseT3Sc7O/qmG90kyXBQC4CMd6RNLT05Wenv6lzztx4oQkyettPuDi9XrV0NDQllsihp3dkPqNwd21bm+VHp88iBACAC7iSI/IiBEjlJqaqrvuukuPPPKIkpKS9PTTT6usrEy33367E7eEizRuULa6pFy2LRpSAcDFHNlHpEuXLlq7dq2CwaDGjh2rYcOG6c0339SLL76oQYMGOXFLuEjhlgNaufVjHbfOqJ3XQ0MqALiYI/uIhAo9IrGlcSTkf949rLqGBiXFt9Ovpw+jIRUAoozxfUSAtiguC+iB1bvUuWO8yj49rlOnz+iK1PZ6LncEm5QBgMsRRGDc/f+1Q4eOndKhao+GXNlJ8V6vFk+8lhACADGAIAJjVm4u06I1e9Q+wSNJSoz36j/+kbNiACCWEEQQdo3Lct/Yd0RnbOm4ZatPegf95JvXEUIAIMYQRBB2PyjapR0Vx3RFSqIqaywtmjBAU3N6mS4LAGAAQQRhU1wW0P3/tUOVx07JliSPR6WP3Wa6LACAQQQROK64LKB7f/+O/lpjqXGteJzXoyfvZE8ZAIh1BBE4KhC0NHtliT47XidJ8kj6Skqi/v07g9mqHQBAEIEzAkFLP1u3T2verVRaxwSdrm+QR9LPvj1YN1/TzXR5AIAIQRBByJVW1WrGird18NhJ2ZI8Hun1eWNYEQMAOA9BBCFVXBbQP/56m6wzn3eDJMV79aup7A0CALgwgghCprEfpDGEJMZ59cd/GclZMQCAiyKI4LIFgpYKNx9QycdH1T0lUXVnGpTsi9e/f2cwIQQA0CKCCC5LaVWtZv7n2/przSmldfSpX7dk/ea7w5mKAQC0CkEElyQQtDT/+V1at+ev8kiyJcW38+qnUwYRQgAArUYQQZsVlwU06zfFqj5V33TtxqxULfkGZ8UAANqGIII2KS4L6DtPb1N9w+cNqR5JP5p4DWfFAAAuCUEErRIIWioqqdDvtn2iM38LId07Japw5o00pAIALhlBBK1SVFKhzfsDuumqLtr04ae66aoumjv+aqZiAACXhSCCFr32/mE9uHqXRvZN0+DMFE0fkaUff+M602UBAFzCa7oARK6CDR9qVmGJjp6s0//srlL7hDhGQAAAIUUQwQUVbPhQS1/d1/T1pMHdNTm7h8GKAABuxNQMmnlhe7nm/+E9naxraLo2/9arlTvmKoNVAQDciiACSZ/3gtz/+x06XndGkhTnlRps6aFbCCEAAOcQRKDX3j+sWYUlza79dPL1mjQ001BFAIBYQY9IjCutqtXs3zYPIT+7kxACAAgPRkRiVCBo6cl1+/T89gqdsb+4TggBAIQTQSQGFZcF9N3Ct1V76oxsW0po51FSQjs9MWWQbr6mm+nyAAAxhCASY4rLAvr2r7Y2jYIkxnm1ctaNGtYrzWxhAICYRBCJES9sL9e8ol06a1WufHFerblvJGfFAACMoVk1BgSClh44J4R0T0nUs7M4sA4AYBYjIi5XXBbQ7JUlivd6dKbBlkfSkzSkAgAiBEHEpQJBSw/81w5t/PBTSVJKUpyye6Zo8cRrGQUBAEQMgogLlVbVKndlifYfOd50bfn0YTSkAgAiDkHEZVZuLtMPX3pfHkkpie1Ua53RQ7dcTQgBAEQkgoiLFJcF9PBL70uSbEl33nClFtx+jdmiAABoAUHEBYrLApr73E79tdaSR5+HkFsGZGjOqD6mSwMAoEUEkShWWlWr+37/jvYcrm26lhTv1W+/ywZlAIDoQBCJUqVVtZrw8006Vf/FQTHdUxJV+F32BgEARA+CSBR6YXu55j63S2edVaeJ13fTwokDldbRZ6wuAADaiiASZZ58dY/+bcNHza79vzG9NffWAYYqAgDg0hFEokRpVa3uKSxWWeBE07V4r/TU1GxOzAUARC2CSBR4YXu55q7eJfusuZicXp31838cylQMACCqEUQiWHFZQFN/vVWnzjS/Pv/Wq5U75iozRQEAEEIEkQi1cnNZ0+ZkjdI7JOj3s7/KqhgAgGsQRCLQhUJIr85JKvrnv2MqBgDgKgSRCFJcFtDdhW/r2MnmczGsigEAuBVBJEK8sL1c33tuV7NrXo/09DRWxQAA3IsgEgHuW/m2/vjeX5tdu7Jzkp656wb6QQAArkYQMaxgw4fnhZDl0xkFAQDEBoKIIQUbPtTSV/edd/3HE68hhAAAYgZBxIALTcWMubqLfnrnYFbFAABiCkEkjAJBS/Of36U/7WkeQu4YmKGfT73BUFUAAJhDEAmTC62K4awYAECsI4iEwYVCyI8nXqOpOb0MVQQAQGQgiDjoYg2pdwzMIIQAACCCiGNWbi47L4QwFQMAQHNep37w9u3bNX78eHXq1ElpaWmaPXu2gsGgU7eLKLP/c9t5Z8V0S07Q1n8dRwgBAOAsjgSRQ4cOady4cerbt6+2bdumtWvXavfu3ZoxY4YTt4sYL2wvV1bey/rT3k+bXb/n73pq64LxLM0FAOAcjkzNrFmzRvHx8frlL38pr/fzrPPUU0/p+uuvV2lpqfr27evEbY0a99P1Kv305HnXp97YQwvuGGigIgAAIp8jQcSyLCUkJDSFEElKSkqSJL355psXDSKWZcmyrKava2pqnCgv5C4UQuI80qrZX9WwXmmGqgIAIPI5MjUzduxYHT58WPn5+Tp9+rSOHj2qvLw8SVJlZeVFX7d06VKlpKQ0PTIzM50oL2Qap2LODSH3/F1PlS69nRACAMCXaFMQycvLk8fjafGxd+9eXXvttfrNb36jxx9/XO3bt1e3bt3Uq1cvde3atdkoybnmz5+v6urqpkd5efll/4JOmfCzjeftDSJ9HkKYigEAoHU8tm3brX3ykSNHFAgEWnxO7969lZCQ0PR1VVWVOnToII/HI7/fr1WrVmnKlCmtul9NTY1SUlJUXV0tv9/f2jId99Uf/0mHg3XNrg3s1kFrvjfaTEEAAESQtnx+t6lHJD09Xenp6W0qpmvXrpKkZ555RomJiRo/fnybXh9JVm4uO29ZriT17ZJECAEA4BI4tqHZL37xC+Xk5Khjx45at26d5s2bp2XLlqlTp05O3dJRQxf/jz472XDedQ6sAwDg0jkWRN566y0tXLhQwWBQ/fv3V0FBgaZNm+bU7RxzoXNiJCle0ofLbg9/QQAAuIhjQaSwsNCpHx02U/7/Jr39yflLiOkHAQAgNDhr5iL6/+vLOnX+TIx+duf1mjQ0spcVAwAQLQgi5yguC2hywdbzrneIl3b/iKkYAABCiSBylgsty5WYigEAwCkEkb/Jynv5gtdvuNKv1f98U5irAQAgNsR8ECnY8KGWvrrvvOseSWWsigEAwFExHUQuNgqS6JX2PkYIAQDAaY4cehcNWpqKIYQAABAeMT0icq4DTMUAABBWMTsici5CCAAA4RezIyIEDwAAzGNEBAAAGEMQAQAAxhBEAACAMQQRAABgDEEEAAAYQxABAADGEEQAAIAxBBEAAGAMQQQAABhDEAEAAMYQRAAAgDEEEQAAYAxBBAAAGBPRp+/ati1JqqmpMVwJAABorcbP7cbP8ZZEdBCpra2VJGVmZhquBAAAtFVtba1SUlJafI7Hbk1cMaShoUGHDh1ScnKyPB6P6XIuqqamRpmZmSovL5ff7zddjuvxfocX73f48F6HF++3c2zbVm1trbp37y6vt+UukIgeEfF6verRo4fpMlrN7/fzxxxGvN/hxfsdPrzX4cX77YwvGwlpRLMqAAAwhiACAACMIYiEgM/n08KFC+Xz+UyXEhN4v8OL9zt8eK/Di/c7MkR0syoAAHA3RkQAAIAxBBEAAGAMQQQAABhDEAEAAMYQREJs+/btGj9+vDp16qS0tDTNnj1bwWDQdFmutW/fPn39619Xly5d5Pf7NXLkSG3YsMF0Wa60ceNGeTyeCz7efvtt0+W51ssvv6zhw4crKSlJqampmjRpkumSXCkrK+u8v+tly5aZLismEERC6NChQxo3bpz69u2rbdu2ae3atdq9e7dmzJhhujTXmjBhgurr67V+/XqVlJRo0KBBmjBhgg4fPmy6NNfJyclRZWVls8fdd9+tXr16adiwYabLc6Xnn39e06ZN08yZM7Vz50795S9/0T/8wz+YLsu1Hn300WZ/3/fdd5/pkmKDjZApKCiwMzIy7DNnzjRd27Vrly3J/vDDDw1W5k5HjhyxJdlvvPFG07Wamhpbkr1u3TqDlcWG06dP2+np6fajjz5quhRXqqurs6+44gr717/+telSYkLPnj3tJ5980nQZMYkRkRCyLEsJCQnNDvhJSkqSJL355pumynKttLQ09evXT4WFhTp+/Ljq6+tVUFCgjIwMZWdnmy7P9V566SUFAgHNnDnTdCmutH37dh08eFBer1dDhgzRV77yFX3ta1/Te++9Z7o011q2bJnS0tI0ZMgQ5efnq76+3nRJMYEgEkJjx47V4cOHlZ+fr9OnT+vo0aPKy8uTJFVWVhquzn08Ho/+/Oc/65133lFycrISExP1xBNPaO3atUpNTTVdnustX75ct956a1QdTBlNPvroI0nSokWL9PDDD2vNmjVKTU3V6NGj9dlnnxmuzn3uv/9+rVq1Shs2bFBubq4ee+wx/eAHPzBdVmwwPSQTDR566CFbUouPPXv22LZt288++6zdtWtXu127dnZCQoL94IMP2l27drWXLVtm+LeIHq19vxsaGuyJEyfaX/va1+w333zTLikpsf/pn/7JvuKKK+xDhw6Z/jWiRlv+vhuVl5fbXq/XLioqMlR19Grt+/3ss8/akuyCgoKm1546dcru0qWL/dRTTxn8DaLHpfxtN1q+fLkdFxdnnzp1KsxVxx62eG+FI0eOKBAItPic3r17KyEhoenrqqoqdejQQR6PR36/X6tWrdKUKVOcLtUVWvt+b9q0SbfccouOHj3a7Ajvq666SrNmzWoajULLLuXv+0c/+pF+/vOf6+DBg4qPj3e6RFdp7fv9l7/8RWPHjtWmTZs0cuTIpu8NHz5c48aN05IlS5wuNepdyt92o927d2vgwIHau3ev+vXr51SJkBRnuoBokJ6ervT09Da9pmvXrpKkZ555RomJiRo/frwTpblSa9/vEydOSFKznpzGrxsaGhypzY3a+vdt27ZWrFih6dOnE0IuQWvf7+zsbPl8Pn3wwQdNQaSurk4HDhxQz549nS7TFS7l/+5GO3bskNfrVUZGRoirwrkIIiH2i1/8Qjk5OerYsaPWrVunefPmadmyZerUqZPp0lxnxIgRSk1N1V133aVHHnlESUlJevrpp1VWVqbbb7/ddHmutX79epWVlenuu+82XYqr+f1+zZkzRwsXLlRmZqZ69uyp/Px8SWJ0NcS2bNmibdu2acyYMUpOTtaWLVs0d+5cTZ06lX6zMCCIhNhbb72lhQsXKhgMqn///iooKNC0adNMl+VKXbp00dq1a7VgwQKNHTtWdXV1uvbaa/Xiiy9q0KBBpstzreXLlysnJ0f9+/c3XYrr5efnKy4uTtOmTdPJkyc1fPhwrV+/ng/HEPP5fFq1apUWLVoky7LUq1cvzZ07V9///vdNlxYT6BEBAADGsHwXAAAYQxABAADGEEQAAIAxBBEAAGAMQQQAABhDEAEAAMYQRAAAgDEEEQAAYAxBBAAAGEMQAQAAxhBEAACAMQQRAABgzP8COS/QeI3C9PwAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter( lncounts.mean(0).cpu().detach(), decoder_centers.cpu().detach(), s=.2 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "b6a73984-42a2-46b2-8146-2bb565ecbb76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.1175e+08, device='cuda:0', grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 156,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "dim1 = 200\n",
    "\n",
    "class Autoencoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Autoencoder, self).__init__()\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(2000, dim1),\n",
    "            nn.BatchNorm1d(dim1),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(dim1, 20),\n",
    "            nn.BatchNorm1d(20),\n",
    "            #nn.ReLU()\n",
    "        )\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(20, dim1),\n",
    "            nn.BatchNorm1d(dim1),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(dim1, 2000),\n",
    "            #nn.BatchNorm1d(2000),\n",
    "            #nn.ReLU()\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        x = self.decoder(x)\n",
    "        return x\n",
    "\n",
    "model = Autoencoder().cuda()\n",
    "\n",
    "counts = torch.tensor( ad.layers[\"counts\"].todense(), dtype=torch.float32 ).cuda()\n",
    "totals = counts.sum(1)\n",
    "lncounts = torch.log(counts / totals[:, None] + 1e-4)\n",
    "gene_means = lncounts.mean(0).detach()\n",
    "\n",
    "eta_hat = model( lncounts - gene_means )\n",
    "lambda_hat = torch.exp( eta_hat + gene_means ) * totals[:,None]\n",
    "\n",
    "alpha = .1**2\n",
    "nbprob = 1. / (lambda_hat * alpha + 1)\n",
    "nbprob = torch.clamp( nbprob, 1e-5, 1-1e-5 )\n",
    "log_likelihood_part = torch.xlogy(counts, 1 - nbprob) + torch.xlogy(1. / alpha, nbprob)\n",
    "loss = -log_likelihood_part.sum()\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebe7db1d-aa38-486d-bc7e-92cd4132f10e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/10000], Loss: 55641060.0000\n",
      "(2.377359628677368, 0.27460822463035583)\n",
      "Epoch [101/10000], Loss: 45340408.0000\n",
      "(1.212778925895691, 0.15660692751407623)\n",
      "Epoch [201/10000], Loss: 41378816.0000\n",
      "(0.7834348082542419, 0.106899693608284)\n",
      "Epoch [301/10000], Loss: 39834776.0000\n",
      "(0.6122679710388184, 0.08111453056335449)\n",
      "Epoch [401/10000], Loss: 39382276.0000\n",
      "(0.5320304036140442, 0.06471727043390274)\n",
      "Epoch [501/10000], Loss: 38694312.0000\n",
      "(0.4885430932044983, 0.05359175428748131)\n",
      "Epoch [601/10000], Loss: 38619464.0000\n",
      "(0.4614768624305725, 0.04572198912501335)\n",
      "Epoch [701/10000], Loss: 38225920.0000\n",
      "(0.4425472915172577, 0.04007025435566902)\n",
      "Epoch [801/10000], Loss: 38140004.0000\n",
      "(0.42818889021873474, 0.03577790781855583)\n",
      "Epoch [901/10000], Loss: 38014160.0000\n",
      "(0.4166598916053772, 0.03245140239596367)\n",
      "Epoch [1001/10000], Loss: 38198444.0000\n",
      "(0.4070424437522888, 0.029804542660713196)\n",
      "Epoch [1101/10000], Loss: 38003936.0000\n",
      "(0.3988170623779297, 0.0276107769459486)\n",
      "Epoch [1201/10000], Loss: 38060884.0000\n",
      "(0.3916913568973541, 0.025798555463552475)\n",
      "Epoch [1301/10000], Loss: 37731856.0000\n",
      "(0.3854571580886841, 0.024263858795166016)\n",
      "Epoch [1401/10000], Loss: 37671012.0000\n",
      "(0.3799617290496826, 0.022961027920246124)\n",
      "Epoch [1501/10000], Loss: 37775160.0000\n",
      "(0.3751099109649658, 0.021816490218043327)\n",
      "Epoch [1601/10000], Loss: 37805012.0000\n",
      "(0.37080007791519165, 0.020791180431842804)\n",
      "Epoch [1701/10000], Loss: 37791608.0000\n",
      "(0.3669433295726776, 0.019885096698999405)\n",
      "Epoch [1801/10000], Loss: 37540480.0000\n",
      "(0.36348432302474976, 0.01909824088215828)\n",
      "Epoch [1901/10000], Loss: 37874628.0000\n",
      "(0.3603724241256714, 0.01838291808962822)\n",
      "Epoch [2001/10000], Loss: 37563600.0000\n",
      "(0.35756102204322815, 0.01771528460085392)\n",
      "Epoch [2101/10000], Loss: 37499528.0000\n",
      "(0.35500863194465637, 0.017119185999035835)\n"
     ]
    }
   ],
   "source": [
    "model = Autoencoder().cuda()\n",
    "\n",
    "optimizer = torch.optim.Adam( model.parameters(), lr=1e-4, eps=1e-3 )\n",
    "num_epochs = 10000\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # Forward pass\n",
    "    eta_hat = model( lncounts - gene_means )\n",
    "    lambda_hat = torch.exp( eta_hat + gene_means ) * totals[:,None]\n",
    "    \n",
    "    # Compute negative binomial loss\n",
    "    alpha = .1**2\n",
    "    nbprob = 1. / (lambda_hat * alpha + 1)\n",
    "    nbprob = torch.clamp( nbprob, 1e-5, 1-1e-5 )\n",
    "    log_likelihood_part = torch.xlogy(counts, 1 - nbprob) + torch.xlogy(1. / alpha, nbprob)\n",
    "    #loss = -log_likelihood_part.sum()\n",
    "\n",
    "    ncells = log_likelihood_part.shape[0]\n",
    "    indices = torch.randperm(ncells)[:int(ncells * 0.5)]\n",
    "    loss = -log_likelihood_part[indices,:].sum()\n",
    "\n",
    "    if torch.isnan(loss).item():\n",
    "        print( \"loss is nan\" )\n",
    "        break\n",
    "    \n",
    "    # Backward pass and optimization\n",
    "    prev_model = copy.deepcopy(model)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    # Print loss for every epoch\n",
    "    if epoch % 100 == 0:\n",
    "        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')\n",
    "\n",
    "        with torch.no_grad():\n",
    "            log_likelihood = torch.lgamma( counts + 1./alpha ) - math.lgamma(1./alpha) - torch.lgamma( counts + 1 ) + \\\n",
    "                torch.xlogy( counts, 1 - nbprob ) + torch.xlogy( 1./alpha, nbprob )\n",
    "            nbprob0 = 1. / ( counts * alpha + 1 )\n",
    "            log_likelihood_saturated = torch.lgamma( counts + 1./alpha ) - math.lgamma(1./alpha) - torch.lgamma( counts + 1 ) + \\\n",
    "                torch.xlogy( counts, 1 - nbprob0 ) + torch.xlogy( 1./alpha, nbprob0 )\n",
    "            deviance = -2* (log_likelihood - log_likelihood_saturated )\n",
    "            print( ( deviance.mean().item(), deviance.median().item() ) )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0671f838-ab40-4370-8738-6cdb5689cbd1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-1.5258789e-05, -1.5258789e-05, -1.5258789e-05, ...,\n",
       "        4.7333188e+02,  5.7289001e+02,  6.2203540e+02], dtype=float32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfl = deviance.cpu().numpy().flatten()\n",
    "dfl.sort()\n",
    "dfl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0464e825-f7af-4c99-a51f-7732682e9dd2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c684629b-06d9-4c85-9945-28ca9aff96da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/10000], Loss: nan\n",
      "(nan, nan)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[12], line 30\u001b[0m\n\u001b[1;32m     27\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39mlog_likelihood_part\u001b[38;5;241m.\u001b[39msum()\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;66;03m# Backward pass and optimization\u001b[39;00m\n\u001b[0;32m---> 30\u001b[0m     \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m     33\u001b[0m \u001b[38;5;66;03m# Print loss\u001b[39;00m\n",
      "File \u001b[0;32m~/pyenv_torch/lib/python3.11/site-packages/torch/_tensor.py:487\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    477\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    478\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    479\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    480\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    485\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    486\u001b[0m     )\n\u001b[0;32m--> 487\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    488\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    489\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/pyenv_torch/lib/python3.11/site-packages/torch/autograd/__init__.py:200\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    195\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    197\u001b[0m \u001b[38;5;66;03m# The reason we repeat same the comment below is that\u001b[39;00m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 200\u001b[0m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    201\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    202\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "model = Autoencoder().cuda()\n",
    "\n",
    "alpha = .1**2\n",
    "\n",
    "optimizer = torch.optim.Adam( model.parameters(), lr=1e-2 )\n",
    "num_epochs = 10000\n",
    "\n",
    "# Create DataLoader\n",
    "dataset = torch.utils.data.TensorDataset( lncounts, counts )\n",
    "dataloader = torch.utils.data.DataLoader( dataset, batch_size=64, shuffle=True )\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "\n",
    "    epoch_loss = 0\n",
    "    for batch_lncounts, batch_counts in dataloader:\n",
    "        optimizer.zero_grad()\n",
    "    \n",
    "        # Forward pass\n",
    "        eta_hat = model( batch_lncounts - gene_means )\n",
    "        lambda_hat = torch.exp( eta_hat + gene_means ) * batch_counts.sum(1)[:,None]\n",
    "    \n",
    "        # Compute negative binomial loss\n",
    "        nbprob = 1. / (lambda_hat * alpha + 1)\n",
    "        log_likelihood_part = torch.xlogy( batch_counts, 1 - nbprob ) + torch.xlogy( 1. / alpha, nbprob )\n",
    "        loss = -log_likelihood_part.sum()\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # Print loss\n",
    "    if epoch % 50 == 0:\n",
    "        model.eval()\n",
    "        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')\n",
    "\n",
    "        with torch.no_grad():\n",
    "            eta_hat = model( lncounts - gene_means )\n",
    "            lambda_hat = torch.exp( eta_hat + gene_means ) * counts.sum(1)[:,None]\n",
    "    \n",
    "            nbprob = 1. / (lambda_hat * alpha + 1)\n",
    "            log_likelihood = torch.lgamma( counts + 1./alpha ) - math.lgamma(1./alpha) - torch.lgamma( counts + 1 ) + \\\n",
    "                torch.xlogy( counts, 1 - nbprob ) + torch.xlogy( 1./alpha, nbprob )\n",
    "            nbprob0 = 1. / ( counts * alpha + 1 )\n",
    "            log_likelihood_saturated = torch.lgamma( counts + 1./alpha ) - math.lgamma(1./alpha) - torch.lgamma( counts + 1 ) + \\\n",
    "                torch.xlogy( counts, 1 - nbprob0 ) + torch.xlogy( 1./alpha, nbprob0 )\n",
    "            deviance = -2* (log_likelihood - log_likelihood_saturated )\n",
    "            print( ( deviance.mean().item(), deviance.median().item() ) )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 559,
   "id": "57e6726d-8913-4d7f-bd1b-267016057eec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64])"
      ]
     },
     "execution_count": 559,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_counts.sum(1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf41eb43-607f-458c-a465-0701ecbe7a9a",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}