xradio 0.0.36__tar.gz → 0.0.38__tar.gz

{xradio-0.0.36/src/xradio.egg-info → xradio-0.0.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.36
+Version: 0.0.38
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License
@@ -37,11 +37,10 @@ License: BSD 3-Clause License
         OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
         OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
         
-Requires-Python: <3.12,>=3.9
+Requires-Python: <3.13,>=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
 Requires-Dist: astropy
-Requires-Dist: casadata
 Requires-Dist: dask
 Requires-Dist: distributed
 Requires-Dist: graphviper
@@ -58,7 +57,7 @@ Requires-Dist: tqdm
 Requires-Dist: xarray
 Requires-Dist: zarr
 Requires-Dist: pyarrow
-Requires-Dist: python_casacore>=3.5.2; sys_platform != "darwin"
+Requires-Dist: python_casacore>=3.6.1; sys_platform != "darwin"
 Requires-Dist: typing_extensions; python_version < "3.10"
 Requires-Dist: typeguard
 Provides-Extra: docs

{xradio-0.0.36 → xradio-0.0.38}/pyproject.toml

@@ -1,17 +1,16 @@
 [project]
 name = "xradio"
-version = "0.0.36"
+version = "0.0.38"
 description = " Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},
 ]
 license = {file = "LICENSE.txt"}
 readme = "README.md"
-requires-python = ">= 3.9, < 3.12"
+requires-python = ">= 3.9, < 3.13"
 
 dependencies = [
     'astropy',
-    'casadata',
     'dask',
     'distributed',
     'graphviper',
@@ -28,7 +27,7 @@ dependencies = [
     'xarray',
     'zarr',
     'pyarrow',
-    'python_casacore>=3.5.2; sys_platform != "darwin" ',
+    'python_casacore>=3.6.1; sys_platform != "darwin" ',
     'typing_extensions; python_version < "3.10"',
     'typeguard'
 ]

{xradio-0.0.36 → xradio-0.0.38}/src/xradio/vis/_vis_utils/_ms/conversion.py

@@ -80,7 +80,7 @@ def check_chunksize(chunksize: dict, xds_type: str) -> None:
         allowed_dims = [
             "time",
             "baseline_id",
-            "antenna_id",
+            "antenna_name",
             "frequency",
             "polarization",
         ]
@@ -138,7 +138,7 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
     It presently relies on the logic of mem_chunksize_to_dict_main_balanced() to find a
     balanced list of dimension sizes for the chunks
 
-    Assumes these relevant dims: (time, antenna_id/baseline_id, frequency,
+    Assumes these relevant dims: (time, antenna_name/baseline_id, frequency,
     polarization).
     """
 
@@ -149,11 +149,11 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
             "Cannot calculate chunk sizes when memory bound ({chunksize}) does not even allow all polarizations in one chunk"
         )
 
-    baseline_or_antenna_id = find_baseline_or_antenna_var(xds)
-    total_size = calc_used_gb(xds.sizes, baseline_or_antenna_id, sizeof_vis)
+    baseline_or_antenna_name = find_baseline_or_antenna_var(xds)
+    total_size = calc_used_gb(xds.sizes, baseline_or_antenna_name, sizeof_vis)
 
     ratio = chunksize / total_size
-    chunked_dims = ["time", baseline_or_antenna_id, "frequency", "polarization"]
+    chunked_dims = ["time", baseline_or_antenna_name, "frequency", "polarization"]
     if ratio >= 1:
         result = {dim: xds.sizes[dim] for dim in chunked_dims}
         logger.debug(
@@ -162,14 +162,17 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
     else:
         xds_dim_sizes = {k: xds.sizes[k] for k in chunked_dims}
         result = mem_chunksize_to_dict_main_balanced(
-            chunksize, xds_dim_sizes, baseline_or_antenna_id, sizeof_vis
+            chunksize, xds_dim_sizes, baseline_or_antenna_name, sizeof_vis
         )
 
     return result
 
 
 def mem_chunksize_to_dict_main_balanced(
-    chunksize: float, xds_dim_sizes: dict, baseline_or_antenna_id: str, sizeof_vis: int
+    chunksize: float,
+    xds_dim_sizes: dict,
+    baseline_or_antenna_name: str,
+    sizeof_vis: int,
 ) -> Dict[str, int]:
     """
     Assumes the ratio is <1 and all pols can fit in memory (from
@@ -236,7 +239,7 @@ def mem_chunksize_to_dict_main_balanced(
             dim_chunksizes[idx] += int_delta
         used = np.prod(dim_chunksizes) * sizeof_vis / GiBYTES_TO_BYTES
 
-    chunked_dim_names = ["time", baseline_or_antenna_id, "frequency", "polarization"]
+    chunked_dim_names = ["time", baseline_or_antenna_name, "frequency", "polarization"]
     dim_chunksizes_int = [int(v) for v in dim_chunksizes]
     result = dict(zip(chunked_dim_names, dim_chunksizes_int))
 
@@ -314,11 +317,11 @@ def mem_chunksize_to_dict_pointing(chunksize: float, xds: xr.Dataset) -> Dict[st
 
 def find_baseline_or_antenna_var(xds: xr.Dataset) -> str:
     if "baseline_id" in xds.coords:
-        baseline_or_antenna_id = "baseline_id"
-    elif "antenna_id" in xds.coords:
-        baseline_or_antenna_id = "antenna_id"
+        baseline_or_antenna_name = "baseline_id"
+    elif "antenna_name" in xds.coords:
+        baseline_or_antenna_name = "antenna_name"
 
-    return baseline_or_antenna_id
+    return baseline_or_antenna_name
 
 
 def itemsize_vis_spec(xds: xr.Dataset) -> int:
@@ -352,11 +355,11 @@ def itemsize_pointing_spec(xds: xr.Dataset) -> int:
 
 
 def calc_used_gb(
-    chunksizes: dict, baseline_or_antenna_id: str, sizeof_vis: int
+    chunksizes: dict, baseline_or_antenna_name: str, sizeof_vis: int
 ) -> float:
     return (
         chunksizes["time"]
-        * chunksizes[baseline_or_antenna_id]
+        * chunksizes[baseline_or_antenna_name]
         * chunksizes["frequency"]
         * chunksizes["polarization"]
         * sizeof_vis
@@ -734,7 +737,6 @@ def convert_and_write_partition(
     """
 
     taql_where = create_taql_query(partition_info)
-    # print("taql_where", taql_where)
     ddi = partition_info["DATA_DESC_ID"][0]
     obs_mode = str(partition_info["OBS_MODE"][0])
 
@@ -877,7 +879,8 @@ def convert_and_write_partition(
             )
 
             # Change antenna_ids to antenna_names
-            xds = antenna_ids_to_names(xds, ant_xds)
+            xds = antenna_ids_to_names(xds, ant_xds, is_single_dish)
+            ant_xds_name_ids = ant_xds["antenna_name"].set_xindex("antenna_id")
             ant_xds = ant_xds.drop_vars(
                 "antenna_id"
             )  # No longer needed after converting to name.
@@ -897,7 +900,7 @@ def convert_and_write_partition(
                 else:
                     pointing_interp_time = None
                 pointing_xds = create_pointing_xds(
-                    in_file, time_min_max, pointing_interp_time
+                    in_file, ant_xds_name_ids, time_min_max, pointing_interp_time
                 )
                 pointing_chunksize = parse_chunksize(
                     pointing_chunksize, "pointing", pointing_xds
@@ -953,7 +956,6 @@ def convert_and_write_partition(
 
             # Fix UVW frame
             # From CASA fixvis docs: clean and the im tool ignore the reference frame claimed by the UVW column (it is often mislabelled as ITRF when it is really FK5 (J2000)) and instead assume the (u, v, w)s are in the same frame as the phase tracking center. calcuvw does not yet force the UVW column and field centers to use the same reference frame! Blank = use the phase tracking frame of vis.
-            # print('##################',field_and_source_xds)
             if is_single_dish:
                 xds.UVW.attrs["frame"] = field_and_source_xds[
                     "FIELD_REFERENCE_CENTER"
@@ -974,11 +976,7 @@ def convert_and_write_partition(
 
             file_name = os.path.join(
                 out_file,
-                pathlib.Path(out_file)
-                .name.replace(".vis.zarr", "")
-                .replace(".zarr", "")
-                + "_"
-                + str(ms_v4_id),
+                pathlib.Path(in_file).name.replace(".ms", "") + "_" + str(ms_v4_id),
             )
 
             if "line_name" in field_and_source_xds.coords:
@@ -1037,17 +1035,19 @@ def convert_and_write_partition(
     # logger.info("Saved ms_v4 " + file_name + " in " + str(time.time() - start_with) + "s")
 
 
-def antenna_ids_to_names(xds, ant_xds):
+def antenna_ids_to_names(
+    xds: xr.Dataset, ant_xds: xr.Dataset, is_single_dish: bool
+) -> xr.Dataset:
     ant_xds = ant_xds.set_xindex(
         "antenna_id"
     )  # Allows for non-dimension coordinate selection.
 
-    if "baseline_antenna1_id" in xds:  # Interferometer
-        xds["baseline_antenna1_id"] = ant_xds["antenna_name"].sel(
-            antenna_id=xds["baseline_antenna1_id"]
+    if not is_single_dish:  # Interferometer
+        xds["baseline_antenna1_id"].data = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna1_id"].data
         )
-        xds["baseline_antenna2_id"] = ant_xds["antenna_name"].sel(
-            antenna_id=xds["baseline_antenna2_id"]
+        xds["baseline_antenna2_id"].data = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna2_id"].data
         )
         xds = xds.rename(
             {
@@ -1055,9 +1055,23 @@ def antenna_ids_to_names(xds, ant_xds):
                 "baseline_antenna2_id": "baseline_antenna2_name",
             }
         )
-    else:  # Single Dish
-        xds["antenna_id"] = ant_xds["antenna_name"].sel(antenna_id=xds["antenna_id"])
-        xds = xds.rename({"antenna_id": "antenna_name"})
+    else:
+        xds["baseline_id"] = ant_xds["antenna_name"].sel(antenna_id=xds["baseline_id"])
+        unwanted_coords_from_ant_xds = [
+            "antenna_id",
+            "antenna_name",
+            "mount",
+            "station",
+        ]
+        for unwanted_coord in unwanted_coords_from_ant_xds:
+            xds = xds.drop_vars(unwanted_coord)
+        xds = xds.rename({"baseline_id": "antenna_name"})
+
+        # 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
+        unwanted_coords_sd = ["baseline_antenna1_id", "baseline_antenna2_id"]
+        for unwanted_coord in unwanted_coords_sd:
+            xds = xds.drop_vars(unwanted_coord)
 
     return xds
 

{xradio-0.0.36 → xradio-0.0.38}/src/xradio/vis/_vis_utils/_ms/create_antenna_xds.py

@@ -127,7 +127,6 @@ def extract_antenna_info(
     )  # Make sure the antenna_id order is correct.
 
     # ['OFFSET', 'POSITION', 'DISH_DIAMETER', 'FLAG_ROW', 'MOUNT', 'NAME', 'STATION']
-    ant_xds = xr.Dataset()
     ant_xds = ant_xds.assign_coords({"cartesian_pos_label": ["x", "y", "z"]})
 
     ant_xds = convert_generic_xds_to_xradio_schema(
@@ -254,9 +253,13 @@ def extract_feed_info(
 
     # print('ant_xds["ANTENNA_FEED_OFFSET"]',ant_xds["ANTENNA_FEED_OFFSET"].data)
     # print('generic_feed_xds["POSITION"].data',generic_feed_xds["POSITION"].data)
+    feed_offset_attrs = ant_xds["ANTENNA_FEED_OFFSET"].attrs
     ant_xds["ANTENNA_FEED_OFFSET"] = (
         ant_xds["ANTENNA_FEED_OFFSET"] + generic_feed_xds["POSITION"].data
     )
+    # recover attrs after arithmetic operation
+    ant_xds["ANTENNA_FEED_OFFSET"].attrs.update(feed_offset_attrs)
+
     coords = {}
     # coords["receptor_label"] = "pol_" + np.arange(ant_xds.sizes["receptor_label"]).astype(str) #Works on laptop but fails in github test runner.
     coords["receptor_label"] = np.array(

{xradio-0.0.36 → xradio-0.0.38}/src/xradio/vis/_vis_utils/_ms/create_field_and_source_xds.py

@@ -30,7 +30,7 @@ def create_field_and_source_xds(
     is_single_dish: bool,
     time_min_max: Tuple[np.float64, np.float64],
     ephemeris_interp_time: Union[xr.DataArray, None] = None,
-):
+) -> tuple[xr.Dataset, int]:
     """
     Create a field and source xarray dataset (xds) from the given input file, field ID, and spectral window ID.
     Data is extracted from the FIELD and SOURCE tables and if there is ephemeris data, it is also extracted.
@@ -57,11 +57,13 @@ def create_field_and_source_xds(
     -------
     field_and_source_xds : xr.Dataset
         The xarray dataset containing the field and source information.
+    num_lines : int
+        Sum of num_lines for all unique sources.
     """
 
     start_time = time.time()
 
-    field_and_source_xds = xr.Dataset()
+    field_and_source_xds = xr.Dataset(attrs={"type": "field_and_source"})
 
     field_and_source_xds, ephemeris_path, ephemeris_table_name, source_id = (
         extract_field_info_and_check_ephemeris(
@@ -433,7 +435,9 @@ def extract_ephemeris_info(
     return xds
 
 
-def extract_source_info(xds, path, source_id, spectral_window_id):
+def extract_source_info(
+    xds: xr.Dataset, path: str, source_id: int, spectral_window_id: int
+) -> tuple[xr.Dataset, int]:
     """
     Extracts source information from the given path and adds it to the xarray dataset.
 

{xradio-0.0.36 → xradio-0.0.38}/src/xradio/vis/_vis_utils/_ms/msv4_sub_xdss.py

@@ -121,11 +121,13 @@ def create_weather_xds(in_file: str):
     # ['ANTENNA_ID', 'TIME', 'INTERVAL', 'H2O', 'IONOS_ELECTRON',
     #  'PRESSURE', 'REL_HUMIDITY', 'TEMPERATURE', 'DEW_POINT',
     #  'WIND_DIRECTION', 'WIND_SPEED']
-    weather_xds = xr.Dataset()
-
+    weather_xds = xr.Dataset(attrs={"type": "weather"})
+    time_attrs = column_description_casacore_to_msv4_measure(
+        weather_column_description["TIME"]
+    )
     coords = {
         "station_id": generic_weather_xds["STATION_ID"].data,
-        "time": generic_weather_xds["TIME"].data,
+        "time": ("time", generic_weather_xds["TIME"].data, time_attrs),
     }
     for key in generic_weather_xds:
         msv4_measure = column_description_casacore_to_msv4_measure(
@@ -181,6 +183,7 @@ def create_weather_xds(in_file: str):
 
 def create_pointing_xds(
     in_file: str,
+    ant_xds_name_ids: xr.DataArray,
     time_min_max: Union[Tuple[np.float64, np.float64], None],
     interp_time: Union[xr.DataArray, None] = None,
 ) -> xr.Dataset:
@@ -194,6 +197,8 @@ def create_pointing_xds(
     ----------
     in_file : str
         Input MS name.
+    ant_xds_name_ids : xr.Dataset
+        antenna_name data array from antenna_xds, with name/id information
     time_min_max : tuple
         min / max times values to constrain loading (from the TIME column)
     interp_time : Union[xr.DataArray, None] (Default value = None)
@@ -220,19 +225,19 @@ def create_pointing_xds(
         # "on_source": "ON_SOURCE",   # removed
         "OVER_THE_TOP": "OVER_THE_TOP",
     }
-    time_ant_dims = ["time", "antenna_id"]
+    time_ant_dims = ["time", "antenna_name"]
     time_ant_dir_dims = time_ant_dims + ["sky_dir_label"]
     data_variable_dims = {
-        # "name": ["time", "antenna_id"],   # removed
-        # "time_origin": ["time", "antenna_id"],   # removed?
+        # "name": ["time", "antenna_name"],   # removed
+        # "time_origin": ["time", "antenna_name"],   # removed?
         "DIRECTION": time_ant_dir_dims,
         "ENCODER": time_ant_dir_dims,
         "TARGET": time_ant_dir_dims,
         "POINTING_OFFSET": time_ant_dir_dims,
         "SOURCE_OFFSET": time_ant_dir_dims,
-        # "pointing_model_id": ["time", "antenna_id"],   # removed
-        # "tracking": ["time", "antenna_id"],   # => attribute
-        # "on_source": ["time", "antenna_id"],  # removed
+        # "pointing_model_id": ["time", "antenna_name"],   # removed
+        # "tracking": ["time", "antenna_name"],   # => attribute
+        # "on_source": ["time", "antenna_name"],  # removed
         "OVER_THE_TOP": time_ant_dims,
     }
     # Unused here
@@ -263,7 +268,7 @@ def create_pointing_xds(
         "ctds_attrs"
     ]["column_descriptions"]
 
-    pointing_xds = xr.Dataset()
+    pointing_xds = xr.Dataset(attrs={"type": "pointing"})
     for key in generic_pointing_xds:
         if key in to_new_data_variable_names:
             data_var_name = to_new_data_variable_names[key]
@@ -312,7 +317,9 @@ def create_pointing_xds(
 
     coords = {
         "time": generic_pointing_xds["TIME"].values,
-        "antenna_id": np.arange(generic_pointing_xds.sizes["ANTENNA_ID"]),
+        "antenna_name": ant_xds_name_ids.sel(
+            antenna_id=generic_pointing_xds["ANTENNA_ID"]
+        ).data,
         "sky_dir_label": ["ra", "dec"],
     }
     pointing_xds = pointing_xds.assign_coords(coords)

{xradio-0.0.36 → xradio-0.0.38}/src/xradio/vis/_vis_utils/_utils/xds_helper.py

@@ -1,4 +1,4 @@
-from importlib_metadata import version
+from importlib.metadata import version
 import graphviper.utils.logger as logger, multiprocessing, psutil
 from typing import Any, Dict, List, Tuple, Union
 

{xradio-0.0.36 → xradio-0.0.38}/src/xradio/vis/read_processing_set.py

@@ -38,7 +38,7 @@ def read_processing_set(
         data_groups = xds.attrs["data_groups"]
 
         if (obs_modes is None) or (
-            xds.attrs["partition_info"]["obs_mode"] in obs_modes
+            bool(set(xds.attrs["partition_info"]["obs_mode"]).intersection(obs_modes))
         ):
             sub_xds_dict, field_and_source_xds_dict = _read_sub_xds(
                 ms_store, file_system=file_system, data_groups=data_groups