{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2b21427-9c55-4e13-88b7-7a1dffb5f02e",
   "metadata": {},
   "source": [
    "# Runing with large scale data (Stereo-seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e0ddddc-d2ad-44a3-8730-5bafbf210969",
   "metadata": {},
   "source": [
    "## import packages and load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d76cb061-90c8-4145-8b9f-127cda66c028",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-25T07:02:51.935271Z",
     "iopub.status.busy": "2024-09-25T07:02:51.935039Z",
     "iopub.status.idle": "2024-09-25T07:02:51.939828Z",
     "shell.execute_reply": "2024-09-25T07:02:51.938885Z",
     "shell.execute_reply.started": "2024-09-25T07:02:51.935249Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scanpy as sc\n",
    "from scipy import sparse, linalg, stats\n",
    "import sys\n",
    "import tqdm\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "sc.settings.set_figure_params(dpi=150)\n",
    "import matplotlib as mpl\n",
    "mpl.rcParams['pdf.fonttype']  = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c847881d-080d-4785-812f-072c0a871435",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-25T07:02:51.941203Z",
     "iopub.status.busy": "2024-09-25T07:02:51.940949Z",
     "iopub.status.idle": "2024-09-25T07:02:53.238200Z",
     "shell.execute_reply": "2024-09-25T07:02:53.236979Z",
     "shell.execute_reply.started": "2024-09-25T07:02:51.941178Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# subsample if necessary\n",
    "scdata_raw = sc.read_h5ad('Path/to/scdata')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7fa9100e-3f8e-4b3e-ad67-3dc43f27d1df",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-25T07:02:53.239593Z",
     "iopub.status.busy": "2024-09-25T07:02:53.239350Z",
     "iopub.status.idle": "2024-09-25T07:03:05.674231Z",
     "shell.execute_reply": "2024-09-25T07:03:05.672800Z",
     "shell.execute_reply.started": "2024-09-25T07:02:53.239570Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "stdata_raw = sc.read_h5ad('Path/to/stdata')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7054f0df-0094-45c1-bbe5-c44e2135814f",
   "metadata": {},
   "source": [
    "## evaluate peak mem, GPU and time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5d64debc-4320-4afa-91a5-5f114619af47",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-25T07:03:05.681735Z",
     "iopub.status.busy": "2024-09-25T07:03:05.681494Z",
     "iopub.status.idle": "2024-09-25T07:03:05.688844Z",
     "shell.execute_reply": "2024-09-25T07:03:05.687803Z",
     "shell.execute_reply.started": "2024-09-25T07:03:05.681710Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from memory_profiler import memory_usage\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f37d4e3d-cb5c-447a-875d-67178e12b60b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-25T07:03:05.689994Z",
     "iopub.status.busy": "2024-09-25T07:03:05.689757Z",
     "iopub.status.idle": "2024-09-25T07:03:05.715924Z",
     "shell.execute_reply": "2024-09-25T07:03:05.714858Z",
     "shell.execute_reply.started": "2024-09-25T07:03:05.689971Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import pynvml\n",
    "\n",
    "# Initialize NVML\n",
    "pynvml.nvmlInit()\n",
    "handle = pynvml.nvmlDeviceGetHandleByIndex(0) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3bfa778",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Polyomino as plmo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f012fd94-f10e-4b11-bf41-0bc083c647d0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-25T07:03:51.306994Z",
     "iopub.status.busy": "2024-09-25T07:03:51.306542Z",
     "iopub.status.idle": "2024-09-25T07:16:09.277826Z",
     "shell.execute_reply": "2024-09-25T07:16:09.276429Z",
     "shell.execute_reply.started": "2024-09-25T07:03:51.306958Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start to run the spot mapping.\n",
      "total_loss: -1.623\n",
      "total_loss: -1.878\n",
      "total_loss: -1.887\n",
      "total_loss: -1.889\n",
      "total_loss: -1.890\n",
      "Finish spot mapping.\n",
      "Finish spot mapping.\n",
      "Mapping Matrix save in zm.spot_matrix\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6437/6437 [10:23<00:00, 10.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Peak memory usage: 60007.22265625 MiB\n",
      "737.9632623195648\n"
     ]
    }
   ],
   "source": [
    "start_time = time()\n",
    "def plmo_mem():\n",
    "    stdata_spot = plmo.generate_grid(stdata_raw,width=5)\n",
    "    plmo_object = plmo.Polyomino(scdata_raw,stdata_spot,device='cuda')\n",
    "    plmo_object.allocate()\n",
    "    cell_alocated_data = plmo.sc2sc(scdata_raw, stdata_raw, zm.spot_matrix, method='max')\n",
    "    pass\n",
    "mem_usage, retval = memory_usage((zmap_mem,), retval=True, max_usage=True)\n",
    "print(f\"Peak memory usage: {mem_usage} MiB\")\n",
    "end_time = time() - start_time\n",
    "print(end_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9caf0fdb-d3b5-4c91-9775-a8f345d6dbae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-25T07:21:08.074250Z",
     "iopub.status.busy": "2024-09-25T07:21:08.073769Z",
     "iopub.status.idle": "2024-09-25T07:21:08.082616Z",
     "shell.execute_reply": "2024-09-25T07:21:08.081196Z",
     "shell.execute_reply.started": "2024-09-25T07:21:08.074208Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Peak GPU memory usage: 5634.22 MiB\n"
     ]
    }
   ],
   "source": [
    "# Get peak GPU memory usage\n",
    "peak_gpu_mem = torch.cuda.max_memory_allocated() / (1024 * 1024)  # Convert to MiB\n",
    "print(f\"Peak GPU memory usage: {peak_gpu_mem:.2f} MiB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dd70ad8-bab3-4a1e-9f21-a941fe28ef1b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b4528ab-daf5-4c10-bcc9-7a2b88a89970",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4a3628c-8c9c-406b-98d6-a246f6bc5d1b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}