xradio 0.0.30__tar.gz → 0.0.33__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.30
+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.30 → xradio-0.0.33}/pyproject.toml

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

xradio-0.0.33/src/xradio/__init__.py

@@ -0,0 +1,13 @@
+import os
+from graphviper.utils.logger import setup_logger
+
+# _logger_name = "xradio"
+# if os.getenv("VIPER_LOGGER_NAME") != _logger_name:
+#     os.environ["VIPER_LOGGER_NAME"] = _logger_name
+#     setup_logger(
+#         logger_name="xradio",
+#         log_to_term=True,
+#         log_to_file=False,  # True
+#         log_file="xradio-logfile",
+#         log_level="DEBUG",
+#     )

xradio-0.0.33/src/xradio/_utils/common.py

@@ -0,0 +1,60 @@
+import numpy as np
+
+_deg_to_rad = np.pi / 180
+
+
+def cast_to_str(x):
+    if isinstance(x, list):
+        return x[0]
+    else:
+        return x
+
+
+def convert_to_si_units(xds):
+    for data_var in xds.data_vars:
+        if "units" in xds[data_var].attrs:
+            for u_i, u in enumerate(xds[data_var].attrs["units"]):
+                if u == "km":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 1e3
+                    xds[data_var].attrs["units"][u_i] = "m"
+                if u == "km/s":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 1e3
+                    xds[data_var].attrs["units"][u_i] = "m/s"
+                if u == "deg":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * np.pi / 180
+                    xds[data_var].attrs["units"][u_i] = "rad"
+                if u == "Au" or u == "AU":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 149597870700
+                    xds[data_var].attrs["units"][u_i] = "m"
+                if u == "Au/d" or u == "AU/d":
+                    xds[data_var][..., u_i] = (
+                        xds[data_var][..., u_i] * 149597870700 / 86400
+                    )
+                    xds[data_var].attrs["units"][u_i] = "m/s"
+                if u == "arcsec":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * np.pi / 648000
+                    xds[data_var].attrs["units"][u_i] = "rad"
+    return xds
+
+
+def add_position_offsets(dv_1, dv_2):
+    # Fun with angles: We are adding angles together. We need to make sure that the results are between -pi and pi.
+    new_pos = dv_1 + dv_2
+
+    while np.any(new_pos[:, 0] > np.pi) or np.any(new_pos[:, 0] < -np.pi):
+        new_pos[:, 0] = np.where(
+            new_pos[:, 0] > np.pi, new_pos[:, 0] - 2 * np.pi, new_pos[:, 0]
+        )
+        new_pos[:, 0] = np.where(
+            new_pos[:, 0] < -np.pi, new_pos[:, 0] + 2 * np.pi, new_pos[:, 0]
+        )
+
+    while np.any(new_pos[:, 1] > np.pi / 2) or np.any(new_pos[:, 1] < -np.pi / 2):
+        new_pos[:, 1] = np.where(
+            new_pos[:, 1] > np.pi / 2, new_pos[:, 1] - np.pi, new_pos[:, 1]
+        )
+        new_pos[:, 1] = np.where(
+            new_pos[:, 1] < -np.pi / 2, new_pos[:, 1] + np.pi, new_pos[:, 1]
+        )
+
+    return new_pos

xradio-0.0.30/src/xradio/_utils/array.py → xradio-0.0.33/src/xradio/_utils/list_and_array.py

@@ -1,6 +1,23 @@
 """Contains optimised functions to be used within other modules."""
 
 import numpy as np
+import xarray as xr
+
+
+def to_list(x):
+    if isinstance(x, (list, np.ndarray)):
+        if x.ndim == 0:
+            return [x.item()]
+        return list(x)  # needed for json serialization
+    return [x]
+
+
+def to_np_array(x):
+    if isinstance(x, (list, np.ndarray)):
+        if x.ndim == 0:
+            return np.array([x.item()])
+        return np.array(x)  # needed for json serialization
+    return np.array([x])
 
 
 def check_if_consistent(array: np.ndarray, array_name: str) -> np.ndarray:
@@ -8,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
@@ -45,7 +62,15 @@ def unique_1d(array: np.ndarray) -> np.ndarray:
         a sorted array of unique values.
 
     """
-    return np.sort(pd.unique(array))
+    if isinstance(array, xr.core.dataarray.DataArray):
+        array = array.values
+
+    if array.ndim == 0:
+        return np.array([array.item()])
+
+    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.30 → 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.30 → xradio-0.0.33}/src/xradio/vis/_processing_set.py

@@ -30,12 +30,14 @@ class processing_set(dict):
     def _summary(self, data_group="base"):
         summary_data = {
             "name": [],
-            "intent": [],
+            "obs_mode": [],
             "shape": [],
             "polarization": [],
-            "spw_id": [],
+            "spw_name": [],
+            # "field_id": [],
             "field_name": [],
-            "field_id": [],
+            # "source_id": [],
+            "source_name": [],
             "field_coords": [],
             "start_frequency": [],
             "end_frequency": [],
@@ -45,9 +47,9 @@ class processing_set(dict):
 
         for key, value in self.items():
             summary_data["name"].append(key)
-            summary_data["intent"].append(value.attrs["partition_info"]["intent"])
-            summary_data["spw_id"].append(
-                value.attrs["partition_info"]["spectral_window_id"]
+            summary_data["obs_mode"].append(value.attrs["partition_info"]["obs_mode"])
+            summary_data["spw_name"].append(
+                value.attrs["partition_info"]["spectral_window_name"]
             )
             summary_data["polarization"].append(value.polarization.values)
 
@@ -61,15 +63,25 @@ class processing_set(dict):
 
             summary_data["shape"].append(value[data_name].shape)
 
-            summary_data["field_id"].append(value.attrs["partition_info"]["field_id"])
+            # summary_data["field_id"].append(value.attrs["partition_info"]["field_id"])
+            # summary_data["source_id"].append(value.attrs["partition_info"]["source_id"])
+
             summary_data["field_name"].append(
-                value[data_name].attrs["field_and_source_xds"].attrs["field_name"]
+                value.attrs["partition_info"]["field_name"]
+            )
+            summary_data["source_name"].append(
+                value.attrs["partition_info"]["source_name"]
             )
             summary_data["start_frequency"].append(value["frequency"].values[0])
             summary_data["end_frequency"].append(value["frequency"].values[-1])
 
             if value[data_name].attrs["field_and_source_xds"].is_ephemeris:
                 summary_data["field_coords"].append("Ephemeris")
+            elif (
+                "time"
+                in value[data_name].attrs["field_and_source_xds"][center_name].coords
+            ):
+                summary_data["field_coords"].append("Multi-Phase-Center")
             else:
                 ra_dec_rad = (
                     value[data_name].attrs["field_and_source_xds"][center_name].values
@@ -106,8 +118,8 @@ class processing_set(dict):
             assert (
                 frame == ms_xds.frequency.attrs["frame"]
             ), "Frequency reference frame not consistent in processing set."
-            if ms_xds.frequency.attrs["spw_id"] not in spw_ids:
-                spw_ids.append(ms_xds.frequency.attrs["spw_id"])
+            if ms_xds.frequency.attrs["spectral_window_id"] not in spw_ids:
+                spw_ids.append(ms_xds.frequency.attrs["spectral_window_id"])
                 freq_axis_list.append(ms_xds.frequency)
 
         freq_axis = xr.concat(freq_axis_list, dim="frequency").sortby("frequency")

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

@@ -16,7 +16,7 @@ from xradio.vis._vis_utils._ms._tables.read_main_table import (
     get_baselines,
     get_baseline_indices,
 )
-from xradio._utils.array import unique_1d
+from xradio._utils.list_and_array import unique_1d
 
 
 def load_expanded_main_table_chunk(
@@ -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.30 → 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):
@@ -747,13 +760,13 @@ def load_generic_cols(
         dict of coordinates and dict of data vars.
     """
 
-    col_cells = find_loadable_filled_cols(tb_tool, ignore)
+    col_types = find_loadable_cols(tb_tool, ignore)
 
     trows = tb_tool.row(ignore, exclude=True)[:]
 
     # Produce coords and data vars from MS columns
     mcoords, mvars = {}, {}
-    for col in col_cells.keys():
+    for col in col_types.keys():
         try:
             # TODO
             # benchmark np.stack() performance
@@ -779,7 +792,7 @@ def load_generic_cols(
             if len(set([isinstance(row[col], dict) for row in trows])) > 1:
                 continue  # can't deal with this case
 
-            data = handle_variable_col_issues(inpath, col, col_cells, trows)
+            data = handle_variable_col_issues(inpath, col, col_types[col], trows)
 
         if len(data) == 0:
             continue
@@ -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
 
@@ -827,7 +840,7 @@ def load_fixed_size_cols(
         dict of coordinates and dict of data vars, ready to construct an xr.Dataset
     """
 
-    loadable_cols = find_loadable_filled_cols(tb_tool, ignore)
+    loadable_cols = find_loadable_cols(tb_tool, ignore)
 
     # Produce coords and data vars from MS columns
     mcoords, mvars = {}, {}
@@ -849,20 +862,23 @@ 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
 
 
-def find_loadable_filled_cols(
+def find_loadable_cols(
     tb_tool: tables.table, ignore: Union[List[str], None]
-) -> Dict:
+) -> Dict[str, str]:
     """
-    For a table, finds the columns that are:
-    - loadable = not of record type, and not to be ignored
-    - filled = the column cells are populated.
+    For a table, finds the columns that are loadable = not of record type,
+    and not to be ignored
+    In extreme cases of variable size columns, it can happen that all the
+    cells are empty (iscelldefined() == false). This is still considered a
+    loadable column, even though all values of the resulting data var will
+    be empty.
 
     Parameters
     ----------
@@ -874,17 +890,15 @@ def find_loadable_filled_cols(
     Returns
     -------
     Dict
-        dict of {column name => first cell} for columns that can/should be loaded
+        dict of {column name: column type} for columns that can/should be loaded
     """
 
     colnames = tb_tool.colnames()
-    # columns that are not populated are skipped. record columns are not supported
+    table_desc = tb_tool.getdesc()
     loadable_cols = {
-        col: tb_tool.getcell(col, 0)
+        col: table_desc[col]["valueType"]
         for col in colnames
-        if (col not in ignore)
-        and (tb_tool.iscelldefined(col, 0))
-        and tb_tool.coldatatype(col) != "record"
+        if (col not in ignore) and tb_tool.coldatatype(col) != "record"
     }
     return loadable_cols
 
@@ -923,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",
@@ -978,7 +991,7 @@ def raw_col_data_to_coords_vars(
 
 
 def handle_variable_col_issues(
-    inpath: str, col: str, col_cells: dict, trows: tables.tablerow
+    inpath: str, col: str, col_type: str, trows: tables.tablerow
 ) -> np.ndarray:
     """
     load variable-size array columns, padding with nans wherever
@@ -992,8 +1005,8 @@ def handle_variable_col_issues(
         path name of the MS
     col : str
         column being loaded
-    col_cells : dict
-        col: cell} values
+    col_type : str
+        type of the column cell values
     trows : tables.tablerow
         rows from a table as loaded by tables.row()
 
@@ -1008,7 +1021,7 @@ def handle_variable_col_issues(
 
     mshape = np.array(max([np.array(row[col]).shape for row in trows]))
     try:
-        pad_nan = get_pad_nan(col_cells[col])
+        pad_nan = get_pad_nan(np.array((), dtype=col_type))
 
         # TODO
         # benchmark np.stack() performance
@@ -1175,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
@@ -1232,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.30 → xradio-0.0.33}/src/xradio/vis/_vis_utils/_ms/_tables/read_main_table.py

@@ -17,23 +17,23 @@ from .read import (
 )
 
 from .table_query import open_table_ro, open_query
-from xradio._utils.array import (
+from xradio._utils.list_and_array import (
     unique_1d,
     pairing_function,
     inverse_pairing_function,
 )
 
 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