def cat_with_some_pols():
# Mock catalog with a couple sources polarized.
Nsrcs = 30
Nfreqs = 10
freqs = np.linspace(100, 130, Nfreqs) * 1e6 * units.Hz
pol_inds = range(12, 15)
stokes = np.zeros((4, Nfreqs, Nsrcs))
stokes[0, :, :] = 1.0
stokes[1, :, pol_inds] = 0.2
stokes[2, :, pol_inds] = 1.2
stokes[3, :, pol_inds] = 0.3
ra = Longitude(np.linspace(0, 2 * np.pi, Nsrcs), 'rad')
dec = Latitude(np.linspace(-np.pi / 2, np.pi / 3, Nsrcs), 'rad')
sky = pyradiosky.SkyModel(
name=np.arange(Nsrcs).astype(str),
ra=ra,
dec=dec,
stokes=stokes * units.Jy,
spectral_type='full',
freq_array=freqs
)
return sky
def test_healpix_catalog():
pytest.importorskip('astropy_healpix')
path = os.path.join(SKY_DATA_PATH, 'healpix_disk.hdf5')
sky = pyradiosky.SkyModel()
sky.read_healpix_hdf5(path)
params = {'sources': {'catalog': path}}
hpx_sky = pyuvsim.simsetup.initialize_catalog_from_params(
params, return_recarray=False
)[0]
assert hpx_sky == sky
def test_skymodeldata_share():
# Test the SkyModelData share method.
sky = pyradiosky.SkyModel('src',
Longitude('10d'),
Latitude('5d'),
np.array([1, 0, 0, 0]) * units.Jy,
'spectral_index',
reference_frequency=np.array([100e6]) * units.Hz,
spectral_index=np.array([-0.74]))
smd = pyuvsim.simsetup.SkyModelData()
if mpi.rank == 0:
smd = pyuvsim.simsetup.SkyModelData(sky)
smd.share() # Shares among processes.
sky2 = smd.get_skymodel()
assert sky2 == sky
def test_mock_catalog_zenith_source(hera_loc):
time = Time(2457458.65410, scale='utc', format='jd')
array_location = hera_loc
source_coord = SkyCoord(
alt=Angle(90 * units.deg), az=Angle(0 * units.deg),
obstime=time, frame='altaz', location=array_location
)
icrs_coord = source_coord.transform_to('icrs')
ra = icrs_coord.ra
dec = icrs_coord.dec
test_source = pyradiosky.SkyModel('src0', ra, dec, units.Quantity([1, 0, 0, 0], 'Jy'), 'flat')
cat, mock_keywords = pyuvsim.create_mock_catalog(time, arrangement='zenith')
assert cat == test_source
def test_mock_catalogs():
time = Time(2458098.27471265, scale='utc', format='jd')
arrangements = ['off-zenith', 'zenith', 'cross', 'triangle', 'long-line', 'random', 'hera_text']
cats = {}
for arr in arrangements:
# rseed is only used by the "random" mock catalog
cat, mock_kwds = pyuvsim.simsetup.create_mock_catalog(time, arr, rseed=2458098)
cats[arr] = cat
# For each mock catalog, verify the Ra/Dec source positions against a text catalog.
text_catalogs = {
'cross': 'mock_cross_2458098.27471.txt',
'hera_text': 'mock_hera_text_2458098.27471.txt',
'long-line': 'mock_long-line_2458098.27471.txt',
'off-zenith': 'mock_off-zenith_2458098.27471.txt',
'triangle': 'mock_triangle_2458098.27471.txt',
'random': 'mock_random_2458098.27471.txt',
'zenith': 'mock_zenith_2458098.27471.txt'
}
with pytest.raises(KeyError, match="Invalid mock catalog arrangement: invalid_catalog_name"):
pyuvsim.create_mock_catalog(time, 'invalid_catalog_name')
radec_catalog = pyradiosky.SkyModel()
for arr in arrangements:
radec_catalog.read_text_catalog(
os.path.join(SIM_DATA_PATH, 'test_catalogs', text_catalogs[arr])
)
assert np.all(radec_catalog == cats[arr])
cat, mock_kwds = pyuvsim.simsetup.create_mock_catalog(time, 'random', save=True)
loc = eval(mock_kwds['array_location'])
loc = EarthLocation.from_geodetic(loc[1], loc[0], loc[2]) # Lon, Lat, alt
fname = 'mock_catalog_random.npz'
alts_reload = np.load(fname)['alts']
cat.update_positions(time, loc)
alt, az = cat.alt_az
assert np.all(alt > np.radians(30.))
assert np.allclose(alts_reload, np.degrees(alt))
os.remove(fname)
def test_skymodeldata(component_type, cat_with_some_pols):
# Test that SkyModelData class can properly recreate a SkyModel and subselect.
if component_type == 'point':
sky = cat_with_some_pols
else:
pytest.importorskip('astropy_healpix')
path = os.path.join(SKY_DATA_PATH, 'healpix_disk.hdf5')
sky = pyradiosky.SkyModel()
sky.read_healpix_hdf5(path)
smd = pyuvsim.simsetup.SkyModelData(sky)
assert (smd.ra == sky.ra.deg).all()
assert (smd.dec == sky.dec.deg).all()
if isinstance(sky.stokes, units.Quantity):
smd.stokes_I *= units.Unit(smd.flux_unit)
if smd.polarized is not None:
smd.stokes_Q *= units.Unit(smd.flux_unit)
smd.stokes_U *= units.Unit(smd.flux_unit)
smd.stokes_V *= units.Unit(smd.flux_unit)
assert (smd.stokes_I == sky.stokes[0]).all()
if smd.polarized is not None:
assert (smd.stokes_Q == sky.stokes[..., smd.polarized][1]).all()
assert (smd.stokes_U == sky.stokes[..., smd.polarized][2]).all()
assert (smd.stokes_V == sky.stokes[..., smd.polarized][3]).all()
# Make skymodel from SkyModelData.
sky1 = smd.get_skymodel()
assert sky1 == sky
# Now try with subselection:
sky1_sub = smd.get_skymodel(range(8, 13))
assert sky1.check()
assert sky1_sub.check()
assert sky1_sub.Ncomponents == 5
if smd.polarized is not None:
assert sky1_sub._n_polarized == 1
def test_mock_catalog_off_zenith_source(hera_loc):
src_az = Angle('90.0d')
src_alt = Angle('85.0d')
time = Time(2457458.65410, scale='utc', format='jd')
array_location = hera_loc
source_coord = SkyCoord(
alt=src_alt, az=src_az, obstime=time, frame='altaz', location=array_location
)
icrs_coord = source_coord.transform_to('icrs')
ra = icrs_coord.ra
dec = icrs_coord.dec
test_source = pyradiosky.SkyModel('src0', ra, dec, units.Quantity([1.0, 0, 0, 0], "Jy"), 'flat')
cat, mock_keywords = pyuvsim.create_mock_catalog(time, arrangement='off-zenith',
alt=src_alt.deg)
assert cat == test_source
def test_skymodeldata_with_quantity_stokes(unit, cat_with_some_pols):
# Support for upcoming pyradiosky change setting SkyModel.stokes
# to an astropy Quantity.
if unit == 'Jy':
sky = cat_with_some_pols
else:
pytest.importorskip('analytic_diffuse')
sky, _ = pyuvsim.simsetup.create_mock_catalog(
Time.now(), arrangement='diffuse', diffuse_model='monopole', map_nside=16
)
if not isinstance(sky.stokes, units.Quantity):
sky.stokes *= units.Unit(unit)
smd = pyuvsim.simsetup.SkyModelData(sky)
assert np.all(sky.stokes.to_value(unit)[0] == smd.stokes_I)
assert smd.flux_unit == unit
sky2 = smd.get_skymodel()
if units.Quantity != pyradiosky.SkyModel()._stokes.expected_type:
sky.stokes = sky.stokes.to_value(unit)
assert sky2 == sky
def fake_tasks():
sky = pyradiosky.SkyModel('src',
Longitude('10d'),
Latitude('5d'),
np.array([1, 0, 0, 0]) * units.Jy,
'spectral_index',
reference_frequency=np.array([100e6]) * units.Hz,
spectral_index=np.array([-0.74]))
n_tasks = 30
t0 = Time.now()
freq = 50 * units.Hz
objs = [
pyuvsim.UVTask(sky, t0, freq, None, None, freq_i=ii)
for ii in range(n_tasks)
]
for ti, task in enumerate(objs):
task.visibility_vector = np.random.uniform(
0, 3) + np.random.uniform(0, 3) * 1j
task.uvdata_index = (ti, 0, 0)
task.sources = 1 # Replace with something easier to compare later.
objs[ti] = task
return objs
def gsm_sky_model(freqs, resolution="hi", nside=None):
"""
Return a pyradiosky SkyModel object populated with a Global Sky Model datacube in
healpix format.
Parameters
----------
freqs : array_like
Frequency array, in Hz.
resolution : str, optional
Whether to use the high or low resolution pygdsm maps. Options are 'hi' or 'low'.
nside : int, optional
Healpix nside to up- or down-sample the GSM sky model to. Default: `None` (use the
default from `pygdsm`, which is 1024).
Returns
-------
sky_model : pyradiosky.SkyModel
SkyModel object.
"""
import pygdsm
# Initialise GSM object
gsm = pygdsm.GlobalSkyModel2016(data_unit="TRJ", resolution=resolution, freq_unit="Hz")
# Construct GSM datacube
hpmap = gsm.generate(freqs=freqs) # FIXME: nside=1024, ring ordering, galactic coords
hpmap_units = "K"
# Set nside or resample
nside_gsm = int(astropy_healpix.npix_to_nside(hpmap.shape[-1]))
if nside is None:
# Use default nside from pygdsm map
nside = nside_gsm
else:
# Transform to a user-selected nside
hpmap_new = np.zeros((hpmap.shape[0], astropy_healpix.nside_to_npix(nside)),
dtype=hpmap.dtype)
for i in range(hpmap.shape[0]):
hpmap_new[i,:] = hp.ud_grade(hpmap[i,:],
nside_out=nside,
order_in="RING",
order_out="RING")
hpmap = hpmap_new
# Get datacube properties
npix = astropy_healpix.nside_to_npix(nside)
indices = np.arange(npix)
history = "pygdsm.GlobalSkyModel2016, data_unit=TRJ, resolution=low, freq_unit=MHz"
freq = units.Quantity(freqs, "hertz")
# hmap is in K
stokes = units.Quantity(np.zeros((4, len(freq), len(indices))), hpmap_units)
stokes[0] = hpmap * units.Unit(hpmap_units)
# Construct pyradiosky SkyModel
sky_model = pyradiosky.SkyModel(
nside=nside,
hpx_inds=indices,
stokes=stokes,
spectral_type="full",
freq_array=freq,
history=history,
frame="galactic",
hpx_order="ring"
)
sky_model.healpix_interp_transform(frame='icrs', full_sky=True, inplace=True) # do coord transform
assert sky_model.component_type == "healpix"
return sky_model
new_map = healpix_utils.load_map(
'{}/Stokes{}_average_map_empirical_rm_in_eor0.fits'.format(
sourcedir, pol_name))
maps.append(new_map)
# Check that the pixel arrays are identical
for pol_ind in range(len(pols) - 1):
if np.max(np.abs(maps[pol_ind].pix_arr - maps[pol_ind + 1].pix_arr)) != 0:
print('ERROR: Mismatched pixel arrays.')
sys.exit(1)
Ncomponents = np.shape(maps[0].signal_arr)[0]
stokes = np.zeros((4, 1, Ncomponents))
for pol_ind in range(len(pols)):
stokes[pol_ind, :, :] = maps[pol_ind].signal_arr
stokes = Quantity(stokes, 'Jy/sr')
freq_array = Quantity([182000000], 'Hz')
if maps[0].nest:
hpx_order = 'nested'
else:
hpx_order = 'ring'
skymodel = pyradiosky.SkyModel(component_type='healpix',
spectral_type='flat',
stokes=stokes,
freq_array=freq_array,
hpx_inds=maps[0].pix_arr,
hpx_order=hpx_order,
nside=maps[0].nside)
skymodel.write_skyh5('/Users/ruby/EoR/diffuse_map.skyh5')