Improve memory usage in uvfits reader

CHANGELOG.md

@@ -6,6 +6,10 @@ All notable changes to this project will be documented in this file.
 ### Changed
 - Added improved handling of a few warning messages from h5py and numpy.
 
+### Fixed
+- Significant reduction in memory when reading UVFITS files, particularly when
+reading the entire file (no selection on read).
+
 ## [3.2.5] - 2025-12-11
 
 ### Changed

src/pyuvdata/parameter.py

@@ -1183,7 +1183,7 @@ def on_moon(self):
         """Detect if this location is on the moon."""
         # figure out if this is on the moon. Doing it this way limits
         # attempted imports of lunarsky
-        if isinstance(self.value, EarthLocation):
+        if self.value is None or isinstance(self.value, EarthLocation):
             return False
         else:
             try:

src/pyuvdata/uvdata/mwa_corr_fits.py

@@ -48,7 +48,7 @@ def read_metafits(
     telescope_info_only=False,
 ):
     # get information from metafits file
-    with fits.open(file, memmap=True) as meta:
+    with fits.open(file, memmap=False) as meta:
         meta_hdr = meta[0].header
 
         telescope_name = meta_hdr.pop("TELESCOP")
@@ -734,6 +734,7 @@ def _read_fits_file(
         # map file number to frequency index
         freq_ind = coarse_ind * num_fine_chans
 
+        need_memmap = False
         if freq_inds is not None:
             # check that we want to read this file
             this_file_f_inds = np.arange(num_fine_chans) + freq_ind
@@ -767,11 +768,17 @@ def _read_fits_file(
             * (len(self.telescope.antenna_numbers) + 1)
             / 2.0
         )
+        if self.Nbls != n_orig_bls:
+            need_memmap = True
         bl_frac = self.Nbls / n_orig_bls
 
         n_orig_pols = 4
+        if self.Npols != n_orig_pols:
+            need_memmap = True
         pol_frac = self.Npols / n_orig_pols
 
+        if n_freq_read != num_fine_chans:
+            need_memmap = True
         freq_frac = n_freq_read / num_fine_chans
 
         blpol_frac = bl_frac * pol_frac
@@ -801,7 +808,7 @@ def _read_fits_file(
                 axis=1,
             ).flatten()
 
-        with fits.open(filename, mode="denywrite") as hdu_list:
+        with fits.open(filename, mode="denywrite", memmap=need_memmap) as hdu_list:
             # if mwax, data is in every other hdu
             if mwax:
                 hdu_list = hdu_list[1::2]
@@ -927,7 +934,7 @@ def _read_flag_file(self, filename, file_nums, num_fine_chans):
         flag_num = int(filename.split("_")[-1][0:2])
         # map file number to frequency index
         freq_ind = np.where(file_nums == flag_num)[0][0] * num_fine_chans
-        with fits.open(filename, mode="denywrite") as aoflags:
+        with fits.open(filename, mode="denywrite", memmap=False) as aoflags:
             flags = aoflags[1].data.field("FLAGS")
         # some flag files are longer than data; crop the end
         flags = flags[: self.Nblts, :]
@@ -1582,7 +1589,7 @@ def read_mwa_corr_fits(
                     raise ValueError("multiple metafits files in filelist")
                 metafits_file = filename
             elif filename.lower().endswith(".fits"):
-                with fits.open(filename, memmap=True) as hdu_list:
+                with fits.open(filename, memmap=False) as hdu_list:
                     hdunames = fits_utils._indexhdus(hdu_list)
                     if "PPDS" in hdunames:
                         ppds_file = filename

src/pyuvdata/uvdata/uvfits.py

@@ -4,6 +4,7 @@
 """Class for reading and writing uvfits files."""
 
 import copy
+import gc
 import os
 import warnings
 
@@ -197,7 +198,7 @@ def _get_parameter_data(
 
     def _get_data(
         self,
-        vis_hdu,
+        filename,
         *,
         antenna_nums,
         antenna_names,
@@ -269,17 +270,38 @@ def _get_data(
         )
 
         if min_frac == 1:
-            # no select, read in all the data
-            if vis_hdu.header["NAXIS"] == 7:
-                raw_data_array = vis_hdu.data.data[:, 0, 0, :, :, :, :]
-                if self.Nspws != raw_data_array.shape[1]:  # pragma: no cover
-                    raise RuntimeError(bad_data_shape_msg)
+            # no select, read in all the data. Don't need memmap
+            with fits.open(filename, memmap=False) as hdu_list:
+                vis_hdu = hdu_list[0]  # assumes the visibilities are in the primary hdu
+                if vis_hdu.header["NAXIS"] == 7:
+                    self.data_array = vis_hdu.data.data[:, 0, 0]
+                    if (
+                        self.Nspws != self.data_array.shape[1]
+                        or len(self.data_array.shape) != 5
+                    ):  # pragma: no cover
+                        raise RuntimeError(bad_data_shape_msg)
+
+                    # Reshape the data array to combine the spw & freq axes
+                    self.data_array = np.reshape(
+                        self.data_array,
+                        (self.Nblts, self.Nfreqs, self.Npols, self.data_array.shape[4]),
+                    )
+
+                else:
+                    self.data_array = vis_hdu.data.data[:, 0, 0, :, :, :]
+
+                    if len(self.data_array.shape) != 4:  # pragma: no cover
+                        raise RuntimeError(bad_data_shape_msg)
+
+            self.flag_array = self.data_array[:, :, :, 2] <= 0
+            self.nsample_array = np.abs(self.data_array[:, :, :, 2])
+
+            # FITS uvw direction convention is opposite ours and Miriad's.
+            # So conjugate the visibilities and flip the uvws:
+            self.data_array = (
+                self.data_array[:, :, :, 0] - 1j * self.data_array[:, :, :, 1]
+            )
 
-            else:
-                # in many uvfits files the spw axis is left out,
-                # here we put it back in so the dimensionality stays the same
-                raw_data_array = vis_hdu.data.data[:, 0, 0, :, :, :]
-                raw_data_array = raw_data_array[:, np.newaxis, :, :]
         else:
             # do select operations on everything except data_array, flag_array
             # and nsample_array
@@ -293,72 +315,80 @@ def _get_data(
             )
 
             # just read in the right portions of the data and flag arrays
-            if blt_frac == min_frac:
-                if vis_hdu.header["NAXIS"] == 7:
-                    raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :, :]
-                    raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
-                    if self.Nspws != raw_data_array.shape[1]:  # pragma: no cover
-                        raise RuntimeError(bad_data_shape_msg)
-                else:
-                    # in many uvfits files the spw axis is left out,
-                    # here we put it back in so the dimensionality stays the same
-                    raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :]
-                    raw_data_array = raw_data_array[:, 0, 0, :, :, :]
-                    raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
-                if freq_frac < 1:
-                    raw_data_array = raw_data_array[:, :, freq_inds, :, :]
-                if pol_frac < 1:
-                    raw_data_array = raw_data_array[:, :, :, pol_inds, :]
-            elif freq_frac == min_frac:
-                if vis_hdu.header["NAXIS"] == 7:
-                    raw_data_array = vis_hdu.data.data[:, :, :, :, freq_inds, :, :]
-                    raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
-                    if self.Nspws != raw_data_array.shape[1]:  # pragma: no cover
-                        raise RuntimeError(bad_data_shape_msg)
-                else:
-                    # in many uvfits files the spw axis is left out,
-                    # here we put it back in so the dimensionality stays the same
-                    raw_data_array = vis_hdu.data.data[:, :, :, freq_inds, :, :]
-                    raw_data_array = raw_data_array[:, 0, 0, :, :, :]
-                    raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
-
-                if blt_frac < 1:
-                    raw_data_array = raw_data_array[blt_inds, :, :, :, :]
-                if pol_frac < 1:
-                    raw_data_array = raw_data_array[:, :, :, pol_inds, :]
-            else:
-                if vis_hdu.header["NAXIS"] == 7:
-                    raw_data_array = vis_hdu.data.data[:, :, :, :, :, pol_inds, :]
-                    raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
-                    if self.Nspws != raw_data_array.shape[1]:  # pragma: no cover
-                        raise RuntimeError(bad_data_shape_msg)
+            # need memmap for this, accept the memory hit
+            with fits.open(filename, memmap=False) as hdu_list:
+                vis_hdu = hdu_list[0]  # assumes the visibilities are in the primary hdu
+                if blt_frac == min_frac:
+                    if vis_hdu.header["NAXIS"] == 7:
+                        raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :, :]
+                        raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
+                        if self.Nspws != raw_data_array.shape[1]:  # pragma: no cover
+                            raise RuntimeError(bad_data_shape_msg)
+                    else:
+                        # in many uvfits files the spw axis is left out,
+                        # here we put it back in so the dimensionality stays the same
+                        raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :]
+                        raw_data_array = raw_data_array[:, 0, 0, :, :, :]
+                        raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
+                    if freq_frac < 1:
+                        raw_data_array = raw_data_array[:, :, freq_inds, :, :]
+                    if pol_frac < 1:
+                        raw_data_array = raw_data_array[:, :, :, pol_inds, :]
+                elif freq_frac == min_frac:
+                    if vis_hdu.header["NAXIS"] == 7:
+                        raw_data_array = vis_hdu.data.data[:, :, :, :, freq_inds, :, :]
+                        raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
+                        if self.Nspws != raw_data_array.shape[1]:  # pragma: no cover
+                            raise RuntimeError(bad_data_shape_msg)
+                    else:
+                        # in many uvfits files the spw axis is left out,
+                        # here we put it back in so the dimensionality stays the same
+                        raw_data_array = vis_hdu.data.data[:, :, :, freq_inds, :, :]
+                        raw_data_array = raw_data_array[:, 0, 0, :, :, :]
+                        raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
+
+                    if blt_frac < 1:
+                        raw_data_array = raw_data_array[blt_inds, :, :, :, :]
+                    if pol_frac < 1:
+                        raw_data_array = raw_data_array[:, :, :, pol_inds, :]
                 else:
-                    # in many uvfits files the spw axis is left out,
-                    # here we put it back in so the dimensionality stays the same
-                    raw_data_array = vis_hdu.data.data[:, :, :, :, pol_inds, :]
-                    raw_data_array = raw_data_array[:, 0, 0, :, :, :]
-                    raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
-
-                if blt_frac < 1:
-                    raw_data_array = raw_data_array[blt_inds, :, :, :, :]
-                if freq_frac < 1:
-                    raw_data_array = raw_data_array[:, :, freq_inds, :, :]
-
-        if len(raw_data_array.shape) != 5:  # pragma: no cover
-            raise RuntimeError(bad_data_shape_msg)
-
-        # Reshape the data array to be the right size if we are working w/ multiple
-        # spectral windows
-        raw_data_array = np.reshape(
-            raw_data_array,
-            (self.Nblts, self.Nfreqs, self.Npols, raw_data_array.shape[4]),
-        )
+                    if vis_hdu.header["NAXIS"] == 7:
+                        raw_data_array = vis_hdu.data.data[:, :, :, :, :, pol_inds, :]
+                        raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
+                        if self.Nspws != raw_data_array.shape[1]:  # pragma: no cover
+                            raise RuntimeError(bad_data_shape_msg)
+                    else:
+                        # in many uvfits files the spw axis is left out,
+                        # here we put it back in so the dimensionality stays the same
+                        raw_data_array = vis_hdu.data.data[:, :, :, :, pol_inds, :]
+                        raw_data_array = raw_data_array[:, 0, 0, :, :, :]
+                        raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
+
+                    if blt_frac < 1:
+                        raw_data_array = raw_data_array[blt_inds, :, :, :, :]
+                    if freq_frac < 1:
+                        raw_data_array = raw_data_array[:, :, freq_inds, :, :]
+
+            if len(raw_data_array.shape) != 5:  # pragma: no cover
+                raise RuntimeError(bad_data_shape_msg)
+
+            # Reshape the data array to be the right size if we are working w/ multiple
+            # spectral windows
+            raw_data_array = np.reshape(
+                raw_data_array,
+                (self.Nblts, self.Nfreqs, self.Npols, raw_data_array.shape[4]),
+            )
 
-        # FITS uvw direction convention is opposite ours and Miriad's.
-        # So conjugate the visibilities and flip the uvws:
-        self.data_array = raw_data_array[:, :, :, 0] - 1j * raw_data_array[:, :, :, 1]
-        self.flag_array = raw_data_array[:, :, :, 2] <= 0
-        self.nsample_array = np.abs(raw_data_array[:, :, :, 2])
+            # FITS uvw direction convention is opposite ours and Miriad's.
+            # So conjugate the visibilities and flip the uvws:
+            self.data_array = (
+                raw_data_array[:, :, :, 0] - 1j * raw_data_array[:, :, :, 1]
+            )
+            self.flag_array = raw_data_array[:, :, :, 2] <= 0
+            self.nsample_array = np.abs(raw_data_array[:, :, :, 2])
+
+            del raw_data_array
+            gc.collect()
 
         if fix_old_proj:
             self.fix_phase(use_ant_pos=fix_use_ant_pos)
@@ -403,7 +433,9 @@ def read_uvfits(
         self.filename = [basename]
         self._filename.form = (1,)
 
-        with fits.open(filename, memmap=True) as hdu_list:
+        # memmap=True uses a lot of memory that doesn't get deallocated quickly
+        # not needed for reading the headers, which is what we're doing here
+        with fits.open(filename, memmap=False) as hdu_list:
             vis_hdu = hdu_list[0]  # assumes the visibilities are in the primary hdu
             vis_hdr = vis_hdu.header.copy()
             hdunames = fits_utils._indexhdus(hdu_list)  # find the rest of the tables
@@ -642,7 +674,7 @@ def read_uvfits(
                 # the array center, but in a rotated ECEF frame so that the x-axis
                 # goes through the local meridian.
                 rot_ecef_positions = ant_hdu.data.field("STABXYZ")
-                _, longitude, altitude = utils.LatLonAlt_from_XYZ(
+                _, longitude, _ = utils.LatLonAlt_from_XYZ(
                     np.array([x_telescope, y_telescope, z_telescope]),
                     frame=telescope_frame,
                     ellipsoid=ellipsoid,
@@ -825,29 +857,29 @@ def read_uvfits(
                     rot_axis=0,
                 )[:, :, 0]
 
-            if read_data:
-                # Now read in the data
-                self._get_data(
-                    vis_hdu,
-                    antenna_nums=antenna_nums,
-                    antenna_names=antenna_names,
-                    ant_str=ant_str,
-                    bls=bls,
-                    frequencies=frequencies,
-                    freq_chans=freq_chans,
-                    spws=spws,
-                    times=times,
-                    time_range=time_range,
-                    lsts=lsts,
-                    lst_range=lst_range,
-                    polarizations=polarizations,
-                    blt_inds=blt_inds,
-                    phase_center_ids=phase_center_ids,
-                    catalog_names=catalog_names,
-                    keep_all_metadata=keep_all_metadata,
-                    fix_old_proj=fix_old_proj,
-                    fix_use_ant_pos=fix_use_ant_pos,
-                )
+        if read_data:
+            # Now read in the data
+            self._get_data(
+                filename,
+                antenna_nums=antenna_nums,
+                antenna_names=antenna_names,
+                ant_str=ant_str,
+                bls=bls,
+                frequencies=frequencies,
+                freq_chans=freq_chans,
+                spws=spws,
+                times=times,
+                time_range=time_range,
+                lsts=lsts,
+                lst_range=lst_range,
+                polarizations=polarizations,
+                blt_inds=blt_inds,
+                phase_center_ids=phase_center_ids,
+                catalog_names=catalog_names,
+                keep_all_metadata=keep_all_metadata,
+                fix_old_proj=fix_old_proj,
+                fix_use_ant_pos=fix_use_ant_pos,
+            )
 
         # check if object has all required UVParameters set
         if run_check: