Skip to content

Commit

Permalink
nbdev setup
Browse files Browse the repository at this point in the history
  • Loading branch information
@cjGO
cjGO committed May 31, 2024
1 parent b849741 commit 67b29f3
Showing 1 changed file with 0 additions and 128 deletions.
128 changes: 0 additions & 128 deletions nbs/03_gametes.ipynb
View file
Original file line number Diff line number Diff line change
@@ -1,133 +1,5 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "c62c0aff",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"def meiosis(haplotypes: torch.Tensor, recombination_map: torch.Tensor, crossover_interference: float = 2.6, p: float = 0):\n",
" \"\"\"Simulates meiosis for an individual.\n",
"\n",
" Args:\n",
" haplotypes (torch.Tensor): A tensor of haplotypes of shape (ploidy, n_chromosomes, n_loci_per_chr).\n",
" recombination_map (torch.Tensor): A tensor of recombination map values for each locus of shape (n_chromosomes, n_loci_per_chr).\n",
" crossover_interference (float, optional): The crossover interference parameter. Defaults to 2.6.\n",
" p (float, optional): The proportion of crossovers coming from a non-interfering pathway. Defaults to 0.\n",
"\n",
" Returns:\n",
" torch.Tensor: A tensor of gametes of shape (ploidy, n_chromosomes, n_loci_per_chr).\n",
" \"\"\"\n",
" \n",
" n_chromosomes = recombination_map.size(0)\n",
" n_loci_per_chr = recombination_map.size(1)\n",
" ploidy = haplotypes.size(0)\n",
" \n",
" gametes = torch.zeros_like(haplotypes)\n",
" \n",
" for chr in range(n_chromosomes):\n",
" \n",
" # Simulate crossovers \n",
" crossover_events = simulate_crossovers(recombination_map[chr], crossover_interference, p)\n",
" \n",
" # Recombine haplotypes\n",
" for ploidy_idx in range(ploidy):\n",
" gametes[ploidy_idx, chr] = recombine_haplotype(haplotypes[ploidy_idx, chr], crossover_events)\n",
" \n",
" return gametes\n",
"\n",
"def simulate_crossovers(recombination_map: torch.Tensor, crossover_interference: float, p: float):\n",
" \"\"\"\n",
" Simulates crossover events along a chromosome based on the recombination map.\n",
"\n",
" Args:\n",
" recombination_map (torch.Tensor): A tensor of recombination map values for each locus.\n",
" crossover_interference (float): Crossover interference parameter.\n",
" p (float): Proportion of crossovers from non-interfering pathway.\n",
"\n",
" Returns:\n",
" torch.Tensor: A tensor of crossover events (1 for crossover, 0 for no crossover) for each locus.\n",
" \"\"\"\n",
" \n",
" # Calculate the probability of crossover at each locus\n",
" crossover_probabilities = 1 - torch.exp(-2 * recombination_map) \n",
"\n",
" # Apply crossover interference using a gamma distribution\n",
" crossover_events = torch.zeros_like(recombination_map, dtype=torch.bool)\n",
"\n",
" # Initialize the last crossover point\n",
" last_crossover = -1\n",
"\n",
" for i in range(recombination_map.size(0)):\n",
" # Probability of crossover at this locus\n",
" crossover_prob = crossover_probabilities[i]\n",
"\n",
" # Probability of crossover from the interfering pathway\n",
" interfering_prob = (1 - p) * crossover_prob\n",
"\n",
" # Probability of crossover from the non-interfering pathway\n",
" non_interfering_prob = p * crossover_prob\n",
"\n",
" # Simulate crossover from the interfering pathway\n",
" if interfering_prob > 0:\n",
" # Sample from a gamma distribution to determine the distance from the last crossover\n",
" distance_from_last_crossover = torch.distributions.Gamma(\n",
" concentration=crossover_interference,\n",
" rate=1.0\n",
" ).sample()\n",
"\n",
" # If the sampled distance is less than the current locus, there is a crossover\n",
" if distance_from_last_crossover <= i - last_crossover:\n",
" crossover_events[i] = True\n",
" last_crossover = i\n",
"\n",
" # Simulate crossover from the non-interfering pathway\n",
" if non_interfering_prob > 0:\n",
" if torch.rand(1) < non_interfering_prob:\n",
" crossover_events[i] = True\n",
" last_crossover = i\n",
"\n",
" return crossover_events\n",
"\n",
"def recombine_haplotype(haplotype: torch.Tensor, crossover_events: torch.Tensor):\n",
" \"\"\"\n",
" Recombines a haplotype based on crossover events.\n",
"\n",
" Args:\n",
" haplotype (torch.Tensor): A tensor of haplotype values.\n",
" crossover_events (torch.Tensor): A tensor of crossover events (1 for crossover, 0 for no crossover).\n",
"\n",
" Returns:\n",
" torch.Tensor: The recombined haplotype.\n",
" \"\"\"\n",
" \n",
" # Split the haplotype at crossover points\n",
" haplotype_segments = []\n",
" start_index = 0\n",
" \n",
" for i in range(crossover_events.size(0)):\n",
" if crossover_events[i]:\n",
" haplotype_segments.append(haplotype[start_index:i])\n",
" start_index = i\n",
" \n",
" haplotype_segments.append(haplotype[start_index:])\n",
" \n",
" # Reverse the segments after the first crossover\n",
" for i in range(len(haplotype_segments) // 2):\n",
" haplotype_segments[i + 1], haplotype_segments[-(i + 1)] = (\n",
" haplotype_segments[-(i + 1)],\n",
" haplotype_segments[i + 1],\n",
" )\n",
" \n",
" # Concatenate the segments\n",
" recombined_haplotype = torch.cat(haplotype_segments)\n",
" \n",
" return recombined_haplotype"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand Down

0 comments on commit 67b29f3

Please sign in to comment.