HERA-Team / hera_cal

Library for HERA data reduction, including redundant calibration, absolute calibration, and LST-binning.
MIT License
13 stars 8 forks source link

Omnical memory improvements #923

Closed jsdillon closed 8 months ago

jsdillon commented 8 months ago

Despite some misleading results from %memit, it appears that omnical's memory usage is the root of some of the problems we've been having with RTP.

This PR improves the memory efficiency of omnical. It does so by studiously avoiding copying data unnecessarily, as well as by updating dictionaries with for loops rather than dictionary comprehension, so that we don't have to make a whole new dictionary before deallocating the old one.

It also changes the default value of wgt_func from lambda x: 1. to None and adds special treatment of the None case so that it doesn't multiply by 1 and create new arrays.

It has a negligible effect on performance (and actually might be a tiny bit faster). In my tests on the file_calibration notebook with zen.2460280.48839.sum.uvh5 locally, runtime improved from 5.07 minutes to 4.63 minutes. It improves peak memory usage from 18.8 GB to 12.3 GB, although it should be noted that ~6.5 GB was the memory usage (mostly raw data) going into omnical.

Here's my memray results before:

image

And after:

image

It's possible we can do even better if we use single precision, which I'll explore next.

This PR relies on https://github.com/HERA-Team/linsolve/pull/51

codecov[bot] commented 8 months ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Comparison is base (790428f) 97.17% compared to head (67d0c29) 97.18%.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #923 +/- ## ======================================= Coverage 97.17% 97.18% ======================================= Files 23 23 Lines 10490 10504 +14 ======================================= + Hits 10194 10208 +14 Misses 296 296 ``` | [Flag](https://app.codecov.io/gh/HERA-Team/hera_cal/pull/923/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HERA-Team) | Coverage Δ | | |---|---|---| | [unittests](https://app.codecov.io/gh/HERA-Team/hera_cal/pull/923/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HERA-Team) | `97.18% <100.00%> (+<0.01%)` | :arrow_up: | Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=HERA-Team#carryforward-flags-in-the-pull-request-comment) to find out more.

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.