added spectral window subtable

daskms/experimental/katdal/tests/test_chunkstore.py

@@ -10,9 +10,10 @@
 import dask.array as da
 import numba
 import numpy as np
+from numpy.testing import assert_array_equal
 import xarray
 
-from daskms.experimental.zarr import xds_to_zarr
+from daskms.experimental.zarr import xds_from_zarr, xds_to_zarr
 
 from katdal.chunkstore_npy import NpyFileChunkStore
 from katdal.dataset import Subarray
@@ -57,8 +58,7 @@ def __init__(
 
         store = NpyFileChunkStore(str(path))
         shape = (len(timestamps), spw.num_chans, len(corr_products))
-        # data, chunk_info = put_fake_dataset(store, "cb1", shape)
-        data, chunk_info = put_fake_dataset(
+        self._test_data, chunk_info = put_fake_dataset(
             store,
             "cb1",
             shape,
@@ -108,7 +108,14 @@ def weights(self):
         return self._weights
 
 
-@pytest.fixture(scope="session")
+NTIME = 20
+NCHAN = 16
+NANT = 4
+DUMP_RATE = 8.0
+DEFAULT_PARAM = {"ntime": NTIME, "nchan": NCHAN, "nant": NANT, "dump_rate": DUMP_RATE}
+
+
+@pytest.fixture(scope="session", params=[DEFAULT_PARAM])
 def dataset(request, tmp_path_factory):
     path = tmp_path_factory.mktemp("chunks")
     targets = [
@@ -121,19 +128,33 @@ def dataset(request, tmp_path_factory):
         Target("J0408-6545 | PKS 0408-65, radec bpcal, 4:08:20.38, -65:45:09.1"),
         Target("J1346-6024 | Cen B, radec, 13:46:49.04, -60:24:29.4"),
     ]
+    ntime = request.param.get("ntime", NTIME)
+    nchan = request.param.get("nchan", NCHAN)
+    nant = request.param.get("nant", NANT)
+    dump_rate = request.param.get("dump_rate", DUMP_RATE)
+
     # Ensure that len(timestamps) is an integer multiple of len(targets)
-    timestamps = 1234667890.0 + 8.0 * np.arange(20)
+    timestamps = 1234667890.0 + dump_rate * np.arange(ntime)
+
+    assert divmod(ntime, len(targets))[-1] == 0
 
     spw = SpectralWindow(
-        centre_freq=1284e6, channel_width=0, num_chans=16, sideband=1, bandwidth=856e6
+        centre_freq=1284e6,
+        channel_width=0,
+        num_chans=nchan,
+        sideband=1,
+        bandwidth=856e6,
     )
 
     return FakeDataset(
-        path, targets, timestamps, antennas=MEERKAT_ANTENNA_DESCRIPTIONS[:4], spw=spw
+        path, targets, timestamps, antennas=MEERKAT_ANTENNA_DESCRIPTIONS[:nant], spw=spw
     )
 
 
-@pytest.mark.parametrize("include_auto_corrs", [False])
+@pytest.mark.parametrize(
+    "dataset", [{"ntime": 20, "nchan": 16, "nant": 4}], indirect=True
+)
+@pytest.mark.parametrize("include_auto_corrs", [True])
 @pytest.mark.parametrize("row_dim", [True, False])
 @pytest.mark.parametrize("out_store", ["output.zarr"])
 def test_chunkstore(tmp_path_factory, dataset, include_auto_corrs, row_dim, out_store):
@@ -269,71 +290,7 @@ def test_chunkstore(tmp_path_factory, dataset, include_auto_corrs, row_dim, out_
             }
         )
     ]
-    spw_xds = [
-        xarray.Dataset(
-            {
-                "CHAN_FREQ": (("row",), np.array([SPW.centre_freq], np.float64)),
-                "CHAN_WIDTH": (("row",), np.array([SPW.channel_width], np.float64)),
-                "EFFECTIVE_BW": (("row",), np.array([SPW.channel_width], np.float64)),
-                "FLAG_ROW": (("row",), np.zeros(1, dtype=np.int32)),
-                "NUM_CHAN": (("row",), np.array([SPW.num_chans], dtype=np.int32)),
-            }
-        )
-    ]
-    # feed_xds = [
-    #     xarray.Dataset(
-    #         {"ANTENNA_ID":(()),
-
-    #         "BEAM_ID":(()),
-
-    #         "BEAM_OFFSET":(()),
-
-    #         "FEED_ID":(()),
-
-    #         "INTERVAL":(()),
-
-    #         "NUM_RECEPTORS":(()),
-
-    #         "POLARIZATION_TYPE":(()),
-
-    #         "POL_RESPONSE":(()),
-
-    #         "POSITION":(()),
-
-    #         "SPECTRAL_WINDOW_ID":(()),
-
-    #         "TIME":(()),
-
-    #         "RECEPTOR_ANGLE":(()),
-
-    #     })
-    # ]
-    # data_desc_xds = [
-    #     xarray.Dataset({
-    #         "FLAG_ROW":(()),
-
-    #         "POLARIZATION_ID":(()),
-
-    #         "SPECTRAL_WINDOW_ID":(()),
-
-    #     })
-    # ]
-    # obs_xds = [
-    #     xarray.Dataset({
-    #         "FLAG_ROW":(()),
-    #         "LOG":(()),
-    #         "OBSERVER":(()),
-    #         "PROJECT":(()),
-    #         "RELEASE_DATE":(()),
-    #         "SCHEDULE":(()),
-    #         "SCHEDULE_TYPE":(()),
-    #         "TELESCOPE_NAME":(()),
-    #         "TIME_RANGE":(()),
-    #     })
-    # ]
-
     print(ant_xds)
-    print(spw_xds)
     print(xds)
 
     # Reintroduce the shutil.rmtree and remote the tmp_path_factory
@@ -343,4 +300,3 @@ def test_chunkstore(tmp_path_factory, dataset, include_auto_corrs, row_dim, out_
 
     dask.compute(xds_to_zarr(xds, out_store))
     dask.compute(xds_to_zarr(ant_xds, f"{out_store}::ANTENNA"))
-    dask.compute(xds_to_zarr(spw_xds, f"{out_store}::SPECTRAL_WINDOW"))