#!/usr/bin/env python2.7 # -*- mode: python; coding: utf-8 -*- # Copyright (c) 2018 Radio Astronomy Software Group # Licensed under the 2-clause BSD License """Test memory usage of read_uvfits.""" import numpy as np from astropy import constants as const from astropy.io import fits from memory_profiler import profile from pyuvdata import UVData @profile def read_uvfits(): """Test memory usage of read_uvfits.""" filename = "/Volumes/Data1/mwa_uvfits/1066571272.uvfits" # first test uvdata.read_uvfits. First read metadata then full data uv_obj = UVData() uv_obj.read_uvfits(filename, read_data=False, read_metadata=False) uv_obj.read_uvfits_metadata(filename) uv_obj.read_uvfits_data(filename) del uv_obj # now test uvdata.read_uvfits with select on read. uv_obj = UVData() uv_obj.read_uvfits(filename, read_data=False, read_metadata=False) uv_obj.read_uvfits_metadata(filename) uv_obj.read_uvfits_data(filename, freq_chans=np.arange(196)) del uv_obj # now test details with astropy hdu_list = fits.open(filename, memmap=True) vis_hdu = hdu_list[0] # only read in times, then uvws, then visibilities time0_array = vis_hdu.data.par("date") uvw_array = ( np.array( np.stack( (vis_hdu.data.par("UU"), vis_hdu.data.par("VV"), vis_hdu.data.par("WW")) ) ) * const.c.to("m/s").value ).T if vis_hdu.header["NAXIS"] == 7: data_array = ( vis_hdu.data.data[:, 0, 0, :, :, :, 0] + 1j * vis_hdu.data.data[:, 0, 0, :, :, :, 1] ) else: data_array = ( vis_hdu.data.data[:, 0, 0, :, :, 0] + 1j * vis_hdu.data.data[:, 0, 0, :, :, 1] ) data_array = data_array[:, np.newaxis, :, :] # test for releasing resources del time0_array del uvw_array del data_array # release file handles del vis_hdu del hdu_list # now test reading a slice of the data hdu_list = fits.open(filename, memmap=True) vis_hdu = hdu_list[0] Nfreqs = vis_hdu.header["NAXIS4"] freq_index = int(Nfreqs // 2) if vis_hdu.header["NAXIS"] == 7: data_slice = ( vis_hdu.data.data[:, 0, 0, :, 0:freq_index, :, 0] + 1j * vis_hdu.data.data[:, 0, 0, :, 0:freq_index, :, 1] ) else: data_slice = ( vis_hdu.data.data[:, 0, 0, 0:freq_index, :, 0] + 1j * vis_hdu.data.data[:, 0, 0, 0:freq_index, :, 1] ) data_slice = data_slice[:, np.newaxis, :, :] del data_slice del vis_hdu del hdu_list del filename return if __name__ == "__main__": read_uvfits()