Restore memory efficiency

.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: |

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
 ====================
 

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

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,15 +175,29 @@ 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):
     # NOTE: THIS WORKS BUT REQUIRES SETTING THESE IN PEDIGREE "PEEL DOWN" ORDER
     # 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,14 +236,15 @@ 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)
         for i in range(nLoci):
             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, :]

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,39 +469,57 @@ 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
         self.genotypes = genotypes
         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 = (
             np.full(nLoci, 0.5, dtype=np.float32),
             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):