xradio 0.0.44__tar.gz → 0.0.46__tar.gz

{xradio-0.0.44/src/xradio.egg-info → xradio-0.0.46}/PKG-INFO

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

{xradio-0.0.44 → xradio-0.0.46}/pyproject.toml

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

xradio-0.0.46/src/xradio/_utils/dict_helpers.py

@@ -0,0 +1,14 @@
+def make_quantity(value, units: str) -> dict:
+    """
+    create a quantity dictionary given value and units
+    Parameters
+    ----------
+    value : numeric or array of numerics
+        Quantity value
+    units: str
+        Quantity units
+    Returns
+    -------
+    dict
+    """
+    return {"value": value, "units": units, "type": "quantity"}

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/image/_util/_casacore/xds_from_casacore.py

@@ -31,6 +31,7 @@ from ..common import (
 )
 from ...._utils._casacore.tables import extract_table_attributes, open_table_ro
 from xradio._utils.coord_math import _deg_to_rad
+from xradio._utils.dict_helpers import make_quantity
 
 """
 def _add_coord_attrs(xds: xr.Dataset, icoords: dict, dir_axes: list) -> xr.Dataset:
@@ -62,13 +63,7 @@ def _add_freq_attrs(xds, coord_dict):
     for k in coord_dict:
         if k.startswith("spectral"):
             sd = coord_dict[k]
-            # meta["native_type"] = _native_types[sd["nativeType"]]
-            meta["rest_frequency"] = {
-                "type": "quantity",
-                "units": "Hz",
-                "value": sd["restfreq"],
-            }
-            # meta["restfreqs"] = {'type': 'quantity', 'units': 'Hz', 'value': list(sd["restfreqs"])}
+            meta["rest_frequency"] = make_quantity(sd["restfreq"], "Hz")
             meta["type"] = "frequency"
             meta["units"] = sd["unit"]
             meta["frame"] = sd["system"]
@@ -184,11 +179,7 @@ def _casa_image_to_xds_attrs(img_full_path: str, history: bool = True) -> dict:
         k = "latpole"
         if k in coord_dir_dict:
             for j in (k, "longpole"):
-                dir_dict[j] = {
-                    "value": coord_dir_dict[j] * _deg_to_rad,
-                    "units": "rad",
-                    "type": "quantity",
-                }
+                dir_dict[j] = make_quantity(coord_dir_dict[j] * _deg_to_rad, "rad")
         for j in ("pc", "projection_parameters", "projection"):
             if j in coord_dir_dict:
                 dir_dict[j] = coord_dir_dict[j]
@@ -518,11 +509,7 @@ def _get_freq_values_attrs(
                     crpix=wcs["crpix"],
                     cdelt=wcs["cdelt"],
                 )
-                attrs["rest_frequency"] = {
-                    "type": "quantity",
-                    "units": "Hz",
-                    "value": sd["restfreq"],
-                }
+                attrs["rest_frequency"] = make_quantity(sd["restfreq"], "Hz")
                 attrs["type"] = "frequency"
                 attrs["units"] = sd["unit"]
                 attrs["frame"] = sd["system"]

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/image/_util/_fits/xds_from_fits.py

@@ -16,6 +16,7 @@ from ..common import (
     _l_m_attr_notes,
 )
 from xradio._utils.coord_math import _deg_to_rad
+from xradio._utils.dict_helpers import make_quantity
 import copy
 import dask
 import dask.array as da
@@ -77,11 +78,7 @@ def _add_freq_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
     freq_coord = xds.coords["frequency"]
     meta = {}
     if helpers["has_freq"]:
-        meta["rest_frequency"] = {
-            "type": "quantity",
-            "units": "Hz",
-            "value": helpers["restfreq"],
-        }
+        meta["rest_frequency"] = make_quantity(helpers["restfreq"], "Hz")
         meta["frame"] = helpers["specsys"]
         meta["units"] = "Hz"
         meta["type"] = "frequency"
@@ -184,16 +181,8 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
         direction["reference"]["value"][i] = x.value
         x = helpers["cdelt"][i] * u.Unit(_get_unit(helpers["cunit"][i]))
         x = x.to("rad")
-    direction["latpole"] = {
-        "value": header["LATPOLE"] * _deg_to_rad,
-        "units": "rad",
-        "type": "quantity",
-    }
-    direction["longpole"] = {
-        "value": header["LONPOLE"] * _deg_to_rad,
-        "units": "rad",
-        "type": "quantity",
-    }
+    direction["latpole"] = make_quantity(header["LATPOLE"] * _deg_to_rad, "rad")
+    direction["longpole"] = make_quantity(header["LONPOLE"] * _deg_to_rad, "rad")
     pc = np.zeros([2, 2])
     for i in (0, 1):
         for j in (0, 1):
@@ -325,9 +314,9 @@ def _beam_attr_from_header(helpers: dict, header) -> Union[dict, str, None]:
     if "BMAJ" in header:
         # single global beam
         beam = {
-            "bmaj": {"type": "quantity", "units": "arcsec", "value": header["BMAJ"]},
-            "bmin": {"type": "quantity", "units": "arcsec", "value": header["BMIN"]},
-            "pa": {"type": "quantity", "units": "arcsec", "value": header["BPA"]},
+            "bmaj": make_quantity(header["BMAJ"], "arcsec"),
+            "bmin": make_quantity(header["BMIN"], "arcsec"),
+            "pa": make_quantity(header["BPA"], "arcsec"),
         }
         return _convert_beam_to_rad(beam)
     elif "CASAMBM" in header and header["CASAMBM"]:

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/image/_util/common.py

@@ -6,6 +6,7 @@ import numpy as np
 from typing import Dict, List
 import xarray as xr
 from xradio._utils.coord_math import _deg_to_rad
+from xradio._utils.dict_helpers import make_quantity
 
 _c = 2.99792458e08 * u.m / u.s
 # OPTICAL = Z
@@ -39,7 +40,7 @@ def _convert_beam_to_rad(beam: dict) -> dict:
             q = u.quantity.Quantity(f"{beam[k]['value']}{beam[k]['units']}")
         q = q.to("rad")
         j = "pa" if k == "positionangle" else k
-        mybeam[j] = {"type": "quantity", "value": q.value, "units": "rad"}
+        mybeam[j] = make_quantity(q.value, "rad")
     return mybeam
 
 
@@ -102,11 +103,7 @@ def _numpy_arrayize_dv(xds: xr.Dataset) -> xr.Dataset:
 
 def _default_freq_info() -> dict:
     return {
-        "rest_frequency": {
-            "value": 1420405751.7860003,
-            "units": "Hz",
-            "type": "quantity",
-        },
+        "rest_frequency": make_quantity(1420405751.7860003, "Hz"),
         "type": "frequency",
         "frame": "LSRK",
         "units": "Hz",

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/image/_util/image_factory.py

@@ -4,6 +4,7 @@ import xarray as xr
 from typing import List, Union
 from .common import _c, _compute_world_sph_dims, _l_m_attr_notes
 from xradio._utils.coord_math import _deg_to_rad
+from xradio._utils.dict_helpers import make_quantity
 
 
 def _input_checks(
@@ -46,11 +47,7 @@ def _add_common_attrs(
     xds.time.attrs = {"format": "MJD", "scale": "UTC", "units": "d"}
     freq_vals = np.array(xds.frequency)
     xds.frequency.attrs = {
-        "rest_frequency": {
-            "type": "quantity",
-            "units": "Hz",
-            "value": restfreq,
-        },
+        "rest_frequency": make_quantity(restfreq, "Hz"),
         "frame": spectral_reference.upper(),
         "units": "Hz",
         "wave_unit": "mm",
@@ -69,8 +66,8 @@ def _add_common_attrs(
                 "value": list(phase_center),
                 "units": ["rad", "rad"],
             },
-            "longpole": {"type": "quantity", "value": np.pi, "units": "rad"},
-            "latpole": {"type": "quantity", "value": 0.0, "units": "rad"},
+            "longpole": make_quantity(np.pi, "rad"),
+            "latpole": make_quantity(0.0, "rad"),
             "pc": np.array([[1.0, 0.0], [0.0, 1.0]]),
             "projection": projection,
             "projection_parameters": [0.0, 0.0],
@@ -289,7 +286,6 @@ def _make_empty_lmuv_image(
         "crval": 0.0,
         "cdelt": -abs(sky_image_cell_size[0]),
         "units": "rad",
-        "type": "quantity",
         "note": attr_note["l"],
     }
     xds.m.attrs = {
@@ -297,7 +293,6 @@ def _make_empty_lmuv_image(
         "crval": 0.0,
         "cdelt": abs(sky_image_cell_size[1]),
         "units": "rad",
-        "type": "quantity",
         "note": attr_note["m"],
     }
     xds.u.attrs = {

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/measurement_set/__init__.py

@@ -1,7 +1,10 @@
 from .processing_set import ProcessingSet
 from .open_processing_set import open_processing_set
 from .load_processing_set import load_processing_set, ProcessingSetIterator
-from .convert_msv2_to_processing_set import convert_msv2_to_processing_set
+from .convert_msv2_to_processing_set import (
+    convert_msv2_to_processing_set,
+    estimate_conversion_memory_and_cores,
+)
 from .measurement_set_xds import MeasurementSetXds
 
 from .schema import SpectrumXds, VisibilityXds
@@ -13,6 +16,7 @@ __all__ = [
     "load_processing_set",
     "ProcessingSetIterator",
     "convert_msv2_to_processing_set",
+    "estimate_conversion_memory_and_cores",
     "SpectrumXds",
     "VisibilityXds",
 ]

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/measurement_set/_utils/_msv2/conversion.py

@@ -679,7 +679,7 @@ def get_weight(
     return xds
 
 
-def create_taql_query(partition_info):
+def create_taql_query_where(partition_info: dict):
     main_par_table_cols = [
         "DATA_DESC_ID",
         "OBSERVATION_ID",
@@ -729,6 +729,192 @@ def fix_uvw_frame(
     return xds
 
 
+def estimate_memory_for_partition(in_file: str, partition: dict) -> float:
+    """
+    Aim: given a partition description, estimates a safe maximum memory value, but avoiding overestimation
+    (at least not adding not well understood factors).
+    """
+
+    def calculate_term_all_data(
+        tb_tool: tables.table, ntimes: float, nbaselines: float
+    ) -> tuple[list[float], bool]:
+        """
+        Size that DATA vars from MS will have in the MSv4, whether this MS has FLOAT_DATA
+        """
+        sizes_all_data_vars = []
+        col_names = tb_tool.colnames()
+        for data_col in ["DATA", "CORRECTED_DATA", "MODEL_DATA", "FLOAT_DATA"]:
+            if data_col in col_names:
+                col_descr = tb_tool.getcoldesc(data_col)
+                if "shape" in col_descr and isinstance(col_descr["shape"], np.ndarray):
+                    # example: "shape": array([15,  4]) => gives pols x channels
+                    cells_in_row = col_descr["shape"].prod()
+                    npols = col_descr["shape"][-1]
+                else:
+                    first_row = np.array(tb_tool.col(data_col)[0])
+                    cells_in_row = np.prod(first_row.shape)
+                    npols = first_row.shape[-1]
+
+                if col_descr["valueType"] == "complex":
+                    # Assume. Otherwise, read first column and get the itemsize:
+                    # col_dtype = np.array(mtable.col(data_col)[0]).dtype
+                    # cell_size = col_dtype.itemsize
+                    cell_size = 4
+                    if data_col != "FLOAT_DATA":
+                        cell_size *= 2
+                elif col_descr["valueType"] == "float":
+                    cell_size = 4
+
+                # cells_in_row should account for the polarization and frequency dims
+                size_data_var = ntimes * nbaselines * cells_in_row * cell_size
+
+                sizes_all_data_vars.append(size_data_var)
+
+        is_float_data = "FLOAT_DATA" in col_names
+
+        return sizes_all_data_vars, is_float_data
+
+    def calculate_term_weight_flag(size_largest_data, is_float_data) -> float:
+        """
+        Size that WEIGHT and FLAG will have in the MSv4, derived from the size of the
+        MSv2 DATA col=> MSv4 VIS/SPECTRUM data var.
+        """
+        # Factors of the relative "cell_size" wrt the DATA var
+        # WEIGHT_SPECTRUM size: DATA (IF), DATA/2 (SD)
+        factor_weight = 1.0 if is_float_data else 0.5
+        factor_flag = 1.0 / 4.0 if is_float_data else 1.0 / 8.0
+
+        return size_largest_data * (factor_weight + factor_flag)
+
+    def calculate_term_other_data_vars(
+        ntimes: int, nbaselines: int, is_float_data: bool
+    ) -> float:
+        """
+        Size all data vars other than the DATA (visibility/spectrum) vars will have in the MSv4
+
+        For the rest of columns, including indices/iteration columns and other
+        scalar columns could say approx ->5% of the (large) data cols
+
+        """
+        # Small ones, but as they are loaded into data arrays, why not including,
+        # For example: UVW (3xscalar), EXPOSURE, TIME_CENTROID
+        # assuming float64 in output MSv4
+        item_size = 8
+        return ntimes * nbaselines * (3 + 1 + 1) * item_size
+
+    def calculate_term_calc_indx_for_row_split(msv2_nrows: int) -> float:
+        """
+        Account for the indices produced in calc_indx_for_row_split():
+        the dominating ones are: tidxs, bidxs, didxs.
+
+        In terms of amount of memory represented by this term relative to the
+        total, it becomes relevant proportionally to the ratio between
+           nrows / (chans x pols)
+        - for example LOFAR long scans/partitions with few channels,
+        but its value is independent from # chans, pols.
+        """
+        item_size = 8
+        # 3 are: tidxs, bidxs, didxs
+        return msv2_nrows * 3 * item_size
+
+    def calculate_term_other_msv2_indices(msv2_nrows: int) -> float:
+        """
+        Account for the allocations to load ID, etc. columns from input MSv2.
+        The converter needs to load: OBSERVATION_ID, INTERVAL, SCAN_NUMBER.
+        These are loaded one after another (allocations do not stack up).
+        Also, in most memory profiles these allocations are released once we
+        get to create_data_variables(). As such, adding this term will most
+        likely lead to overestimation (but adding it for safety).
+
+        Simlarly as with calculate_term_calc_indx_for_row_split() this term
+        becomes relevant when the ratio 'nrows / (chans x pols)' is high.
+        """
+        # assuming float64/int64 in input MSv2, which seems to be the case,
+        # except for OBSERVATION_ID (int32)
+        item_size = 8
+        return msv2_nrows * item_size
+
+    def calculate_term_attrs(size_estimate_main_xds: float) -> float:
+        """Rough guess which seems to be more than enough"""
+        # could also account for info_dicts (which seem to require typically ~1 MB)
+        return 10 * 1024 * 1024
+
+    def calculate_term_sub_xds(size_estimate_main_xds: float) -> float:
+        """
+        This is still very rough. Just seemingly working for now. Not taking into account the dims
+        of the sub-xdss, interpolation options used, etc.
+        """
+        # Most cases so far 1% seems enough
+        return 0.015 * size_estimate_main_xds
+
+    def calculate_term_to_zarr(size_estimate_main_xds: float) -> float:
+        """
+        The to_zarr call on the main_xds seems to allocate 10s or 100s of MBs, presumably for buffers.
+        That adds on top of the expected main_xds size.
+        This is currently a very rough extrapolation and is being (mis)used to give a safe up to 5-6%
+        overestimation. Perhaps we should drop this term once other sub-xdss are accounted for (and
+        this term could be replaced by a similar, smaller but still safe over-estimation percentage).
+        """
+        return 0.05 * size_estimate_main_xds
+
+    taql_partition = create_taql_query_where(partition)
+    taql_main = f"select * from $mtable {taql_partition}"
+    with open_table_ro(in_file) as mtable:
+        col_names = mtable.colnames()
+        with open_query(mtable, taql_main) as tb_tool:
+            # Do not feel tempted to rely on nrows. nrows tends to underestimate memory when baselines are missing.
+            # For some EVN datasets that can easily underestimate by a 50%
+            utimes, _tol = get_utimes_tol(mtable, taql_partition)
+            ntimes = len(utimes)
+            nbaselines = len(get_baselines(tb_tool))
+
+            # Still, use nrwos for estimations related to sizes of input (MSv2)
+            # columns, not sizes of output (MSv4) data vars
+            msv2_nrows = tb_tool.nrows()
+
+            sizes_all_data, is_float_data = calculate_term_all_data(
+                tb_tool, ntimes, nbaselines
+            )
+
+    size_largest_data = np.max(sizes_all_data)
+    sum_sizes_data = np.sum(sizes_all_data)
+    estimate_main_xds = (
+        sum_sizes_data
+        + calculate_term_weight_flag(size_largest_data, is_float_data)
+        + calculate_term_other_data_vars(ntimes, nbaselines, is_float_data)
+    )
+    estimate = (
+        estimate_main_xds
+        + calculate_term_calc_indx_for_row_split(msv2_nrows)
+        + calculate_term_other_msv2_indices(msv2_nrows)
+        + calculate_term_sub_xds(estimate_main_xds)
+        + calculate_term_to_zarr(estimate_main_xds)
+    )
+    estimate /= GiBYTES_TO_BYTES
+
+    return estimate
+
+
+def estimate_memory_and_cores_for_partitions(
+    in_file: str, partitions: list
+) -> tuple[float, int, int]:
+    """
+    Estimates approximate memory required to convert an MSv2 to MSv4, given
+    a predefined set of partitions.
+    """
+    max_cores = len(partitions)
+
+    size_estimates = [
+        estimate_memory_for_partition(in_file, part_description)
+        for part_description in partitions
+    ]
+    max_estimate = np.max(size_estimates) if size_estimates else 0.0
+
+    recommended_cores = np.ceil(max_cores / 4).astype("int")
+
+    return float(max_estimate), int(max_cores), int(recommended_cores)
+
+
 def convert_and_write_partition(
     in_file: str,
     out_file: str,
@@ -790,7 +976,7 @@ def convert_and_write_partition(
         _description_
     """
 
-    taql_where = create_taql_query(partition_info)
+    taql_where = create_taql_query_where(partition_info)
     ddi = partition_info["DATA_DESC_ID"][0]
     intents = str(partition_info["OBS_MODE"][0])
 
@@ -839,9 +1025,11 @@ def convert_and_write_partition(
             start = time.time()
             xds = xr.Dataset(
                 attrs={
-                    "creation_date": datetime.datetime.utcnow().isoformat(),
+                    "creation_date": datetime.datetime.now(
+                        datetime.timezone.utc
+                    ).isoformat(),
                     "xradio_version": importlib.metadata.version("xradio"),
-                    "schema_version": "4.0.-9994",
+                    "schema_version": "4.0.-9991",
                     "type": "visibility",
                 }
             )
@@ -1085,6 +1273,8 @@ def convert_and_write_partition(
                 else:
                     xds.attrs["type"] = "visibility"
 
+            import sys
+
             start = time.time()
             if storage_backend == "zarr":
                 xds.to_zarr(store=os.path.join(file_name, "correlated_xds"), mode=mode)
@@ -1193,7 +1383,12 @@ def antenna_ids_to_names(
         ]
         for unwanted_coord in unwanted_coords_from_ant_xds:
             xds = xds.drop_vars(unwanted_coord)
-        xds = xds.rename({"baseline_id": "antenna_name"})
+
+        # Rename a dim coord started generating warnings (index not re-created). Swap dims, create coord
+        # https://github.com/pydata/xarray/pull/6999
+        xds = xds.swap_dims({"baseline_id": "antenna_name"})
+        xds = xds.assign_coords({"antenna_name": xds["baseline_id"].data})
+        xds = xds.drop_vars("baseline_id")
 
         # drop more vars that seem unwanted in main_sd_xds, but there shouuld be a better way
         # of not creating them in the first place
@@ -1204,58 +1399,58 @@ def antenna_ids_to_names(
     return xds
 
 
+def add_group_to_data_groups(
+    data_groups: dict, what_group: str, correlated_data_name: str, uvw: bool = True
+):
+    """
+    Adds one correlated_data variable to the data_groups dict.
+    A utility function to use when creating/updating data_groups from MSv2 data columns
+    / data variables.
+
+    Parameters
+    ----------
+    data_groups: str
+        The data_groups dict of an MSv4 xds. It is updated in-place
+    what_group: str
+        Name of the data group: "base", "corrected", "model", etc.
+    correlated_data_name: str
+        Name of the correlated_data var: "VISIBILITY", "VISIBILITY_CORRECTED", "SPECTRUM", etc.
+    uvw: bool
+        Whether to add a uvw field to the data group (assume True = interferometric data).
+    """
+    data_groups[what_group] = {
+        "correlated_data": correlated_data_name,
+        "flag": "FLAG",
+        "weight": "WEIGHT",
+    }
+    if uvw:
+        data_groups[what_group]["uvw"] = "UVW"
+
+
 def add_data_groups(xds):
     xds.attrs["data_groups"] = {}
+
+    data_groups = xds.attrs["data_groups"]
     if "VISIBILITY" in xds:
-        xds.attrs["data_groups"]["base"] = {
-            "correlated_data": "VISIBILITY",
-            "flag": "FLAG",
-            "weight": "WEIGHT",
-            "uvw": "UVW",
-        }
+        add_group_to_data_groups(data_groups, "base", "VISIBILITY")
 
     if "VISIBILITY_CORRECTED" in xds:
-        xds.attrs["data_groups"]["corrected"] = {
-            "correlated_data": "VISIBILITY_CORRECTED",
-            "flag": "FLAG",
-            "weight": "WEIGHT",
-            "uvw": "UVW",
-        }
+        add_group_to_data_groups(data_groups, "corrected", "VISIBILITY_CORRECTED")
 
     if "VISIBILITY_MODEL" in xds:
-        xds.attrs["data_groups"]["model"] = {
-            "correlated_data": "VISIBILITY_MODEL",
-            "flag": "FLAG",
-            "weight": "WEIGHT",
-            "uvw": "UVW",
-        }
+        add_group_to_data_groups(data_groups, "model", "VISIBILITY_MODEL")
 
     is_single_dish = False
     if "SPECTRUM" in xds:
-        xds.attrs["data_groups"]["base"] = {
-            "correlated_data": "SPECTRUM",
-            "flag": "FLAG",
-            "weight": "WEIGHT",
-            "uvw": "UVW",
-        }
+        add_group_to_data_groups(data_groups, "base", "SPECTRUM", False)
         is_single_dish = True
 
     if "SPECTRUM_MODEL" in xds:
-        xds.attrs["data_groups"]["model"] = {
-            "correlated_data": "SPECTRUM_MODEL",
-            "flag": "FLAG",
-            "weight": "WEIGHT",
-            "uvw": "UVW",
-        }
+        add_group_to_data_groups(data_groups, "model", "SPECTRUM_MODEL", False)
         is_single_dish = True
 
     if "SPECTRUM_CORRECTED" in xds:
-        xds.attrs["data_groups"]["corrected"] = {
-            "correlated_data": "SPECTRUM_CORRECTED",
-            "flag": "FLAG",
-            "weight": "WEIGHT",
-            "uvw": "UVW",
-        }
+        add_group_to_data_groups(data_groups, "corrected", "SPECTRUM_CORRECTED", False)
         is_single_dish = True
 
     return xds, is_single_dish

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/measurement_set/_utils/_msv2/create_antenna_xds.py

@@ -15,7 +15,9 @@ from xradio.measurement_set._utils._msv2._tables.read import (
     table_exists,
 )
 from xradio._utils.schema import convert_generic_xds_to_xradio_schema
-from xradio.measurement_set._utils._msv2.msv4_sub_xdss import interpolate_to_time
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
+    rename_and_interpolate_to_time,
+)
 
 from xradio._utils.list_and_array import (
     check_if_consistent,
@@ -509,27 +511,8 @@ def create_phase_calibration_xds(
         phase_cal_xds.time_phase_cal.astype("float64").astype("float64") / 10**9
     )
 
-    phase_cal_xds = interpolate_to_time(
-        phase_cal_xds,
-        phase_cal_interp_time,
-        "antenna_xds",
-        time_name="time_phase_cal",
+    phase_cal_xds = rename_and_interpolate_to_time(
+        phase_cal_xds, "time_phase_cal", phase_cal_interp_time, "phase_cal_xds"
     )
 
-    time_coord_attrs = {
-        "type": "time",
-        "units": ["s"],
-        "scale": "utc",
-        "format": "unix",
-    }
-
-    # If we interpolate rename the time_phase_cal axis to time.
-    if phase_cal_interp_time is not None:
-        time_coord = {"time": ("time_phase_cal", phase_cal_interp_time.data)}
-        phase_cal_xds = phase_cal_xds.assign_coords(time_coord)
-        phase_cal_xds.coords["time"].attrs.update(time_coord_attrs)
-        phase_cal_xds = phase_cal_xds.swap_dims({"time_phase_cal": "time"}).drop_vars(
-            "time_phase_cal"
-        )
-
     return phase_cal_xds

{xradio-0.0.44 → xradio-0.0.46}/src/xradio/measurement_set/_utils/_msv2/create_field_and_source_xds.py

@@ -6,7 +6,9 @@ import numpy as np
 import xarray as xr
 
 import toolviper.utils.logger as logger
-from xradio.measurement_set._utils._msv2.msv4_sub_xdss import interpolate_to_time
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
+    rename_and_interpolate_to_time,
+)
 from xradio.measurement_set._utils._msv2.subtables import subt_rename_ids
 from xradio.measurement_set._utils._msv2._tables.read import (
     convert_casacore_time_to_mjd,
@@ -363,20 +365,13 @@ def extract_ephemeris_info(
     }
     temp_xds["time_ephemeris"].attrs.update(time_coord_attrs)
 
-    # Convert to si units and interpolate if ephemeris_interpolate=True:
+    # Convert to si units
     temp_xds = convert_to_si_units(temp_xds)
-    temp_xds = interpolate_to_time(
-        temp_xds, interp_time, "field_and_source_xds", time_name="time_ephemeris"
-    )
 
-    # If we interpolate rename the time_ephemeris axis to time.
-    if interp_time is not None:
-        time_coord = {"time": ("time_ephemeris", interp_time.data)}
-        temp_xds = temp_xds.assign_coords(time_coord)
-        temp_xds.coords["time"].attrs.update(time_coord_attrs)
-        temp_xds = temp_xds.swap_dims({"time_ephemeris": "time"}).drop_vars(
-            "time_ephemeris"
-        )
+    # interpolate if ephemeris_interpolate/interp_time=True, and rename time_ephemeris=>time
+    temp_xds = rename_and_interpolate_to_time(
+        temp_xds, "time_ephemeris", interp_time, "field_and_source_xds"
+    )
 
     xds = xr.merge([xds, temp_xds])