dev.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91534e22-616c-4fac-8c5c-b0962e58a1c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from swan_vis import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f334058-7005-4dd2-a111-5e0bb15ea431",
   "metadata": {},
   "source": [
    "## Creating tss / tes / ic adatas from cerberus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "cca6064b-9b1b-40a2-a038-13c412c85d6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_cerberus_adata(sg, mode):\n",
    "    \"\"\"\n",
    "    Parameters:\n",
    "        mode (str): {'ic', 'tss', 'tes'}\n",
    "    \"\"\"\n",
    "    id_col = '{}_id'.format(mode)\n",
    "    name_col = '{}_name'.format(mode)\n",
    "    \n",
    "    # merge ic, tss, tes info with transcript counts\n",
    "    gb_cols = ['gid', 'gname', id_col]\n",
    "    subset_cols = copy.deepcopy(gb_cols)\n",
    "    if mode in ['tss', 'tes']:\n",
    "        end = True\n",
    "        vert_col = '{}_vertex'.format(mode)\n",
    "        subset_cols.append('loc_path')\n",
    "        if mode == 'tss':\n",
    "            ind = 0\n",
    "        elif mode == 'tes':\n",
    "            ind = -1\n",
    "    else:\n",
    "        end = False\n",
    "    \n",
    "    df = self.t_df[subset_cols].reset_index()\n",
    "    \n",
    "    # merge with vertex id\n",
    "    if end:\n",
    "        df[vert_col] = [path[ind] for path in df.loc_path.values.tolist()]\n",
    "        df.drop('loc_path', axis=1, inplace=True)\n",
    "        gb_cols.append(vert_col)\n",
    "     \n",
    "    t_counts = self.get_transcript_abundance(kind='counts')\n",
    "    t_counts = t_counts.merge(df, how='left', on='tid')\n",
    "    # if mode == 'tes':\n",
    "    #     print(t_counts.loc[t_counts.tes_id == 'ENCODEHG000058837_1', 'hl60_m2_72hr_1_1'])\n",
    "    \n",
    "    # gb and sum counts over the different features\n",
    "    t_counts.drop('tid', axis=1, inplace=True)\n",
    "    # print(mode)\n",
    "    # print(gb_cols)\n",
    "    t_counts = t_counts.groupby(gb_cols).sum()\n",
    "    \n",
    "    # get separate components of the table\n",
    "    X = sparse.csr_matrix(t_counts.transpose().values)\n",
    "    obs = self.adata.obs\n",
    "    \n",
    "    # add name of thing\n",
    "    t_counts.reset_index(inplace=True)\n",
    "    t_counts[name_col] = t_counts.gname+'_'+t_counts[id_col].str.split('_', expand=True)[1]    \n",
    "    gb_cols.append(name_col)\n",
    "    var = t_counts[gb_cols]\n",
    "    var.set_index(id_col, inplace=True)\n",
    "    \n",
    "    if end:\n",
    "        var.rename({vert_col: 'vertex'}, axis=1, inplace=True)\n",
    "    \n",
    "    return obs, var, X\n",
    "\n",
    "def create_feat_adata(self, kind):\n",
    "    \"\"\"\n",
    "    Create a tss / tes / ic-level adata object. Enables calculating tss / tes\n",
    "    usage across samples.\n",
    "\n",
    "    Parameters:\n",
    "        kind (str): {'tss', 'ic', 'tes'}\n",
    "    \"\"\"\n",
    "    \n",
    "    # check to see if end information is already present in swangraph\n",
    "    id_col = '{}_id'.format(kind)\n",
    "    if id_col in self.t_df.columns:\n",
    "        obs, var, X = create_cerberus_adata(sg, kind)\n",
    "        \n",
    "    else:\n",
    "        # limit to only expresed transcripts\n",
    "        t_df = self.t_df.copy(deep=True)\n",
    "        t_df = t_df.loc[self.adata.var.index.tolist()]\n",
    "        df = get_ends(t_df, kind)\n",
    "\n",
    "        # get a mergeable transcript expression df\n",
    "        tid = self.adata.var.index.tolist()\n",
    "        obs = self.adata.obs.index.tolist()\n",
    "        data = self.adata.layers['counts'].transpose()\n",
    "        t_exp_df = pd.DataFrame.sparse.from_spmatrix(columns=obs, data=data, index=tid)\n",
    "        t_exp_df = t_exp_df.merge(t_df, how='left',\n",
    "            left_index=True, right_index=True)\n",
    "\n",
    "        # merge counts per transcript with end expression\n",
    "        df = df.merge(t_exp_df, how='left',\n",
    "            left_index=True, right_index=True)\n",
    "\n",
    "        # sort based on vertex id\n",
    "        df.sort_index(inplace=True, ascending=True)\n",
    "\n",
    "        # set index to gene ID, gene name, and vertex id\n",
    "        df.reset_index(drop=True, inplace=True)\n",
    "        df.set_index(['gid', 'gname', 'vertex_id'], inplace=True)\n",
    "        df = df[self.datasets]\n",
    "\n",
    "        # groupby on gene and assign each unique TSS / gene combo an ID\n",
    "        # use dense representation b/c I've already removed 0 counts and sparse\n",
    "        # gb operations are known to be slow in pandas\n",
    "        # https://github.com/pandas-dev/pandas/issues/36123\n",
    "        # maybe try this? :\n",
    "        # https://cmdlinetips.com/2019/03/how-to-write-pandas-groupby-function-using-sparse-matrix/\n",
    "        id_col = '{}_id'.format(kind)\n",
    "        name_col = '{}_name'.format(kind)\n",
    "        df = df.copy()\n",
    "        df.reset_index(inplace=True)\n",
    "        df[self.datasets] = df[self.datasets].sparse.to_dense()\n",
    "        df = df.groupby(['gid', 'gname', 'vertex_id']).sum().reset_index()\n",
    "\n",
    "        df['end_gene_num'] = df.sort_values(['gid', 'vertex_id'],\n",
    "                        ascending=[True, True])\\\n",
    "                        .groupby(['gid']) \\\n",
    "                        .cumcount() + 1\n",
    "        df[id_col] = df['gid']+'_'+df['end_gene_num'].astype(str)\n",
    "        df[name_col] = df['gname']+'_'+df['end_gene_num'].astype(str)\n",
    "        df.drop('end_gene_num', axis=1, inplace=True)\n",
    "\n",
    "        # obs, var, and X tables for new data\n",
    "        var_cols = ['gid', 'gname', 'vertex_id', id_col, name_col]\n",
    "        var = df[var_cols]\n",
    "        var.set_index('{}_id'.format(kind), inplace=True)\n",
    "        df.drop(var_cols, axis=1, inplace=True)\n",
    "        df = df[self.adata.obs.index.tolist()]\n",
    "        X = sparse.csr_matrix(df.transpose().values)\n",
    "        obs = self.adata.obs\n",
    "\n",
    "    # create anndata\n",
    "    adata = anndata.AnnData(var=var, obs=obs, X=X)\n",
    "\n",
    "    # add counts and tpm as layers\n",
    "    adata.layers['counts'] = adata.X\n",
    "    adata.layers['tpm'] = sparse.csr_matrix(calc_tpm(adata, recalc=True).to_numpy())\n",
    "    if not self.sc:\n",
    "        adata.layers['pi'] = sparse.csr_matrix(calc_pi(adata,\n",
    "                adata.var)[0].to_numpy())\n",
    "\n",
    "    # assign adata and clean up unstructured data if needed\n",
    "    if kind == 'tss':\n",
    "        if self.has_abundance():\n",
    "            adata.uns = self.tss_adata.uns\n",
    "        self.tss_adata = adata\n",
    "\n",
    "    elif kind == 'tes':\n",
    "        if self.has_abundance():\n",
    "            adata.uns = self.tss_adata.uns\n",
    "        self.tes_adata = adata\n",
    "    \n",
    "    elif kind == 'ic':\n",
    "        if self.has_abundance():\n",
    "            adata.uns = self.ic_adata.uns\n",
    "        self.ic_adata = adata\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "546503ac-6033-4918-b39c-1c3f62425c55",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read in graph from /Users/fairliereese/Documents/programming/mortazavi_lab/data/rnawg/lr_bulk/cerberus/swan.p\n"
     ]
    }
   ],
   "source": [
    "fname = '/Users/fairliereese/Documents/programming/mortazavi_lab/data/rnawg/lr_bulk/cerberus/swan.p'\n",
    "sg = read(fname)\n",
    "sg.ic_adata = anndata.AnnData()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "e559106f-4f3e-483f-b402-cf00d552ef61",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/fairliereese/Documents/programming/mortazavi_lab/bin/swan_vis/swan_vis/utils.py:410: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  df.reset_index(inplace=True)\n",
      "/Users/fairliereese/Documents/programming/mortazavi_lab/bin/swan_vis/swan_vis/utils.py:410: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  df.reset_index(inplace=True)\n",
      "/Users/fairliereese/Documents/programming/mortazavi_lab/bin/swan_vis/swan_vis/utils.py:410: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  df.reset_index(inplace=True)\n"
     ]
    }
   ],
   "source": [
    "create_feat_adata(sg, kind='tss')\n",
    "create_feat_adata(sg, kind='tes')\n",
    "create_feat_adata(sg, kind='ic')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "70c99574-074c-4b36-aa26-7f146d2a2975",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 138 × 71958\n",
       "    obs: 'dataset', 'total_counts', 'classification', 'sample', 'health_status'\n",
       "    var: 'gid', 'gname', 'vertex', 'tss_name'\n",
       "    uns: 'sample_colors', '{}_dict', 'health_status_colors'\n",
       "    layers: 'counts', 'tpm', 'pi'"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sg.tss_adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "a8337f19-1161-4d8d-ac97-813caa759368",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 138 × 86085\n",
       "    obs: 'dataset', 'total_counts', 'classification', 'sample', 'health_status'\n",
       "    var: 'gid', 'gname', 'vertex', 'tes_name'\n",
       "    uns: 'sample_colors', '{}_dict', 'health_status_colors'\n",
       "    layers: 'counts', 'tpm', 'pi'"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sg.tes_adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "34a8eec0-f2ba-4899-909b-575b05401400",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 138 × 132640\n",
       "    obs: 'dataset', 'total_counts', 'classification', 'sample', 'health_status'\n",
       "    var: 'gid', 'gname', 'ic_name'\n",
       "    layers: 'counts', 'tpm', 'pi'"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sg.ic_adata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b754f1b6-71ca-402c-b004-5c5f544dac5a",
   "metadata": {},
   "source": [
    "## Adding tss / tes / ic from cerberus "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f6039da1-3bcd-4e4f-b6ba-01b9438e12e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def abundance_to_adata(sg, counts_file, how='iso'):\n",
    "    # read in abundance file\n",
    "    check_file_loc(counts_file, 'abundance matrix')\n",
    "    try:\n",
    "        df = pd.read_csv(counts_file, sep='\\t')\n",
    "    except:\n",
    "        raise ValueError('Problem reading expression matrix {}'.format(counts_file))\n",
    "\n",
    "    # check if abundance matrix is a talon abundance matrix\n",
    "    cols = ['gene_ID', 'transcript_ID', 'annot_gene_id', 'annot_transcript_id',\n",
    "        'annot_gene_name', 'annot_transcript_name', 'n_exons', 'length',\n",
    "        'gene_novelty', 'transcript_novelty', 'ISM_subtype']\n",
    "    if df.columns.tolist()[:11] == cols:\n",
    "        df = reformat_talon_abundance(df, how=how)\n",
    "        \n",
    "    # rename id ID column\n",
    "    col = df.columns[0]\n",
    "    if how == 'gene':\n",
    "        id_col = 'gid'\n",
    "    elif how == 'iso':\n",
    "        id_col = 'tid'\n",
    "\n",
    "    df.rename({col: id_col}, axis=1, inplace=True)\n",
    "    \n",
    "    # sum for gene level\n",
    "    if how == 'gene':\n",
    "        df = df.groupby(id_col).sum().reset_index()\n",
    "    \n",
    "    # limit to just the transcripts already in the graph\n",
    "    if how == 'iso':\n",
    "        sg_tids = sg.t_df.tid.tolist()\n",
    "        ab_tids = df.tid.tolist()\n",
    "        tids = list(set(sg_tids)&set(ab_tids))\n",
    "        df = df.loc[df.tid.isin(tids)]\n",
    "        \n",
    "    # transpose to get adata format\n",
    "    df.set_index(id_col, inplace=True)\n",
    "    df = df.T\n",
    "    \n",
    "    # get adata components - obs, var, and X\n",
    "    var = df.columns.to_frame()\n",
    "    var.columns = [id_col]\n",
    "    obs = df.index.to_frame()\n",
    "    obs.columns = ['dataset']\n",
    "    X = sparse.csr_matrix(df.to_numpy())\n",
    "    \n",
    "    # create transcript-level adata object and filter out unexpressed transcripts\n",
    "    adata = anndata.AnnData(var=var, obs=obs, X=X)\n",
    "    genes, _  = sc.pp.filter_genes(adata, min_counts=1, inplace=False)\n",
    "    adata = adata[:, genes]\n",
    "    adata.layers['counts'] = adata.X\n",
    "    \n",
    "    return adata    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "08598eef-f82c-40b8-9148-1ddd4ff5b329",
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_adata_abundance(sg, adata, how='iso'):\n",
    "    \n",
    "    if how == 'gene':\n",
    "        dataset_list = sg.gene_datasets\n",
    "        ab_bool = sg.has_gene_abundance()\n",
    "        sg_adata = sg.gene_adata\n",
    "    elif how == 'iso':\n",
    "        dataset_list = sg.datasets\n",
    "        ab_bool = sg.has_abundance()\n",
    "        sg_adata = sg.adata\n",
    "        \n",
    "    print(adata)\n",
    "    \n",
    "    # add each dataset to list of \"datasets\", check if any are already there!\n",
    "    datasets = adata.obs.dataset.tolist()\n",
    "    for d in datasets:\n",
    "        if d in dataset_list:\n",
    "            raise ValueError('Dataset {} already present in the SwanGraph.'.format(d))\n",
    "    dataset_list.extend(datasets)\n",
    "\n",
    "    print()\n",
    "    if len(datasets) <= 5:\n",
    "        print('Adding abundance for datasets {} to SwanGraph.'.format(', '.join(datasets)))\n",
    "    else:\n",
    "        mini_datasets = datasets[:5]\n",
    "        n = len(datasets) - len(mini_datasets)\n",
    "        print('Adding abundance for datasets {}... (and {} more) to SwanGraph'.format(', '.join(mini_datasets), n))\n",
    "\n",
    "    # if there is preexisting abundance data in the SwanGraph, concatenate\n",
    "    # otherwise, adata is the new transcript level adata\n",
    "    if not ab_bool:\n",
    "\n",
    "        # create transcript-level adata object\n",
    "        sg_adata = adata\n",
    "\n",
    "        # add counts as layers\n",
    "        sg_adata.layers['counts'] = sg_adata.X\n",
    "        print('Calculating TPM...')\n",
    "        sg_adata.layers['tpm'] = sparse.csr_matrix(calc_tpm(sg_adata, recalc=True).to_numpy())\n",
    "\n",
    "        if not sg.sc and how == 'iso':\n",
    "            print('Calculating PI...')\n",
    "            sg_adata.layers['pi'] = sparse.csr_matrix(calc_pi(sg_adata, sg.t_df)[0].to_numpy())\n",
    "    else:\n",
    "\n",
    "        # first set current layer to be counts\n",
    "        sg_adata.X = sg_adata.layers['counts']\n",
    "\n",
    "        # concatenate existing adata with new one\n",
    "        # outer join to add all new transcripts (that are from added\n",
    "        # annotation or transcriptome) to the abundance\n",
    "        uns = sg_adata.uns\n",
    "        sg_adata = sg_adata.concatenate(adata, join='outer', index_unique=None)\n",
    "        sg_adata.uns = uns\n",
    "\n",
    "        # recalculate pi and tpm\n",
    "        print('Calculating TPM...')\n",
    "        sg_adata.layers['tpm'] = sparse.csr_matrix(calc_tpm(sg_adata, recalc=True).to_numpy())\n",
    "\n",
    "        if not sg.sc and how == 'iso':\n",
    "            print('Calculating PI...')\n",
    "            sg_adata.layers['pi'] = sparse.csr_matrix(calc_pi(sg_adata, sg.t_df)[0].to_numpy())\n",
    "\n",
    "    # # add abundance for edges, TSS per gene, and TES per gene\n",
    "    # if how == 'iso':\n",
    "    #     print('Calculating edge usage...')\n",
    "    #     sg.create_edge_adata()\n",
    "    #     print('Calculating TSS usage...')\n",
    "    #     sg.create_end_adata(kind='tss')\n",
    "    #     print('Calculating TES usage...')\n",
    "    #     sg.create_end_adata(kind='tes')\n",
    "\n",
    "    # set abundance flag to true\n",
    "    # and make adata object\n",
    "    if how == 'iso':\n",
    "        sg.abundance = True  \n",
    "        sg.adata = sg_adata\n",
    "    elif how == 'gene':\n",
    "        sg.gene_abundance = True\n",
    "        sg.gene_adata = sg_adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2274854c-0a04-4337-9e2a-b8ccfe0d3aca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_end_adata(self, kind):\n",
    "    \"\"\"\n",
    "    Create a tss / tes-level adata object. Enables calculating tss / tes\n",
    "    usage across samples.\n",
    "\n",
    "    Parameters:\n",
    "        kind (str): Choose from 'tss' or 'tes'\n",
    "    \"\"\"\n",
    "\n",
    "    # limit to only expresed transcripts\n",
    "    t_df = sg.t_df.copy(deep=True)\n",
    "    t_df = t_df.loc[sg.adata.var.index.tolist()]\n",
    "    df = get_ends(t_df, kind)\n",
    "\n",
    "    # get a mergeable transcript expression df\n",
    "    tid = sg.adata.var.index.tolist()\n",
    "    obs = sg.adata.obs.index.tolist()\n",
    "    data = sg.adata.layers['counts'].transpose()\n",
    "    t_exp_df = pd.DataFrame.sparse.from_spmatrix(columns=obs, data=data, index=tid)\n",
    "    t_exp_df = t_exp_df.merge(t_df, how='left',\n",
    "        left_index=True, right_index=True)\n",
    "\n",
    "    # merge counts per transcript with end expression\n",
    "    df = df.merge(t_exp_df, how='left',\n",
    "        left_index=True, right_index=True)\n",
    "\n",
    "    # sort based on vertex id\n",
    "    df.sort_index(inplace=True, ascending=True)\n",
    "\n",
    "    # set index to gene ID, gene name, and vertex id\n",
    "    df.reset_index(drop=True, inplace=True)\n",
    "    df.set_index(['gid', 'gname', 'vertex_id'], inplace=True)\n",
    "    df = df[sg.datasets]\n",
    "\n",
    "    # groupby on gene and assign each unique TSS / gene combo an ID\n",
    "    # use dense representation b/c I've already removed 0 counts and sparse\n",
    "    # gb operations are known to be slow in pandas\n",
    "    # https://github.com/pandas-dev/pandas/issues/36123\n",
    "    # maybe try this? :\n",
    "    # https://cmdlinetips.com/2019/03/how-to-write-pandas-groupby-function-using-sparse-matrix/\n",
    "    id_col = '{}_id'.format(kind)\n",
    "    name_col = '{}_name'.format(kind)\n",
    "    df = df.copy()\n",
    "    df.reset_index(inplace=True)\n",
    "    df[sg.datasets] = df[sg.datasets].sparse.to_dense()\n",
    "    df = df.groupby(['gid', 'gname', 'vertex_id']).sum().reset_index()\n",
    "\n",
    "    df['end_gene_num'] = df.sort_values(['gid', 'vertex_id'],\n",
    "                    ascending=[True, True])\\\n",
    "                    .groupby(['gid']) \\\n",
    "                    .cumcount() + 1\n",
    "    df[id_col] = df['gid']+'_'+df['end_gene_num'].astype(str)\n",
    "    df[name_col] = df['gname']+'_'+df['end_gene_num'].astype(str)\n",
    "    df.drop('end_gene_num', axis=1, inplace=True)\n",
    "\n",
    "    # obs, var, and X tables for new data\n",
    "    var_cols = ['gid', 'gname', 'vertex_id', id_col, name_col]\n",
    "    var = df[var_cols]\n",
    "    var.set_index('{}_id'.format(kind), inplace=True)\n",
    "    df.drop(var_cols, axis=1, inplace=True)\n",
    "    df = df[sg.adata.obs.index.tolist()]\n",
    "    X = sparse.csr_matrix(df.transpose().values)\n",
    "    obs = sg.adata.obs\n",
    "\n",
    "    # create anndata\n",
    "    adata = anndata.AnnData(var=var, obs=obs, X=X)\n",
    "\n",
    "    # add counts and tpm as layers\n",
    "    adata.layers['counts'] = adata.X\n",
    "    adata.layers['tpm'] = sparse.csr_matrix(calc_tpm(adata, recalc=True).to_numpy())\n",
    "    if not sg.sc:\n",
    "        adata.layers['pi'] = sparse.csr_matrix(calc_pi(adata,\n",
    "                adata.var)[0].to_numpy())\n",
    "\n",
    "    # assign adata and clean up unstructured data if needed\n",
    "    if kind == 'tss':\n",
    "        if sg.has_abundance():\n",
    "            adata.uns = sg.tss_adata.uns\n",
    "        sg.tss_adata = adata\n",
    "\n",
    "    elif kind == 'tes':\n",
    "        if sg.has_abundance():\n",
    "            adata.uns = sg.tss_adata.uns\n",
    "        sg.tes_adata = adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ec6dbcab-cb17-46df-a61a-461891d56730",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_abundance(sg, counts_file, how='iso'):\n",
    "    adata = abundance_to_adata(sg, counts_file, how=how)\n",
    "    merge_adata_abundance(sg, adata, how=how)\n",
    "    \n",
    "    return sg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "60ad6ad9-3770-4bd9-b30f-698159fb854c",
   "metadata": {},
   "outputs": [],
   "source": [
    "s_pref = '/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/swan_transcriptome'\n",
    "annot = '/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/v29_cerberus.gtf'\n",
    "ab = '/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/human_cerberus_abundance.tsv'\n",
    "gene_ab = '/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/talon/human_talon_abundance.tsv'\n",
    "gtf = '/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/cerberus.gtf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "45becab8-778e-4d28-995a-63e43c63dbd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Adding annotation to the SwanGraph\n",
      "\n",
      "Adding transcriptome to the SwanGraph\n",
      "Saving graph as /Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/swan_transcriptome.p\n"
     ]
    }
   ],
   "source": [
    "sg = SwanGraph()\n",
    "sg.add_annotation(annot)\n",
    "sg.add_transcriptome(gtf)\n",
    "sg.save_graph(s_pref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9022104b-9b94-412b-a3a2-cb6fd0fc3418",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_file = '/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/human_cerberus_abundance.tsv'\n",
    "how = 'iso'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "223060f7-ce14-4107-8501-2146e3f5d421",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AnnData object with n_obs × n_vars = 138 × 157691\n",
      "    obs: 'dataset'\n",
      "    var: 'tid'\n",
      "    layers: 'counts'\n",
      "\n",
      "Adding abundance for datasets hepg2_1_1, imr90_1_1, mesenteric_fat_pad_1_1, ovary_1_1, hl60_1_1... (and 133 more) to SwanGraph\n",
      "Calculating TPM...\n",
      "Calculating PI...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/fairliereese/Documents/programming/mortazavi_lab/bin/swan_vis/swan_vis/utils.py:410: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead.  To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  df.reset_index(inplace=True)\n"
     ]
    }
   ],
   "source": [
    "sg = add_abundance(sg, counts_file, how=how)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "78ac36cc-39bf-4b69-a435-73b841085e61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def add_abundance(sg, counts_file, how='iso'):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b70a81aa-e15c-4c4e-86b8-e27b53deefa1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6616cbf6-5462-4b4f-bc94-53457a57872e",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'swan' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-16-6c57c74e1612>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0msg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mswan\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/swan.p'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'swan' is not defined"
     ]
    }
   ],
   "source": [
    "sg = swan.read('/Users/fairliereese/mortazavi_lab/data/rnawg/lr_bulk/cerberus/swan.p')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7aa719e1-ca89-44a5-b5c6-5676a7897b70",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3a24ff11-879c-4a2d-98b4-8a520e845f34",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "99b28cfd-fbd5-45d3-99f6-74cd803be901",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8473861d-8824-43b9-8d84-b23b8fbb26d2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f902502-b9aa-48d9-abd6-544e5f02724f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# add each dataset to list of \"datasets\", check if any are already there!\n",
    "datasets = adata.obs.dataset.tolist()\n",
    "for d in datasets:\n",
    "    if d in sg.datasets:\n",
    "        raise ValueError('Dataset {} already present in the SwanGraph.'.format(d))\n",
    "sg.datasets.extend(datasets)\n",
    "\n",
    "print()\n",
    "if len(datasets) <= 5:\n",
    "    print('Adding abundance for datasets {} to SwanGraph.'.format(', '.join(datasets)))\n",
    "else:\n",
    "    mini_datasets = datasets[:5]\n",
    "    n = len(datasets) - len(mini_datasets)\n",
    "    print('Adding abundance for datasets {}... (and {} more) to SwanGraph'.format(', '.join(mini_datasets), n))\n",
    "\n",
    "# if there is preexisting abundance data in the SwanGraph, concatenate\n",
    "# otherwise, adata is the new transcript level adata\n",
    "if not sg.has_abundance():\n",
    "\n",
    "    # create transcript-level adata object\n",
    "    sg.adata = adata\n",
    "\n",
    "    # add counts as layers\n",
    "    sg.adata.layers['counts'] = sg.adata.X\n",
    "    print('Calculating transcript TPM...')\n",
    "    sg.adata.layers['tpm'] = sparse.csr_matrix(calc_tpm(sg.adata, recalc=True).to_numpy())\n",
    "\n",
    "    if not sg.sc:\n",
    "        print('Calculating PI...')\n",
    "        sg.adata.layers['pi'] = sparse.csr_matrix(calc_pi(sg.adata, sg.t_df)[0].to_numpy())\n",
    "else:\n",
    "\n",
    "    # first set current layer to be counts\n",
    "    sg.adata.X = sg.adata.layers['counts']\n",
    "\n",
    "    # concatenate existing adata with new one\n",
    "    # outer join to add all new transcripts (that are from added\n",
    "    # annotation or transcriptome) to the abundance\n",
    "    uns = sg.adata.uns\n",
    "    sg.adata = sg.adata.concatenate(adata, join='outer', index_unique=None)\n",
    "    sg.adata.uns = uns\n",
    "\n",
    "    # recalculate pi and tpm\n",
    "    print('Calculating transcript TPM...')\n",
    "    sg.adata.layers['tpm'] = sparse.csr_matrix(calc_tpm(sg.adata, recalc=True).to_numpy())\n",
    "\n",
    "    if not sg.sc:\n",
    "        print('Calculating PI...')\n",
    "        sg.adata.layers['pi'] = sparse.csr_matrix(calc_pi(sg.adata, sg.t_df)[0].to_numpy())\n",
    "\n",
    "# add abundance for edges, TSS per gene, and TES per gene\n",
    "print('Calculating edge usage...')\n",
    "sg.create_edge_adata()\n",
    "print('Calculating TSS usage...')\n",
    "sg.create_end_adata(kind='tss')\n",
    "print('Calculating TES usage...')\n",
    "sg.create_end_adata(kind='tes')\n",
    "\n",
    "# set abundance flag to true\n",
    "sg.abundance = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3237c77c-df48-4be5-8832-46feef1b9015",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}