xradio 0.0.49__tar.gz → 0.0.51__tar.gz

{xradio-0.0.49/src/xradio.egg-info → xradio-0.0.51}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: xradio
-Version: 0.0.49
+Version: 0.0.51
 Summary:  Xarray Radio Astronomy Data IO
-Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
+Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>, Federico Montesino Pouzols <pouzols@eso.edu>, Dave Mehringer <dmehring@nrao.edu>, Peter Wortmann <peter.wortmann@skao.int>
 License: BSD 3-Clause License
         
         All works in this repository are copyrighted 2024.  
@@ -43,7 +43,7 @@ License-File: LICENSE.txt
 Requires-Dist: astropy
 Requires-Dist: dask
 Requires-Dist: distributed
-Requires-Dist: toolviper
+Requires-Dist: toolviper>=0.0.11
 Requires-Dist: numba>=0.57.0
 Requires-Dist: numpy
 Requires-Dist: pytest
@@ -56,6 +56,7 @@ Requires-Dist: zarr<3,>=2
 Requires-Dist: pyarrow
 Requires-Dist: python_casacore>=3.6.1; sys_platform != "darwin"
 Requires-Dist: typeguard
+Requires-Dist: numcodecs<0.16
 Provides-Extra: interactive
 Requires-Dist: matplotlib; extra == "interactive"
 Requires-Dist: prettytable; extra == "interactive"
@@ -78,7 +79,13 @@ Dynamic: license-file
 # xradio
 Xarray Radio Astronomy Data IO is still in development.
 
-[![Python 3.9 3.10 3.11 3.12](https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11%203.12-blue)](https://www.python.org/downloads/release/python-380/)
+[![Python 3.11 3.12 3.13](https://img.shields.io/badge/python-3.11%20%7C%203.12%20%7C%203.13-blue)](https://www.python.org/downloads/release/python-3130/)
+![Linux Tests](https://github.com/casangi/xradio/actions/workflows/python-testing-linux.yml/badge.svg?branch=main)
+![macOS Tests](https://github.com/casangi/xradio/actions/workflows/python-testing-macos.yml/badge.svg?branch=main)
+![ipynb Tests](https://github.com/casangi/xradio/actions/workflows/run-ipynb.yml/badge.svg?branch=main)
+[![Coverage](https://codecov.io/gh/casangi/xradio/branch/main/graph/badge.svg)](https://codecov.io/gh/casangi/xradio/branch/main/xradio)
+[![Documentation Status](https://readthedocs.org/projects/xradio/badge/?version=latest)](https://xradio.readthedocs.io)
+[![Version Status](https://img.shields.io/pypi/v/xradio.svg)](https://pypi.python.org/pypi/xradio/)
 
 # Installing
 It is recommended to use the conda environment manager from [miniforge](https://github.com/conda-forge/miniforge) to create a clean, self-contained runtime where XRADIO and all its dependencies can be installed:

xradio-0.0.51/README.md

@@ -0,0 +1,27 @@
+# xradio
+Xarray Radio Astronomy Data IO is still in development.
+
+[![Python 3.11 3.12 3.13](https://img.shields.io/badge/python-3.11%20%7C%203.12%20%7C%203.13-blue)](https://www.python.org/downloads/release/python-3130/)
+![Linux Tests](https://github.com/casangi/xradio/actions/workflows/python-testing-linux.yml/badge.svg?branch=main)
+![macOS Tests](https://github.com/casangi/xradio/actions/workflows/python-testing-macos.yml/badge.svg?branch=main)
+![ipynb Tests](https://github.com/casangi/xradio/actions/workflows/run-ipynb.yml/badge.svg?branch=main)
+[![Coverage](https://codecov.io/gh/casangi/xradio/branch/main/graph/badge.svg)](https://codecov.io/gh/casangi/xradio/branch/main/xradio)
+[![Documentation Status](https://readthedocs.org/projects/xradio/badge/?version=latest)](https://xradio.readthedocs.io)
+[![Version Status](https://img.shields.io/pypi/v/xradio.svg)](https://pypi.python.org/pypi/xradio/)
+
+# Installing
+It is recommended to use the conda environment manager from [miniforge](https://github.com/conda-forge/miniforge) to create a clean, self-contained runtime where XRADIO and all its dependencies can be installed:
+```sh
+conda create --name xradio python=3.12 --no-default-packages
+conda activate xradio
+```
+> 📝 On macOS it is required to pre-install `python-casacore` using `conda install -c conda-forge python-casacore`.
+
+XRADIO can now be installed using:
+```sh
+pip install xradio
+```
+This will also install the minimal dependencies for XRADIO. To install the minimal dependencies and the interactive components (JupyterLab) use:
+```sh
+pip install "xradio[interactive]"
+```

{xradio-0.0.49 → xradio-0.0.51}/pyproject.toml

@@ -1,9 +1,12 @@
 [project]
 name = "xradio"
-version = "0.0.49"
+version = "0.0.51"
 description = " Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},
+    {name = "Federico Montesino Pouzols", email="pouzols@eso.edu"},
+    {name = "Dave Mehringer", email="dmehring@nrao.edu"},
+    {name = "Peter Wortmann", email="peter.wortmann@skao.int"}
 ]
 license = {file = "LICENSE.txt"}
 readme = "README.md"
@@ -13,7 +16,7 @@ dependencies = [
     'astropy',
     'dask',
     'distributed',
-    'toolviper',
+    'toolviper>=0.0.11',
     'numba>=0.57.0',
     'numpy',
     'pytest',
@@ -25,7 +28,8 @@ dependencies = [
     'zarr>=2,<3',
     'pyarrow',
     'python_casacore>=3.6.1; sys_platform != "darwin" ',
-    'typeguard'
+    'typeguard',
+    'numcodecs<0.16'
 ]
 
 [project.optional-dependencies]

{xradio-0.0.49 → xradio-0.0.51}/src/xradio/_utils/schema.py

@@ -190,14 +190,14 @@ casacore_to_msv4_measure_type = {
             "LSRK": "lsrk",
             "LSRD": "lsrd",
             "BARY": "BARY",
-            "GEO": "GEO",
+            "GEO": "gcrs",
             "TOPO": "TOPO",
         },  # The frames/observer we are not sure if/how to translate to astropy are uppercase
     },
     "position": {
         "type": "location",
         "Ref": "frame",
-        "Ref_map": {"ITRF": "GRS80"},
+        "Ref_map": {"ITRF": "ITRS"},
     },
     "uvw": {
         "type": "uvw",

{xradio-0.0.49 → xradio-0.0.51}/src/xradio/measurement_set/_utils/_msv2/conversion.py

@@ -1214,13 +1214,13 @@ def convert_and_write_partition(
         )
         # but before, keep the name-id arrays, we need them for the pointing and weather xds
         ant_xds_name_ids = ant_xds["antenna_name"].set_xindex("antenna_id")
-        ant_xds_station_name_ids = ant_xds["station"].set_xindex("antenna_id")
+        ant_position_xds_with_ids = ant_xds["ANTENNA_POSITION"].set_xindex("antenna_id")
         # No longer needed after converting to name.
         ant_xds = ant_xds.drop_vars("antenna_id")
 
         # Create weather_xds
         start = time.time()
-        weather_xds = create_weather_xds(in_file, ant_xds_station_name_ids)
+        weather_xds = create_weather_xds(in_file, ant_position_xds_with_ids)
         logger.debug("Time weather " + str(time.time() - start))
 
         # Create pointing_xds
@@ -1341,7 +1341,7 @@ def convert_and_write_partition(
 
         ms_xdt["/antenna_xds"] = ant_xds
         for group_name in xds.attrs["data_groups"]:
-            ms_xdt["/" + f"field_and_source_xds_{group_name}"] = field_and_source_xds
+            ms_xdt["/" + f"field_and_source_{group_name}_xds"] = field_and_source_xds
 
         if with_pointing and len(pointing_xds.data_vars) > 0:
             ms_xdt["/pointing_xds"] = pointing_xds
@@ -1425,7 +1425,7 @@ def antenna_ids_to_names(
             "antenna_id",
             "antenna_name",
             "mount",
-            "station",
+            "station_name",
         ]
         for unwanted_coord in unwanted_coords_from_ant_xds:
             xds = xds.drop_vars(unwanted_coord)
@@ -1468,6 +1468,7 @@ def add_group_to_data_groups(
         "correlated_data": correlated_data_name,
         "flag": "FLAG",
         "weight": "WEIGHT",
+        "field_and_source": f"field_and_source_{what_group}_xds",
         "description": f"Data group derived from the data column '{correlated_data_name}' of an MSv2 converted to MSv4",
         "date": datetime.datetime.now(datetime.timezone.utc).isoformat(),
     }

{xradio-0.0.49 → xradio-0.0.51}/src/xradio/measurement_set/_utils/_msv2/create_antenna_xds.py

@@ -109,7 +109,7 @@ def extract_antenna_info(
 
     to_new_coords = {
         "NAME": ["antenna_name", ["antenna_name"]],
-        "STATION": ["station", ["antenna_name"]],
+        "STATION": ["station_name", ["antenna_name"]],
         "MOUNT": ["mount", ["antenna_name"]],
         # "PHASED_ARRAY_ID": ["phased_array_id", ["antenna_name"]],
         "antenna_id": ["antenna_id", ["antenna_name"]],
@@ -158,9 +158,9 @@ def extract_antenna_info(
 
         # None of the native numpy functions work on the github test runner.
         antenna_name = ant_xds["antenna_name"].values
-        station = ant_xds["station"].values
+        station_name = ant_xds["station_name"].values
         antenna_name = np.array(
-            list(map(lambda x, y: x + "_" + y, antenna_name, station))
+            list(map(lambda x, y: x + "_" + y, antenna_name, station_name))
         )
 
         ant_xds["antenna_name"] = xr.DataArray(antenna_name, dims=["antenna_name"])
@@ -376,7 +376,7 @@ def create_gain_curve_xds(
 
     ant_borrowed_coords = {
         "antenna_name": ant_xds.coords["antenna_name"],
-        "station": ant_xds.coords["station"],
+        "station_name": ant_xds.coords["station_name"],
         "mount": ant_xds.coords["mount"],
         "telescope_name": ant_xds.coords["telescope_name"],
         "receptor_label": ant_xds.coords["receptor_label"],
@@ -486,7 +486,7 @@ def create_phase_calibration_xds(
 
     ant_borrowed_coords = {
         "antenna_name": ant_xds.coords["antenna_name"],
-        "station": ant_xds.coords["station"],
+        "station_name": ant_xds.coords["station_name"],
         "mount": ant_xds.coords["mount"],
         "telescope_name": ant_xds.coords["telescope_name"],
         "receptor_label": ant_xds.coords["receptor_label"],

{xradio-0.0.49 → xradio-0.0.51}/src/xradio/measurement_set/_utils/_msv2/create_field_and_source_xds.py

@@ -226,9 +226,10 @@ def extract_ephemeris_info(
             "type": "location",
             "units": ["deg", "deg", "m"],
             "data": observer_position,
-            "frame": "WGS84",
+            "frame": "ITRS",
             "origin_object_name": "Earth",
             "coordinate_system": ephemeris_meta["obsloc"].lower(),
+            "ellipsoid": "WGS84",
         }
     )  # I think the units are ['deg','deg','m'] and 'WGS84'.
 

{xradio-0.0.49 → xradio-0.0.51}/src/xradio/measurement_set/_utils/_msv2/msv4_sub_xdss.py

@@ -172,7 +172,130 @@ def make_taql_where_weather(
     return taql_where
 
 
-def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
+def prepare_generic_weather_xds_and_station_name(
+    generic_weather_xds: xr.Dataset,
+    in_file: str,
+    ant_position_with_ids: xr.DataArray,
+    has_asdm_station_position: bool,
+) -> tuple[xr.Dataset, np.ndarray]:
+    """
+    A generic_weather_xds loaded with load_generic_table() might still need to be reloaded
+    with an additional WHERE condition to constrain the indices of antennas. But this depends on whether
+    ASDM/importasdm extension columns are present or not.
+
+    This also prepares the station_name values:
+    - if has_asdm_station_ids:
+       - tries to find from ASDM_STATION the station names,
+       - otherwise, takes ids (antenna_ids in generic_weather were actually the ASDM_STATION_IDs
+    - else: get the values from antenna_xds (the stations present)
+
+
+    Parameters
+    ----------
+    generic_weather_xds : xr.Dataset
+        generic dataset read from an MSv2 WEATHER subtable
+    in_file : str
+        Input MS name.
+    ant_position_with_ids : xr.DataArray
+        antenna_position data var from the antenna_xds (expected to still include the initial ANTENNA_ID
+        coordinate as well as other coordinates from the antenna_xds)
+    has_asdm_station_position : bool
+        wHether this generic weatehr_xds should be treated as including the nonstandard extensions
+        NS_WX_STATION_ID and NS_WX_STATION_POSITION as created by CASA/importasdm (ALMA and VLA).
+
+    Returns
+    -------
+    (generic_weather_xds, station_name): tuple[[xarray.Dataset, numpy.ndarray]
+        Weather Xarray Dataset prepared for generic conversion to MSv4, values for the station_name coordinate
+    """
+
+    if has_asdm_station_position:
+        asdm_station_path = os.path.join(in_file, "ASDM_STATION")
+        if table_exists(asdm_station_path):
+            asdm_station_xds = load_generic_table(in_file, "ASDM_STATION")
+            station_name = asdm_station_xds.name.values[
+                generic_weather_xds["ANTENNA_ID"].values
+            ]
+        else:
+            # if no info from ASDM_STATION, use the indices from antenna_id which was actually the NS_WX_STATION_ID
+            len_antenna_id = generic_weather_xds.sizes["ANTENNA_ID"]
+            station_name = list(
+                map(
+                    lambda x, y: x + "_" + y,
+                    ["Station"] * len_antenna_id,
+                    generic_weather_xds["ANTENNA_ID"].values.astype(str),
+                )
+            )
+
+    else:
+        taql_where = make_taql_where_weather(in_file, ant_position_with_ids)
+        generic_weather_xds = load_generic_table(
+            in_file,
+            "WEATHER",
+            rename_ids=subt_rename_ids["WEATHER"],
+            taql_where=taql_where,
+        )
+
+        if not generic_weather_xds.data_vars:
+            # for example when the weather subtable only has info for antennas/stations
+            # not present in the MSv4 (no overlap between antennas loaded in ant_xds and weather)
+            return None, None
+
+        stations_present = ant_position_with_ids.sel(
+            antenna_id=generic_weather_xds["ANTENNA_ID"]
+        ).station_name
+        station_name = stations_present.values
+
+    return generic_weather_xds, station_name
+
+
+def finalize_station_position(
+    weather_xds: xr.Dataset, ant_position_with_ids, has_asdm_station_position: bool
+) -> xr.Dataset:
+    """
+    For a STATION_POSITION data var being added to a weather_xds, make sure coordinates and dimensions
+    are conforming to the schema.
+
+    Parameters
+    ----------
+    weather_xds : xr.Dataset
+        weather_xds where we still need to ensure the right coordinates and attributes
+    ant_position_with_ids : xr.DataArray
+        antenna_position data var from the antenna_xds (expected to still include the initial ANTENNA_ID
+        coordinate as well as other coordinates from the antenna_xds)
+    has_asdm_station_position : bool
+        Whether this generic weatehr_xds should be treated as including the nonstandard extensions
+        NS_WX_STATION_ID and NS_WX_STATION_POSITION as created by CASA/importasdm (ALMA and VLA).
+
+    Returns
+    -------
+    weather_xds: xarray.Dataset
+        Weather Xarray Dataset with all coordinates and attributes in STATION_POSITION
+    """
+    if has_asdm_station_position:
+        # STATION_POSITION has been created but needs prooper dimensions and attrs
+        # Drop the time dim
+        weather_xds["STATION_POSITION"] = weather_xds["STATION_POSITION"].sel(
+            time_weather=0, drop=True, method="nearest"
+        )
+        # borrow location frame attributes from antenna position
+        weather_xds["STATION_POSITION"].attrs = ant_position_with_ids.attrs
+    else:
+        # borrow from ant_posision_with_ids but without carrying over other coords
+        weather_xds = weather_xds.assign(
+            {
+                "STATION_POSITION": (
+                    ["station_name", "cartesian_pos_label"],
+                    ant_position_with_ids.values,
+                    ant_position_with_ids.attrs,
+                )
+            }
+        )
+
+    return weather_xds
+
+
+def create_weather_xds(in_file: str, ant_position_with_ids: xr.DataArray):
     """
     Creates a Weather Xarray Dataset from a MS v2 WEATHER table.
 
@@ -180,8 +303,9 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     ----------
     in_file : str
         Input MS name.
-    ant_xds_station_name_ids : xr.DataArray
-        station name data array from antenna_xds, with name/id information
+    ant_position_with_ids : xr.DataArray
+        antenna_position data var from the antenna_xds (expected to still including the initial ANTENNA_ID coordinate
+        as wellas other coordinates from the antenna_xds)
 
     Returns
     -------
@@ -190,32 +314,32 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     """
 
     try:
-        taql_where = make_taql_where_weather(in_file, ant_xds_station_name_ids)
         generic_weather_xds = load_generic_table(
             in_file,
             "WEATHER",
             rename_ids=subt_rename_ids["WEATHER"],
-            taql_where=taql_where,
         )
     except ValueError as _exc:
         return None
 
-    if not generic_weather_xds.data_vars:
-        # for example when the weather subtable only has info for antennas/stations
-        # not present in the MSv4 (no overlap between antennas loaded in ant_xds and weather)
+    has_asdm_station_position = (
+        "NS_WX_STATION_POSITION" in generic_weather_xds.data_vars
+    )
+    generic_weather_xds, station_name = prepare_generic_weather_xds_and_station_name(
+        generic_weather_xds, in_file, ant_position_with_ids, has_asdm_station_position
+    )
+    if not generic_weather_xds:
         return None
 
     weather_xds = xr.Dataset(attrs={"type": "weather"})
-    stations_present = ant_xds_station_name_ids.sel(
-        antenna_id=generic_weather_xds["ANTENNA_ID"]
-    )
     coords = {
-        "station_name": stations_present.data,
-        "antenna_name": stations_present.coords["antenna_name"].data,
+        "station_name": station_name,
+        "cartesian_pos_label": ["x", "y", "z"],
     }
     weather_xds = weather_xds.assign_coords(coords)
 
     dims_station_time = ["station_name", "time_weather"]
+    dims_station_time_position = dims_station_time + ["cartesian_pos_label"]
     to_new_data_variables = {
         "H20": ["H2O", dims_station_time],
         "IONOS_ELECTRON": ["IONOS_ELECTRON", dims_station_time],
@@ -226,6 +350,15 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
         "WIND_DIRECTION": ["WIND_DIRECTION", dims_station_time],
         "WIND_SPEED": ["WIND_SPEED", dims_station_time],
     }
+    if has_asdm_station_position:
+        to_new_data_variables.update(
+            {
+                "NS_WX_STATION_POSITION": [
+                    "STATION_POSITION",
+                    dims_station_time_position,
+                ],
+            }
+        )
 
     to_new_coords = {
         "TIME": ["time_weather", ["time_weather"]],
@@ -234,6 +367,9 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     weather_xds = convert_generic_xds_to_xradio_schema(
         generic_weather_xds, weather_xds, to_new_data_variables, to_new_coords
     )
+    weather_xds = finalize_station_position(
+        weather_xds, ant_position_with_ids, has_asdm_station_position
+    )
 
     # TODO: option to interpolate to main time
 
@@ -256,6 +392,7 @@ def correct_generic_pointing_xds(
     and tries to correct several deviations from the MSv2 specs seen in
     common test data.
     The problems fixed here include wrong dimensions:
+
     - for example transposed dimensions with respect to the MSv2 specs (output
     from CASA simulator),
     - missing/additional unexpected dimensions when some of the columns are
@@ -423,6 +560,7 @@ def prepare_generic_sys_cal_xds(generic_sys_cal_xds: xr.Dataset) -> xr.Dataset:
     sys_cal_xds dataset, as their structure differs in dimensions and order
     of dimensions.
     This function performs various prepareation steps, such as:
+
     - filter out dimensions not neeed for an individual MSv4 (SPW, FEED),
     - drop variables loaded from columns with all items set to empty array,
     - transpose the dimensions frequency,receptor

{xradio-0.0.49 → xradio-0.0.51}/src/xradio/measurement_set/measurement_set_xdt.py

@@ -93,17 +93,34 @@ class MeasurementSetXdt:
         if data_group_name is not None:
             sel_data_group_set = set(
                 self._xdt.attrs["data_groups"][data_group_name].values()
-            )
+            ) - set(["date", "description"])
+
+            sel_field_and_source_xds = self._xdt.attrs["data_groups"][data_group_name][
+                "field_and_source"
+            ]
 
             data_variables_to_drop = []
-            for dg in self._xdt.attrs["data_groups"].values():
-                temp_set = set(dg.values()) - sel_data_group_set
+            field_and_source_to_drop = []
+            for dg_name, dg in self._xdt.attrs["data_groups"].items():
+                print(f"Data group: {dg_name}", dg)
+                f_and_s = dg["field_and_source"]
+                dg_copy = dg.copy()
+                dg_copy.pop("date", None)
+                dg_copy.pop("description", None)
+                dg_copy.pop("field_and_source", None)
+                temp_set = set(dg_copy.values()) - sel_data_group_set
                 data_variables_to_drop.extend(list(temp_set))
 
+                if f_and_s != sel_field_and_source_xds:
+                    field_and_source_to_drop.append(f_and_s)
+
             data_variables_to_drop = list(set(data_variables_to_drop))
 
             sel_ms_xdt = self._xdt
 
+            print("Data variables to drop: ", data_variables_to_drop)
+            print("Field and source to drop: ", field_and_source_to_drop)
+
             sel_corr_xds = self._xdt.ds.sel(
                 indexers, method, tolerance, drop, **indexers_kwargs
             ).drop_vars(data_variables_to_drop)
@@ -140,7 +157,10 @@ class MeasurementSetXdt:
             else:
                 data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
 
-        return self._xdt[f"field_and_source_xds_{data_group_name}"].ds
+        field_and_source_xds_name = self._xdt.attrs["data_groups"][data_group_name][
+            "field_and_source"
+        ]
+        return self._xdt[field_and_source_xds_name].ds
 
     def get_partition_info(self, data_group_name: str = None) -> dict:
         """
@@ -195,3 +215,105 @@ class MeasurementSetXdt:
         }
 
         return partition_info
+
+    def add_data_group(
+        self,
+        new_data_group_name: str,
+        correlated_data: str = None,
+        weight: str = None,
+        flag: str = None,
+        uvw: str = None,
+        field_and_source_xds: str = None,
+        date_time: str = None,
+        description: str = None,
+        data_group_dv_shared_with: str = None,
+    ) -> xr.DataTree:
+        """_summary_
+
+        Parameters
+        ----------
+        new_data_group_name : str
+            _description_
+        correlated_data : str, optional
+            _description_, by default None
+        weights : str, optional
+            _description_, by default None
+        flag : str, optional
+            _description_, by default None
+        uvw : str, optional
+            _description_, by default None
+        field_and_source_xds : str, optional
+            _description_, by default None
+        date_time : str, optional
+            _description_, by default None
+        description : str, optional
+            _description_, by default None
+        data_group_dv_shared_with : str, optional
+            _description_, by default "base"
+
+        Returns
+        -------
+        xr.DataTree
+            _description_
+        """
+
+        if data_group_dv_shared_with is None:
+            data_group_dv_shared_with = self._xdt.xr_ms._get_default_data_group_name()
+        default_data_group = self._xdt.attrs["data_groups"][data_group_dv_shared_with]
+
+        new_data_group = {}
+
+        if correlated_data is None:
+            correlated_data = default_data_group["correlated_data"]
+        new_data_group["correlated_data"] = correlated_data
+        assert (
+            correlated_data in self._xdt.ds.data_vars
+        ), f"Data variable {correlated_data} not found in dataset."
+
+        if weight is None:
+            weight = default_data_group["weight"]
+        new_data_group["weight"] = weight
+        assert (
+            weight in self._xdt.ds.data_vars
+        ), f"Data variable {weight} not found in dataset."
+
+        if flag is None:
+            flag = default_data_group["flag"]
+        new_data_group["flag"] = flag
+        assert (
+            flag in self._xdt.ds.data_vars
+        ), f"Data variable {flag} not found in dataset."
+
+        if self._xdt.attrs["type"] == "visibility":
+            if uvw is None:
+                uvw = default_data_group["uvw"]
+            new_data_group["uvw"] = uvw
+            assert (
+                uvw in self._xdt.ds.data_vars
+            ), f"Data variable {uvw} not found in dataset."
+
+        if field_and_source_xds is None:
+            field_and_source_xds = default_data_group["field_and_source_xds"]
+        new_data_group["field_and_source"] = field_and_source_xds
+        assert (
+            field_and_source_xds in self._xdt.children
+        ), f"Data variable {field_and_source_xds} not found in dataset."
+
+        if date_time is None:
+            date_time = datetime.now().isoformat()
+        new_data_group["date"] = date_time
+
+        if description is None:
+            description = ""
+        new_data_group["description"] = description
+
+        self._xdt.attrs["data_groups"][new_data_group_name] = new_data_group
+
+        return self._xdt
+
+    def _get_default_data_group_name(self):
+        if "base" in self._xdt.attrs["data_groups"].keys():
+            data_group_name = "base"
+        else:
+            data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
+        return data_group_name