import mwalib
import numpy as np
import pytest
MWAX_CORRELATOR_METAFITS = "test_files/1244973688_1_timestep/1244973688.metafits"
MWAX_CORRELATOR_GPUBOX_FILES = [
"test_files/1244973688_1_timestep/1244973688_20190619100110_ch114_000.fits",
"test_files/1244973688_1_timestep/1244973688_20190619100110_ch114_001.fits",
]
LEGACY_CORRELATOR_METAFITS = "test_files/1101503312_1_timestep/1101503312.metafits"
LEGACY_CORRELATOR_GPUBOX_FILES = [
"test_files/1101503312_1_timestep/1101503312_20141201210818_gpubox01_00.fits",
]
@pytest.fixture
def mwax_cc() -> mwalib.CorrelatorContext:
return mwalib.CorrelatorContext(MWAX_CORRELATOR_METAFITS, MWAX_CORRELATOR_GPUBOX_FILES)
@pytest.fixture
def legacy_cc() -> mwalib.CorrelatorContext:
return mwalib.CorrelatorContext(LEGACY_CORRELATOR_METAFITS, LEGACY_CORRELATOR_GPUBOX_FILES)
def test_legacy_corr_context_attributes(legacy_cc: mwalib.CorrelatorContext):
assert legacy_cc.mwa_version == mwalib.MWAVersion.CorrLegacy
assert legacy_cc.bscale == 0.5
def test_mwax_corr_context_attributes(mwax_cc: mwalib.CorrelatorContext):
assert mwax_cc.mwa_version == mwalib.MWAVersion.CorrMWAXv2
assert mwax_cc.bscale == 1.0
def test_mwax_corr_context_read_visibilities(mwax_cc: mwalib.CorrelatorContext):
timestep = 0
coarse_chan = 10
data_by_bl = mwax_cc.read_by_baseline(timestep, coarse_chan)
assert np.sum(data_by_bl, dtype=np.float64) == 1389140766690.4983
assert data_by_bl.shape == (
mwax_cc.metafits_context.num_baselines,
mwax_cc.metafits_context.num_corr_fine_chans_per_coarse,
mwax_cc.metafits_context.num_visibility_pols * 2,
)
data_by_freq = mwax_cc.read_by_frequency(timestep, coarse_chan)
assert data_by_freq.shape == (
mwax_cc.metafits_context.num_corr_fine_chans_per_coarse,
mwax_cc.metafits_context.num_baselines,
mwax_cc.metafits_context.num_visibility_pols * 2,
)
assert np.sum(data_by_bl, dtype=np.float64) == np.sum(data_by_freq, dtype=np.float64)
def test_mwax_corr_context_read_weights_by_baseline(mwax_cc: mwalib.CorrelatorContext):
timestep = 0
coarse_chan = 10
data = mwax_cc.read_weights_by_baseline(timestep, coarse_chan)
assert data.shape == (
mwax_cc.metafits_context.num_baselines,
mwax_cc.metafits_context.num_visibility_pols,
)
def test_legacy_corr_context_get_fine_chan_freqs_hz_array(
legacy_cc: mwalib.CorrelatorContext,
):
chans = [10, 20]
fine_chan_freqs = legacy_cc.get_fine_chan_freqs_hz_array(chans)
assert len(fine_chan_freqs) == 256
assert fine_chan_freqs[0] == 151680000.0
assert fine_chan_freqs[128] == 164480000.0