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
