From 1d01bf5f8caa94414fe8e7377ce08d3229942b0b Mon Sep 17 00:00:00 2001 From: cjgo Date: Tue, 11 Jun 2024 05:04:24 -0500 Subject: [PATCH] packaged multi_trait generator into function --- README.md | 73 ++++-------- chewc/_modidx.py | 5 +- chewc/cross.py | 10 +- chewc/trait.py | 42 ++++++- nbs/02_trait.ipynb | 128 ++++++++++------------ nbs/{04_crossing.ipynb => 04_cross.ipynb} | 21 ++-- nbs/index.ipynb | 106 ++++++++---------- 7 files changed, 178 insertions(+), 207 deletions(-) rename nbs/{04_crossing.ipynb => 04_cross.ipynb} (96%) diff --git a/README.md b/README.md index dcae4cb..ffedc62 100644 --- a/README.md +++ b/README.md @@ -33,64 +33,31 @@ pip install chewc First, define the genome of your crop ``` python -# import random +import torch -# ploidy = 2 -# number_chromosomes = 10 -# loci_per_chromosome = 100 -# genetic_map = create_random_genetic_map(number_chromosomes,loci_per_chromosome) -# crop_genome = Genome(ploidy, number_chromosomes, loci_per_chromosome, genetic_map) +ploidy = 2 +n_chr = 10 +n_loci = 100 +n_Ind = 333 +g = Genome(ploidy, n_chr, n_loci) +population = Population() +population.create_random_founder_population(g, n_founders=n_Ind) +init_pop = population.get_dosages().float() # gets allele dosage for calculating trait values -# n_founders = 500 -# founder_pop = create_random_founder_pop(crop_genome , n_founders) -# sim_param = SimParam -# sim_param.founder_pop = founder_pop -# sim_param.genome = crop_genome +# multi_traits -# #add a single additive trait -# qtl_loci = 20 -# qtl_map = select_qtl_loci(qtl_loci,sim_param.genome) +target_means = torch.tensor([0, 5]) +target_vars = torch.tensor([1, 1]) # Note: I'm assuming you want a variance of 1 for the second trait +correlation_values = [ + [1.0, 0.2], + [0.2, 1.0], + ] -# ta = TraitA(qtl_map,sim_param,0, 1) -# ta.sample_initial_effects() -# ta.scale_genetic_effects() -# ta.calculate_intercept() +correlated_traits = corr_traits(g, init_pop, target_means, target_vars, correlation_values) +``` + Created genetic map - - - -# # Ensure sim_param.device is defined and correct -# device = sim_param.device - -# years = 20 -# current_pop = founder_pop.to(device) -# pmean = [] -# pvar = [] - -# for _ in range(years): -# # phenotype current pop -# TOPK = 10 -# new_pop = [] -# pheno = ta.phenotype(current_pop, h2=0.14).to(device) -# topk = torch.topk(pheno, TOPK).indices.to(device) - -# for _ in range(200): -# sampled_indices = torch.multinomial(torch.ones(topk.size(0), device=device), 2, replacement=False) -# sampled_parents = topk[sampled_indices] -# m, f = current_pop[sampled_parents[0]], current_pop[sampled_parents[1]] -# new_pop.append(make_cross(sim_param, m, f).to(device)) - -# current_pop = torch.stack(new_pop).to(device) -# pmean.append(ta.calculate_genetic_values(current_pop).mean().item()) -# pvar.append(ta.calculate_genetic_values(current_pop).var().item()) - -# pmean_normalized = torch.tensor(pmean, device=device) / max(pmean) -# pvar_normalized = torch.tensor(pvar, device=device) / max(pvar) - -# plt.scatter(range(len(pmean_normalized)), pmean_normalized.cpu()) -# plt.scatter(range(len(pvar_normalized)), pvar_normalized.cpu()) -# plt.show() -``` + NameError: name 'init_pop' is not defined diff --git a/chewc/_modidx.py b/chewc/_modidx.py index 583b744..4b87bf1 100644 --- a/chewc/_modidx.py +++ b/chewc/_modidx.py @@ -33,8 +33,8 @@ 'chewc.core.PopulationDataset.__init__': ('core.html#populationdataset.__init__', 'chewc/core.py'), 'chewc.core.PopulationDataset.__len__': ('core.html#populationdataset.__len__', 'chewc/core.py'), 'chewc.core.create_population_dataloader': ('core.html#create_population_dataloader', 'chewc/core.py')}, - 'chewc.cross': { 'chewc.cross.double_haploid': ('crossing.html#double_haploid', 'chewc/cross.py'), - 'chewc.cross.random_crosses': ('crossing.html#random_crosses', 'chewc/cross.py')}, + 'chewc.cross': { 'chewc.cross.double_haploid': ('cross.html#double_haploid', 'chewc/cross.py'), + 'chewc.cross.random_crosses': ('cross.html#random_crosses', 'chewc/cross.py')}, 'chewc.crossing': { 'chewc.crossing.double_haploid': ('crossing.html#double_haploid', 'chewc/crossing.py'), 'chewc.crossing.random_crosses': ('crossing.html#random_crosses', 'chewc/crossing.py')}, 'chewc.meiosis': { 'chewc.meiosis.poisson_crossing_over': ('meiosis.html#poisson_crossing_over', 'chewc/meiosis.py'), @@ -42,4 +42,5 @@ 'chewc.trait': { 'chewc.trait.TraitA': ('trait.html#traita', 'chewc/trait.py'), 'chewc.trait.TraitA.__init__': ('trait.html#traita.__init__', 'chewc/trait.py'), 'chewc.trait.TraitA.forward': ('trait.html#traita.forward', 'chewc/trait.py'), + 'chewc.trait.corr_traits': ('trait.html#corr_traits', 'chewc/trait.py'), 'chewc.trait.select_qtl_loci': ('trait.html#select_qtl_loci', 'chewc/trait.py')}}} diff --git a/chewc/cross.py b/chewc/cross.py index 637582e..df1cd08 100644 --- a/chewc/cross.py +++ b/chewc/cross.py @@ -1,9 +1,13 @@ -# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_crossing.ipynb. +# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_cross.ipynb. # %% auto 0 __all__ = ['random_crosses', 'double_haploid'] -# %% ../nbs/04_crossing.ipynb 4 +# %% ../nbs/04_cross.ipynb 3 +from .core import * +from .trait import * +from .meiosis import * +import torch def random_crosses( genome: Genome, parent_haplotypes: torch.Tensor, n_crosses: int) -> torch.Tensor: """ Generate random crosses from a set of parent haplotypes. @@ -41,7 +45,7 @@ def random_crosses( genome: Genome, parent_haplotypes: torch.Tensor, n_crosses: return progeny_haplotypes -# %% ../nbs/04_crossing.ipynb 7 +# %% ../nbs/04_cross.ipynb 6 def double_haploid(genome: Genome, parent_haplotypes: torch.Tensor) -> torch.Tensor: """ Generate doubled haploid individuals from a set of parent haplotypes. diff --git a/chewc/trait.py b/chewc/trait.py index 503be43..51f0919 100644 --- a/chewc/trait.py +++ b/chewc/trait.py @@ -1,7 +1,7 @@ # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_trait.ipynb. # %% auto 0 -__all__ = ['select_qtl_loci', 'TraitA'] +__all__ = ['select_qtl_loci', 'TraitA', 'corr_traits'] # %% ../nbs/02_trait.ipynb 3 from .core import * @@ -88,3 +88,43 @@ def forward(self, dosages: torch.Tensor, h2: Optional[float] = None, varE: Optio breeding_values += env_noise return breeding_values + +# %% ../nbs/02_trait.ipynb 4 +def corr_traits(genome, init_pop, target_means, target_vars, correlation_matrix): + + n_chr, n_loci = genome.genetic_map.shape + n_traits = target_means.shape[0] + corA = torch.tensor(correlation_matrix) + L = torch.linalg.cholesky(corA) + + # Sample initial additive effects from a standard normal distribution + uncorrelated_effects = torch.randn(n_chr, n_loci, n_traits) + + # Reshape for proper multiplication with Cholesky factor + uncorrelated_effects = uncorrelated_effects.reshape(n_chr * n_loci, n_traits) + + # Introduce correlation FIRST + correlated_effects = torch.matmul(L, uncorrelated_effects.T).T + correlated_effects = correlated_effects.reshape(n_chr, n_loci, n_traits) + + # Calculate unscaled breeding values using CORRELATED effects + unscaled_bvs = torch.einsum('ijk,lij->lk', correlated_effects, init_pop) + unscaled_var = unscaled_bvs.var(dim=0) + unscaled_mean = unscaled_bvs.mean(dim=0) + trait_intercepts = [] + + # Scale correlated effects and calculate intercepts + for i in range(n_traits): + scaling_factor = torch.sqrt(target_vars[i] / unscaled_var[i]) + correlated_effects[:, :, i] *= scaling_factor # Scale the CORRELATED effects + trait_intercepts.append(target_means[i] - (unscaled_mean[i] * scaling_factor)) + + # Now we have the additive marker effects and intercepts to calculate breeding values + trait_intercepts = torch.tensor(trait_intercepts) +# scaled_bvs = torch.einsum('ijk,lij->lk', correlated_effects, init_pop) + + + Traits = [] + for i in range(len(target_means)): + Traits.append(TraitA(genome, correlated_effects[:,:,i], trait_intercepts[i])) + return Traits diff --git a/nbs/02_trait.ipynb b/nbs/02_trait.ipynb index 16fa6ac..97b3939 100644 --- a/nbs/02_trait.ipynb +++ b/nbs/02_trait.ipynb @@ -125,6 +125,54 @@ " return breeding_values" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "2f4b77c3", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def corr_traits(genome, init_pop, target_means, target_vars, correlation_matrix):\n", + " \n", + " n_chr, n_loci = genome.genetic_map.shape\n", + " n_traits = target_means.shape[0]\n", + " corA = torch.tensor(correlation_matrix)\n", + " L = torch.linalg.cholesky(corA)\n", + "\n", + " # Sample initial additive effects from a standard normal distribution\n", + " uncorrelated_effects = torch.randn(n_chr, n_loci, n_traits)\n", + "\n", + " # Reshape for proper multiplication with Cholesky factor\n", + " uncorrelated_effects = uncorrelated_effects.reshape(n_chr * n_loci, n_traits)\n", + "\n", + " # Introduce correlation FIRST\n", + " correlated_effects = torch.matmul(L, uncorrelated_effects.T).T\n", + " correlated_effects = correlated_effects.reshape(n_chr, n_loci, n_traits)\n", + "\n", + " # Calculate unscaled breeding values using CORRELATED effects\n", + " unscaled_bvs = torch.einsum('ijk,lij->lk', correlated_effects, init_pop)\n", + " unscaled_var = unscaled_bvs.var(dim=0)\n", + " unscaled_mean = unscaled_bvs.mean(dim=0)\n", + " trait_intercepts = []\n", + "\n", + " # Scale correlated effects and calculate intercepts\n", + " for i in range(n_traits):\n", + " scaling_factor = torch.sqrt(target_vars[i] / unscaled_var[i])\n", + " correlated_effects[:, :, i] *= scaling_factor # Scale the CORRELATED effects\n", + " trait_intercepts.append(target_means[i] - (unscaled_mean[i] * scaling_factor))\n", + "\n", + " # Now we have the additive marker effects and intercepts to calculate breeding values\n", + " trait_intercepts = torch.tensor(trait_intercepts)\n", + "# scaled_bvs = torch.einsum('ijk,lij->lk', correlated_effects, init_pop)\n", + " \n", + " \n", + " Traits = []\n", + " for i in range(len(target_means)):\n", + " Traits.append(TraitA(genome, correlated_effects[:,:,i], trait_intercepts[i]))\n", + " return Traits" + ] + }, { "cell_type": "code", "execution_count": null, @@ -151,61 +199,19 @@ "population.create_random_founder_population(g, n_founders=n_Ind)\n", "init_pop = population.get_dosages().float() # gets allele dosage for calculating trait values\n", "\n", - "# Trait A Class\n", + "# multi_traits\n", + "\n", + "\n", "target_means = torch.tensor([0, 5, 20])\n", "target_vars = torch.tensor([1, 1, 0.5]) # Note: I'm assuming you want a variance of 1 for the second trait\n", - "n_traits = target_means.shape[0]\n", "correlation_values = [\n", - " [1.0, 0.2, 0.58],\n", - " [0.2, 1.0, -0.37],\n", - " [0.58, -0.37, 1.0],\n", - "]\n", - "\n", - "corA = torch.tensor(correlation_values)\n", - "L = torch.linalg.cholesky(corA)\n", - "\n", - "# Sample initial additive effects from a standard normal distribution\n", - "uncorrelated_effects = torch.randn(n_chr, n_loci, n_traits)\n", - "\n", - "# Reshape for proper multiplication with Cholesky factor\n", - "uncorrelated_effects = uncorrelated_effects.reshape(n_chr * n_loci, n_traits)\n", - "\n", - "# Introduce correlation FIRST\n", - "correlated_effects = torch.matmul(L, uncorrelated_effects.T).T\n", - "correlated_effects = correlated_effects.reshape(n_chr, n_loci, n_traits)\n", - "\n", - "# Calculate unscaled breeding values using CORRELATED effects\n", - "unscaled_bvs = torch.einsum('ijk,lij->lk', correlated_effects, init_pop)\n", - "unscaled_var = unscaled_bvs.var(dim=0)\n", - "unscaled_mean = unscaled_bvs.mean(dim=0)\n", - "trait_intercepts = []\n", - "\n", - "# Scale correlated effects and calculate intercepts\n", - "for i in range(n_traits):\n", - " scaling_factor = torch.sqrt(target_vars[i] / unscaled_var[i])\n", - " correlated_effects[:, :, i] *= scaling_factor # Scale the CORRELATED effects\n", - " trait_intercepts.append(target_means[i] - (unscaled_mean[i] * scaling_factor))\n", - "\n", - "# Now we have the additive marker effects and intercepts to calculate breeding values\n", - "trait_intercepts = torch.tensor(trait_intercepts)\n", - "scaled_bvs = torch.einsum('ijk,lij->lk', correlated_effects, init_pop)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6108ca7c", - "metadata": {}, - "outputs": [], - "source": [ + " [1.0, 0.2, 0.58],\n", + " [0.2, 1.0, -0.37],\n", + " [0.58, -0.37, 1.0],\n", + " ]\n", "\n", - "ta = TraitA(g, correlated_effects[:,:,0], trait_intercepts[0])\n", - "tb = TraitA(g, correlated_effects[:,:,1], trait_intercepts[1])\n", - "tc = TraitA(g, correlated_effects[:,:,2], trait_intercepts[2])\n", "\n", - "ta_values = ta(init_pop,h2=1)\n", - "tb_values = tb(init_pop,h2=1)\n", - "tc_values = tc(init_pop,h2=.2)" + "correlated_traits = corr_traits(g, init_pop, target_means, target_vars, correlation_values)" ] }, { @@ -213,27 +219,7 @@ "execution_count": null, "id": "f836f5b1", "metadata": {}, - "outputs": [ - { - "ename": "JSONDecodeError", - "evalue": "Extra data: line 396 column 1 (char 159937)", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mJSONDecodeError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[4], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m#| hide\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnbdev\u001b[39;00m; nbdev\u001b[38;5;241m.\u001b[39mnbdev_export()\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/site-packages/fastcore/script.py:110\u001b[0m, in \u001b[0;36mcall_parse.._f\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 107\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(func)\n\u001b[1;32m 108\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_f\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m 109\u001b[0m mod \u001b[38;5;241m=\u001b[39m inspect\u001b[38;5;241m.\u001b[39mgetmodule(inspect\u001b[38;5;241m.\u001b[39mcurrentframe()\u001b[38;5;241m.\u001b[39mf_back)\n\u001b[0;32m--> 110\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m mod: \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 111\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m SCRIPT_INFO\u001b[38;5;241m.\u001b[39mfunc \u001b[38;5;129;01mand\u001b[39;00m mod\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;241m==\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m: SCRIPT_INFO\u001b[38;5;241m.\u001b[39mfunc \u001b[38;5;241m=\u001b[39m func\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 112\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(sys\u001b[38;5;241m.\u001b[39margv)\u001b[38;5;241m>\u001b[39m\u001b[38;5;241m1\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m sys\u001b[38;5;241m.\u001b[39margv[\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m==\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m: sys\u001b[38;5;241m.\u001b[39margv\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;241m1\u001b[39m)\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/site-packages/nbdev/doclinks.py:142\u001b[0m, in \u001b[0;36mnbdev_export\u001b[0;34m(path, procs, **kwargs)\u001b[0m\n\u001b[1;32m 140\u001b[0m procs \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28mgetattr\u001b[39m(nbdev\u001b[38;5;241m.\u001b[39mexport, p) \u001b[38;5;28;01mfor\u001b[39;00m p \u001b[38;5;129;01min\u001b[39;00m L(procs)]\n\u001b[1;32m 141\u001b[0m files \u001b[38;5;241m=\u001b[39m nbglob(path\u001b[38;5;241m=\u001b[39mpath, as_path\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\u001b[38;5;241m.\u001b[39msorted(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mname\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m--> 142\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m files: \u001b[43mnb_export\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprocs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprocs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 143\u001b[0m add_init(get_config()\u001b[38;5;241m.\u001b[39mlib_path)\n\u001b[1;32m 144\u001b[0m _build_modidx()\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/site-packages/nbdev/export.py:67\u001b[0m, in \u001b[0;36mnb_export\u001b[0;34m(nbname, lib_path, procs, debug, mod_maker, name)\u001b[0m\n\u001b[1;32m 65\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m lib_path \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m: lib_path \u001b[38;5;241m=\u001b[39m get_config()\u001b[38;5;241m.\u001b[39mlib_path\n\u001b[1;32m 66\u001b[0m exp \u001b[38;5;241m=\u001b[39m ExportModuleProc()\n\u001b[0;32m---> 67\u001b[0m nb \u001b[38;5;241m=\u001b[39m \u001b[43mNBProcessor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnbname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[43mexp\u001b[49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43mL\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprocs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdebug\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdebug\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 68\u001b[0m nb\u001b[38;5;241m.\u001b[39mprocess()\n\u001b[1;32m 69\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m mod,cells \u001b[38;5;129;01min\u001b[39;00m exp\u001b[38;5;241m.\u001b[39mmodules\u001b[38;5;241m.\u001b[39mitems():\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/site-packages/nbdev/process.py:93\u001b[0m, in \u001b[0;36mNBProcessor.__init__\u001b[0;34m(self, path, procs, nb, debug, rm_directives, process)\u001b[0m\n\u001b[1;32m 92\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, path\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, procs\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, nb\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, debug\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, rm_directives\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, process\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m):\n\u001b[0;32m---> 93\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnb \u001b[38;5;241m=\u001b[39m \u001b[43mread_nb\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mif\u001b[39;00m nb \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m nb\n\u001b[1;32m 94\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlang \u001b[38;5;241m=\u001b[39m nb_lang(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnb)\n\u001b[1;32m 95\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m cell \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnb\u001b[38;5;241m.\u001b[39mcells: cell\u001b[38;5;241m.\u001b[39mdirectives_ \u001b[38;5;241m=\u001b[39m extract_directives(cell, remove\u001b[38;5;241m=\u001b[39mrm_directives, lang\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlang)\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/site-packages/execnb/nbio.py:57\u001b[0m, in \u001b[0;36mread_nb\u001b[0;34m(path)\u001b[0m\n\u001b[1;32m 55\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mread_nb\u001b[39m(path):\n\u001b[1;32m 56\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mReturn notebook at `path`\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m---> 57\u001b[0m res \u001b[38;5;241m=\u001b[39m dict2nb(\u001b[43m_read_json\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mencoding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mutf-8\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m)\n\u001b[1;32m 58\u001b[0m res[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpath_\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mstr\u001b[39m(path)\n\u001b[1;32m 59\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m res\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/site-packages/execnb/nbio.py:16\u001b[0m, in \u001b[0;36m_read_json\u001b[0;34m(self, encoding, errors)\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_read_json\u001b[39m(\u001b[38;5;28mself\u001b[39m, encoding\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, errors\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mloads\u001b[49m\u001b[43m(\u001b[49m\u001b[43mPath\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread_text\u001b[49m\u001b[43m(\u001b[49m\u001b[43mencoding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mencoding\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43merrors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43merrors\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/json/__init__.py:346\u001b[0m, in \u001b[0;36mloads\u001b[0;34m(s, cls, object_hook, parse_float, parse_int, parse_constant, object_pairs_hook, **kw)\u001b[0m\n\u001b[1;32m 341\u001b[0m s \u001b[38;5;241m=\u001b[39m s\u001b[38;5;241m.\u001b[39mdecode(detect_encoding(s), \u001b[38;5;124m'\u001b[39m\u001b[38;5;124msurrogatepass\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m 343\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\u001b[38;5;28mcls\u001b[39m \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m object_hook \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m\n\u001b[1;32m 344\u001b[0m parse_int \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m parse_float \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m\n\u001b[1;32m 345\u001b[0m parse_constant \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m object_pairs_hook \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m kw):\n\u001b[0;32m--> 346\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_default_decoder\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdecode\u001b[49m\u001b[43m(\u001b[49m\u001b[43ms\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 347\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mcls\u001b[39m \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 348\u001b[0m \u001b[38;5;28mcls\u001b[39m \u001b[38;5;241m=\u001b[39m JSONDecoder\n", - "File \u001b[0;32m~/miniconda3/envs/breeding/lib/python3.11/json/decoder.py:340\u001b[0m, in \u001b[0;36mJSONDecoder.decode\u001b[0;34m(self, s, _w)\u001b[0m\n\u001b[1;32m 338\u001b[0m end \u001b[38;5;241m=\u001b[39m _w(s, end)\u001b[38;5;241m.\u001b[39mend()\n\u001b[1;32m 339\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m end \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mlen\u001b[39m(s):\n\u001b[0;32m--> 340\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m JSONDecodeError(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExtra data\u001b[39m\u001b[38;5;124m\"\u001b[39m, s, end)\n\u001b[1;32m 341\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m obj\n", - "\u001b[0;31mJSONDecodeError\u001b[0m: Extra data: line 396 column 1 (char 159937)" - ] - } - ], + "outputs": [], "source": [ "#| hide\n", "import nbdev; nbdev.nbdev_export()" diff --git a/nbs/04_crossing.ipynb b/nbs/04_cross.ipynb similarity index 96% rename from nbs/04_crossing.ipynb rename to nbs/04_cross.ipynb index a899205..aee19e9 100644 --- a/nbs/04_crossing.ipynb +++ b/nbs/04_cross.ipynb @@ -21,25 +21,12 @@ "#| default_exp cross" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "42554857", - "metadata": {}, - "outputs": [], - "source": [ - "from chewc.core import *\n", - "from chewc.trait import *\n", - "from chewc.meiosis import *\n", - "import torch" - ] - }, { "cell_type": "markdown", "id": "df6bef44", "metadata": {}, "source": [ - "## Crossing\n", + "## Cross\n", "> Using gamete operations to produce new individuals" ] }, @@ -51,6 +38,12 @@ "outputs": [], "source": [ "#| export\n", + "\n", + "\n", + "from chewc.core import *\n", + "from chewc.trait import *\n", + "from chewc.meiosis import *\n", + "import torch\n", "def random_crosses( genome: Genome, parent_haplotypes: torch.Tensor, n_crosses: int) -> torch.Tensor:\n", " \"\"\"\n", " Generate random crosses from a set of parent haplotypes.\n", diff --git a/nbs/index.ipynb b/nbs/index.ipynb index fdde9ac..2fe0b0f 100644 --- a/nbs/index.ipynb +++ b/nbs/index.ipynb @@ -17,10 +17,10 @@ "outputs": [], "source": [ "#| hide\n", - "# from chewc.core import *\n", - "# from chewc.trait import *\n", - "# from chewc.meiosis import *\n", - "# from chewc.crossing import *\n", + "from chewc.core import *\n", + "from chewc.trait import *\n", + "from chewc.meiosis import *\n", + "from chewc.cross import *\n", "\n", "\n", "import torch\n", @@ -90,68 +90,51 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Created genetic map\n" + ] + }, + { + "ename": "NameError", + "evalue": "name 'init_pop' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[3], line 23\u001b[0m\n\u001b[1;32m 16\u001b[0m target_vars \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor([\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m]) \u001b[38;5;66;03m# Note: I'm assuming you want a variance of 1 for the second trait\u001b[39;00m\n\u001b[1;32m 17\u001b[0m correlation_values \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 18\u001b[0m [\u001b[38;5;241m1.0\u001b[39m, \u001b[38;5;241m0.2\u001b[39m],\n\u001b[1;32m 19\u001b[0m [\u001b[38;5;241m0.2\u001b[39m, \u001b[38;5;241m1.0\u001b[39m],\n\u001b[1;32m 20\u001b[0m ]\n\u001b[0;32m---> 23\u001b[0m correlated_traits \u001b[38;5;241m=\u001b[39m \u001b[43mcorr_traits\u001b[49m\u001b[43m(\u001b[49m\u001b[43mg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget_means\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcorrelation_values\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m/mnt/c/Users/cltng/OneDrive/chatgpt/chewc/trait.py:111\u001b[0m, in \u001b[0;36mcorr_traits\u001b[0;34m(genome, target_means, target_vars, correlation_matrix)\u001b[0m\n\u001b[1;32m 108\u001b[0m correlated_effects \u001b[38;5;241m=\u001b[39m correlated_effects\u001b[38;5;241m.\u001b[39mreshape(n_chr, n_loci, n_traits)\n\u001b[1;32m 110\u001b[0m \u001b[38;5;66;03m# Calculate unscaled breeding values using CORRELATED effects\u001b[39;00m\n\u001b[0;32m--> 111\u001b[0m unscaled_bvs \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39meinsum(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mijk,lij->lk\u001b[39m\u001b[38;5;124m'\u001b[39m, correlated_effects, \u001b[43minit_pop\u001b[49m)\n\u001b[1;32m 112\u001b[0m unscaled_var \u001b[38;5;241m=\u001b[39m unscaled_bvs\u001b[38;5;241m.\u001b[39mvar(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m 113\u001b[0m unscaled_mean \u001b[38;5;241m=\u001b[39m unscaled_bvs\u001b[38;5;241m.\u001b[39mmean(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\n", + "\u001b[0;31mNameError\u001b[0m: name 'init_pop' is not defined" + ] + } + ], "source": [ - "# import random\n", - "\n", - "# ploidy = 2\n", - "# number_chromosomes = 10\n", - "# loci_per_chromosome = 100\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", - "\n", - "# n_founders = 500\n", - "# founder_pop = create_random_founder_pop(crop_genome , n_founders)\n", - "# sim_param = SimParam\n", - "# sim_param.founder_pop = founder_pop\n", - "# sim_param.genome = crop_genome\n", - "\n", - "\n", - "# #add a single additive trait\n", - "# qtl_loci = 20\n", - "# qtl_map = select_qtl_loci(qtl_loci,sim_param.genome)\n", - "\n", - "# ta = TraitA(qtl_map,sim_param,0, 1)\n", - "# ta.sample_initial_effects()\n", - "# ta.scale_genetic_effects()\n", - "# ta.calculate_intercept()\n", - "\n", - "\n", - "\n", - "\n", - "\n", + "import torch\n", "\n", - "# # Ensure sim_param.device is defined and correct\n", - "# device = sim_param.device\n", + "ploidy = 2\n", + "n_chr = 10\n", + "n_loci = 100\n", + "n_Ind = 333\n", + "g = Genome(ploidy, n_chr, n_loci)\n", + "population = Population()\n", + "population.create_random_founder_population(g, n_founders=n_Ind)\n", + "init_pop = population.get_dosages().float() # gets allele dosage for calculating trait values\n", "\n", - "# years = 20\n", - "# current_pop = founder_pop.to(device)\n", - "# pmean = []\n", - "# pvar = []\n", + "# multi_traits\n", "\n", - "# for _ in range(years):\n", - "# # phenotype current pop\n", - "# TOPK = 10\n", - "# new_pop = []\n", - "# pheno = ta.phenotype(current_pop, h2=0.14).to(device)\n", - "# topk = torch.topk(pheno, TOPK).indices.to(device)\n", "\n", - "# for _ in range(200):\n", - "# sampled_indices = torch.multinomial(torch.ones(topk.size(0), device=device), 2, replacement=False)\n", - "# sampled_parents = topk[sampled_indices]\n", - "# m, f = current_pop[sampled_parents[0]], current_pop[sampled_parents[1]]\n", - "# new_pop.append(make_cross(sim_param, m, f).to(device))\n", - " \n", - "# current_pop = torch.stack(new_pop).to(device)\n", - "# pmean.append(ta.calculate_genetic_values(current_pop).mean().item())\n", - "# pvar.append(ta.calculate_genetic_values(current_pop).var().item())\n", + "target_means = torch.tensor([0, 5])\n", + "target_vars = torch.tensor([1, 1]) # Note: I'm assuming you want a variance of 1 for the second trait\n", + "correlation_values = [\n", + " [1.0, 0.2],\n", + " [0.2, 1.0],\n", + " ]\n", "\n", - "# pmean_normalized = torch.tensor(pmean, device=device) / max(pmean)\n", - "# pvar_normalized = torch.tensor(pvar, device=device) / max(pvar)\n", "\n", - "# plt.scatter(range(len(pmean_normalized)), pmean_normalized.cpu())\n", - "# plt.scatter(range(len(pvar_normalized)), pvar_normalized.cpu())\n", - "# plt.show()" + "correlated_traits = corr_traits(g, init_pop, target_means, target_vars, correlation_values)" ] }, { @@ -159,10 +142,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "\n", - "\n" - ] + "source": [] } ], "metadata": {