Restore memory efficiency

[XingerTang]

What changed

• segregation probability would store in memory only if segregation outpus is required

Why this change

• To reduce unnecessary memory usage

Notes / Risks

• Small changes only

[XingerTang]

After the change, the memory usage for the combined mode accuracy test dropped from

Allocation results for tests/accuracy_tests/run_accu_test.py::test_accu[combined] at the high watermark

         📦 Total memory allocated: 158.0MiB
         📏 Total allocations: 11
         📊 Histogram of allocation sizes: |▅▂█ ▂|
         🥇 Biggest allocating functions:
                - _read:/Users/xtang3/anaconda3/envs/phase/lib/python3.11/site-packages/numpy/lib/_npyio_impl.py:1046 -> 96.9MiB
                - _read:/Users/xtang3/anaconda3/envs/phase/lib/python3.11/site-packages/numpy/lib/_npyio_impl.py:1046 -> 61.0MiB
                - decode:<frozen codecs>:322 -> 32.1KiB
                - _read:/Users/xtang3/anaconda3/envs/phase/lib/python3.11/site-packages/numpy/lib/_npyio_impl.py:962 -> 15.6KiB
                - assess_peeling:/Users/xtang3/AlphaImpute2/tests/accuracy_tests/run_accu_test.py:140 -> 15.6KiB

to

Allocation results for tests/accuracy_tests/run_accu_test.py::test_accu[combined] at the high watermark

         📦 Total memory allocated: 79.9MiB
         📏 Total allocations: 10
         📊 Histogram of allocation sizes: |▂▂█ ▅|
         🥇 Biggest allocating functions:
                - _read:/Users/xtang3/anaconda3/envs/phase/lib/python3.11/site-packages/numpy/lib/_npyio_impl.py:1046 -> 49.3MiB
                - _read:/Users/xtang3/anaconda3/envs/phase/lib/python3.11/site-packages/numpy/lib/_npyio_impl.py:1046 -> 30.5MiB
                - decode:<frozen codecs>:322 -> 32.1KiB
                - _read:/Users/xtang3/anaconda3/envs/phase/lib/python3.11/site-packages/numpy/lib/_npyio_impl.py:962 -> 15.6KiB
                - assess_peeling:/Users/xtang3/AlphaImpute2/tests/accuracy_tests/run_accu_test.py:134 -> 15.6KiB

[gregorgorjanc]

I personally like to add arg keyword as it makes code clearer to me, but I don't know what's the preferred/popular Python style;) Your call @XingerTang!

[XingerTang]

The whole args might be unnecessary here. I can make seg_outputa variable that is assigned earlier separately; it would be like:

seg_output = bool(args.seg_output)
for ind in pedigree:
  ind.setPeelingView(seg_output)

[gregorgorjanc]

I agree it's not needed. I personally prefer code like ind.setPeelingView(store_out_segregation = bool(args.seg_output)) or ind.setPeelingView(store_out_segregation = seg_output), but that's just my taste. You decide what works best for this codebase!

[XingerTang]

I updated the code to the latter way. Python requires all positional arguments to come before any keyword arguments; I usually do not specify the keyword if it can be positional during development to avoid the ordering issue. But this one only has one argument, so it doesn't matter.

[gregorgorjanc]

Looks like a simple change! Great.

[XingerTang]

In 7ac3f4f, the memory is no longer allocated to the segregation prob output vector if that is not required. While the pytest memray benchmark results are promising, the experiment with the real data is not as effective as expected.

In the early version v0.0.1, the memory usage is:

 Max vmem         = 3.917G
 Max rss          = 2.821G

Before this change (v0.0.4), the memory usage is:

 Max vmem         = 5.503G
 Max rss          = 5.284G

After the commit, without seg_output requested, the memory usage is:

 Max vmem         = 4.792G
 Max rss          = 4.562G

There is a significant decrease after the change, but it is not as much as the difference between v0.0.1 and v0.0.4.

Therefore, in 46d2ca1, I cleared the peeling_view when seg_output is not requested.
This change does not change the memory usage of pytest, but it might be that the data size is not enough, as the code would just make some of the used memory space available for a new memory requirement.

[XingerTang]

After 46d2ca1, the memory usage dropped to

 Max vmem         = 3.722G
 Max rss          = 3.487G

which the Max vmem is smaller than the v0.0.1 usage. However, the Max rss is still higher.

[XingerTang]

A possible source of the difference of Max vmem might be led by the full storage of genotype probabilities of all individuals instead of the individuals with progeny, which was introduced in the v0.0.3 to correct the haplotypes when the haplotypes and genotypes do not match.

dc5595f restored the original probability storage pattern without affecting the haplotype and genotype matching.

The memray profile on pytest still stay the same, but the real data result might be different.

[XingerTang]

The memory usage for dc5595f on the same data gives:

 Max vmem         = 2.850G
 Max rss          = 2.615G

which I'm quite happy to see.

diff against output files shows that the outputs after the changes are exactly the same as the outputs before the changes.

[gregorgorjanc]

@XingerTang this is great testing! I assume all accuracy metrics and runtime are unaffected - good to check. I am asking because I see you are not storing or you are clearing some objects, which is what we want so I guess this is all fine and gives you lower memory footprint, but you seem to have to do some extra work with phase_probabilities too (will that change runtime, probably not). Looking at the code it seems that accuracies should not change.

[XingerTang]

@XingerTang this is great testing! I assume all accuracy metrics and runtime are unaffected - good to check. I am asking because I see you are not storing or you are clearing some objects, which is what we want so I guess this is all fine and gives you lower memory footprint, but you seem to have to do some extra work with phase_probabilities too (will that change runtime, probably not). Looking at the code it seems that accuracies should not change.

All the metrics are the same, and essentially the output files are exactly the same. There are more temporary holders (such as phase_probabilities) instead of big blocks of probabilities, but essentially it does not change anything.

[gregorgorjanc]

All the metrics are the same, and essentially the output files are exactly the same. There are more temporary holders (such as phase_probabilities) instead of big blocks of probabilities, but essentially it does not change anything.

Excellent - all is well then. Well done on cracking this memory issue!

[XingerTang]

There are some issues with GitHub Action runners: https://www.githubstatus.com/

[XingerTang]

I think it's okay we ignore the failing tests (all caused by the GitHub runner issue) and merge

[gregorgorjanc]

@XingerTang I suggest we also add maintainers here and later co-authors.

[gregorgorjanc]

I agree - all the failures are due to outage on GitHub side. Add the maintainers and merge;)