From 55cb265b461dec5a1234e0e6c7ad739226fd30fd Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Tue, 26 Mar 2024 17:58:57 +0200 Subject: [PATCH] Fix imports --- daskms/experimental/katdal/__init__.py | 4 - daskms/experimental/katdal/mock_dataset.py | 88 ++++++++++++++++ daskms/experimental/katdal/tests/conftest.py | 102 ++----------------- 3 files changed, 96 insertions(+), 98 deletions(-) create mode 100644 daskms/experimental/katdal/mock_dataset.py diff --git a/daskms/experimental/katdal/__init__.py b/daskms/experimental/katdal/__init__.py index e9781f3e..e69de29b 100644 --- a/daskms/experimental/katdal/__init__.py +++ b/daskms/experimental/katdal/__init__.py @@ -1,4 +0,0 @@ -try: - import katdal -except ImportError as e: - raise ImportError("katdal is not installed\n" "pip install dask-ms[katdal]") from e diff --git a/daskms/experimental/katdal/mock_dataset.py b/daskms/experimental/katdal/mock_dataset.py new file mode 100644 index 00000000..3b163500 --- /dev/null +++ b/daskms/experimental/katdal/mock_dataset.py @@ -0,0 +1,88 @@ +from katdal.lazy_indexer import DaskLazyIndexer +from katdal.chunkstore_npy import NpyFileChunkStore +from katdal.dataset import Subarray +from katdal.spectral_window import SpectralWindow +from katdal.vis_flags_weights import ChunkStoreVisFlagsWeights +from katdal.test.test_vis_flags_weights import put_fake_dataset +from katdal.test.test_dataset import MinimalDataSet +from katpoint import Antenna + + +from daskms.experimental.katdal.meerkat_antennas import MEERKAT_ANTENNA_DESCRIPTIONS + +SPW = SpectralWindow( + centre_freq=1284e6, channel_width=0, num_chans=16, sideband=1, bandwidth=856e6 +) + + +class MockDataset(MinimalDataSet): + def __init__( + self, + path, + targets, + timestamps, + antennas=MEERKAT_ANTENNA_DESCRIPTIONS, + spw=SPW, + ): + antennas = list(map(Antenna, antennas)) + corr_products = [ + (a1.name + c1, a2.name + c2) + for i, a1 in enumerate(antennas) + for a2 in antennas[i:] + for c1 in ("h", "v") + for c2 in ("h", "v") + ] + + subarray = Subarray(antennas, corr_products) + assert len(subarray.ants) > 0 + + store = NpyFileChunkStore(str(path)) + shape = (len(timestamps), spw.num_chans, len(corr_products)) + self._test_data, chunk_info = put_fake_dataset( + store, + "cb1", + shape, + chunk_overrides={ + "correlator_data": (1, spw.num_chans, len(corr_products)), + "flags": (1, spw.num_chans, len(corr_products)), + "weights": (1, spw.num_chans, len(corr_products)), + }, + ) + self._vfw = ChunkStoreVisFlagsWeights(store, chunk_info) + self._vis = None + self._weights = None + self._flags = None + super().__init__(targets, timestamps, subarray, spw) + + def _set_keep( + self, + time_keep=None, + freq_keep=None, + corrprod_keep=None, + weights_keep=None, + flags_keep=None, + ): + super()._set_keep(time_keep, freq_keep, corrprod_keep, weights_keep, flags_keep) + stage1 = (time_keep, freq_keep, corrprod_keep) + self._vis = DaskLazyIndexer(self._vfw.vis, stage1) + self._weights = DaskLazyIndexer(self._vfw.weights, stage1) + self._flags = DaskLazyIndexer(self._vfw.flags, stage1) + + @property + def vis(self): + if self._vis is None: + raise ValueError("Selection has not yet been performed") + return self._vis + + @property + def flags(self): + if self._flags is None: + raise ValueError("Selection has not yet been performed") + return self._flags + + @property + def weights(self): + if self._weights is None: + raise ValueError("Selection has not yet been performed") + + return self._weights diff --git a/daskms/experimental/katdal/tests/conftest.py b/daskms/experimental/katdal/tests/conftest.py index 9a82f4f3..589a6361 100644 --- a/daskms/experimental/katdal/tests/conftest.py +++ b/daskms/experimental/katdal/tests/conftest.py @@ -1,98 +1,7 @@ import pytest -katdal = pytest.importorskip("katdal") - -import numpy as np - -from katdal.lazy_indexer import DaskLazyIndexer -from katdal.chunkstore_npy import NpyFileChunkStore -from katdal.dataset import Subarray -from katdal.spectral_window import SpectralWindow -from katdal.vis_flags_weights import ChunkStoreVisFlagsWeights -from katdal.test.test_vis_flags_weights import put_fake_dataset -from katdal.test.test_dataset import MinimalDataSet -from katpoint import Antenna, Target, Timestamp - - from daskms.experimental.katdal.meerkat_antennas import MEERKAT_ANTENNA_DESCRIPTIONS - -SPW = SpectralWindow( - centre_freq=1284e6, channel_width=0, num_chans=16, sideband=1, bandwidth=856e6 -) - - -class FakeDataset(MinimalDataSet): - def __init__( - self, - path, - targets, - timestamps, - antennas=MEERKAT_ANTENNA_DESCRIPTIONS, - spw=SPW, - ): - antennas = list(map(Antenna, antennas)) - corr_products = [ - (a1.name + c1, a2.name + c2) - for i, a1 in enumerate(antennas) - for a2 in antennas[i:] - for c1 in ("h", "v") - for c2 in ("h", "v") - ] - - subarray = Subarray(antennas, corr_products) - assert len(subarray.ants) > 0 - - store = NpyFileChunkStore(str(path)) - shape = (len(timestamps), spw.num_chans, len(corr_products)) - self._test_data, chunk_info = put_fake_dataset( - store, - "cb1", - shape, - chunk_overrides={ - "correlator_data": (1, spw.num_chans, len(corr_products)), - "flags": (1, spw.num_chans, len(corr_products)), - "weights": (1, spw.num_chans, len(corr_products)), - }, - ) - self._vfw = ChunkStoreVisFlagsWeights(store, chunk_info) - self._vis = None - self._weights = None - self._flags = None - super().__init__(targets, timestamps, subarray, spw) - - def _set_keep( - self, - time_keep=None, - freq_keep=None, - corrprod_keep=None, - weights_keep=None, - flags_keep=None, - ): - super()._set_keep(time_keep, freq_keep, corrprod_keep, weights_keep, flags_keep) - stage1 = (time_keep, freq_keep, corrprod_keep) - self._vis = DaskLazyIndexer(self._vfw.vis, stage1) - self._weights = DaskLazyIndexer(self._vfw.weights, stage1) - self._flags = DaskLazyIndexer(self._vfw.flags, stage1) - - @property - def vis(self): - if self._vis is None: - raise ValueError("Selection has not yet been performed") - return self._vis - - @property - def flags(self): - if self._flags is None: - raise ValueError("Selection has not yet been performed") - return self._flags - - @property - def weights(self): - if self._weights is None: - raise ValueError("Selection has not yet been performed") - - return self._weights - +import numpy as np NTIME = 20 NCHAN = 16 @@ -103,6 +12,12 @@ def weights(self): @pytest.fixture(scope="session", params=[DEFAULT_PARAM]) def dataset(request, tmp_path_factory): + MockDataset = pytest.importorskip( + "daskms.experimental.katdal.mock_dataset" + ).MockDataset + SpectralWindow = pytest.importorskip("katdal.spectral_window").SpectralWindow + Target = pytest.importorskip("katpoint").Target + path = tmp_path_factory.mktemp("chunks") targets = [ # It would have been nice to have radec = 19:39, -63:42 but then @@ -121,7 +36,6 @@ def dataset(request, tmp_path_factory): # Ensure that len(timestamps) is an integer multiple of len(targets) timestamps = 1234667890.0 + dump_rate * np.arange(ntime) - assert divmod(ntime, len(targets))[-1] == 0 spw = SpectralWindow( @@ -132,6 +46,6 @@ def dataset(request, tmp_path_factory): bandwidth=856e6, ) - return FakeDataset( + return MockDataset( path, targets, timestamps, antennas=MEERKAT_ANTENNA_DESCRIPTIONS[:nant], spw=spw )