11_atac_apply.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dff3bf90-87ac-491e-b109-955dbf931d67",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 0\n",
      "/projects/b1038/Pulmonary/nmarkov/2022-multimodal/code/venv/lib/python3.9/site-packages/pytorch_lightning/utilities/warnings.py:53: LightningDeprecationWarning: pytorch_lightning.utilities.warnings.rank_zero_deprecation has been deprecated in v1.6 and will be removed in v1.8. Use the equivalent function from the pytorch_lightning.utilities.rank_zero module instead.\n",
      "  new_rank_zero_deprecation(\n",
      "/projects/b1038/Pulmonary/nmarkov/2022-multimodal/code/venv/lib/python3.9/site-packages/pytorch_lightning/utilities/warnings.py:58: LightningDeprecationWarning: The `pytorch_lightning.loggers.base.rank_zero_experiment` is deprecated in v1.7 and will be removed in v1.9. Please use `pytorch_lightning.loggers.logger.rank_zero_experiment` instead.\n",
      "  return new_rank_zero_deprecation(*args, **kwargs)\n"
     ]
    }
   ],
   "source": [
    "import scanpy as sc\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import anndata\n",
    "import sc_utils\n",
    "import scvi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7675dbd7-5d37-4035-a5e1-2ff3b9ef4dbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import hdf5plugin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dee2bdcd-00b8-4474-9a2e-6a2eae459f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = h5py.File(\"../data/test_multi_inputs.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "10afb5a4-2a78-4a0f-ab48-8264aa95d943",
   "metadata": {},
   "outputs": [],
   "source": [
    "names = f[\"test_multi_inputs\"][\"axis0\"].asstr()[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aa1610b6-7bf5-4689-bc6a-74b410ab2fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "names = pd.Series(names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b2937eed-bbf4-4715-b462-cbe04883b3ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "chromosomes = pd.Series(names.str.split(\":\").str[0].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "806926dd-d434-4001-aa58-c2aa1e348584",
   "metadata": {},
   "outputs": [],
   "source": [
    "chromosomes = chromosomes[chromosomes.str.startswith(\"chr\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "705320b9-59b7-42bc-8ca4-6b77635f5631",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterator, Mapping, Tuple, NamedTuple, Sequence, Union\n",
    "\n",
    "from absl import app\n",
    "from absl import flags\n",
    "from absl import logging\n",
    "import haiku as hk\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import optax\n",
    "import matplotlib.pyplot as plt\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ba75e90c-cf18-4ed4-a5e6-ff2ce4fddb52",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128 #\"Size of the batch to train on.\")\n",
    "learning_rate = 0.001 # \"Learning rate for the optimizer.\")\n",
    "training_steps = 1000 #, \"Number of training steps to run.\")\n",
    "eval_frequency = 10 #, \"How often to evaluate the model.\")\n",
    "random_seed = 42 #, \"Random seed.\")\n",
    "l1_coef = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e4adbc06-f9f4-4008-9016-dc1096ca6a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "PRNGKey = jnp.ndarray\n",
    "Batch = Mapping[str, np.ndarray]\n",
    "HiddenSize = Union[int, tuple]\n",
    "\n",
    "MNIST_IMAGE_SHAPE: Sequence[int] = (28, 28, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5fc3443f-6e6b-42b4-bdff-0f1c6dd97085",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(hk.Module):\n",
    "    \"\"\"Encoder model.\"\"\"\n",
    "\n",
    "    def __init__(self, hidden_size: HiddenSize = 512, latent_size: int = 10):\n",
    "        super().__init__()\n",
    "        if not isinstance(hidden_size, tuple):\n",
    "            hidden_size = (hidden_size,)\n",
    "        self._hidden_size = hidden_size\n",
    "        self._latent_size = latent_size\n",
    "\n",
    "    def __call__(self, x: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:\n",
    "        x = hk.Flatten()(x)\n",
    "        activation_sum = 0\n",
    "        for layer_size in self._hidden_size:\n",
    "            x = hk.Linear(layer_size)(x)\n",
    "            activation_sum += jnp.mean(jnp.abs(x))\n",
    "            x = jax.nn.leaky_relu(x)\n",
    "\n",
    "        return hk.Linear(self._latent_size)(x), activation_sum\n",
    "\n",
    "\n",
    "class Decoder(hk.Module):\n",
    "    \"\"\"Decoder model.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        hidden_size: HiddenSize = 512,\n",
    "        output_shape: Sequence[int] = MNIST_IMAGE_SHAPE,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        if not isinstance(hidden_size, tuple):\n",
    "            hidden_size = (hidden_size,)\n",
    "        self._hidden_size = hidden_size\n",
    "        self._output_shape = output_shape\n",
    "\n",
    "    def __call__(self, z: jnp.ndarray) -> jnp.ndarray:\n",
    "        activation_sum = 0\n",
    "        for layer_size in reversed(self._hidden_size):\n",
    "            z = hk.Linear(layer_size)(z)\n",
    "            activation_sum += jnp.mean(jnp.abs(z))\n",
    "            z = jax.nn.leaky_relu(z)\n",
    "\n",
    "        return hk.Linear(np.prod(self._output_shape))(z), activation_sum\n",
    "\n",
    "\n",
    "class AEOutput(NamedTuple):\n",
    "    data: jnp.ndarray\n",
    "    latent: jnp.ndarray\n",
    "    act_sum: float\n",
    "\n",
    "\n",
    "class AutoEncoder(hk.Module):\n",
    "    \"\"\"Main VAE model class, uses Encoder & Decoder under the hood.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        hidden_size: HiddenSize = 512,\n",
    "        latent_size: int = 10,\n",
    "        output_shape: Sequence[int] = MNIST_IMAGE_SHAPE,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self._hidden_size = hidden_size\n",
    "        self._latent_size = latent_size\n",
    "        self._output_shape = output_shape\n",
    "\n",
    "    def __call__(self, x: jnp.ndarray) -> AEOutput:\n",
    "        x = x.astype(jnp.float32)\n",
    "        z, act_sum_en = Encoder(self._hidden_size, self._latent_size)(x)\n",
    "        data, act_sum_de = Decoder(self._hidden_size, self._output_shape)(z)\n",
    "        act_sum_total = act_sum_en + act_sum_de + jnp.mean(jnp.abs(z))\n",
    "        return AEOutput(data, z, act_sum_total)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0d00c229-0e33-4c16-945b-8f3d5e898dd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def loss_fn(params: hk.Params, batch) -> jnp.ndarray:\n",
    "    output: AEOutput = model.apply(params, batch)\n",
    "    return jnp.mean(jnp.square(batch - output.data)) + l1_coef * output.act_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "922a090d-db85-4af9-818a-9d9eed3602c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def update(\n",
    "    params: hk.Params,\n",
    "    opt_state: optax.OptState,\n",
    "    batch: Batch,\n",
    ") -> Tuple[hk.Params, optax.OptState]:\n",
    "    \"\"\"Single SGD update step.\"\"\"\n",
    "    grads = jax.grad(loss_fn)(params, batch)\n",
    "    updates, new_opt_state = optimizer.update(grads, opt_state)\n",
    "    new_params = optax.apply_updates(params, updates)\n",
    "    return new_params, new_opt_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d22818d-31e7-4c04-8c85-9778f9f9d0a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Applying for chromosome chr10: 11009\n",
      "Applying for chromosome chr11: 10889\n",
      "Applying for chromosome chr12: 11354\n",
      "Applying for chromosome chr13: 6168\n",
      "Applying for chromosome chr14: 7299\n",
      "Applying for chromosome chr15: 7541\n",
      "Applying for chromosome chr16: 6996\n",
      "Applying for chromosome chr17: 9674\n",
      "Applying for chromosome chr18: 4969\n",
      "Applying for chromosome chr19: 7391\n",
      "Applying for chromosome chr1: 21706\n",
      "Applying for chromosome chr20: 6223\n",
      "Applying for chromosome chr21: 2751\n",
      "Applying for chromosome chr22: 4506\n",
      "Applying for chromosome chr2: 19071\n",
      "Applying for chromosome chr3: 16018\n",
      "Applying for chromosome chr4: 11120\n",
      "Applying for chromosome chr5: 12306\n",
      "Applying for chromosome chr6: 13826\n",
      "Applying for chromosome chr7: 11976\n",
      "Applying for chromosome chr8: 10361\n"
     ]
    }
   ],
   "source": [
    "for c in chromosomes:\n",
    "    c_names = names.str.startswith(c + \":\")\n",
    "    c_data = f[\"test_multi_inputs\"][\"block0_values\"][:, c_names.values]\n",
    "    c_shape = c_names.sum()\n",
    "    print(f\"Applying for chromosome {c}: {c_shape}\")\n",
    "    model = hk.transform(\n",
    "        lambda x: AutoEncoder(\n",
    "            hidden_size=(c_shape, c_shape // 16),\n",
    "            latent_size=64, # TODO: try different sizes, including 1\n",
    "            output_shape=(c_shape,)\n",
    "        )(x)\n",
    "    )  # pylint: disable=unnecessary-lambda\n",
    "    model = hk.without_apply_rng(model)\n",
    "    optimizer = optax.adam(learning_rate)\n",
    "    rng = jax.random.PRNGKey(1066)\n",
    "    params = jnp.load(f\"10_models/{c}.model.npy\", allow_pickle=True).ravel()[0]\n",
    "    output: AEOutput = model.apply(params, c_data)\n",
    "    jnp.save(f\"11_latent/{c}.latent\", output.latent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5156e451-49f2-4d84-8273-9d9c1ac41707",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}