LIA-Independent_alleles.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Independent alleles, from [rosalind.info](http://rosalind.info/)\n",
    "\n",
    "(text copied from http://rosalind.info/problems/lia/)\n",
    "\n",
    "> **Mendel's Second Law**\n",
    ">\n",
    "> Recall that Mendel's first law states that for any factor, an individual randomly assigns one of its two alleles to its offspring.\n",
    "> Yet this law does not state anything regarding the relationship with which alleles for different factors will be inherited.\n",
    ">\n",
    "> After recording the results of crossing thousands of pea plants for seven years, Mendel surmised that alleles for different factors are inherited with no dependence on each other.\n",
    "> This statement has become his second law, also known as the law of independent assortment.\n",
    ">\n",
    "> What does it mean for factors to be \"assorted independently?\"\n",
    "> If we cross two organisms, then a shortened form of independent assortment states that if we look only at organisms having the same alleles for one factor, then the inheritance of another factor should not change.\n",
    ">\n",
    "> For example, Mendel's first law states that if we cross two $Aa$ organisms, then 1/4 of their offspring will be $aa$, 1/4 will be $AA$, and 1/2 will be $Aa$.\n",
    "> Now, say that we cross plants that are both heterozygous for two factors, so that both of their genotypes may be written as $Aa Bb$.\n",
    "> Next, examine only $Bb$ offspring: Mendel's second law states that the same proportions of $AA$, $Aa$, and $aa$ individuals will be observed in these offspring.\n",
    "> The same fact holds for $BB$ and $bb$ offspring.\n",
    ">\n",
    "> As a result, independence will allow us to say that the probability of an $aa BB$ offspring is simply equal to the probability of an $aa$ offspring times the probability of a $BB$ organism, i.e., 1/16.\n",
    ">\n",
    "> Because of independence, we can also extend the idea of Punnett squares to multiple factors, as shown in Figure 1. We now wish to quantify Mendel's notion of independence using probability.\n",
    ">\n",
    "> [Figure 1](http://rosalind.info/media/problems/lia/dihybrid_cross.png).\n",
    "> Mendel's second law dictates that every one of the 16 possible assignments of parental alleles is equally likely.\n",
    "> The Punnett square for two factors therefore places each of these assignments in a cell of a 4 X 4 table.\n",
    "> The probability of an offspring's genome is equal to the number of times it appears in the table, divided by 16.\n",
    "\n",
    "> **Problem**\n",
    ">\n",
    "> Two events $A$ and $B$ are independent if $Pr(A and B)$ is equal to $Pr(A)×Pr(B)$.\n",
    "> In other words, the events do not influence each other, so that we may simply calculate each of the individual probabilities separately and then multiply.\n",
    ">\n",
    "> More generally, random variables $X$ and $Y$ are independent if whenever $A$ and $B$ are respective events for $X$ and $Y$, $A$ and $B$ are independent (i.e., $Pr(A and B)=Pr(A)×Pr(B)$).\n",
    ">\n",
    "> As an example of how helpful independence can be for calculating probabilities, let $X$ and $Y$ represent the numbers showing on two six-sided dice.\n",
    "> Intuitively, the number of pips showing on one die should not affect the number showing on the other die.\n",
    "> If we want to find the probability that $X+Y$ is odd, then we don't need to draw a tree diagram and consider all possibilities.\n",
    "> We simply first note that for $X+Y$ to be odd, either $X$ is even and $Y$ is odd or $X$ is odd and $Y$ is even.\n",
    "> In terms of probability, $Pr(X+Y is odd) = Pr(X is even and Y is odd) + Pr(X is odd and Y is even)$.\n",
    "> Using independence, this becomes $[Pr(X is even) × Pr(Y is odd)] + [Pr(X is odd) × Pr(Y is even)]$, or $(\\frac{1}{2})^2 + (\\frac{1}{2})^2 = \\frac{1}{2}$.\n",
    "> You can verify this result in Figure 2, which shows all 36 outcomes for rolling two dice.\n",
    ">\n",
    "> [Figure 2](http://rosalind.info/media/problems/lia/two_dice.png). The probability of each outcome for the sum of the values on two rolled dice (black and white), broken down depending on the number of pips showing on each die.\n",
    "> You can verify that 18 of the 36 equally probable possibilities result in an odd sum.\n",
    ">\n",
    "> **Given**: Two positive integers $k (k≤7)$ and $N (N≤2k)$.\n",
    "> In this problem, we begin with Tom, who in the 0th generation has genotype Aa Bb.\n",
    "> Tom has two children in the 1st generation, each of whom has two children, and so on.\n",
    "> Each organism always mates with an organism having genotype Aa Bb.\n",
    ">\n",
    "> **Return**: The probability that at least $N$ Aa Bb organisms will belong to the $k$-th generation of Tom's family tree (don't count the Aa Bb mates at each level). \n",
    "> Assume that Mendel's second law holds for the factors.\n",
    ">\n",
    "> **Sample Dataset**\n",
    "> ```\n",
    "> 2 1\n",
    "> ```\n",
    "> **Sample Output**\n",
    "> ```\n",
    "> 0.684\n",
    "> ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## My interpretation/reasoning\n",
    "\n",
    "1. This is an exercise of genetics and probabilities.\n",
    "\n",
    "2. The exercise focuses on double heterozygous organisms ($AaBb$).\n",
    "\n",
    "3. We start from one individual, Tom, with genotype AaBb.\n",
    "\n",
    "4. Tom mates with an AaBb and has two children in generation 1 (k = 1).\n",
    "\n",
    "5. In the next generation, those children each mate with an AaBb and have two children of their own, until generation k (<=7).\n",
    "\n",
    "6. For a given number N, I am going to calculate the probability that at least that many organisms have genotype AaBb. (Assuming independent inheritence.)\n",
    "\n",
    "7. The example reads: After two (k = 2) generations, there will be 2 * 2 = 4 organisms. The probability that at least 1 (N = 1) has genotype AaBb is 0.684.\n",
    "\n",
    "8. For the first generation, there are two organisms. Their genotypes can be any of the 16 shown in [figure 1](http://rosalind.info/media/problems/lia/dihybrid_cross.png). 4/16 possibilities (1/4) are AaBb. Then for next generations, probabilities have to be calculated for each possible genotype that mate with an AaBb. It is probably best to calculate all and finally return only the probability for AaBb in the script.\n",
    "\n",
    "9. Since we assume independent inheritence, we may calculate probabilities for Aa and Bb separately and multiply in the end.\n",
    "\n",
    "10. Chances for Aa and Bb should be identical given the situation, so only one needs to be calculated and can be squared in the end.\n",
    "\n",
    "11. The difficulty for me now seems to in the 'at least N organisms have genotype AaBb' part. \n",
    "    - The number of organisms doubles per generation.\n",
    "    - Probabilities are equal for all organisms\n",
    "    - The probability for 'at least N organisms' should be the sum of probabilities for N organisms in that generation.\n",
    "    - N cannot be larger than 2 * k, because that is the total number of organisms.\n",
    "    \n",
    "The difficulty seems to be mostly in the fact that organisms in generation 2 and up may have a parent from 'the Tom line' that was not AaBb. In generation 2, Aa may be the offspring of AA x Aa, Aa x Aa or aa x Aa. All of these are viable possibilities. Then for generation 3, the 'path of ancestors' becomes even more complicated, because in the generation before organisms may have had any genotype.\n",
    "\n",
    "However, this may not really matter in the end. For each generation, there are three possible genotypes: AA, Aa and aa. Each with their respective probability.\n",
    "In the first generation, probabilities are: AA = 1/4, Aa = 1/2, aa = 1/4. Then for the second generation, the probabilities are like this:\n",
    "\n",
    "AA = 1/2 * (AA first) + 1/4 * (Aa first) + 0 * (aa first). Where the 'first' refers to the probabilities observed in the previous generation. This same principle should apply to each generation and each genotype. Because the 'mate' has a set genotype (heterozygous), the chances for each genotype are always the same. From the second generation onwards, the only added difficulty is summing the probabilities from all 3 previous genotypes. Let's see how that can work in code...\n",
    "\n",
    "Now practically, this means the code is to:\n",
    "\n",
    "  - Open and read a text file with two numbers: k (generations) and N (number of individuals with genotype AaBb)\n",
    "  - For the first generation, the probabilities per organism are known: AA = 1/4, Aa = 1/2, aa = 1/4\n",
    "  - For each next generation, the probabilities are calculated based on previous probabilities (so the probabilities have to be stored, at least temporarily)\n",
    "  - This calculation has to be repeated k times.\n",
    "  - The final answer will be the probability of Aa (which is stored) at timepoint k squared (to incorporate Bb for which probabilities should be identical), to the power N (for at least N organisms have genotype Aa) PLUS that same chance for each value greater than N up to 2 * k!! (**at least N means including any number greater than N!**)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_numbers(input_file):\n",
    "    \"\"\"\n",
    "Read numbers k (generations) and N (number of organisms) from a file.\n",
    "Confirm that k is no greater than 7 and N is no greater than 2 * k.\n",
    "    \"\"\"\n",
    "    with open(input_file, 'r') as read_file:\n",
    "        for line in read_file:\n",
    "            numbers = line.split()\n",
    "            generation = int(numbers[0])\n",
    "            number = int(numbers[1])\n",
    "            \n",
    "    if generation > 7 or generation < 1:\n",
    "        print(\"Generation is outside range 1-7: %i\" % generation)\n",
    "        return(None)\n",
    "    else:\n",
    "        pass\n",
    "    \n",
    "    if number > (generation ** 2) or number < 0:\n",
    "        print(\"Number 'N' is too great, or below zero: %i\" % number)\n",
    "        return(None)\n",
    "    else:\n",
    "        pass\n",
    "    \n",
    "    return(generation, number)\n",
    "\n",
    "def calculate_probabilities(previous_probabilities):\n",
    "    \"\"\"\n",
    "Calculate the probabilities for each genotype (AA, Aa, aa)\n",
    "based on the probabilities of the previous generation.\n",
    "    \"\"\"\n",
    "    probability_AA = (previous_probabilities[\"AA\"] * 1/2 +\n",
    "                      previous_probabilities[\"Aa\"] * 1/4 +\n",
    "                      previous_probabilities[\"aa\"] * 0)\n",
    "    \n",
    "    probability_Aa = (previous_probabilities[\"AA\"] * 1/2 +\n",
    "                      previous_probabilities[\"Aa\"] * 1/2 +\n",
    "                      previous_probabilities[\"aa\"] * 1/2)\n",
    "            \n",
    "    probability_aa = (previous_probabilities[\"AA\"] * 0 +\n",
    "                      previous_probabilities[\"Aa\"] * 1/4 +\n",
    "                      previous_probabilities[\"AA\"] * 1/2)\n",
    "            \n",
    "    new_probabilities = { \"AA\": probability_AA,\n",
    "                          \"Aa\": probability_Aa,\n",
    "                          \"aa\": probability_aa }\n",
    "    \n",
    "    return(new_probabilities)\n",
    "\n",
    "def solve_independent_alleles(input_file):\n",
    "    \"\"\"\n",
    "Function that reads an input file with numbers for k (generations) \n",
    "and N (number of organisms) to solve the exercise and return\n",
    "the probability that at least N organisms in generation k have\n",
    "the heterozygous genotype Aa.\n",
    "    \"\"\"\n",
    "    (generation, number) = read_numbers(input_file)\n",
    "    \n",
    "    #First calculate the probabilities for the 3 genotypes in generation k\n",
    "    initial_probabilities = { \"AA\": 1/4, \"Aa\": 1/2, \"aa\": 1/4 }\n",
    "    \n",
    "    if generation == 1:\n",
    "        probabilities = initial_probabilities\n",
    "        \n",
    "        probability_Aa = probabilities[\"Aa\"]\n",
    "        \n",
    "        print(\"The probability that one organism has genotype Aa is: %f\" % probability_Aa)\n",
    "        \n",
    "    elif generation > 1:\n",
    "        probabilities = initial_probabilities\n",
    "        \n",
    "        for n in range(2, generation):\n",
    "            #From the second generation onwards,\n",
    "            # calculate the probabilities for each genotype.\n",
    "            new_probabilities = calculate_probabilities(probabilities)\n",
    "            #Overwrite the 'probabilities' variable so that it is recalculated\n",
    "            # for each step in this loop (number of generations from 2).\n",
    "            probabilities = new_probabilities\n",
    "            \n",
    "            print(new_probabilities)\n",
    "        \n",
    "    else:\n",
    "        return(\"Generation seems to be a number < 1: %i\" % generation)\n",
    "        \n",
    "    if number < generation * 2:\n",
    "        #If there are multiple solutions to 'at least N in k*2', e.g.\n",
    "        # if N = 2 and k = 4, there are 8 organisms. At least 2/8 means it can\n",
    "        # be 3, 4, 5 ... 8 as well, so add those probabilities.\n",
    "        total_probability = 0\n",
    "\n",
    "        for n in range(number, generation * 2):\n",
    "            probability = probabilities[\"Aa\"] ** 2 ** n\n",
    "            #The probability is the probability for Aa squared,\n",
    "            # to incorporate the identical probability of Bb,\n",
    "            # then times the number of organisms that have to have\n",
    "            # that exact genotype.\n",
    "            total_probability += probability\n",
    "            \n",
    "            print(total_probability)\n",
    "\n",
    "    else:\n",
    "        #N is equal to 2*k, i.e. there is only one solution: N is all organisms in generation k.\n",
    "        total_probability = probabilities[\"Aa\"] ** 2 ** number\n",
    "        #The probability is the probability for Aa squared,\n",
    "        # to incorporate the identical probability of Bb,\n",
    "        # then times the number of organisms that have to have\n",
    "        # that exact genotype.\n",
    "\n",
    "    return(total_probability)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.25\n",
      "0.3125\n",
      "0.31640625\n",
      "0.31640625\n"
     ]
    }
   ],
   "source": [
    "print(solve_independent_alleles(\"data/Example_independent_alleles.txt\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like I'm underestimating the probabilities. Something is going wrong. Does it have to do with possible combinations of organisms? That if at least one out of two has to have that genotype, that is may be number one, number two, or both (a total of three possibilities)? I don't know. I will look at it later."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These combinations can be calculated with the  `comb` function from the `math` module. Also see https://stackoverflow.com/questions/4941753/is-there-a-math-ncr-function-in-python for more information on the topic. I think this is what I need to incorporate into my function to solve the problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import comb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_independent_alleles2(input_file):\n",
    "    \"\"\"\n",
    "    -UPDATED-\n",
    "Function that reads an input file with numbers for k (generations) \n",
    "and N (number of organisms) to solve the exercise and return\n",
    "the probability that at least N organisms in generation k have\n",
    "the heterozygous genotype Aa.\n",
    "    \"\"\"\n",
    "    (generation, number) = read_numbers(input_file)\n",
    "    \n",
    "    #First calculate the probabilities for the 3 genotypes in generation k\n",
    "    initial_probabilities = { \"AA\": 1/4, \"Aa\": 1/2, \"aa\": 1/4 }\n",
    "    \n",
    "    if generation == 1:\n",
    "        probabilities = initial_probabilities\n",
    "        \n",
    "        probability_Aa = probabilities[\"Aa\"]\n",
    "        \n",
    "        print(\"The probability that one organism has genotype Aa is: %f\" % probability_Aa)\n",
    "        \n",
    "    elif generation > 1:\n",
    "        probabilities = initial_probabilities\n",
    "        \n",
    "        for n in range(1, generation):\n",
    "            #From the second generation onwards,\n",
    "            # calculate the probabilities for each genotype.\n",
    "            new_probabilities = calculate_probabilities(probabilities)\n",
    "            #Overwrite the 'probabilities' variable so that it is recalculated\n",
    "            # for each step in this loop (number of generations from 2).\n",
    "            probabilities = new_probabilities\n",
    "            \n",
    "            print(new_probabilities)\n",
    "        \n",
    "    else:\n",
    "        return(\"Generation seems to be a number < 1: %i\" % generation)\n",
    "        \n",
    "    if number < generation * 2:\n",
    "        #If there are multiple solutions to 'at least N in k*2', e.g.\n",
    "        # if N = 2 and k = 4, there are 8 organisms. At least 2/8 means it can\n",
    "        # be 3, 4, 5 ... 8 as well, so add those probabilities.\n",
    "        total_probability = 0\n",
    "\n",
    "        for n in range(number, generation * 2):\n",
    "            print(probabilities)\n",
    "            probability = (probabilities[\"Aa\"] * comb(generation * 2, n)) ** 2\n",
    "            #The probability is the probability for Aa squared,\n",
    "            # to incorporate the identical probability of Bb,\n",
    "            # then times the number of organisms that have to have\n",
    "            # that exact genotype.\n",
    "            total_probability += probability\n",
    "            \n",
    "            print(total_probability)\n",
    "\n",
    "    else:\n",
    "        #N is equal to 2*k, i.e. there is only one solution: N is all organisms in generation k.\n",
    "        total_probability = probabilities[\"Aa\"] ** 2 ** number * comb(generation * 2, number)\n",
    "        #The probability is the probability for Aa squared,\n",
    "        # to incorporate the identical probability of Bb,\n",
    "        # then times the number of organisms that have to have\n",
    "        # that exact genotype.\n",
    "\n",
    "    return(total_probability)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'AA': 0.25, 'Aa': 0.5, 'aa': 0.25}\n",
      "{'AA': 0.25, 'Aa': 0.5, 'aa': 0.25}\n",
      "4.0\n",
      "{'AA': 0.25, 'Aa': 0.5, 'aa': 0.25}\n",
      "13.0\n",
      "{'AA': 0.25, 'Aa': 0.5, 'aa': 0.25}\n",
      "17.0\n",
      "17.0\n"
     ]
    }
   ],
   "source": [
    "print(solve_independent_alleles2(\"data/Example_independent_alleles.txt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6953125"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1.390625 / 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It seems that the chances of each genotype are identical each generation. The most important thing to take into account then is that when at least 1 organism in generation 2 is to be Aa(Bb), then that can happen in 4 possible ways (like 1000, 0100, 0010, 0001). AND it also counts if 2, 3 or 4 organisms have that genotype, so the chances for each of those outcomes should be added."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n"
     ]
    }
   ],
   "source": [
    "chance_of_1 = 0.5 ** 2\n",
    "combinations_of_1 = comb(4, 1)\n",
    "\n",
    "total_probability_1 = chance_of_1 * combinations_of_1\n",
    "\n",
    "print(total_probability_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But this is too simple. This does not take into account that when the 1 does have genotype Aa, the other three do NOT. (Right?)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.75\n"
     ]
    }
   ],
   "source": [
    "chance_of_1 = 0.5 ** 2 * (3 * 0.5 ** 2)\n",
    "combinations_of_1 = comb(4, 1)\n",
    "\n",
    "total_probability_1 = chance_of_1 * combinations_of_1\n",
    "\n",
    "print(total_probability_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.16796875"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/4 * 3/4 ** 3 * comb(4,1) + 1/4 ** 2 * 3/4 ** 2 * comb(4,2) + 1/4 ** 3 * 3/4 ** 1 * comb(4,3) + 1/4 ** 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n"
     ]
    }
   ],
   "source": [
    "total = 4\n",
    "number = 1\n",
    "\n",
    "chance_of_1 = 1/4\n",
    "chance_not = 3/4\n",
    "combinations_of_1 = comb(total, number)\n",
    "\n",
    "print(combinations_of_1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.421875"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/4 * (3/4) ** 3 * 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.421875"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "27/64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.2109375"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/16 * 9/16 * 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.046875"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/4 ** 3 * 3/4 * 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00390625"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/4 ** 4 * (3/4) ** 0 * 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.68359375"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "0.421875 + 0.2109375 + 0.046875 + 0.00390625"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.0"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "3/4 ** 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import pow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pow(3/4, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(3/4) ** 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With some experimentation and calculating on paper, I have come to the conclusion that:\n",
    " - chances for each genotype remain the same over generations (so do not need to be recalculated)\n",
    " - with numbers k and N it is actually relatively simple to calculate the chance of 'at least N organisms have a heterozygous genotype'\n",
    "\n",
    "The calculation is as follows:\n",
    "The number of organisms per generation (k) = 2 * k\n",
    "The number of organisms that match the 'at least N' criterion is a list ranging from N  to 2 * k - for each of these outcomes you need to calculate the probability and sum all these probabilities\n",
    "For each of these numbers, there can be some combinations of organisms that give this outcome, e.g. when exactly one organism is heterozygous in a population of four, this can be represented as 1000, 0100, 0010 and 0001. So 4 possibilities: multiply the probability by 4.\n",
    "In each of those situations, the assumption is that the other organisms are NOT heterozygous.\n",
    "The chance of being heterozygous is 1/4, the chance of being not heterozygous is, therefore, 3/4. (Or 1 - 1/4)\n",
    "We then get: probability(heterozygous) to the power [number of organisms] times probability(not heterozygous) to the power [number of organisms] times the number of combinations\n",
    "For example in case of 1 organism out of a total of 4 (generation 2):\n",
    "  1/4 ** 1 * 3/4 ** 3 * 4 = 27/64 or 0.421875\n",
    "This is the chance that exactly 1 organism is heterozygous. Then add the chances for exactly 2, 3 or 4 (N+1 - 2 * k), and that should give the right answer!\n",
    "\n",
    "Now to put this into a function..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_independent_alleles3(generation, number):\n",
    "    \"\"\"\n",
    "    -UPDATED AGAIN-\n",
    "Function that reads an input file with numbers for k (generations) \n",
    "and N (number of organisms) to solve the exercise and return\n",
    "the probability that at least N organisms in generation k have\n",
    "the heterozygous genotype Aa.\n",
    "    \"\"\"\n",
    "    #The probabilities for each genotype are the same throughout generations,\n",
    "    # because each organisms mates with a heterozygous organism.\n",
    "    probabilities = { \"AA\": 1/4, \"Aa\": 1/2, \"aa\": 1/4 }\n",
    "    \n",
    "    if number < generation * 2:\n",
    "        #If there are multiple solutions to 'at least N in k*2', e.g.\n",
    "        # if N = 2 and k = 4, there are 8 organisms. At least 2/8 means it can\n",
    "        # be 3, 4, 5 ... 8 as well, so add those probabilities.\n",
    "        total_probability = 0\n",
    "\n",
    "        for n in range(number, generation * 2 + 1):\n",
    "            probability = (probabilities[\"Aa\"] ** 2) ** n * (1 - probabilities[\"Aa\"] ** 2) ** (2 * generation - n) * comb(generation * 2, n)\n",
    "            #The probability is the probability for Aa squared,\n",
    "            # to incorporate the identical probability of Bb,\n",
    "            # then to the power of the number of organisms that have to have\n",
    "            # that exact genotype (n). Then this is multiplied by the chance\n",
    "            # that an organism has another genotype (1 - probabilities[\"Aa\"] ** 2),\n",
    "            # to the power of the number of organisms that have this genotype\n",
    "            # (2 * generation - n) = (2k - n).\n",
    "            # And finally, the outcome is multiplied by the number of combinations\n",
    "            # that are possible for the given scenario.\n",
    "            total_probability += probability\n",
    "\n",
    "    else:\n",
    "        #N is equal to 2*k, i.e. there is only one solution: N is all organisms in generation k.\n",
    "        total_probability = (probabilities[\"Aa\"] ** 2) ** number * (1 - probabilities[\"Aa\"] ** 2) ** 0 * comb(generation * 2, number)\n",
    "        #The probability is the probability for Aa squared,\n",
    "        # to incorporate the identical probability of Bb,\n",
    "        # then times the number of organisms that have to have\n",
    "        # that exact genotype.\n",
    "\n",
    "    return(total_probability)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.68359375"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(generation, number) = read_numbers(\"data/Example_independent_alleles.txt\")\n",
    "solve_independent_alleles3(generation, number)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.03827075660228729"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve_independent_alleles3(7, 7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For some reason this feels like I have created a very roundabout way of saying something like this formula:\n",
    "\n",
    "$E(X) = \\sum_{k=N}^{2g} (\\frac{1}{4})^N × (\\frac{3}{4}) ^ {2g - N} × \\binom{N}{2g}$\n",
    "\n",
    "where 'g' = the generation (k in the example)\n",
    "\n",
    "Now I am not entirely sure if this mathematical notation is correct, but let's try writing this as a 'simpler' function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_independent_alleles_short(k, N):\n",
    "    \"\"\"\n",
    "The same formula as above, written shorter.\n",
    "    \"\"\"\n",
    "    return(sum([(1/4) ** n * (3/4) ** (2*k - n) * comb(2*k, n) for n in range(N, 2*k + 1)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.68359375"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve_independent_alleles_short(2, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.684"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "round(solve_independent_alleles_short(2, 1), 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just in case it needs to be rounded to 3 decimals, I can use the round function.\n",
    "\n",
    "---\n",
    "\n",
    "Now that this seems to work, let's try the real exercise!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00032343652002547274"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(generation, number) = read_numbers(\"data/rosalind_lia.txt\")\n",
    "\n",
    "solve_independent_alleles_short(generation, number)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This clearly didn't work... I think I got it wrong with the number of organisms per generation. I made it 2 * generation, but it should be $generation ^ 2$.\n",
    "\n",
    "No! Perhaps I mixed up the number, put them the wrong way round? That also doesn't seem to be the case. First generation, then number.\n",
    "\n",
    "Now what could have gone wrong...?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I just spotted a major overlook in the number of organisms per generation! I used the calculation: organisms = 2 * k, which coincidentally works for generations 1 and 2 (which works well enough with the example).\n",
    "However, the number of organisms should be $2^k$. So the function then becomes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_independent_alleles_short2(k, N):\n",
    "    \"\"\"\n",
    "The same formula as above, written shorter.\n",
    "    \"\"\"\n",
    "    return(sum([(1/4) ** n * (3/4) ** (2**k - n) * comb(2**k, n) for n in range(N, 2**k + 1)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the formula more like:\n",
    "\n",
    "$E(X) = \\sum_{k=N}^{2^g} (\\frac{1}{4})^N × (\\frac{3}{4}) ^ {2^g - N} × \\binom{N}{2^g}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.23176258105462272"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(generation, number) = read_numbers(\"data/rosalind_lia.txt\")\n",
    "\n",
    "solve_independent_alleles_short2(generation, number)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now I think that has solved the issue I had. Let's see if I can pass the test now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.45268544440216346"
      ]
     },
     "execution_count": 107,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(generation, number) = read_numbers(\"data/rosalind_lia2.txt\")\n",
    "\n",
    "solve_independent_alleles_short2(generation, number)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Success!!\n",
    "\n",
    "Yes, that worked. Now let's save all of this and commit to git."
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}