From ae67a857195f2a4b1156abfb6eec369dd7a8a982 Mon Sep 17 00:00:00 2001 From: XingerTang Date: Fri, 7 Aug 2026 12:30:12 +0100 Subject: [PATCH] Restore the memory efficiency --- .github/workflows/tests.yml | 2 +- docs/source/changelog.rst | 12 +++++ pyproject.toml | 12 ++++- .../Imputation/Heuristic_Peeling.py | 32 +++++++++---- .../Imputation/ImputationIndividual.py | 46 +++++++++++++------ 5 files changed, 79 insertions(+), 25 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index d490e5c..e38fab2 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -37,7 +37,7 @@ jobs: run: python3 -m build - name: Install AlphaImpute2 - run: pip install dist/alphaimpute2-0.0.4-py3-none-any.whl + run: pip install dist/alphaimpute2-0.0.5-py3-none-any.whl - name: Install pytest run: | diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 5245404..886dfa6 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -2,6 +2,18 @@ Changelog ========= +[0.0.5] - 2026-08-07 +==================== + +Maintenance +----------- + +* Restore the memory efficiency with + * Optional segregation probabilities storage and + * Reduced memory usage for genotype probabilities + (:pr:`75`, :user:`XingerTang`, :user:`gregorgorjanc`). + + [0.0.4] - 2026-06-12 ==================== diff --git a/pyproject.toml b/pyproject.toml index 4ff890a..2c979a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,9 +4,17 @@ build-backend = "setuptools.build_meta" [project] name = "AlphaImpute2" -version = "0.0.4" +version = "0.0.5" authors = [ - { name="Andrew Whalen", email="alphagenes.dev@gmail.com" }, + { name="Andrew Whalen" }, + { name="Xinger Tang", email="alphagenes.dev@gmail.com" }, + { name="Rosalind Craddock", email="alphagenes.dev@gmail.com" }, + { name="Yu Zhang", email="alphagenes.dev@gmail.com" }, + { name="Gregor Gorjanc", email="alphagenes.dev@gmail.com" }, +] +maintainers = [ + { name="Xinger Tang", email="alphagenes.dev@gmail.com" }, + { name="Gregor Gorjanc", email="alphagenes.dev@gmail.com" }, ] description = "An imputation software for massive livestock populations." # readme corresponds to the long_description diff --git a/src/alphaimpute2/Imputation/Heuristic_Peeling.py b/src/alphaimpute2/Imputation/Heuristic_Peeling.py index 71d8efa..b1f1c5e 100644 --- a/src/alphaimpute2/Imputation/Heuristic_Peeling.py +++ b/src/alphaimpute2/Imputation/Heuristic_Peeling.py @@ -56,9 +56,7 @@ def runHeuristicPeeling(pedigree, args, final_cutoff=0.3): for ind in pedigree: call_genotypes(ind, final_cutoff, args.error) - # Clear peeling view from all individuals to reduce memory impact. - # for ind in pedigree: - # ind.peeling_view = None + clearPeelingViewsIfNotNeeded(pedigree, args) @profile @@ -127,9 +125,7 @@ def run_integrated_peeling(pedigree, args, final_cutoff=0.3, arrays=None): for ind in ld_for_ped_imputation: ind.restore_original_genotypes() - # Clear peeling view from all individuals to reduce memory impact. - # for ind in pedigree: - # ind.peeling_view = None + clearPeelingViewsIfNotNeeded(pedigree, args) return hd_individuals, ld_for_pop_imputation, ld_for_ped_imputation @@ -179,8 +175,20 @@ def mask_genotypes(mat, mask): def setupHeuristicPeeling(pedigree, args): # Sets the founder anterior values and penetrance value for Heuristic peeling. + + seg_output = bool(args.seg_output) for ind in pedigree: - ind.setPeelingView() + ind.setPeelingView(store_out_segregation=seg_output) + + +def clearPeelingViewsIfNotNeeded(pedigree, args): + if args.seg_output: + return + + # Segregation output is written from the peeling view, so only release it when + # the output will not be requested later. + for ind in pedigree: + ind.peeling_view = None def call_genotypes(ind, final_cutoff, error_rate): @@ -188,6 +196,8 @@ def call_genotypes(ind, final_cutoff, error_rate): # IF NOT, PARENT'S ANTERIOR VALUES MAY NOT BE CORRECTLY SET. # FIX: RUN FINAL ROUND OF PEEL DOWN AT THE END. + phase_probabilities = None + if ind.peeling_view.has_offspring: if ind.sire is not None and ind.dam is not None: # For individuals with offspring, recalculate the anterior value, and then set the genotypes. @@ -203,6 +213,8 @@ def call_genotypes(ind, final_cutoff, error_rate): # If no parents, directly set genotypes from the penetrance field. ind.peeling_view.setGenotypesAll(final_cutoff) + phase_probabilities = ind.peeling_view.genotypeProbabilities + else: nLoci = len(ind.genotypes) @@ -224,6 +236,7 @@ def call_genotypes(ind, final_cutoff, error_rate): ind.peeling_view.setGenotypesFromGenotypeProbabilities( genotypeProbabilities, final_cutoff ) + phase_probabilities = genotypeProbabilities if final_cutoff < 0.5: nLoci = len(ind.genotypes) @@ -231,7 +244,7 @@ def call_genotypes(ind, final_cutoff, error_rate): if ind.genotypes[i] == 1: if ind.haplotypes[0][i] + ind.haplotypes[1][i] != 1: # correct the genotype-haplotype mismatch - phase_probs = ind.peeling_view.genotypeProbabilities[1:3, i] + phase_probs = phase_probabilities[1:3, i] phase = np.argmax(phase_probs) ind.haplotypes[0][i] = phase ind.haplotypes[1][i] = 1 - phase @@ -524,7 +537,8 @@ def setSegregation(ind, sire, dam): pointEstimates, 1.0 / nLoci * ind.map_length ) # This is where different map lengths could be added. - ind.out_segregation = smoothedEstimates.copy() + if ind.store_out_segregation: + ind.out_segregation = smoothedEstimates.copy() # Then set the segregation values for the individual. ind.segregation[0][:] = smoothedEstimates[2, :] + smoothedEstimates[3, :] diff --git a/src/alphaimpute2/Imputation/ImputationIndividual.py b/src/alphaimpute2/Imputation/ImputationIndividual.py index 05e4c39..038408f 100644 --- a/src/alphaimpute2/Imputation/ImputationIndividual.py +++ b/src/alphaimpute2/Imputation/ImputationIndividual.py @@ -189,8 +189,7 @@ def setupIndividual(self): # self.setPhasingView() # self.setPeelingView() - def setPeelingView(self): - # Set the + def setPeelingView(self, store_out_segregation=False): if self.genotypes is None or self.haplotypes is None: raise ValueError( "In order to create a jit_Peeling_Individual both the genotypes and haplotypes need to be created." @@ -214,6 +213,7 @@ def setPeelingView(self): has_parents, nLoci, self.map_length, + store_out_segregation, ) def setPhasingView(self): @@ -446,6 +446,7 @@ def get_example_phasing_individual(): (np.array([0, 1], dtype=np.float32), np.array([0], dtype=np.float32)) ) spec["out_segregation"] = optional(float32[:, :]) +spec["store_out_segregation"] = boolean spec["newPosterior"] = optional(numba.typeof([np.full((4, 100), 0, dtype=np.float32)])) @@ -468,7 +469,15 @@ class jit_Peeling_Individual(object): """ def __init__( - self, idn, genotypes, haplotypes, has_offspring, has_parents, nLoci, map_length + self, + idn, + genotypes, + haplotypes, + has_offspring, + has_parents, + nLoci, + map_length, + store_out_segregation, ): self.nLoci = nLoci self.idn = idn @@ -476,6 +485,7 @@ def __init__( self.haplotypes = haplotypes self.map_length = map_length + self.store_out_segregation = store_out_segregation # Initial value for segregation is .5 to represent uncertainty between haplotype inheritance. self.segregation = ( @@ -483,24 +493,33 @@ def __init__( np.full(nLoci, 0.5, dtype=np.float32), ) - self.out_segregation = np.full((4, nLoci), 0.25, dtype=np.float32) + if self.store_out_segregation: + self.out_segregation = np.full((4, nLoci), 0.25, dtype=np.float32) + else: + self.out_segregation = None # Create the posterior terms. self.has_offspring = has_offspring self.has_parents = has_parents - self.original_genotypes = self.genotypes.copy() - self.original_haplotypes = ( - self.haplotypes[0].copy(), - self.haplotypes[1].copy(), - ) + if self.has_offspring: + self.original_genotypes = self.genotypes.copy() + self.original_haplotypes = ( + self.haplotypes[0].copy(), + self.haplotypes[1].copy(), + ) + self.posterior = np.full((4, nLoci), 1, dtype=np.float32) + self.genotypeProbabilities = np.full((4, nLoci), 1, dtype=np.float32) + else: + self.original_genotypes = self.genotypes + self.original_haplotypes = self.haplotypes + self.posterior = np.full((0, 0), 1, dtype=np.float32) + self.genotypeProbabilities = np.full((0, 0), 1, dtype=np.float32) + self.anterior_availible = False self.posterior_open = False - self.posterior = np.full((4, nLoci), 1, dtype=np.float32) self.anterior = np.full((0, 0), 1, dtype=np.float32) - self.genotypeProbabilities = np.full((4, nLoci), 1, dtype=np.float32) - self.newPosterior = None # Current state indicates which genotype combination an individual is set to, and the cutoff value used for calling their genotypes. @@ -706,7 +725,8 @@ def setGenotypesFromPeelingData( def setGenotypesFromGenotypeProbabilities(self, finalGenotypes, cutoff): nLoci = self.nLoci normalize(finalGenotypes) - self.genotypeProbabilities[:, :] = finalGenotypes + if self.has_offspring: + self.genotypeProbabilities[:, :] = finalGenotypes # set genotypes/haplotypes from this value. for i in range(nLoci):