SRS_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "02a1507d",
   "metadata": {},
   "source": [
    "#### 1. generating example data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "id": "a7de6935",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s1</th>\n",
       "      <th>s2</th>\n",
       "      <th>s3</th>\n",
       "      <th>s4</th>\n",
       "      <th>s5</th>\n",
       "      <th>s6</th>\n",
       "      <th>s7</th>\n",
       "      <th>s8</th>\n",
       "      <th>s9</th>\n",
       "      <th>s10</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ASV_1</th>\n",
       "      <td>0</td>\n",
       "      <td>86</td>\n",
       "      <td>16</td>\n",
       "      <td>19</td>\n",
       "      <td>9</td>\n",
       "      <td>92</td>\n",
       "      <td>51</td>\n",
       "      <td>10</td>\n",
       "      <td>68</td>\n",
       "      <td>23</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_2</th>\n",
       "      <td>14</td>\n",
       "      <td>63</td>\n",
       "      <td>21</td>\n",
       "      <td>46</td>\n",
       "      <td>3</td>\n",
       "      <td>56</td>\n",
       "      <td>88</td>\n",
       "      <td>80</td>\n",
       "      <td>46</td>\n",
       "      <td>54</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_3</th>\n",
       "      <td>79</td>\n",
       "      <td>71</td>\n",
       "      <td>14</td>\n",
       "      <td>77</td>\n",
       "      <td>15</td>\n",
       "      <td>25</td>\n",
       "      <td>53</td>\n",
       "      <td>84</td>\n",
       "      <td>58</td>\n",
       "      <td>85</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_4</th>\n",
       "      <td>29</td>\n",
       "      <td>44</td>\n",
       "      <td>37</td>\n",
       "      <td>22</td>\n",
       "      <td>91</td>\n",
       "      <td>54</td>\n",
       "      <td>76</td>\n",
       "      <td>12</td>\n",
       "      <td>59</td>\n",
       "      <td>26</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_5</th>\n",
       "      <td>76</td>\n",
       "      <td>71</td>\n",
       "      <td>87</td>\n",
       "      <td>39</td>\n",
       "      <td>43</td>\n",
       "      <td>76</td>\n",
       "      <td>38</td>\n",
       "      <td>91</td>\n",
       "      <td>69</td>\n",
       "      <td>98</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_6</th>\n",
       "      <td>33</td>\n",
       "      <td>43</td>\n",
       "      <td>26</td>\n",
       "      <td>56</td>\n",
       "      <td>69</td>\n",
       "      <td>73</td>\n",
       "      <td>52</td>\n",
       "      <td>89</td>\n",
       "      <td>27</td>\n",
       "      <td>43</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_7</th>\n",
       "      <td>1</td>\n",
       "      <td>82</td>\n",
       "      <td>94</td>\n",
       "      <td>96</td>\n",
       "      <td>26</td>\n",
       "      <td>84</td>\n",
       "      <td>30</td>\n",
       "      <td>64</td>\n",
       "      <td>22</td>\n",
       "      <td>98</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_8</th>\n",
       "      <td>52</td>\n",
       "      <td>3</td>\n",
       "      <td>70</td>\n",
       "      <td>12</td>\n",
       "      <td>39</td>\n",
       "      <td>83</td>\n",
       "      <td>48</td>\n",
       "      <td>61</td>\n",
       "      <td>70</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_9</th>\n",
       "      <td>36</td>\n",
       "      <td>23</td>\n",
       "      <td>22</td>\n",
       "      <td>66</td>\n",
       "      <td>92</td>\n",
       "      <td>95</td>\n",
       "      <td>53</td>\n",
       "      <td>9</td>\n",
       "      <td>41</td>\n",
       "      <td>57</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_10</th>\n",
       "      <td>88</td>\n",
       "      <td>93</td>\n",
       "      <td>76</td>\n",
       "      <td>82</td>\n",
       "      <td>11</td>\n",
       "      <td>36</td>\n",
       "      <td>16</td>\n",
       "      <td>30</td>\n",
       "      <td>84</td>\n",
       "      <td>57</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        s1  s2  s3  s4  s5  s6  s7  s8  s9  s10\n",
       "ASV_1    0  86  16  19   9  92  51  10  68   23\n",
       "ASV_2   14  63  21  46   3  56  88  80  46   54\n",
       "ASV_3   79  71  14  77  15  25  53  84  58   85\n",
       "ASV_4   29  44  37  22  91  54  76  12  59   26\n",
       "ASV_5   76  71  87  39  43  76  38  91  69   98\n",
       "ASV_6   33  43  26  56  69  73  52  89  27   43\n",
       "ASV_7    1  82  94  96  26  84  30  64  22   98\n",
       "ASV_8   52   3  70  12  39  83  48  61  70   13\n",
       "ASV_9   36  23  22  66  92  95  53   9  41   57\n",
       "ASV_10  88  93  76  82  11  36  16  30  84   57"
      ]
     },
     "execution_count": 220,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import random as rd\n",
    "\n",
    "#df with 10 ASVs x 10 samples populated by random integers between 0 and 99\n",
    "df = pd.DataFrame(np.random.randint(0,100,size=(10, 10)), \n",
    "                  columns=['s1','s2','s3','s4','s5','s6','s7','s8','s9','s10'], \n",
    "                  index = [\"{}_{}\".format('ASV', i) for i in range(1, 11)])\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0a1deb4",
   "metadata": {},
   "source": [
    "#### 2. defining SRS function "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 233,
   "id": "b3e9bbcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def SRS(data, Cmin, set_seed = True, seed = 1):\n",
    "    \n",
    "    #set the seed that allows reproducible results\n",
    "    if set_seed is True:\n",
    "        np.random.seed(seed)\n",
    "    \n",
    "    #check whether samples will have to be discarded due to low library size (sequencing depth)\n",
    "    if Cmin > min(data.sum(axis = 0)):\n",
    "        max_lib_size = max(data.sum(axis = 0))\n",
    "        if Cmin > max_lib_size:\n",
    "            print(\"Please select a Cmin <= \", max_lib_size)\n",
    "            raise ValueError(\"All samples discarded due to low number of counts.\")\n",
    "        else:\n",
    "            discarded = data.columns.values[data.sum() < Cmin]\n",
    "            print(\"Sample(s) discarded due to low number of counts (number of counts < Cmin): \", end='')\n",
    "            print(*discarded, sep=\", \")\n",
    "            data = data[data.columns[data.sum() >= Cmin]]\n",
    "            samples = data.columns\n",
    "            SRS(data, Cmin, set_seed, seed)\n",
    "            \n",
    "    #proceed with samples with library size >= Cmin\n",
    "    else:\n",
    "        \n",
    "        #throw error in case of invalid Cmin\n",
    "        if Cmin < 0:\n",
    "            raise ValueError(\"Cmin < 0. Please select Cmin >= 0.\")\n",
    "        if float(Cmin).is_integer() is False:\n",
    "            raise ValueError(\"SRS accepts only integers for Cmin.\")\n",
    "        \n",
    "        #start normalization\n",
    "        else:\n",
    "            #scaling\n",
    "            data_scaled = data.div(data.sum(axis = 0) / Cmin, axis = 1)\n",
    "            #picking the integer part of the scaled data\n",
    "            data_scaled_floor = np.floor(data_scaled)\n",
    "            #picking the fractional part of the scaled data\n",
    "            data_scaled_frac = data_scaled - data_scaled_floor\n",
    "            #n of counts missing for each sample in the integer\n",
    "            #part of the scaled data in order to reach Cmin\n",
    "            counts_missing = data_scaled_frac.sum(axis = 0)\n",
    "            data_normalized = data_scaled_floor\n",
    "            \n",
    "            #start distribution of missing counts per sample\n",
    "            for sample in data.columns:\n",
    "                counts_missing_s = int(counts_missing[sample])\n",
    "                data_scaled_frac_s = data_scaled_frac[sample]\n",
    "                data_scaled_frac_s_rank = data_scaled_frac_s.rank(method = 'min').sort_values(ascending = False)\n",
    "                data_scaled_floor_s = data_scaled_floor[sample]\n",
    "                while counts_missing_s > 0:\n",
    "                    for asv in data_scaled_frac_s_rank.index:\n",
    "                        if counts_missing_s == 0:\n",
    "                            break\n",
    "                        rank_current_asv = data_scaled_frac_s_rank.loc[asv]\n",
    "                        n_asvs_sharing_frac_rank = data_scaled_frac_s_rank.value_counts()[rank_current_asv]\n",
    "                        \n",
    "                        #distribute counts based on the frac rank (descending)\n",
    "                        if n_asvs_sharing_frac_rank <= counts_missing_s:\n",
    "                            data_normalized.loc[asv, sample] = data_normalized.loc[asv, sample] + 1\n",
    "                            counts_missing_s -= 1\n",
    "                            data_scaled_frac_s_rank = data_scaled_frac_s_rank.drop(asv)\n",
    "                        \n",
    "                        #if there are asvs sharing the same frac rank and there are less\n",
    "                        #missing counts than the number of asvs sharing the same frac rank\n",
    "                        else:\n",
    "                            asvs_sharing_frac_rank = data_scaled_frac_s_rank.head(n_asvs_sharing_frac_rank).index\n",
    "                            data_scaled_floor_s_asfr = data_scaled_floor_s.filter(asvs_sharing_frac_rank)\n",
    "                            data_scaled_floor_s_asfr_rank = data_scaled_floor_s_asfr.rank(method = 'min').sort_values(ascending = False)\n",
    "                            for asv in data_scaled_floor_s_asfr_rank.index:\n",
    "                                if counts_missing_s == 0:\n",
    "                                    break\n",
    "                                rank_current_asv = data_scaled_floor_s_asfr_rank.loc[asv]\n",
    "                                n_asvs_sharing_floor_rank = data_scaled_floor_s_asfr_rank.value_counts()[rank_current_asv]\n",
    "                                \n",
    "                                #distribute counts based on the int rank (descending)\n",
    "                                if n_asvs_sharing_floor_rank <= counts_missing_s:\n",
    "                                    data_normalized.loc[asv, sample] = data_normalized.loc[asv, sample] + 1\n",
    "                                    counts_missing_s -= 1\n",
    "                                    data_scaled_floor_s_asfr_rank = data_scaled_floor_s_asfr_rank.drop(asv)\n",
    "                                \n",
    "                                #if these asvs also share the same int rank, then\n",
    "                                #distribute missing counts randomly (without replacement)\n",
    "                                else:\n",
    "                                    while counts_missing_s > 0:\n",
    "                                        random_asv = rd.choice(data_scaled_floor_s_asfr_rank.index)\n",
    "                                        data_normalized.loc[random_asv, sample] = data_normalized.loc[random_asv, sample] + 1\n",
    "                                        counts_missing_s -= 1\n",
    "                                        data_scaled_floor_s_asfr_rank = data_scaled_floor_s_asfr_rank.drop(random_asv)\n",
    "    \n",
    "    #return df with SRS-normalized data (counts)\n",
    "    return(data_normalized.astype(int))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e029f94",
   "metadata": {},
   "source": [
    "#### 3. running SRS function "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 235,
   "id": "25844be9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s1</th>\n",
       "      <th>s2</th>\n",
       "      <th>s3</th>\n",
       "      <th>s4</th>\n",
       "      <th>s5</th>\n",
       "      <th>s6</th>\n",
       "      <th>s7</th>\n",
       "      <th>s8</th>\n",
       "      <th>s9</th>\n",
       "      <th>s10</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ASV_1</th>\n",
       "      <td>0</td>\n",
       "      <td>18</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>2</td>\n",
       "      <td>16</td>\n",
       "      <td>12</td>\n",
       "      <td>2</td>\n",
       "      <td>15</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_2</th>\n",
       "      <td>4</td>\n",
       "      <td>13</td>\n",
       "      <td>5</td>\n",
       "      <td>11</td>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "      <td>21</td>\n",
       "      <td>18</td>\n",
       "      <td>10</td>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_3</th>\n",
       "      <td>23</td>\n",
       "      <td>14</td>\n",
       "      <td>4</td>\n",
       "      <td>18</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>13</td>\n",
       "      <td>19</td>\n",
       "      <td>13</td>\n",
       "      <td>19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_4</th>\n",
       "      <td>9</td>\n",
       "      <td>9</td>\n",
       "      <td>9</td>\n",
       "      <td>5</td>\n",
       "      <td>27</td>\n",
       "      <td>10</td>\n",
       "      <td>18</td>\n",
       "      <td>3</td>\n",
       "      <td>13</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_5</th>\n",
       "      <td>22</td>\n",
       "      <td>15</td>\n",
       "      <td>22</td>\n",
       "      <td>9</td>\n",
       "      <td>13</td>\n",
       "      <td>14</td>\n",
       "      <td>9</td>\n",
       "      <td>20</td>\n",
       "      <td>15</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_6</th>\n",
       "      <td>10</td>\n",
       "      <td>9</td>\n",
       "      <td>7</td>\n",
       "      <td>13</td>\n",
       "      <td>21</td>\n",
       "      <td>13</td>\n",
       "      <td>12</td>\n",
       "      <td>20</td>\n",
       "      <td>6</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_7</th>\n",
       "      <td>0</td>\n",
       "      <td>17</td>\n",
       "      <td>24</td>\n",
       "      <td>22</td>\n",
       "      <td>8</td>\n",
       "      <td>15</td>\n",
       "      <td>7</td>\n",
       "      <td>14</td>\n",
       "      <td>5</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_8</th>\n",
       "      <td>15</td>\n",
       "      <td>0</td>\n",
       "      <td>18</td>\n",
       "      <td>3</td>\n",
       "      <td>12</td>\n",
       "      <td>15</td>\n",
       "      <td>11</td>\n",
       "      <td>14</td>\n",
       "      <td>15</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_9</th>\n",
       "      <td>11</td>\n",
       "      <td>5</td>\n",
       "      <td>6</td>\n",
       "      <td>15</td>\n",
       "      <td>28</td>\n",
       "      <td>17</td>\n",
       "      <td>13</td>\n",
       "      <td>2</td>\n",
       "      <td>9</td>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ASV_10</th>\n",
       "      <td>26</td>\n",
       "      <td>19</td>\n",
       "      <td>20</td>\n",
       "      <td>19</td>\n",
       "      <td>3</td>\n",
       "      <td>6</td>\n",
       "      <td>4</td>\n",
       "      <td>7</td>\n",
       "      <td>19</td>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        s1  s2  s3  s4  s5  s6  s7  s8  s9  s10\n",
       "ASV_1    0  18   4   4   2  16  12   2  15    5\n",
       "ASV_2    4  13   5  11   1  10  21  18  10   12\n",
       "ASV_3   23  14   4  18   4   4  13  19  13   19\n",
       "ASV_4    9   9   9   5  27  10  18   3  13    6\n",
       "ASV_5   22  15  22   9  13  14   9  20  15   21\n",
       "ASV_6   10   9   7  13  21  13  12  20   6    9\n",
       "ASV_7    0  17  24  22   8  15   7  14   5   21\n",
       "ASV_8   15   0  18   3  12  15  11  14  15    3\n",
       "ASV_9   11   5   6  15  28  17  13   2   9   12\n",
       "ASV_10  26  19  20  19   3   6   4   7  19   12"
      ]
     },
     "execution_count": 235,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SRS(df, Cmin = 120)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f179493",
   "metadata": {},
   "source": [
    "check https://cran.r-project.org/web/packages/SRS/SRS.pdf (page 2) for details on SRS function params"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}