Skip to content

Commit

Permalink
meiosis with more tensor ops + moved intercept in trait
Browse files Browse the repository at this point in the history
  • Loading branch information
@cjGO
cjGO committed Jun 8, 2024
1 parent c35132f commit 54f7806
Show file tree
Hide file tree
Showing 7 changed files with 123 additions and 91 deletions.
2 changes: 1 addition & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,4 +85,4 @@ plt.tight_layout() # Adjust layout to prevent overlap
plt.show()
```

AttributeError: 'int' object has no attribute 'float'
![](index_files/figure-commonmark/cell-4-output-1.png)
49 changes: 18 additions & 31 deletions chewc/meiosis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,31 +25,19 @@ def gamma_interference_model(length, rate, shape):
return crossover_positions


def simulate_meiosis(num_chromosomes, chromosome_length, crossover_rate, interference_strength):
import torch

def simulate_meiosis(num_chromosomes, map_length, num_individuals, num_crossovers):
"""
Simulate meiosis with crossover events and interference.
Parameters:
num_chromosomes (int): Number of chromatids.
chromosome_length (float): Length of the chromosome.
crossover_rate (float): Rate of crossover events.
interference_strength (float): Strength of interference.
Returns:
list of torch.Tensor: List of crossover positions for each chromatid.
This function simulates random crossover events across chromosomes.
"""
chromatid_crossovers = []
for _ in range(num_chromosomes):
crossovers = gamma_interference_model(chromosome_length, crossover_rate, interference_strength)
chromatid_crossovers.append(crossovers)
return chromatid_crossovers
return [torch.sort(torch.rand((num_individuals, num_crossovers)) * map_length, dim=-1)[0] for _ in range(num_chromosomes)]

def simulate_gametes(genetic_map, parent_genome):
"""
Simulate the formation of a single gamete given crossover positions, genetic map, and parent genomes.
Parameters:
crossover_positions (list of torch.Tensor): List of positions of crossover events for each chromosome.
genetic_map (list of torch.Tensor): List of positions of genetic markers on the chromosomes.
parent_genome (torch.Tensor): Genomes of the parents. SHAPE: (ploidy, num_chromosomes, num_loci)
Expand All @@ -58,21 +46,20 @@ def simulate_gametes(genetic_map, parent_genome):
"""
ploidy, num_chromosomes, num_loci = parent_genome.shape
gamete_genome = torch.zeros((ploidy // 2, num_chromosomes, num_loci), dtype=parent_genome.dtype)
crossover_positions = simulate_meiosis(num_chromosomes,100.0,1,1)
crossover_positions = simulate_meiosis(num_chromosomes, 100.0, 1, 1)

for chrom in range(num_chromosomes):
parent_1 = parent_genome[0, chrom]
parent_2 = parent_genome[1, chrom]

crossover_sites = crossover_positions[chrom]
chromatid = parent_1.clone()

if len(crossover_sites) > 0:
crossover_index = 0
for i, marker_position in enumerate(genetic_map[chrom]):
if crossover_index < len(crossover_sites) and marker_position >= crossover_sites[crossover_index]:
chromatid = parent_2 if chromatid is parent_1 else parent_1
crossover_index += 1
gamete_genome[0, chrom, i] = chromatid[i]
# Create a crossover mask for each chromosome
crossover_mask = torch.zeros(num_loci, dtype=torch.bool)
crossover_sites = crossover_positions[chrom].flatten()
for position in crossover_sites:
# Find the nearest marker index for each crossover position
index = torch.argmin(torch.abs(genetic_map[chrom] - position))
crossover_mask[index:] = ~crossover_mask[index:] # Flip values at and after the crossover index

# Use the mask to select genes from parent 1 or parent 2
gamete_genome[0, chrom] = torch.where(crossover_mask, parent_genome[1, chrom], parent_genome[0, chrom])

return gamete_genome


8 changes: 6 additions & 2 deletions chewc/trait.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ class TraitA:
target_mean: float
target_variance: float
distribution: str ='normal'
intercept: None = 0

def calculate_scaled_additive_dosages(self, genotypes: torch.Tensor) -> torch.Tensor:
"""
Expand Down Expand Up @@ -160,7 +161,10 @@ def calculate_genetic_values(self, pop):
# pop.shape = [50,2,10,10]
# self.scaled_effects.shape = [10,10]
# output [50]
return torch.einsum('ijkl,kl->i', self.calculate_scaled_additive_dosages(pop.float()), self.scaled_effects)
if self.intercept:
return torch.einsum('ijkl,kl->i', self.calculate_scaled_additive_dosages(pop.float()), self.scaled_effects) + self.intercept
else:
return torch.einsum('ijkl,kl->i', self.calculate_scaled_additive_dosages(pop.float()), self.scaled_effects)

def calculate_intercept(self):
"""
Expand Down Expand Up @@ -192,7 +196,7 @@ def phenotype(self, pop, h2=None, varE=None):
if h2 is None and varE is None:
raise ValueError("Provide either h2 or varE.")

genetic_values = self.calculate_genetic_values(pop) + self.intercept
genetic_values = self.calculate_genetic_values(pop)

if h2 is not None:
# Calculate varE based on target heritability
Expand Down
Binary file added index_files/figure-commonmark/cell-4-output-1.png
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
12 changes: 8 additions & 4 deletions nbs/02_trait.ipynb
View file

Large diffs are not rendered by default.

116 changes: 82 additions & 34 deletions nbs/03_meiosis.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,31 +70,19 @@
" return crossover_positions\n",
"\n",
"\n",
"def simulate_meiosis(num_chromosomes, chromosome_length, crossover_rate, interference_strength):\n",
"import torch\n",
"\n",
"def simulate_meiosis(num_chromosomes, map_length, num_individuals, num_crossovers):\n",
" \"\"\"\n",
" Simulate meiosis with crossover events and interference.\n",
" \n",
" Parameters:\n",
" num_chromosomes (int): Number of chromatids.\n",
" chromosome_length (float): Length of the chromosome.\n",
" crossover_rate (float): Rate of crossover events.\n",
" interference_strength (float): Strength of interference.\n",
" \n",
" Returns:\n",
" list of torch.Tensor: List of crossover positions for each chromatid.\n",
" This function simulates random crossover events across chromosomes.\n",
" \"\"\"\n",
" chromatid_crossovers = []\n",
" for _ in range(num_chromosomes):\n",
" crossovers = gamma_interference_model(chromosome_length, crossover_rate, interference_strength)\n",
" chromatid_crossovers.append(crossovers)\n",
" return chromatid_crossovers\n",
" return [torch.sort(torch.rand((num_individuals, num_crossovers)) * map_length, dim=-1)[0] for _ in range(num_chromosomes)]\n",
"\n",
"def simulate_gametes(genetic_map, parent_genome):\n",
" \"\"\"\n",
" Simulate the formation of a single gamete given crossover positions, genetic map, and parent genomes.\n",
"\n",
" Parameters:\n",
" crossover_positions (list of torch.Tensor): List of positions of crossover events for each chromosome.\n",
" genetic_map (list of torch.Tensor): List of positions of genetic markers on the chromosomes.\n",
" parent_genome (torch.Tensor): Genomes of the parents. SHAPE: (ploidy, num_chromosomes, num_loci)\n",
"\n",
Expand All @@ -103,39 +91,99 @@
" \"\"\"\n",
" ploidy, num_chromosomes, num_loci = parent_genome.shape\n",
" gamete_genome = torch.zeros((ploidy // 2, num_chromosomes, num_loci), dtype=parent_genome.dtype)\n",
" crossover_positions = simulate_meiosis(num_chromosomes,100.0,1,1)\n",
" crossover_positions = simulate_meiosis(num_chromosomes, 100.0, 1, 1)\n",
"\n",
" for chrom in range(num_chromosomes):\n",
" parent_1 = parent_genome[0, chrom]\n",
" parent_2 = parent_genome[1, chrom]\n",
" \n",
" crossover_sites = crossover_positions[chrom]\n",
" chromatid = parent_1.clone()\n",
" \n",
" if len(crossover_sites) > 0:\n",
" crossover_index = 0\n",
" for i, marker_position in enumerate(genetic_map[chrom]):\n",
" if crossover_index < len(crossover_sites) and marker_position >= crossover_sites[crossover_index]:\n",
" chromatid = parent_2 if chromatid is parent_1 else parent_1\n",
" crossover_index += 1\n",
" gamete_genome[0, chrom, i] = chromatid[i]\n",
" # Create a crossover mask for each chromosome\n",
" crossover_mask = torch.zeros(num_loci, dtype=torch.bool)\n",
" crossover_sites = crossover_positions[chrom].flatten()\n",
" for position in crossover_sites:\n",
" # Find the nearest marker index for each crossover position\n",
" index = torch.argmin(torch.abs(genetic_map[chrom] - position))\n",
" crossover_mask[index:] = ~crossover_mask[index:] # Flip values at and after the crossover index\n",
"\n",
" return gamete_genome\n"
" # Use the mask to select genes from parent 1 or parent 2\n",
" gamete_genome[0, chrom] = torch.where(crossover_mask, parent_genome[1, chrom], parent_genome[0, chrom])\n",
"\n",
" return gamete_genome\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1ca41486",
"metadata": {},
"outputs": [],
"outputs": [
{
"data": {
"text/plain": [
"tensor([[[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1,\n",
" 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1,\n",
" 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 1,\n",
" 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1,\n",
" 0, 0, 1, 1, 1, 0, 1, 1],\n",
" [1, 1, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1,\n",
" 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1,\n",
" 0, 1, 1, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1,\n",
" 1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 1, 1, 0,\n",
" 0, 0, 1, 0, 0, 1, 0, 0],\n",
" [0, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0,\n",
" 1, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0,\n",
" 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0,\n",
" 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1,\n",
" 1, 0, 0, 0, 1, 0, 1, 1],\n",
" [1, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1,\n",
" 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0,\n",
" 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0,\n",
" 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0,\n",
" 1, 1, 0, 1, 0, 0, 0, 0],\n",
" [1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0,\n",
" 0, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 0,\n",
" 0, 1, 0, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1,\n",
" 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0,\n",
" 1, 1, 0, 0, 1, 1, 1, 1],\n",
" [0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 0, 1,\n",
" 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0,\n",
" 1, 0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 0, 1, 1,\n",
" 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0,\n",
" 1, 0, 1, 1, 0, 0, 1, 0],\n",
" [1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
" 1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1,\n",
" 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0,\n",
" 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 0,\n",
" 0, 0, 0, 0, 0, 1, 0, 1],\n",
" [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0,\n",
" 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1,\n",
" 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0,\n",
" 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0,\n",
" 1, 1, 1, 0, 0, 0, 1, 1],\n",
" [1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1,\n",
" 0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1,\n",
" 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0,\n",
" 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0,\n",
" 1, 1, 0, 1, 0, 0, 1, 1],\n",
" [1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0,\n",
" 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0,\n",
" 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1,\n",
" 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0,\n",
" 1, 1, 1, 0, 1, 0, 1, 1]]])"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ploidy = 2\n",
"number_chromosomes = 10\n",
"loci_per_chromosome = 100\n",
"n_founders = 50\n",
"genetic_map = create_random_genetic_map(number_chromosomes,loci_per_chromosome)\n",
"crop_genome = Genome(ploidy, number_chromosomes, loci_per_chromosome, genetic_map)\n",
"founder_pop = create_random_founder_pop(crop_genome , n_founders)"
"founder_pop = create_random_founder_pop(crop_genome , n_founders)\n",
"simulate_gametes(genetic_map, founder_pop[0])"
]
},
{
Expand Down
27 changes: 8 additions & 19 deletions nbs/index.ipynb
View file

Large diffs are not rendered by default.

0 comments on commit 54f7806

Please sign in to comment.