Skip to content

Commit 106a7b6

Browse files
bhazeltonbhazelton
committed
Only use memmap when required
1 parent 0e6f7aa commit 106a7b6

File tree

1 file changed

+94
-83
lines changed

1 file changed

+94
-83
lines changed

‎src/pyuvdata/uvdata/uvfits.py‎

Lines changed: 94 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
"""Class for reading and writing uvfits files."""
55

66
import copy
7+
import gc
78
import os
89
import warnings
910

@@ -197,7 +198,7 @@ def _get_parameter_data(
197198

198199
def _get_data(
199200
self,
200-
vis_hdu,
201+
filename,
201202
*,
202203
antenna_nums,
203204
antenna_names,
@@ -269,17 +270,19 @@ def _get_data(
269270
)
270271

271272
if min_frac == 1:
272-
# no select, read in all the data
273-
if vis_hdu.header["NAXIS"] == 7:
274-
raw_data_array = vis_hdu.data.data[:, 0, 0, :, :, :, :]
275-
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
276-
raise RuntimeError(bad_data_shape_msg)
273+
# no select, read in all the data. Don't need memmap
274+
with fits.open(filename, memmap=False) as hdu_list:
275+
vis_hdu = hdu_list[0] # assumes the visibilities are in the primary hdu
276+
if vis_hdu.header["NAXIS"] == 7:
277+
raw_data_array = vis_hdu.data.data[:, 0, 0, :, :, :, :]
278+
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
279+
raise RuntimeError(bad_data_shape_msg)
277280

278-
else:
279-
# in many uvfits files the spw axis is left out,
280-
# here we put it back in so the dimensionality stays the same
281-
raw_data_array = vis_hdu.data.data[:, 0, 0, :, :, :]
282-
raw_data_array = raw_data_array[:, np.newaxis, :, :]
281+
else:
282+
# in many uvfits files the spw axis is left out,
283+
# here we put it back in so the dimensionality stays the same
284+
raw_data_array = vis_hdu.data.data[:, 0, 0, :, :, :]
285+
raw_data_array = raw_data_array[:, np.newaxis, :, :]
283286
else:
284287
# do select operations on everything except data_array, flag_array
285288
# and nsample_array
@@ -293,56 +296,59 @@ def _get_data(
293296
)
294297

295298
# just read in the right portions of the data and flag arrays
296-
if blt_frac == min_frac:
297-
if vis_hdu.header["NAXIS"] == 7:
298-
raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :, :]
299-
raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
300-
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
301-
raise RuntimeError(bad_data_shape_msg)
302-
else:
303-
# in many uvfits files the spw axis is left out,
304-
# here we put it back in so the dimensionality stays the same
305-
raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :]
306-
raw_data_array = raw_data_array[:, 0, 0, :, :, :]
307-
raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
308-
if freq_frac < 1:
309-
raw_data_array = raw_data_array[:, :, freq_inds, :, :]
310-
if pol_frac < 1:
311-
raw_data_array = raw_data_array[:, :, :, pol_inds, :]
312-
elif freq_frac == min_frac:
313-
if vis_hdu.header["NAXIS"] == 7:
314-
raw_data_array = vis_hdu.data.data[:, :, :, :, freq_inds, :, :]
315-
raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
316-
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
317-
raise RuntimeError(bad_data_shape_msg)
318-
else:
319-
# in many uvfits files the spw axis is left out,
320-
# here we put it back in so the dimensionality stays the same
321-
raw_data_array = vis_hdu.data.data[:, :, :, freq_inds, :, :]
322-
raw_data_array = raw_data_array[:, 0, 0, :, :, :]
323-
raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
324-
325-
if blt_frac < 1:
326-
raw_data_array = raw_data_array[blt_inds, :, :, :, :]
327-
if pol_frac < 1:
328-
raw_data_array = raw_data_array[:, :, :, pol_inds, :]
329-
else:
330-
if vis_hdu.header["NAXIS"] == 7:
331-
raw_data_array = vis_hdu.data.data[:, :, :, :, :, pol_inds, :]
332-
raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
333-
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
334-
raise RuntimeError(bad_data_shape_msg)
299+
# need memmap for this, accept the memory hit
300+
with fits.open(filename, memmap=False) as hdu_list:
301+
vis_hdu = hdu_list[0] # assumes the visibilities are in the primary hdu
302+
if blt_frac == min_frac:
303+
if vis_hdu.header["NAXIS"] == 7:
304+
raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :, :]
305+
raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
306+
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
307+
raise RuntimeError(bad_data_shape_msg)
308+
else:
309+
# in many uvfits files the spw axis is left out,
310+
# here we put it back in so the dimensionality stays the same
311+
raw_data_array = vis_hdu.data.data[blt_inds, :, :, :, :, :]
312+
raw_data_array = raw_data_array[:, 0, 0, :, :, :]
313+
raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
314+
if freq_frac < 1:
315+
raw_data_array = raw_data_array[:, :, freq_inds, :, :]
316+
if pol_frac < 1:
317+
raw_data_array = raw_data_array[:, :, :, pol_inds, :]
318+
elif freq_frac == min_frac:
319+
if vis_hdu.header["NAXIS"] == 7:
320+
raw_data_array = vis_hdu.data.data[:, :, :, :, freq_inds, :, :]
321+
raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
322+
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
323+
raise RuntimeError(bad_data_shape_msg)
324+
else:
325+
# in many uvfits files the spw axis is left out,
326+
# here we put it back in so the dimensionality stays the same
327+
raw_data_array = vis_hdu.data.data[:, :, :, freq_inds, :, :]
328+
raw_data_array = raw_data_array[:, 0, 0, :, :, :]
329+
raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
330+
331+
if blt_frac < 1:
332+
raw_data_array = raw_data_array[blt_inds, :, :, :, :]
333+
if pol_frac < 1:
334+
raw_data_array = raw_data_array[:, :, :, pol_inds, :]
335335
else:
336-
# in many uvfits files the spw axis is left out,
337-
# here we put it back in so the dimensionality stays the same
338-
raw_data_array = vis_hdu.data.data[:, :, :, :, pol_inds, :]
339-
raw_data_array = raw_data_array[:, 0, 0, :, :, :]
340-
raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
336+
if vis_hdu.header["NAXIS"] == 7:
337+
raw_data_array = vis_hdu.data.data[:, :, :, :, :, pol_inds, :]
338+
raw_data_array = raw_data_array[:, 0, 0, :, :, :, :]
339+
if self.Nspws != raw_data_array.shape[1]: # pragma: no cover
340+
raise RuntimeError(bad_data_shape_msg)
341+
else:
342+
# in many uvfits files the spw axis is left out,
343+
# here we put it back in so the dimensionality stays the same
344+
raw_data_array = vis_hdu.data.data[:, :, :, :, pol_inds, :]
345+
raw_data_array = raw_data_array[:, 0, 0, :, :, :]
346+
raw_data_array = raw_data_array[:, np.newaxis, :, :, :]
341347

342-
if blt_frac < 1:
343-
raw_data_array = raw_data_array[blt_inds, :, :, :, :]
344-
if freq_frac < 1:
345-
raw_data_array = raw_data_array[:, :, freq_inds, :, :]
348+
if blt_frac < 1:
349+
raw_data_array = raw_data_array[blt_inds, :, :, :, :]
350+
if freq_frac < 1:
351+
raw_data_array = raw_data_array[:, :, freq_inds, :, :]
346352

347353
if len(raw_data_array.shape) != 5: # pragma: no cover
348354
raise RuntimeError(bad_data_shape_msg)
@@ -360,6 +366,9 @@ def _get_data(
360366
self.flag_array = raw_data_array[:, :, :, 2] <= 0
361367
self.nsample_array = np.abs(raw_data_array[:, :, :, 2])
362368

369+
del raw_data_array
370+
gc.collect()
371+
363372
if fix_old_proj:
364373
self.fix_phase(use_ant_pos=fix_use_ant_pos)
365374

@@ -403,7 +412,9 @@ def read_uvfits(
403412
self.filename = [basename]
404413
self._filename.form = (1,)
405414

406-
with fits.open(filename, memmap=True) as hdu_list:
415+
# memmap=True uses a lot of memory that doesn't get deallocated quickly
416+
# not needed for reading the headers, which is what we're doing here
417+
with fits.open(filename, memmap=False) as hdu_list:
407418
vis_hdu = hdu_list[0] # assumes the visibilities are in the primary hdu
408419
vis_hdr = vis_hdu.header.copy()
409420
hdunames = fits_utils._indexhdus(hdu_list) # find the rest of the tables
@@ -825,29 +836,29 @@ def read_uvfits(
825836
rot_axis=0,
826837
)[:, :, 0]
827838

828-
if read_data:
829-
# Now read in the data
830-
self._get_data(
831-
vis_hdu,
832-
antenna_nums=antenna_nums,
833-
antenna_names=antenna_names,
834-
ant_str=ant_str,
835-
bls=bls,
836-
frequencies=frequencies,
837-
freq_chans=freq_chans,
838-
spws=spws,
839-
times=times,
840-
time_range=time_range,
841-
lsts=lsts,
842-
lst_range=lst_range,
843-
polarizations=polarizations,
844-
blt_inds=blt_inds,
845-
phase_center_ids=phase_center_ids,
846-
catalog_names=catalog_names,
847-
keep_all_metadata=keep_all_metadata,
848-
fix_old_proj=fix_old_proj,
849-
fix_use_ant_pos=fix_use_ant_pos,
850-
)
839+
if read_data:
840+
# Now read in the data
841+
self._get_data(
842+
filename,
843+
antenna_nums=antenna_nums,
844+
antenna_names=antenna_names,
845+
ant_str=ant_str,
846+
bls=bls,
847+
frequencies=frequencies,
848+
freq_chans=freq_chans,
849+
spws=spws,
850+
times=times,
851+
time_range=time_range,
852+
lsts=lsts,
853+
lst_range=lst_range,
854+
polarizations=polarizations,
855+
blt_inds=blt_inds,
856+
phase_center_ids=phase_center_ids,
857+
catalog_names=catalog_names,
858+
keep_all_metadata=keep_all_metadata,
859+
fix_old_proj=fix_old_proj,
860+
fix_use_ant_pos=fix_use_ant_pos,
861+
)
851862

852863
# check if object has all required UVParameters set
853864
if run_check:

0 commit comments

Comments
 (0)