xradio 0.0.31__tar.gz → 0.0.33__tar.gz

{xradio-0.0.31/src/xradio.egg-info → xradio-0.0.33}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.31
+Version: 0.0.33
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License

{xradio-0.0.31 → xradio-0.0.33}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xradio"
-version = "0.0.31"
+version = "0.0.33"
 description = " Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/_utils/list_and_array.py

@@ -25,9 +25,9 @@ def check_if_consistent(array: np.ndarray, array_name: str) -> np.ndarray:
 
     Parameters
     ----------
-    col : _type_
+    array : _type_
         _description_
-    col_name : _type_
+    array_name : _type_
         _description_
 
     Returns
@@ -68,7 +68,9 @@ def unique_1d(array: np.ndarray) -> np.ndarray:
     if array.ndim == 0:
         return np.array([array.item()])
 
-    return np.sort(pd.unique(array))
+    return np.sort(
+        pd.unique(array)
+    )  # Don't remove the sort! It will cause errors that are very difficult to detect. Specifically create_field_info_and_check_ephemeris has a TaQL query that requires this.
 
 
 def pairing_function(antenna_pairs: np.ndarray) -> np.ndarray:

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/__init__.py

@@ -2,14 +2,12 @@ from .read_processing_set import read_processing_set
 from .load_processing_set import load_processing_set
 from .convert_msv2_to_processing_set import convert_msv2_to_processing_set
 
-from .vis_io import read_vis, load_vis_block, write_vis
-
 from .schema import VisibilityXds
 
 __all__ = [
-    "read_vis",
-    "load_vis_block",
-    "write_vis",
+    "read_processing_set",
+    "load_processing_set",
+    "convert_msv2_to_processing_set",
     "VisibilityXds",
     "PointingXds",
     "AntennaXds",

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_processing_set.py

@@ -33,7 +33,7 @@ class processing_set(dict):
             "obs_mode": [],
             "shape": [],
             "polarization": [],
-            "spw_id": [],
+            "spw_name": [],
             # "field_id": [],
             "field_name": [],
             # "source_id": [],
@@ -48,8 +48,8 @@ class processing_set(dict):
         for key, value in self.items():
             summary_data["name"].append(key)
             summary_data["obs_mode"].append(value.attrs["partition_info"]["obs_mode"])
-            summary_data["spw_id"].append(
-                value.attrs["partition_info"]["spectral_window_id"]
+            summary_data["spw_name"].append(
+                value.attrs["partition_info"]["spectral_window_name"]
             )
             summary_data["polarization"].append(value.polarization.values)
 

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/_tables/load_main_table.py

@@ -187,7 +187,7 @@ def load_ddi_cols_chunk(
         cell_shape = cdata.shape
         if len(cell_shape) == 0:
             col_dims = dims[:2]
-            mvars[col.lower()] = xr.DataArray(
+            mvars[col] = xr.DataArray(
                 load_col_chunk(
                     tb_tool, col, (ctlen, cblen), tidxs, bidxs, didxs, None, None
                 ),
@@ -196,7 +196,7 @@ def load_ddi_cols_chunk(
 
         elif col == "UVW":
             col_dims = dims[:2] + ["uvw_coords"]
-            mvars[col.lower()] = xr.DataArray(
+            mvars[col] = xr.DataArray(
                 load_col_chunk(
                     tb_tool, col, (ctlen, cblen, 3), tidxs, bidxs, didxs, None, None
                 ),
@@ -206,7 +206,7 @@ def load_ddi_cols_chunk(
         elif len(cell_shape) == 1:
             pols, col_dims = get_col_1d_pols(cell_shape, dims, chan_cnt, pol_cnt, chunk)
             cshape = (ctlen, cblen) + (pols[1] - pols[0] + 1,)
-            mvars[col.lower()] = xr.DataArray(
+            mvars[col] = xr.DataArray(
                 load_col_chunk(tb_tool, col, cshape, tidxs, bidxs, didxs, pols, None),
                 dims=col_dims,
             )
@@ -215,7 +215,7 @@ def load_ddi_cols_chunk(
             chans, pols = get_col_2d_chans_pols(cell_shape, chan_cnt, pol_cnt, chunk)
             cshape = (ctlen, cblen) + (chans[1] - chans[0] + 1, pols[1] - pols[0] + 1)
             col_dims = dims
-            mvars[col.lower()] = xr.DataArray(
+            mvars[col] = xr.DataArray(
                 load_col_chunk(tb_tool, col, cshape, tidxs, bidxs, didxs, chans, pols),
                 dims=col_dims,
             )

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/_tables/read.py

@@ -279,7 +279,10 @@ def add_units_measures(
     col_descrs = cc_attrs["column_descriptions"]
     # TODO: Should probably loop the other way around, over mvars
     for col in col_descrs:
-        var_name = col.lower()
+        if col == "TIME":
+            var_name = "time"
+        else:
+            var_name = col
         if var_name in mvars and "keywords" in col_descrs[col]:
             if "QuantumUnits" in col_descrs[col]["keywords"]:
                 cc_units = col_descrs[col]["keywords"]["QuantumUnits"]
@@ -364,12 +367,12 @@ def make_freq_attrs(spw_xds: xr.Dataset, spw_id: int) -> Dict[str, Any]:
     ctds_cols = spw_xds.attrs["other"]["msv2"]["ctds_attrs"]["column_descriptions"]
     cfreq = ctds_cols["CHAN_FREQ"]
 
-    cf_attrs = spw_xds.chan_freq.attrs
+    cf_attrs = spw_xds.data_vars["CHAN_FREQ"].attrs
     if "MEASINFO" in cfreq["keywords"] and "VarRefCol" in cfreq["keywords"]["MEASINFO"]:
         fattrs = cfreq["keywords"]["MEASINFO"]
         var_ref_col = fattrs["VarRefCol"]
         # This should point to the SPW/MEAS_FREQ_REF col
-        meas_freq_ref_idx = spw_xds.data_vars[var_ref_col.lower()].values[spw_id]
+        meas_freq_ref_idx = spw_xds.data_vars[var_ref_col].values[spw_id]
 
         if "TabRefCodes" not in fattrs or "TabRefTypes" not in fattrs:
             # Datasets like vla/ic2233_1.ms say "VarRefCol" but "TabRefTypes" is missing
@@ -384,8 +387,8 @@ def make_freq_attrs(spw_xds: xr.Dataset, spw_id: int) -> Dict[str, Any]:
         }
 
         # Also set the 'VarRefCol' for CHAN_FREQ and REF_FREQUENCEY
-        spw_xds.data_vars["chan_freq"].attrs.update(cf_attrs)
-        spw_xds.data_vars["ref_frequency"].attrs.update(cf_attrs)
+        spw_xds.data_vars["CHAN_FREQ"].attrs.update(cf_attrs)
+        spw_xds.data_vars["REF_FREQUENCY"].attrs.update(cf_attrs)
 
     return cf_attrs
 
@@ -440,18 +443,18 @@ def redimension_ms_subtable(xds: xr.Dataset, subt_name: str) -> xr.Dataset:
         (one dimension for every columns)
     """
     subt_key_cols = {
-        "DOPPLER": ["doppler_id", "source_id"],
+        "DOPPLER": ["DOPPLER_ID", "SOURCE_ID"],
         "FREQ_OFFSET": [
-            "antenna1",
-            "antenna2",
-            "feed_id",
-            "spectral_window_id",
-            "time",
+            "ANTENNA1",
+            "ANTENNA2",
+            "FEED_ID",
+            "SPECTRAL_WINDOW_ID",
+            "TIME",
         ],
-        "POINTING": ["time", "antenna_id"],
-        "SOURCE": ["source_id", "time", "spectral_window_id"],
-        "SYSCAL": ["antenna_id", "feed_id", "spectral_window_id", "time"],
-        "WEATHER": ["antenna_id", "time"],
+        "POINTING": ["TIME", "ANTENNA_ID"],
+        "SOURCE": ["SOURCE_ID", "TIME", "SPECTRAL_WINDOW_ID"],
+        "SYSCAL": ["ANTENNA_ID", "FEED_ID", "SPECTRAL_WINDOW_ID", "TIME"],
+        "WEATHER": ["ANTENNA_ID", "TIME"],
         # added tables (MSv3 but not preent in MSv2). Build it from "EPHEMi_... tables
         # Not clear what to do about 'time' var/dim:  , "time"],
         "EPHEMERIDES": ["ephemeris_row_id", "ephemeris_id"],
@@ -476,10 +479,13 @@ def redimension_ms_subtable(xds: xr.Dataset, subt_name: str) -> xr.Dataset:
         # we need to reset to the original type.
         for var in rxds.data_vars:
             if rxds[var].dtype != xds[var].dtype:
-                rxds[var] = rxds[var].astype(xds[var].dtype)
+                # beware of gaps/empty==nan values when redimensioning
+                with np.errstate(invalid="ignore"):
+                    rxds[var] = rxds[var].astype(xds[var].dtype)
     except Exception as exc:
         logger.warning(
-            f"Cannot expand rows to {key_dims}, possibly duplicate values in those coordinates. Exception: {exc}"
+            f"Cannot expand rows in table {subt_name} to {key_dims}, possibly duplicate values in those coordinates. "
+            f"Exception: {exc}"
         )
         rxds = xds.copy()
 
@@ -500,9 +506,9 @@ def add_ephemeris_vars(tname: str, xds: xr.Dataset) -> xr.Dataset:
         ephem_id = 0
 
     xds["ephemeris_id"] = np.uint32(ephem_id) * xr.ones_like(
-        xds["mjd"], dtype=np.uint32
+        xds["MJD"], dtype=np.uint32
     )
-    xds = xds.rename({"mjd": "time"})
+    xds = xds.rename({"MJD": "time"})
     xds["ephemeris_row_id"] = (
         xr.zeros_like(xds["time"], dtype=np.uint32) + xds["row"].values
     )
@@ -529,7 +535,7 @@ def is_nested_ms(attrs: Dict) -> bool:
     )
 
 
-def read_generic_table(
+def load_generic_table(
     inpath: str,
     tname: str,
     timecols: Union[List[str], None] = None,
@@ -574,7 +580,7 @@ def read_generic_table(
     infile = str(infile.expanduser())
     if not os.path.isdir(infile):
         raise ValueError(
-            f"invalid input filename to read_generic_table: {infile} table {tname}"
+            f"invalid input filename to load_generic_table: {infile} table {tname}"
         )
 
     cc_attrs = extract_table_attributes(infile)
@@ -632,7 +638,14 @@ def read_generic_table(
         )
     )
 
-    if tname in ["DOPPLER", "FREQ_OFFSET", "POINTING", "SOURCE", "SYSCAL", "WEATHER"]:
+    if tname in [
+        "DOPPLER",
+        "FREQ_OFFSET",
+        "POINTING",
+        "SOURCE",
+        "SYSCAL",
+        "WEATHER",
+    ]:
         xds = redimension_ms_subtable(xds, tname)
 
     if is_ephem_subtable(tname):
@@ -788,9 +801,9 @@ def load_generic_cols(
             inpath, tb_tool, col, data, timecols
         )
         if array_type == "coord":
-            mcoords[col.lower()] = array_data
+            mcoords[col] = array_data
         elif array_type == "data_var":
-            mvars[col.lower()] = array_data
+            mvars[col] = array_data
 
     return mcoords, mvars
 
@@ -849,9 +862,9 @@ def load_fixed_size_cols(
             inpath, tb_tool, col, data, timecols
         )
         if array_type == "coord":
-            mcoords[col.lower()] = array_data
+            mcoords[col] = array_data
         elif array_type == "data_var":
-            mvars[col.lower()] = array_data
+            mvars[col] = array_data
 
     return mcoords, mvars
 
@@ -924,7 +937,6 @@ def raw_col_data_to_coords_vars(
     # Almost sure that when TIME is present (in a standard MS subt) it
     # is part of the key. But what about non-std subtables, ASDM subts?
     subts_with_time_key = (
-        "FEED",
         "FLAG_CMD",
         "FREQ_OFFSET",
         "HISTORY",
@@ -1176,6 +1188,7 @@ def read_col_conversion(
     cshape: Tuple[int],
     tidxs: np.ndarray,
     bidxs: np.ndarray,
+    use_table_iter: bool,
 ) -> np.ndarray:
     """
     Function to perform delayed reads from table columns when converting
@@ -1233,23 +1246,26 @@ def read_col_conversion(
     data = np.full(cshape + extra_dimensions, np.nan, dtype=col_dtype)
 
     # Use built-in casacore table iterator to populate the data column by unique times.
-    start_row = 0
-    for ts in tb_tool.iter("TIME", sort=False):
-        num_rows = ts.nrows()
+    if use_table_iter:
+        start_row = 0
+        for ts in tb_tool.iter("TIME", sort=False):
+            num_rows = ts.nrows()
 
-        # Create small temporary array to store the partial column
-        tmp_arr = np.full((num_rows,) + extra_dimensions, np.nan, dtype=col_dtype)
+            # Create small temporary array to store the partial column
+            tmp_arr = np.full((num_rows,) + extra_dimensions, np.nan, dtype=col_dtype)
 
-        # Note we don't use `getcol()` because it's less safe. See:
-        # https://github.com/casacore/python-casacore/issues/130#issuecomment-463202373
-        ts.getcolnp(col, tmp_arr)
+            # Note we don't use `getcol()` because it's less safe. See:
+            # https://github.com/casacore/python-casacore/issues/130#issuecomment-463202373
+            ts.getcolnp(col, tmp_arr)
 
-        # Get the slice of rows contained in `tmp_arr`.
-        # Used to get the relevant integer indexes from `tidxs` and `bidxs`
-        tmp_slice = slice(start_row, start_row + num_rows)
+            # Get the slice of rows contained in `tmp_arr`.
+            # Used to get the relevant integer indexes from `tidxs` and `bidxs`
+            tmp_slice = slice(start_row, start_row + num_rows)
 
-        # Copy `tmp_arr` into correct elements of `tmp_arr`
-        data[tidxs[tmp_slice], bidxs[tmp_slice]] = tmp_arr
-        start_row += num_rows
+            # Copy `tmp_arr` into correct elements of `tmp_arr`
+            data[tidxs[tmp_slice], bidxs[tmp_slice]] = tmp_arr
+            start_row += num_rows
+    else:
+        data[tidxs, bidxs] = tb_tool.getcol(col)
 
     return data

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/_tables/read_main_table.py

@@ -24,16 +24,16 @@ from xradio._utils.list_and_array import (
 )
 
 rename_msv2_cols = {
-    "antenna1": "antenna1_id",
-    "antenna2": "antenna2_id",
-    "feed1": "feed1_id",
-    "feed2": "feed2_id",
+    "ANTENNA1": "antenna1_id",
+    "ANTENNA2": "antenna2_id",
+    "FEED1": "feed1_id",
+    "FEED2": "feed2_id",
     # optional cols:
-    "weight_spectrum": "weight",
-    "corrected_data": "vis_corrected",
-    "data": "vis",
-    "model_data": "vis_model",
-    "float_data": "autocorr",
+    "WEIGHT_SPECTRUM": "WEIGHT",
+    "CORRECTED_DATA": "VIS_CORRECTED",
+    "DATA": "VIS",
+    "MODEL_DATA": "VIS_MODEL",
+    "FLOAT_DATA": "AUTOCORR",
 }
 
 
@@ -83,11 +83,11 @@ def redim_id_data_vars(mvars: Dict[str, xr.DataArray]) -> Dict[str, xr.DataArray
     """
     # Vars to drop baseline dim
     var_names = [
-        "array_id",
-        "observation_id",
-        "processor_id",
-        "scan_number",
-        "state_id",
+        "ARRAY_ID",
+        "OBSERVATION_ID",
+        "PROCESSOR_ID",
+        "SCAN_NUMBER",
+        "STATE_ID",
     ]
     for vname in var_names:
         if "baseline" in mvars[vname].coords:
@@ -566,7 +566,7 @@ def concat_tvars_to_mvars(
 
     mvars = {}
     for tvr in tvars.keys():
-        data_var = tvr.lower()
+        data_var = tvr
         if tvr == "UVW":
             mvars[data_var] = xr.DataArray(
                 dask.array.concatenate(tvars[tvr], axis=0),
@@ -742,8 +742,7 @@ def read_flat_main_table(
     # now concat all the dask chunks from each time to make the xds
     mvars = {}
     for kk in bvars.keys():
-        # from uppercase MS col names to lowercase xds var names:
-        data_var = kk.lower()
+        data_var = kk
         if len(bvars[kk]) == 0:
             ignore += [kk]
             continue
@@ -766,7 +765,7 @@ def read_flat_main_table(
             )
 
     mvars["time"] = xr.DataArray(
-        convert_casacore_time(mvars["time"].values), dims=["row"]
+        convert_casacore_time(mvars["TIME"].values), dims=["row"]
     ).chunk({"row": chunks[0]})
 
     # add xds global attributes

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/_tables/read_subtables.py

@@ -16,7 +16,7 @@ from .read import (
     extract_table_attributes,
     add_units_measures,
     table_exists,
-    read_generic_table,
+    load_generic_table,
 )
 from .write import revert_time
 from xradio._utils.list_and_array import unique_1d
@@ -49,7 +49,7 @@ def read_ephemerides(
         logger.debug(f"Reading ephemerides info from: FIELD / {sdir.name}")
         # One "EPHEM_*.tab" (each with a difference ephemeris_id) to concatenate
         ephem.append(
-            read_generic_table(infile, str(Path(*sdir.parts[-2:])), timecols=["MJD"])
+            load_generic_table(infile, str(Path(*sdir.parts[-2:])), timecols=["MJD"])
         )
 
     if ephem:
@@ -339,7 +339,7 @@ def read_delayed_pointing_chunks(
                 ):
                     continue
                 if col not in bvars:
-                    bvars[col.lower()] = []
+                    bvars[col] = []
 
                 cdata = tb_tool.getcol(col, 0, 1)[0]
                 if isinstance(cdata, str):
@@ -356,7 +356,7 @@ def read_delayed_pointing_chunks(
                         None,
                         None,
                     )
-                    bvars[col.lower()] += [
+                    bvars[col] += [
                         dask.array.from_delayed(
                             delayed_array, (ctlen, cblen), cdata.dtype
                         )
@@ -390,6 +390,6 @@ def read_delayed_pointing_chunks(
                                 )
                             ]
                         d1_list += [dask.array.concatenate(d2_list, axis=3)]
-                    bvars[col.lower()] += [dask.array.concatenate(d1_list, axis=2)]
+                    bvars[col] += [dask.array.concatenate(d1_list, axis=2)]
 
     return bvars

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/_tables/write.py

@@ -6,8 +6,6 @@ import xarray as xr
 
 from casacore import tables
 
-from ..msv2_msv3 import ignore_msv2_cols
-
 
 def revert_time(datetimes: np.ndarray) -> np.ndarray:
     """
@@ -72,9 +70,9 @@ def create_table(
 
     if cols is None:
         if ctds_attrs and "column_descriptions" in ctds_attrs:
-            cols = {col: col.lower() for col in ctds_attrs["column_descriptions"]}
+            cols = {col: col for col in ctds_attrs["column_descriptions"]}
         else:
-            cols = {var.upper(): var for var in xds.data_vars}
+            cols = {var: var for var in xds.data_vars}
             # Would add all xds data vars regardless of description availability
             # +
             # list(xds.data_vars) +

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/_tables/write_exp_api.py

@@ -15,16 +15,16 @@ from casacore import tables
 # TODO: this should be consolidated with the equivalent in read_main_table,
 # if we keep this mapping
 rename_to_msv2_cols = {
-    "antenna1_id": "antenna1",
-    "antenna2_id": "antenna2",
-    "feed1_id": "feed1",
-    "feed2_id": "feed2",
+    "antenna1_id": "ANTENNA1",
+    "antenna2_id": "ANTENNA2",
+    "feed1_id": "FEED1",
+    "feed2_id": "FEED2",
     # optional cols:
-    # "weight": "weight_spectrum",
-    "vis_corrected": "corrected_data",
-    "vis": "data",
-    "vis_model": "model_data",
-    "autocorr": "float_data",
+    # "WEIGHT": "WEIGHT_SPECTRUM",
+    "VIS_CORRECTED": "CORRECTED_DATA",
+    "VIS": "DATA",
+    "VIS_MODEL": "MODEL_DATA",
+    "AUTOCORR": "FLOAT_DATA",
 }
 # cols added in xds not in MSv2
 cols_not_in_msv2 = ["baseline_ant1_id", "baseline_ant2_id"]
@@ -214,9 +214,15 @@ def write_ms(
                 continue
 
             col_chunk_size = np.prod([kk[0] for kk in txds[col].chunks])
-            col_rows = (
-                int(np.ceil(max_chunk_size / col_chunk_size)) * txds[col].chunks[0][0]
-            )
+            if max_chunk_size <= 0:
+                max_chunk_size = 19200
+            if col_chunk_size <= 0:
+                col_rows = max_chunk_size
+            else:
+                col_rows = (
+                    int(np.ceil(max_chunk_size / col_chunk_size))
+                    * txds[col].chunks[0][0]
+                )
             for rr in range(0, txds[col].row.shape[0], col_rows):
                 txda = txds[col].isel(row=slice(rr, rr + col_rows))
                 delayed_writes += [
@@ -338,7 +344,7 @@ def write_ms_serial(
                 print(f"Exception writing subtable {subtable}: {exc}")
 
     part_key0 = next(iter(mxds.partitions))
-    vis_data_shape = mxds.partitions[part_key0].vis.shape
+    vis_data_shape = mxds.partitions[part_key0].VIS.shape
     rows_chunk_size = calc_optimal_ms_chunk_shape(
         memory_available_in_bytes, vis_data_shape, 16, "DATA"
     )

{xradio-0.0.31 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/chunks.py

@@ -2,18 +2,11 @@ from pathlib import Path
 from typing import Dict, Tuple
 
 import xarray as xr
-import numpy as np
 
 
-from .msv2_msv3 import ignore_msv2_cols
-from .partition_queries import make_partition_ids_by_ddi_scan
 from .subtables import subt_rename_ids
-from ._tables.load_main_table import load_expanded_main_table_chunk
-from ._tables.read import read_generic_table, make_freq_attrs
+from ._tables.read import load_generic_table
 from ._tables.read_subtables import read_delayed_pointing_table
-from .._utils.partition_attrs import add_partition_attrs
-from .._utils.xds_helper import make_coords
-from xradio._utils.list_and_array import unique_1d
 
 
 def read_spw_ddi_ant_pol(inpath: str) -> Tuple[xr.Dataset]:
@@ -30,81 +23,21 @@ def read_spw_ddi_ant_pol(inpath: str) -> Tuple[xr.Dataset]:
     Tuple[xr.Dataset]
         tuple with antenna, ddi, spw, and polarization setup subtables info
     """
-    spw_xds = read_generic_table(
+    spw_xds = load_generic_table(
         inpath,
         "SPECTRAL_WINDOW",
         rename_ids=subt_rename_ids["SPECTRAL_WINDOW"],
     )
-    ddi_xds = read_generic_table(inpath, "DATA_DESCRIPTION")
-    ant_xds = read_generic_table(
+    ddi_xds = load_generic_table(inpath, "DATA_DESCRIPTION")
+    ant_xds = load_generic_table(
         inpath, "ANTENNA", rename_ids=subt_rename_ids["ANTENNA"]
     )
-    pol_xds = read_generic_table(
+    pol_xds = load_generic_table(
         inpath, "POLARIZATION", rename_ids=subt_rename_ids["POLARIZATION"]
     )
     return ant_xds, ddi_xds, spw_xds, pol_xds
 
 
-def load_main_chunk(
-    infile: str, chunk: Dict[str, slice]
-) -> Dict[Tuple[int, int], xr.Dataset]:
-    """
-    Loads a chunk of visibility data. For every DDI, a separate
-    dataset is produced.
-    This is very loosely equivalent to the
-    partitions.read_*_partitions functions, but in a load (not lazy)
-    fashion and with an implicit single partition wrt. anything but
-    DDIs.
-    Metainfo (sub)tables) are not loaded, and the result is one or more
-    Xarray datasets. It produces one dataset per DDI found within the
-    chunk slice of time/baseline.
-
-    Parameters
-    ----------
-    infile : str
-        MS path (main table)
-    chunk : Dict[str, slice]
-        specification of chunk to load
-
-    Returns
-    -------
-    Dict[Tuple[int, int], xr.Dataset]
-        dictionary of chunk datasets (keys are spw and pol_setup IDs)
-    """
-
-    chunk_dims = ["time", "baseline", "freq", "pol"]
-    if not all(key in chunk_dims for key in chunk):
-        raise ValueError(f"chunks dict has unknown keys. Accepted ones: {chunk_dims}")
-
-    ant_xds, ddi_xds, spw_xds, pol_xds = read_spw_ddi_ant_pol(infile)
-
-    # TODO: constrain this better/ properly
-    data_desc_id, scan_number, state_id = make_partition_ids_by_ddi_scan(infile, False)
-
-    all_xdss = {}
-    data_desc_id = unique_1d(data_desc_id)
-    for ddi in data_desc_id:
-        xds, part_ids, attrs = load_expanded_main_table_chunk(
-            infile, ddi, chunk, ignore_msv2_cols=ignore_msv2_cols
-        )
-
-        coords = make_coords(xds, ddi, (ant_xds, ddi_xds, spw_xds, pol_xds))
-        xds = xds.assign_coords(coords)
-        xds = add_partition_attrs(xds, ddi, ddi_xds, part_ids, other_attrs={})
-
-        # freq dim needs to pull its units/measure info from the SPW subtable
-        spw_id = xds.attrs["partition_ids"]["spw_id"]
-        xds.freq.attrs.update(make_freq_attrs(spw_xds, spw_id))
-        pol_setup_id = ddi_xds.polarization_id.values[ddi]
-
-        chunk_ddi_key = (spw_id, pol_setup_id)
-        all_xdss[chunk_ddi_key] = xds
-
-    chunk_xdss = finalize_chunks(infile, all_xdss, chunk)
-
-    return chunk_xdss
-
-
 def finalize_chunks(
     infile: str, chunks: Dict[str, xr.Dataset], chunk_spec: Dict[str, slice]
 ) -> Dict[Tuple[int, int], xr.Dataset]: