def _comply_vispol(pol): '''Maps an input visibility polarization string onto a string compliant with pyuvdata and hera_cal.''' if _is_cardinal(pol): return polnum2str(polstr2num(pol, x_orientation='north'), x_orientation='north') else: return polnum2str(polstr2num(pol))
def test_pol_funcs_x_orientation(): """ Test utility functions to convert between polarization strings and numbers with x_orientation """ pol_nums = [-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4] x_orient1 = 'e' pol_str = [ 'ne', 'en', 'nn', 'ee', 'lr', 'rl', 'll', 'rr', 'pI', 'pQ', 'pU', 'pV' ] assert pol_nums == uvutils.polstr2num(pol_str, x_orientation=x_orient1) assert pol_str == uvutils.polnum2str(pol_nums, x_orientation=x_orient1) # Check individuals assert -6 == uvutils.polstr2num('NN', x_orientation=x_orient1) assert 'pV' == uvutils.polnum2str(4) # Check errors pytest.raises(KeyError, uvutils.polstr2num, 'foo', x_orientation=x_orient1) pytest.raises(ValueError, uvutils.polstr2num, 1, x_orientation=x_orient1) pytest.raises(ValueError, uvutils.polnum2str, 7.3, x_orientation=x_orient1) # Check parse assert uvutils.parse_polstr("eE", x_orientation=x_orient1) == 'ee' assert uvutils.parse_polstr("xx", x_orientation=x_orient1) == 'ee' assert uvutils.parse_polstr("NN", x_orientation=x_orient1) == 'nn' assert uvutils.parse_polstr("yy", x_orientation=x_orient1) == 'nn' assert uvutils.parse_polstr('i', x_orientation=x_orient1) == 'pI' x_orient2 = 'n' pol_str = [ 'en', 'ne', 'ee', 'nn', 'lr', 'rl', 'll', 'rr', 'pI', 'pQ', 'pU', 'pV' ] assert pol_nums == uvutils.polstr2num(pol_str, x_orientation=x_orient2) assert pol_str == uvutils.polnum2str(pol_nums, x_orientation=x_orient2) # Check individuals assert -6 == uvutils.polstr2num('EE', x_orientation=x_orient2) assert 'pV' == uvutils.polnum2str(4) # Check errors pytest.raises(KeyError, uvutils.polstr2num, 'foo', x_orientation=x_orient2) pytest.raises(ValueError, uvutils.polstr2num, 1, x_orientation=x_orient2) pytest.raises(ValueError, uvutils.polnum2str, 7.3, x_orientation=x_orient2) # Check parse assert uvutils.parse_polstr("nN", x_orientation=x_orient2) == 'nn' assert uvutils.parse_polstr("xx", x_orientation=x_orient2) == 'nn' assert uvutils.parse_polstr("EE", x_orientation=x_orient2) == 'ee' assert uvutils.parse_polstr("yy", x_orientation=x_orient2) == 'ee' assert uvutils.parse_polstr('i', x_orientation=x_orient2) == 'pI' # check warnings for non-recognized x_orientation assert uvtest.checkWarnings(uvutils.polstr2num, ['xx'], {'x_orientation': 'foo'}, message='x_orientation not recognized') == -5 assert uvtest.checkWarnings(uvutils.polnum2str, [-6], {'x_orientation': 'foo'}, message='x_orientation not recognized') == 'yy'
def power_beam_sq_int(self, pol='pI'): """ Computes the integral of the beam**2 over solid angle to give a beam**2 area (in str) as a function of frequency. Parameters ---------- pol: str, optional Which polarization to compute the beam scalar for. 'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX' Default: pI. Returns ------- primary_beam_area: array_like Array of floats containing the primary beam-squared area. """ # type check if isinstance(pol, str): pol = uvutils.polstr2num(pol, x_orientation=self.x_orientation) if pol in self.OmegaPP.keys(): return self.OmegaPP[pol] else: raise KeyError("OmegaPP not specified for polarization '%s'. " "Available polarizations are: %s" \ % (pol, self.available_pols))
def test_pol_funcs(): """ Test utility functions to convert between polarization strings and numbers """ pol_nums = [-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4] pol_str = [ 'YX', 'XY', 'YY', 'XX', 'LR', 'RL', 'LL', 'RR', 'pI', 'pQ', 'pU', 'pV' ] nt.assert_equal(pol_nums, uvutils.polstr2num(pol_str)) nt.assert_equal(pol_str, uvutils.polnum2str(pol_nums)) # Check individuals nt.assert_equal(-6, uvutils.polstr2num('YY')) nt.assert_equal('pV', uvutils.polnum2str(4)) # Check errors nt.assert_raises(KeyError, uvutils.polstr2num, 'foo') nt.assert_raises(ValueError, uvutils.polstr2num, 1) nt.assert_raises(ValueError, uvutils.polnum2str, 7.3)
def efield_to_power(self): """ Tell interp to return values corresponding with a power beam. """ self.beam_type = 'power' pol_strings = ['XX', 'XY', 'YX', 'YY'] self.polarization_array = np.array( [uvutils.polstr2num(ps.upper()) for ps in pol_strings])
def polpair_tuple2int(polpair, x_orientation=None): """ Convert a tuple pair of polarization strings/integers into an pol-pair integer. The polpair integer is formed by adding 20 to each standardized polarization integer (see polstr2num and AIPS memo 117) and then concatenating them. For example, polarization pair ('pI', 'pQ') == (1, 2) == 2122. Parameters ---------- polpair : tuple, length 2 A length-2 tuple containing a pair of polarization strings or integers, e.g. ('XX', 'YY') or (-5, -5). x_orientation: str, optional Orientation in cardinal direction east or north of X dipole. Default keeps polarization in X and Y basis. Returns ------- polpair : int Integer representation of polarization pair. """ # Recursive evaluation if isinstance(polpair, (list, np.ndarray)): return [polpair_tuple2int(p) for p in polpair] # Check types assert type(polpair) in (tuple, ), "pol must be a tuple" assert len(polpair) == 2, "polpair tuple must have 2 elements" # Convert strings to ints if necessary pol1, pol2 = polpair if type(pol1) in (str, np.str): pol1 = polstr2num(pol1, x_orientation=x_orientation) if type(pol2) in (str, np.str): pol2 = polstr2num(pol2, x_orientation=x_orientation) # Convert to polpair integer ppint = (20 + pol1) * 100 + (20 + pol2) return ppint
def test_pol_funcs(): """ Test utility functions to convert between polarization strings and numbers """ pol_nums = [-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4] pol_str = [ 'yx', 'xy', 'yy', 'xx', 'lr', 'rl', 'll', 'rr', 'pI', 'pQ', 'pU', 'pV' ] assert pol_nums == uvutils.polstr2num(pol_str) assert pol_str == uvutils.polnum2str(pol_nums) # Check individuals assert -6 == uvutils.polstr2num('YY') assert 'pV' == uvutils.polnum2str(4) # Check errors pytest.raises(KeyError, uvutils.polstr2num, 'foo') pytest.raises(ValueError, uvutils.polstr2num, 1) pytest.raises(ValueError, uvutils.polnum2str, 7.3) # Check parse assert uvutils.parse_polstr("xX") == 'xx' assert uvutils.parse_polstr("XX") == 'xx' assert uvutils.parse_polstr('i') == 'pI'
def add_pol(self, pol, OmegaP, OmegaPP): """ Add OmegaP and OmegaPP for a new polarization. Parameters ---------- pol: str Which polarization to add beam solid angle arrays for. Valid options are: 'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX' If the arrays already exist for the specified polarization, they will be overwritten. OmegaP : array_like of float Integral over beam solid angle, as a function of frequency. Must have the same shape as self.beam_freqs. OmegaPP : array_like of float Integral over beam solid angle squared, as a function of frequency. Must have the same shape as self.beam_freqs. """ # Type check if isinstance(pol, str): pol = uvutils.polstr2num(pol, x_orientation=self.x_orientation) # Check for allowed polarization if pol not in self.allowed_pols: raise KeyError("Polarization '%s' is not valid." % pol) # Make sure OmegaP and OmegaPP are arrays try: OmegaP = np.array(OmegaP).astype(np.float) OmegaPP = np.array(OmegaPP).astype(np.float) except: raise TypeError("OmegaP and OmegaPP must both be array_like.") # Check that array dimensions are consistent if OmegaP.shape != self.beam_freqs.shape \ or OmegaPP.shape != self.beam_freqs.shape: raise ValueError("OmegaP and OmegaPP should both " "have the same shape as beam_freqs.") # Store arrays self.OmegaP[pol] = OmegaP self.OmegaPP[pol] = OmegaPP # get available pols self.available_pols = ", ".join( map(uvutils.polnum2str, self.OmegaP.keys()))
def gen_model(**kwargs): p = Dict2Obj(**kwargs) utils.log("\n{}\n...Generating a Flux Model", f=p.lf, verbose=p.verbose) # compile complist_gleam.py command cmd = casa + ["-c", "{}/complist_gleam.py".format(casa_scripts)] cmd += ['--point_ra', p.source_ra, '--point_dec', p.latitude, '--outdir', p.out_dir, '--gleamfile', p.gleamfile, '--radius', p.radius, '--min_flux', p.min_flux, '--freqs', p.freqs, '--cell', p.cell, '--imsize', p.imsize] if p.image: cmd += ['--image'] if p.use_peak: cmd += ['--use_peak'] if p.overwrite: cmd += ['--overwrite'] if hasattr(p, 'regions'): cmd += ['--regions', p.regions, '--exclude', '--region_radius', p.region_radius] if hasattr(p, 'file_ext'): cmd += ['--ext', p.file_ext] else: p.file_ext = '' cmd = map(str, cmd) ecode = subprocess.check_call(cmd) modelstem = os.path.join(p.out_dir, "gleam{}.cl".format(p.file_ext)) model = modelstem if p.image: model += ".image" # pbcorrect if p.pbcorr: utils.log("...applying PB to model", f=p.lf, verbose=p.verbose) assert p.image, "Cannot pbcorrect flux model without image == True" cmd = ["pbcorr.py", "--lon", p.longitude, "--lat", p.latitude, "--time", p.time, "--pols"] \ + [uvutils.polstr2num(pol) for pol in p.pols] \ + ["--outdir", p.out_dir, "--multiply", "--beamfile", p.beamfile] if p.overwrite: cmd.append("--overwrite") cmd.append(modelstem + '.fits') # generate component list and / or image cube flux model cmd = map(str, cmd) ecode = subprocess.check_call(cmd) modelstem = os.path.join(p.out_dir, modelstem) # importfits cmd = p.casa + ["-c", "importfits('{}', '{}', overwrite={})".format(modelstem + '.pbcorr.fits', modelstem + '.pbcorr.image', p.overwrite)] ecode = subprocess.check_call(cmd) model = modelstem + ".pbcorr.image" return model
def test_pol_funcs(): """ Test utility functions to convert between polarization strings and numbers """ pol_nums = [-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4] pol_str = [ 'yx', 'xy', 'yy', 'xx', 'lr', 'rl', 'll', 'rr', 'pI', 'pQ', 'pU', 'pV' ] nt.assert_equal(pol_nums, uvutils.polstr2num(pol_str)) nt.assert_equal(pol_str, uvutils.polnum2str(pol_nums)) # Check individuals nt.assert_equal(-6, uvutils.polstr2num('YY')) nt.assert_equal('pV', uvutils.polnum2str(4)) # Check errors nt.assert_raises(KeyError, uvutils.polstr2num, 'foo') nt.assert_raises(ValueError, uvutils.polstr2num, 1) nt.assert_raises(ValueError, uvutils.polnum2str, 7.3) # Check parse nt.assert_equal(uvutils.parse_polstr("xX"), 'xx') nt.assert_equal(uvutils.parse_polstr("XX"), 'xx') nt.assert_equal(uvutils.parse_polstr('i'), 'pI') nt.assert_equal(uvutils.parse_jpolstr('x'), 'Jxx') nt.assert_equal(uvutils.parse_jpolstr('xy'), 'Jxy') nt.assert_equal(uvutils.parse_jpolstr('XY'), 'Jxy')
def apply_gains(uvdata, gains, inverse=False): """apply gains to a uvdata object. Parameters ---------- uvdata: UVData object. UVData for data to have gains applied. gains: UVCal object. UVCal object containing gains to be applied. inverse: bool, optional Multiply gains instead of dividing. Returns ------- calibrated: UVData object. UVData object containing calibrated data. """ calibrated = copy.deepcopy(uvdata) for pnum, pol in enumerate(uvdata.get_pols()): for ap in calibrated.get_antpairs(): dinds = calibrated.antpair2ind(ap) gindp = np.where(gains.jones_array == uvutils.polstr2num(pol, x_orientation=gains.x_orientation))[0][0] for time in calibrated.time_array[dinds]: dindt = np.where(calibrated.time_array[dinds] == time)[0][0] gindt = np.where(np.isclose(gains.time_array, time, rtol=0.0, atol=1e-7))[0][0] aind0 = np.where(gains.ant_array == ap[0])[0][0] aind1 = np.where(gains.ant_array == ap[1])[0][0] if not inverse: calibrated.data_array[dinds[dindt], 0, :, pnum] = calibrated.data_array[ dinds[dindt], 0, :, pnum ] / ( gains.gain_array[aind0, 0, :, gindt, gindp] * np.conj(gains.gain_array[aind1, 0, :, gindt, gindp]) ) else: calibrated.data_array[dinds[dindt], 0, :, pnum] = calibrated.data_array[ dinds[dindt], 0, :, pnum ] * ( gains.gain_array[aind0, 0, :, gindt, gindp] * np.conj(gains.gain_array[aind1, 0, :, gindt, gindp]) ) calibrated.flag_array[dinds[dindt], 0, :, pnum] = calibrated.flag_array[dinds[dindt], 0, :, pnum] | ( gains.flag_array[aind0, 0, :, gindt, gindp] | gains.flag_array[aind1, 0, :, gindt, gindp] ) return calibrated
def setup_uvdata( array_layout=None, telescope_location=None, telescope_name=None, Nfreqs=None, start_freq=None, bandwidth=None, freq_array=None, Ntimes=None, time_cadence=None, start_time=None, time_array=None, bls=None, anchor_ant=None, antenna_nums=None, no_autos=True, pols=["xx"], make_full=False, redundancy=None, run_check=True, ): """ Setup a UVData object for simulating. Args: array_layout : str Filepath to array layout in ENU coordinates [meters] telescope_location : len-3 tuple Telescope location on Earth in LatLonAlt coordinates [deg, deg, meters] telescope_name : str Name of telescope Nfreqs : int Number of frequency channels start_freq : float Starting frequency [Hz] bandwidth : float Total frequency bandwidth of spectral window [Hz] freq_array : ndarray frequency array [Hz], cannot be specified if start_freq, Nfreqs or bandwidth is specified Ntimes : int Number of integration bins time_cadence : float Cadence of time bins [seconds] start_time : float Time of the first integration bin [Julian Date] time_array : ndarray time array [Julian Date], cannot be specified if start_time, Ntimes and time_cadence is specified bls : list List of antenna-pair tuples for baseline selection anchor_ant: int Selects baselines such that one of the pair is a specified antenna number redundancy: float Redundant baseline selection tolerance for selection antenna_nums : list List of antenna numbers to keep in array make_full : Generate the full UVData object, includes arrays of length Nblts. Default behavior creates an invalid UVData object, where baseline_array has length Nbls, etc. This is to avoid excessive memory usage up front when it's not necessary. Returns: UVData object with zeroed data_array """ # get antenna information tele_dict = parse_telescope_params( dict( array_layout=array_layout, telescope_location=telescope_location, telescope_name=telescope_name, )) lat, lon, alt = uvutils.LatLonAlt_from_XYZ(tele_dict["telescope_location"]) antpos = tele_dict["antenna_positions"] enu = uvutils.ENU_from_ECEF( tele_dict["antenna_positions"] + tele_dict["telescope_location"], lat, lon, alt) anums = tele_dict["antenna_numbers"] antnames = tele_dict["antenna_names"] Nants = tele_dict["Nants_data"] # setup object and fill in basic info uv_obj = UVData() uv_obj.telescope_location = tele_dict["telescope_location"] uv_obj.telescope_location_lat_lon_alt = (lat, lon, alt) uv_obj.telescope_location_lat_lon_alt_degrees = ( np.degrees(lat), np.degrees(lon), alt, ) uv_obj.antenna_numbers = anums uv_obj.antenna_names = antnames uv_obj.antenna_positions = antpos uv_obj.Nants_telescope = Nants # fill in frequency and time info: wait to fill in len-Nblts arrays until later if freq_array is not None: if Nfreqs is not None or start_freq is not None or bandwidth is not None: raise ValueError( "Cannot specify freq_array as well as Nfreqs, start_freq or bandwidth" ) if freq_array.ndim == 1: freq_array = freq_array.reshape(1, -1) Nfreqs = freq_array.size else: freq_dict = parse_frequency_params( dict(Nfreqs=Nfreqs, start_freq=start_freq, bandwidth=bandwidth)) freq_array = freq_dict["freq_array"] if time_array is not None: if Ntimes is not None or start_time is not None or time_cadence is not None: raise ValueError( "Cannot specify time_array as well as Ntimes, start_time or time_cadence" ) Ntimes = time_array.size else: time_dict = parse_time_params( dict(Ntimes=Ntimes, start_time=start_time, time_cadence=time_cadence)) time_array = time_dict["time_array"] uv_obj.freq_array = freq_array uv_obj.Nfreqs = Nfreqs uv_obj.Ntimes = Ntimes # fill in other attributes uv_obj.spw_array = np.array([0], dtype=int) uv_obj.Nspws = 1 uv_obj.polarization_array = np.array( [uvutils.polstr2num(pol) for pol in pols], dtype=int) uv_obj.Npols = uv_obj.polarization_array.size if uv_obj.Nfreqs > 1: uv_obj.channel_width = np.diff(uv_obj.freq_array[0])[0] else: uv_obj.channel_width = 1.0 uv_obj._set_drift() uv_obj.telescope_name = tele_dict["telescope_name"] uv_obj.instrument = "simulator" uv_obj.object_name = "zenith" uv_obj.vis_units = "Jy" uv_obj.history = "" if redundancy is not None: red_tol = redundancy reds, vec_bin_centers, lengths = uvutils.get_antenna_redundancies( anums, enu, tol=red_tol, include_autos=not bool(no_autos)) bls = [r[0] for r in reds] bls = [uvutils.baseline_to_antnums(bl_ind, Nants) for bl_ind in bls] # Setup array and implement antenna/baseline selections. bl_array = [] _bls = [(a1, a2) for a1 in anums for a2 in anums if a1 <= a2] if bls is not None: if isinstance(bls, str): bls = ast.literal_eval(bls) bls = [bl for bl in _bls if bl in bls] else: bls = _bls if anchor_ant is not None: bls = [bl for bl in bls if anchor_ant in bl] if bool(no_autos): bls = [bl for bl in bls if bl[0] != bl[1]] if antenna_nums is not None: if isinstance(antenna_nums, str): antenna_nums = ast.literal_eval(antenna_nums) if isinstance(antenna_nums, int): antenna_nums = [antenna_nums] bls = [(a1, a2) for (a1, a2) in bls if a1 in antenna_nums or a2 in antenna_nums] bls = sorted(bls) for (a1, a2) in bls: bl_array.append(uvutils.antnums_to_baseline(a1, a2, 1)) bl_array = np.asarray(bl_array) print("Nbls: {}".format(bl_array.size)) if bl_array.size == 0: raise ValueError("No baselines selected.") uv_obj.time_array = time_array # Keep length Ntimes uv_obj.baseline_array = bl_array # Length Nbls if make_full: uv_obj = complete_uvdata(uv_obj, run_check=run_check) return uv_obj
# Create frequency array data.freq_array = np.zeros((data.Nspws, data.Nfreqs)) data.freq_array[0, :] = fqs # Set spw array data.spw_array = np.zeros(data.Nspws).astype(int) # Channel width data.channel_width = data.freq_array[0, 1] - data.freq_array[0, 0] # Convert to Nblt ordering and copy data data.data_array = np.zeros( (data.Nblts, data.Nspws, data.Nfreqs, data.Npols), dtype=complex) data.data_array[:, 0, :, 0] = np.transpose(poco[0]) data.polarization_array = np.array([polstr2num(pol)]).astype(int) # Integration time (including dead time) data.integration_time = times[1] - times[0] # Antenna numbers data.ant_1_array = [ants[pair[0]] for pair in bls] * args.num_spectra data.ant_2_array = [ants[pair[1]] for pair in bls] * args.num_spectra data.ant_1_array = np.asarray( [str(a)[2:-1] for a in data.ant_1_array]).astype(int) data.ant_2_array = np.asarray( [str(a)[2:-1] for a in data.ant_2_array]).astype(int)
def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None, x_orientation=None): """ Primary beam model built from user-defined arrays for the integrals over beam solid angle and beam solid angle squared. Allowed polarizations are: pI, pQ, pU, pV, XX, YY, XY, YX Other polarizations will be ignored. Parameters ---------- OmegaP : array_like of float (or dict of array_like) Integral over beam solid angle, as a function of frequency. If only one array is specified, this will be assumed to be for the I polarization. If a dict is specified, an OmegaP array for several polarizations can be specified. OmegaPP : array_like of float (or dict of array_like) Integral over beam solid angle squared, as a function of frequency. If only one array is specified, this will be assumed to be for the I polarization. If a dict is specified, an OmegaP array for several polarizations can be specified. beam_freqs : array_like of float Frequencies at which beam solid angles OmegaP and OmegaPP are evaluated, in Hz. This should be specified as a single array, not as a dict. cosmo : conversions.Cosmo_Conversions object, optional Cosmology object. Uses the default cosmology object if not specified. Default: None. x_orientation : str, optional Orientation in cardinal direction east or north of X dipole. Default keeps polarization in X and Y basis. """ self.OmegaP = {} self.OmegaPP = {} self.x_orientation = x_orientation # these are allowed pols in AIPS polarization integer convention # see pyuvdata.utils.polstr2num() for details self.allowed_pols = [1, 2, 3, 4, -5, -6, -7, -8] # Set beam_freqs self.beam_freqs = np.asarray(beam_freqs) if isinstance(OmegaP, np.ndarray) and isinstance(OmegaPP, np.ndarray): # Only single arrays were specified; assume I OmegaP = {1: OmegaP} OmegaPP = {1: OmegaPP} elif isinstance(OmegaP, np.ndarray) or isinstance(OmegaPP, np.ndarray): # Mixed dict and array types are not allowed raise TypeError("OmegaP and OmegaPP must both be either dicts " "or arrays. Mixing dicts and arrays is not " "allowed.") else: pass # Should now have two dicts if everything is OK if not isinstance(OmegaP, (odict, dict)) or not isinstance( OmegaPP, (odict, dict)): raise TypeError("OmegaP and OmegaPP must both be either dicts or " "arrays.") # Check for disallowed polarizations for key in list(OmegaP.keys()): # turn into pol integer if a pol string if isinstance(key, str): new_key = uvutils.polstr2num(key, x_orientation=self.x_orientation) OmegaP[new_key] = OmegaP.pop(key) key = new_key # check its an allowed pol if key not in self.allowed_pols: raise KeyError("Unrecognized polarization '%s' in OmegaP." % key) for key in list(OmegaPP.keys()): # turn into pol integer if a pol string if isinstance(key, str): new_key = uvutils.polstr2num(key, x_orientation=self.x_orientation) OmegaPP[new_key] = OmegaPP.pop(key) key = new_key # check its an allowed pol if key not in self.allowed_pols: raise KeyError("Unrecognized polarization '%s' in OmegaPP." % key) # Check for available polarizations for pol in self.allowed_pols: if pol in OmegaP.keys() or pol in OmegaPP.keys(): if pol not in OmegaP.keys() or pol not in OmegaPP.keys(): raise KeyError("Polarization '%s' must be specified for" " both OmegaP and OmegaPP." % pol) # Add arrays for this polarization self.add_pol(pol, OmegaP[pol], OmegaPP[pol]) # Set cosmology if cosmo is None: self.cosmo = conversions.Cosmo_Conversions() else: self.cosmo = cosmo
**kwargs) except: print(sys.exc_info()) continue peak_flux.append(output[0]) peak_flux_err.append(output[2]) peak_gauss_flux.append(output[3]) int_gauss_flux.append(output[4]) freqs.append(output[5]) if len(peak_flux) == 0: raise ValueError("Couldn't get source flux from any input files") freqs = np.array(freqs) pols = np.asarray([ pol if isinstance(pol, (int, np.integer)) else uvutils.polstr2num(pol) for pol in args.pols ]) peak_flux = np.array(peak_flux) peak_flux_err = np.array(peak_flux_err) peak_gauss_flux = np.array(peak_gauss_flux) int_gauss_flux = np.array(int_gauss_flux) # save spectrum print("...saving {}".format(output_fname)) notes = "Freqs [MHz], Peak Flux [Jy/beam], Peak Gauss Flux [Jy/beam], Integrated Gauss Flux [Jy]" np.savez(output_fname, frequencies=freqs, polarizations=pols, peak_flux=peak_flux, peak_flux_err=peak_flux_err,
def source_extract(imfile, source, source_ra, source_dec, source_ext='', radius=1, gaussfit_mult=1.5, rms_max_r=None, rms_min_r=None, pols=1, plot_fit=False): # open fits file hdu = fits.open(imfile) # get header head = hdu[0].header # get info RA, DEC, pol_arr, freqs, stok_ax, freq_ax = casa_utils.get_hdu_info(hdu) dra, ddec = head['CDELT1'], head['CDELT2'] # get axes info npix1 = head["NAXIS1"] npix2 = head["NAXIS2"] nstok = head["NAXIS{}".format(stok_ax)] nfreq = head["NAXIS{}".format(freq_ax)] # get frequency of image freq = head["CRVAL{}".format(freq_ax)] # get radius coordinates: flat-sky approx R = np.sqrt((RA - source_ra)**2 + (DEC - source_dec)**2) # select pixels select = R < radius # polarization check if isinstance(pols, (int, np.integer, str, np.str)): pols = [pols] # iterate over polarizations peak, peak_err, rms, peak_gauss_flux, int_gauss_flux = [], [], [], [], [] for pol in pols: # get polstr if isinstance(pol, (int, np.integer)): polint = pol polstr = uvutils.polnum2str(polint) elif isinstance(pol, (str, np.str)): polstr = pol polint = uvutils.polstr2num(polstr) if polint not in pol_arr: raise ValueError( "Requested polarization {} not found in pol_arr {}".format( polint, pol_arr)) pol_ind = pol_arr.tolist().index(polint) # get data if stok_ax == 3: data = hdu[0].data[0, pol_ind, :, :] elif stok_ax == 4: data = hdu[0].data[pol_ind, 0, :, :] # get beam info for this polarization bmaj, bmin, bpa = casa_utils.get_beam_info(hdu, pol_ind=pol_ind) # check for tclean failed PSF if np.isclose(bmaj, bmin, 1e-6): raise ValueError( "The PSF is not defined for pol {}.".format(polstr)) # relate FWHM of major and minor axes to standard deviation maj_std = bmaj / 2.35 min_std = bmin / 2.35 # calculate beam area in degrees^2 # https://casa.nrao.edu/docs/CasaRef/image.fitcomponents.html beam_area = (bmaj * bmin * np.pi / 4 / np.log(2)) # calculate pixel area in degrees^2 pixel_area = np.abs(dra * ddec) Npix_beam = beam_area / pixel_area # get peak brightness within pixel radius _peak = np.nanmax(data[select]) # get rms outside of source radius if rms_max_r is not None and rms_max_r is not None: rms_select = (R < rms_max_r) & (R > rms_min_r) _rms = np.sqrt(np.mean(data[rms_select]**2)) else: _rms = np.sqrt(np.mean(data[~select]**2)) ## fit a 2D gaussian and get integrated and peak flux statistics ## # recenter R array by peak flux point and get thata T array peak_ind = np.argmax(data[select]) peak_ra = RA[select][peak_ind] peak_dec = DEC[select][peak_ind] X = (RA - peak_ra) Y = (DEC - peak_dec) R = np.sqrt(X**2 + Y**2) X[np.where(np.isclose(X, 0.0))] = 1e-5 T = np.arctan(Y / X) # use synthesized beam as data mask ecc = maj_std / min_std beam_theta = bpa * np.pi / 180 + np.pi / 2 EMAJ = R * np.sqrt( np.cos(T + beam_theta)**2 + ecc**2 * np.sin(T + beam_theta)**2) fit_mask = EMAJ < (maj_std * gaussfit_mult) masked_data = data.copy() masked_data[~fit_mask] = 0.0 # fit 2d gaussian gauss_init = mod.functional_models.Gaussian2D(_peak, peak_ra, peak_dec, x_stddev=maj_std, y_stddev=min_std) fitter = mod.fitting.LevMarLSQFitter() gauss_fit = fitter(gauss_init, RA[fit_mask], DEC[fit_mask], data[fit_mask]) # get gaussian fit properties _peak_gauss_flux = gauss_fit.amplitude.value P = np.array([X, Y]).T beam_theta -= np.pi / 2 # correct for previous + np.pi/2 Prot = P.dot( np.array([[np.cos(beam_theta), -np.sin(beam_theta)], [np.sin(beam_theta), np.cos(beam_theta)]])) gauss_cov = np.array([[gauss_fit.x_stddev.value**2, 0], [0, gauss_fit.y_stddev.value**2]]) # try to get integrated flux try: model_gauss = stats.multivariate_normal.pdf(Prot, mean=np.array([0, 0]), cov=gauss_cov) model_gauss *= gauss_fit.amplitude.value / model_gauss.max() nanmask = ~np.isnan(model_gauss) _int_gauss_flux = np.nansum(model_gauss) / Npix_beam except: model_gauss = np.zeros_like(data) _int_gauss_flux = 0 # get peak error # http://www.gb.nrao.edu/~bmason/pubs/m2mapspeed.pdf beam = np.exp(-((X / maj_std)**2 + (Y / min_std)**2)) _peak_err = _rms / np.sqrt(np.sum(beam**2)) # append peak.append(_peak) peak_err.append(_peak_err) rms.append(_rms) peak_gauss_flux.append(_peak_gauss_flux) int_gauss_flux.append(_int_gauss_flux) # plot if plot_fit: # get postage cutout ra_axis = RA[npix1 // 2] dec_axis = DEC[:, npix2 // 2] ra_select = np.where(np.abs(ra_axis - source_ra) < radius)[0] dec_select = np.where(np.abs(dec_axis - source_dec) < radius)[0] d = data[ra_select[0]:ra_select[-1] + 1, dec_select[0]:dec_select[-1] + 1] m = model_gauss[ra_select[0]:ra_select[-1] + 1, dec_select[0]:dec_select[-1] + 1] # setup wcs and figure wcs = WCS(head, naxis=2) fig = plt.figure(figsize=(14, 5)) fig.subplots_adjust(wspace=0.2) fig.suptitle("Source {} from {}\n{:.2f} MHz".format( source, imfile, freq / 1e6), fontsize=10) # make 3D plot if mplot: ax = fig.add_subplot(131, projection='3d') ax.axis('off') x, y = np.meshgrid(ra_select, dec_select) ax.plot_wireframe(x, y, m, color='steelblue', lw=2, rcount=20, ccount=20, alpha=0.75) ax.plot_surface(x, y, d, rcount=40, ccount=40, cmap='magma', alpha=0.5) # plot cut-out ax = fig.add_subplot(132, projection=wcs) cax = ax.imshow(data, origin='lower', cmap='magma') ax.contour(fit_mask, origin='lower', colors='lime', levels=[0.5]) ax.contour(model_gauss, origin='lower', colors='snow', levels=np.array([0.5, 0.9]) * np.nanmax(m)) ax.grid(color='w') cbar = fig.colorbar(cax, ax=ax) [tl.set_size(8) for tl in cbar.ax.yaxis.get_ticklabels()] [tl.set_size(10) for tl in ax.get_xticklabels()] [tl.set_size(10) for tl in ax.get_yticklabels()] ax.set_xlim(ra_select[0], ra_select[-1] + 1) ax.set_ylim(dec_select[0], dec_select[-1] + 1) ax.set_xlabel('Right Ascension', fontsize=12) ax.set_ylabel('Declination', fontsize=12) ax.set_title("Source Flux and Gaussian Fit", fontsize=10) # plot residual ax = fig.add_subplot(133, projection=wcs) resid = data - model_gauss vlim = np.abs(resid[fit_mask]).max() cax = ax.imshow(resid, origin='lower', cmap='magma', vmin=-vlim, vmax=vlim) ax.contour(fit_mask, origin='lower', colors='lime', levels=[0.5]) ax.grid(color='w') ax.set_xlabel('Right Ascension', fontsize=12) cbar = fig.colorbar(cax, ax=ax) cbar.set_label(head['BUNIT'], fontsize=10) [tl.set_size(8) for tl in cbar.ax.yaxis.get_ticklabels()] [tl.set_size(10) for tl in ax.get_xticklabels()] [tl.set_size(10) for tl in ax.get_yticklabels()] ax.set_xlim(ra_select[0], ra_select[-1] + 1) ax.set_ylim(dec_select[0], dec_select[-1] + 1) ax.set_title("Residual", fontsize=10) fig.savefig('{}.{}.png'.format( os.path.splitext(imfile)[0], source + source_ext)) plt.close() peak = np.asarray(peak) peak_err = np.asarray(peak_err) rms = np.asarray(rms) peak_gauss_flux = np.asarray(peak_gauss_flux) int_gauss_flux = np.asarray(int_gauss_flux) return peak, peak_err, rms, peak_gauss_flux, int_gauss_flux, freq
def empty_uvdata(nfreq, ntimes, ants, antpairs=None, pols=[ 'xx', ], time_per_integ=10.7, min_freq=0.1, channel_bw=0.1 / 1024., instrument='hera_sim', telescope_location=HERA_LOCATION, telescope_lat_lon_alt=HERA_LAT_LON_ALT, object_name='sim_data', start_jd=2458119.5, vis_units='uncalib'): """ Create an empty UVData object with valid metadata and zeroed data arrays with the correct dimensions. Args: nfreq (int) : number of frequency channels. ntimes (int): number of LST bins. ant (dict): antenna positions. The key should be an integer antenna ID, and the value should be a tuple of (x, y, z) positions in the units required by UVData.antenna_positions (meters, position relative to telescope_location). Example:: ants = {0 : (20., 20., 0.)} antpairs (list of len-2 tuples): List of baselines as antenna pair tuples, e.g. ``bls = [(1,2), (3,4)]``. All antennas must be in the ants dict. pols (list of str, optional): polarization strings. time_per_integ (float, optional): Time per integration. min_freq (float, optional): minimum frequency of the frequency array [GHz] channel_bw (float, optional): frequency channel bandwidth [GHz]. instrument (str, optional): name of the instrument. telescope_location (list of float, optional): location of the telescope, in default UVData coordinate system. Expects a list of length 3. telescope_lat_lon_alt (tuple of float, optional): Latitude, longitude, and altitude of telescope, corresponding to the coordinates in telescope_location. Default: HERA_LAT_LON_ALT. object_name (str, optional): name of UVData object start_jd (float, optional): Julian date of the first time sample in the dataset. vis_units (str, optional): assumed units of the visibility data. Returns: :class:`pyuvdata.UVData`: A new UVData object containing valid metadata and blank (zeroed) arrays. """ # Generate empty UVData object uvd = uv.UVData() # Basic time and freq. specs sim_freq = (min_freq + np.arange(nfreq) * channel_bw) * 1e9 # Hz sim_times = start_jd + np.arange(ntimes) * time_per_integ / SEC_PER_SDAY sim_pols = pols lat, lon, alt = telescope_lat_lon_alt sim_lsts = get_lst_for_time(sim_times, lat, lon, alt) # Basic telescope metadata uvd.instrument = instrument uvd.telescope_name = uvd.instrument uvd.telescope_location = np.array(telescope_location) uvd.telescope_lat_lon_alt = telescope_lat_lon_alt uvd.history = "Generated by hera_sim" uvd.object_name = object_name uvd.vis_units = vis_units # Fill-in array layout using dish positions nants = len(ants.keys()) uvd.antenna_numbers = np.array([int(antid) for antid in ants.keys()], dtype=np.int) uvd.antenna_names = [str(antid) for antid in uvd.antenna_numbers] uvd.antenna_positions = np.zeros((nants, 3)) uvd.Nants_data = nants uvd.Nants_telescope = nants # Populate antenna position table for i, antid in enumerate(ants.keys()): uvd.antenna_positions[i] = np.array(ants[antid]) # Generate the antpairs if they are not given explicitly. antpairs, ant1, ant2 = _get_antpairs(ants, antpairs) defined_ants = ants.keys() ants_not_found = [] for _ant in np.unique((ant1, ant2)): if _ant not in defined_ants: ants_not_found.append(_ant) if len(ants_not_found) > 0: raise KeyError("Baseline list contains antennas that were not " "defined in the 'ants' dict: %s" % ants_not_found) # Convert to baseline integers bls = [uvd.antnums_to_baseline(*_antpair) for _antpair in antpairs] bls = np.unique(bls) # Convert back to ant1 and ant2 lists ant1, ant2 = list(zip(*[uvd.baseline_to_antnums(_bl) for _bl in bls])) # Add frequency and polarization arrays uvd.freq_array = sim_freq.reshape((1, sim_freq.size)) uvd.polarization_array = np.array([polstr2num(_pol) for _pol in sim_pols], dtype=np.int) uvd.channel_width = sim_freq[1] - sim_freq[0] uvd.Nfreqs = sim_freq.size uvd.Nspws = 1 uvd.Npols = len(sim_pols) # Generate LST array (for each LST: Nbls copy of LST) # and bls array (repeat bls list Ntimes times) bl_arr, lst_arr = np.meshgrid(np.array(bls), sim_lsts) uvd.baseline_array = bl_arr.flatten() uvd.lst_array = lst_arr.flatten() # Time array _, time_arr = np.meshgrid(np.array(bls), sim_times) uvd.time_array = time_arr.flatten() # Set antenna arrays (same shape as baseline_array) ant1_arr, _ = np.meshgrid(np.array(ant1), sim_lsts) ant2_arr, _ = np.meshgrid(np.array(ant2), sim_lsts) uvd.ant_1_array = ant1_arr.flatten() uvd.ant_2_array = ant2_arr.flatten() # Sets UVWs uvd.set_uvws_from_antenna_positions() # Populate array lengths uvd.Nbls = len(bls) uvd.Ntimes = sim_lsts.size uvd.Nblts = bl_arr.size # Initialise data, flag, and integration arrays uvd.data_array = np.zeros((uvd.Nblts, uvd.Nspws, uvd.Nfreqs, uvd.Npols), dtype=np.complex64) uvd.flag_array = np.zeros((uvd.Nblts, uvd.Nspws, uvd.Nfreqs, uvd.Npols), dtype=bool) uvd.nsample_array = np.ones((uvd.Nblts, uvd.Nspws, uvd.Nfreqs, uvd.Npols), dtype=np.float32) uvd.spw_array = np.ones(1, dtype=np.int) uvd.integration_time = time_per_integ * np.ones(uvd.Nblts) # per bl-time uvd.phase_type = 'drift' # Check validity and return uvd.check() return uvd
def uvd_downselect(uvd, use_ants=None, use_bls=None, use_pols='linear', use_autos=True, use_cross=False, use_times=None, use_freqs=None): """Select only a subset of the data in ``uvd``. XXX refer to numpydoc style (does UVData etc receive ``...``?) Parameters ---------- uvd : UVData or list of UVData UVData object, or path to a file that may be read by a UVData object, on which to perform the data reduction. A list of UVData objects or strings may also be passed, but the list must be of uniform type. If a list is passed, then *all* UVData objects are loaded in a single UVData object. use_ants : array-like of int, optional List of antenna numbers whose data should be kept. Default is to keep data for all antennas. use_bls : array-like of 2- or 3-tuples, optional List of antenna pairs or baseline tuples specifying which baselines (and possibly polarizations) to keep. Default is to keep all baselines. use_pols : str or array-like, optional If passing a string, then it must be one of the following: 'linear', 'cross', or 'all' (these refer to which visibility polarizations to keep). If passing an array-like object, then the entries must either be polarization strings or polarization integers. Polarization strings are automatically converted to polarization integers. Default is to use only the linear polarizations ('xx' and 'yy' or 'ee' and 'nn'). use_autos : bool, optional Whether to keep the autocorrelations. Default is to keep the autos. use_cross : bool, optional Whether to keep the cross-correlations. Default is to discard the cross-correlations. use_times : array-like, optional Times to keep. If length-2, then it is interpreted as a range of time values and must be specified in Julian Date. Otherwise, the times must exist in the ``time_array`` attribute of the ``UVData`` object corresponding to ``uvd``. Default is to use all times. use_freqs : array-like, optional Frequencies or frequency channels to keep. If each entry is an integer, then it is interpreted as frequency channels. Default is to use all frequencies. Returns ------- uvd : UVData UVData object downselected according to the parameters chosen. """ # handle different types for ``uvd`` if isinstance(uvd, str): uvd_ = uvd uvd = UVData() uvd = uvd.read(uvd_) elif isinstance(uvd, UVData): pass elif isinstance(uvd, (list, tuple)): if all([isinstance(uvd_, str) for uvd_ in uvd]): uvd_ = uvd uvd = UVData() uvd = uvd.read(uvd_) elif all([isinstance(uvd_, UVData) for uvd_ in uvd]): _uvd = uvd[0] for uvd_ in uvd[1:]: _uvd += uvd_ uvd = _uvd else: raise ValueError( "If you pass a list or tuple for ``uvd``, then every entry " "in the list must be of the same type (str or UVData)." ) else: raise ValueError( "``uvd`` must be either a string, UVData object, or list/tuple " "of strings or UVData objects (no mixing of types allowed)." ) # first downselect: polarization # XXX can make a helper function for this pol_strings = uvd.get_pols() pol_array = uvd.polarization_array if use_pols == 'linear': use_pols = [ polstr2num(pol) for pol in ('xx', 'yy') # polstr2num(pol) for pol in pol_strings # if pol[0] == pol[1] ] elif use_pols == 'cross': use_pols = [ polstr2num(pol) for pol in ('xy', 'yx') # polstr2num(pol) for pol in pol_strings # if pol[0] != pol[1] ] elif use_pols == 'all': use_pols = pol_array else: try: _ = iter(use_pols) except TypeError: raise ValueError( "``use_pols`` must be one of the following:\n" "'linear' : use only linear polarizations\n" "'cross' : use only cross polarizations\n" "'all' : use all polarizations\n" "iterable of polarization numbers" ) try: use_pols = [ polstr2num(pol) if isinstance(pol, str) else pol for pol in use_pols ] if not all([type(pol) is int for pol in use_pols]): raise KeyError # in case polarizations aren't recognized except KeyError: warnings.warn( "Polarizations not recognized. Skipping polarization downselect." ) use_pols = pol_array # actually downselect along polarization if use_pols != pol_array: uvd.select(polarizations=use_pols, keep_all_metadata=False) # next downselect: visibility type if use_autos and not use_cross: ant_str = 'auto' elif not use_autos and use_cross: ant_str = 'cross' else: ant_str = 'all' if ant_str != 'all': uvd.select(ant_str=ant_str, keep_all_metadata=False) # next downselect: antennas if use_ants is not None: uvd.select(antenna_nums=use_ants, keep_all_metadata=False) # next downselect: baselines if use_bls is not None: uvd.select(bls=use_bls, keep_all_metadata=False) # next downselect: frequency if use_freqs is not None: if all([isinstance(freq, int) for freq in use_freqs]): uvd.select(freq_chans=use_freqs, keep_all_metadata=False) else: uvd.select(frequencies=use_freqs, keep_all_metadata=False) # next downselect: time if use_times is not None: if len(use_times) == 2: uvd.select(time_range=use_times, keep_all_metadata=False) else: uvd.select(times=use_times, keep_all_metadata=False) # all the downselecting should be done at this point return uvd
def uvdata_from_sim_data( array_lat, array_lon, array_height, jd_axis, nu_axis, r_axis, ant_pairs, V_sim, integration_time="derived", channel_width="derived", antenna_numbers="derived", antenna_names="derived", instrument="left blank by user", telescope_name="left blank by user", history="left blank by user", object_name="left blank by user", ): HERA_LAT = array_lat HERA_LON = array_lon HERA_HEIGHT = array_height HERA_LAT_LON_ALT = (HERA_LAT, HERA_LON, HERA_HEIGHT) HERA_LOC = coord.EarthLocation( lat=HERA_LAT * units.rad, lon=HERA_LON * units.rad, height=HERA_HEIGHT * units.meter, ) jd_obj = Time(jd_axis, format="jd", location=HERA_LOC) lst_axis = jd_obj.sidereal_time("apparent").radian if integration_time == "derived": del_jd = jd_obj[1] - jd_obj[0] integration_time = del_jd.sec uvd = UVData() uvd.telescope_location = HERA_LOC.value uvd.telescope_location_lat_lon_alt = HERA_LAT_LON_ALT if antenna_numbers == "derived": uvd.antenna_numbers = np.arange(r_axis.shape[0], dtype=np.int64) ant_num_pairs = ant_pairs else: uvd.antenna_numbers = antenna_numbers ant_num_pairs = np.array([(antenna_numbers[ap[0]], antenna_numbers[ap[1]]) for ap in ant_pairs]) if antenna_names == "derived": uvd.antenna_names = [str(ant_ind) for ant_ind in uvd.antenna_numbers] else: uvd.antenna_names = antenna_names ant_pos_ECEF = pyuvdata.utils.ECEF_from_ENU(r_axis, HERA_LOC.lat.rad, HERA_LOC.lon.rad, HERA_LOC.height.value) uvd.antenna_positions = ant_pos_ECEF - uvd.telescope_location bls = [ uvd.antnums_to_baseline(ant_num_pair[0], ant_num_pair[1]) for ant_num_pair in ant_num_pairs ] uvd.freq_array = nu_axis.reshape((1, nu_axis.size)) pols = ["xx", "yy", "xy", "yx"] uvd.polarization_array = np.array( [polstr2num(pol_str) for pol_str in pols]) uvd.x_orientation = "east" if channel_width == "derived": uvd.channel_width = nu_axis[1] - nu_axis[0] else: uvd.channel_width = channel_width uvd.Nfreqs = nu_axis.size uvd.Nspws = 1 uvd.Npols = len(pols) bl_arr, lst_arr = np.meshgrid(np.array(bls), lst_axis) uvd.baseline_array = bl_arr.flatten() uvd.lst_array = lst_arr.flatten() _, time_arr = np.meshgrid(np.array(bls), jd_axis) uvd.time_array = time_arr.flatten() ant1_arr, _ = np.meshgrid(ant_num_pairs[:, 0], lst_axis) ant2_arr, _ = np.meshgrid(ant_num_pairs[:, 1], lst_axis) uvd.ant_1_array = ant1_arr.flatten() uvd.ant_2_array = ant2_arr.flatten() # Nants_data might be less than r_axis.shape[0] for redundant arrays uvd.Nants_data = len( np.unique(np.r_[ant1_arr.flatten(), ant2_arr.flatten()])) uvd.Nants_telescope = r_axis.shape[0] uvd.set_uvws_from_antenna_positions() uvd.Nbls = len(bls) uvd.Ntimes = lst_axis.size uvd.Nblts = bl_arr.size uvd.data_array = np.zeros((uvd.Nblts, uvd.Nspws, uvd.Nfreqs, uvd.Npols), dtype=np.complex128) uvd.flag_array = np.zeros_like(uvd.data_array, dtype=np.bool) uvd.nsample_array = np.ones((uvd.Nblts, uvd.Nspws, uvd.Nfreqs, uvd.Npols), dtype=np.float64) uvd.spw_array = np.ones(1, dtype=np.int64) uvd.integration_time = integration_time * np.ones(uvd.Nblts) uvd.phase_type = "drift" # (0,0) <-> 'xx' <-> East-East, etc. # matches order of uvd.polarization_array pol_map = {(0, 0): 0, (1, 1): 1, (0, 1): 2, (1, 0): 3} for i_a in range(2): for i_b in range(2): i_p = pol_map[(i_a, i_b)] for k in range(ant_num_pairs.shape[0]): a_i, a_j = ant_num_pairs[k] bl_num = uvd.antnums_to_baseline(a_i, a_j) bl_ind = np.where(uvd.baseline_array == bl_num)[0] uvd.data_array[bl_ind, 0, :, i_p] = np.copy(V_sim[:, :, k, i_a, i_b]) uvd.vis_units = "Jy" uvd.instrument = instrument uvd.telescope_name = telescope_name uvd.history = history uvd.object_name = object_name uvd.check() return uvd
def chunk_files(filenames, inputfile, outputfile, chunk_size, type="data", polarizations=None, spw_range=None, throw_away_flagged_ants=False, clobber=False, ant_flag_yaml=None): """Chunk a list of data or cal files together into a single file. Parameters ---------- filenames: list of strings list of filenames to chunk. Should be homogenous in blt. inputfile: string, name of the file within filenames to use for the start of the chunk. data between the index of input file and the index of inputfile + chunk_size will be chunked together. outpufile: str name of outputfile to write time-concatenated data too. chunk_size: int number of files to chunk after the index of the input file. type : str specify whether "data", "gains" polarizations: list of strs, optional Limit output to polarizations listed. Default None selects all polarizations. spw_range: 2-list or 2-tuple of integers optional lower and upper channel range to select throw_away_flagged_ants: bool, optional if true, throw away baselines that are fully flagged. default is False. clobber: bool, optional if true, overwrite any preexisting output files. defualt is false. flag_yaml : str, optional yaml file with list of antennas to flag and throw away if throw_away_flagged_ants is True Returns ------- None """ filenames = sorted(filenames) start = filenames.index(inputfile) end = start + chunk_size if type == 'data': chunked_files = io.HERAData(filenames[start:end]) elif type == 'gains': chunked_files = io.HERACal(filenames[start:end]) else: raise ValueError("Invalid type provided. Must be in ['data', 'gains']") read_args = {} if type == 'data': if polarizations is None: if len(chunked_files.filepaths) > 1: polarizations = list(chunked_files.pols.values())[0] else: polarizations = chunked_files.pols if spw_range is None: spw_range = (0, chunked_files.Nfreqs) data, flags, nsamples = chunked_files.read(axis='blt', polarizations=polarizations, freq_chans=range( spw_range[0], spw_range[1])) elif type == 'gains': chunked_files.read() if polarizations is None: polarizations = [pol[1:] for pol in chunked_files.pols] if spw_range is None: spw_range = (0, chunked_files.Nfreqs) # convert polarizations to jones integers. jones = [ uvutils.polstr2num(pol, x_orientation=chunked_files.x_orientation) for pol in polarizations ] chunked_files.select(freq_chans=np.arange(spw_range[0], spw_range[1]).astype(int), jones=jones) # throw away fully flagged baselines. if throw_away_flagged_ants: from hera_qm.utils import apply_yaml_flags chunked_files = apply_yaml_flags(chunked_files, ant_flag_yaml, flag_freqs=False, flag_times=False, flag_ants=True, ant_indices_only=True, throw_away_flagged_ants=True) if type == 'data': chunked_files.write_uvh5(outputfile, clobber=clobber) elif type == 'gains': chunked_files.write_calfits(outputfile, clobber=clobber)
def beam_normalized_response(self, pol='pI', freq=None, x_orientation=None): """ Outputs beam response for given polarization as a function of pixels on the sky and input frequencies. The response needs to be peak normalized, and is read in from Healpix coordinates. Uses interp_freq function from uvbeam for interpolation of beam response over given frequency values. Parameters ---------- pol: str, optional Which polarization to compute the beam response for. 'pI', 'pQ', 'pU', 'pV', 'XX', 'YY', 'XY', 'YX' The output shape is (Nfreq, Npixels) Default: 'pI' freq: array, optional Frequencies [Hz] to interpolate onto. x_orientation: str, optional Orientation in cardinal direction east or north of X dipole. Default keeps polarization in X and Y basis. Returns ------- beam_res : float, array-like Beam response as a function healpix indices and frequency. omega : float, array-like Beam solid angle as a function of frequency nside : int, scalar used to compute resolution """ if self.primary_beam.beam_type != 'power': raise ValueError('beam_type must be power') if self.primary_beam.Naxes_vec > 1: raise ValueError('Expect scalar for power beam, found vector') if self.primary_beam._data_normalization.value != 'peak': raise ValueError('beam must be peak normalized') if self.primary_beam.pixel_coordinate_system != 'healpix': raise ValueError('Currently only healpix format supported') nside = self.primary_beam.nside beam_res = self.primary_beam._interp_freq( freq ) # interpolate beam in frequency, based on the data frequencies beam_res = beam_res[0] if isinstance(pol, (str, np.str)): pol = uvutils.polstr2num(pol, x_orientation=x_orientation) pol_array = self.primary_beam.polarization_array if pol in pol_array: stokes_p_ind = np.where(np.isin(pol_array, pol))[0][0] beam_res = beam_res[ 0, 0, stokes_p_ind] # extract the beam with the correct polarization, dim (nfreq X npix) else: raise ValueError('Do not have the right polarization information') omega = np.sum(beam_res, axis=-1) * np.pi / ( 3. * nside**2 ) #compute beam solid angle as a function of frequency return beam_res, omega, nside
if _is_cardinal(antpol): return jnum2str(jstr2num(antpol, x_orientation='north'), x_orientation='north') else: return jnum2str(jstr2num(antpol)) def _comply_vispol(pol): '''Maps an input visibility polarization string onto a string compliant with pyuvdata and hera_cal.''' if _is_cardinal(pol): return polnum2str(polstr2num(pol, x_orientation='north'), x_orientation='north') else: return polnum2str(polstr2num(pol)) _VISPOLS = set([pol for pol in list(POL_STR2NUM_DICT.keys()) if polstr2num(pol) < 0]) # Add east/north polarizations to _VISPOLS while relying only on pyuvdata definitions for pol in copy.deepcopy(_VISPOLS): try: _VISPOLS.add(polnum2str(polstr2num(pol), x_orientation='north')) except KeyError: pass SPLIT_POL = {pol: (_comply_antpol(pol[0]), _comply_antpol(pol[1])) for pol in _VISPOLS} JOIN_POL = {v: k for k, v in SPLIT_POL.items()} def split_pol(pol): '''Splits visibility polarization string (pyuvdata's polstr) into antenna polarization strings (pyuvdata's jstr).''' return SPLIT_POL[_comply_vispol(pol)]
from pyuvdata import UVData from pyuvdata.utils import polstr2num # figure out which directory the data lives in filename = sys.argv[1] dirname, file_basename = os.path.split(filename) jd_pattern = re.compile("[0-9]{7}") jd = jd_pattern.findall(file_basename)[0] file_glob = sorted( glob.glob(os.path.join(dirname, "zen.{jd}.*.uvh5".format(jd=jd)))) file_glob = list([fname for fname in file_glob if "diff" not in fname]) if len(file_glob) == 0: raise FileNotFoundError("Something went wrong--no files were found.") # load in the data, downselect to autos and linear pols use_pols = [polstr2num(pol) for pol in ('xx', 'yy')] uvd = UVData() uvd.read(file_glob, ant_str='auto', polarizations=use_pols) # just do everything in this script; first isolate the rfi rfi_data = np.zeros_like(uvd.data_array, dtype=np.float) normalized_rfi_data = np.zeros_like(rfi_data, dtype=np.float) rfi_flags = np.zeros_like(rfi_data, dtype=np.bool) for antpairpol in uvd.get_antpairpols(): # get indices for properly slicing through data array blt_inds, conj_blt_inds, pol_inds = uvd._key2inds(antpairpol) this_slice = slice(blt_inds, 0, None, pol_inds[0]) # approximately remove all non-rfi signal this_data = uvd.get_data(antpairpol).real