Bug: interpolate_data_coords_onto_parallel_coords failed when ps_partition=["field_name"]

[tnakazato]

interpolate_data_coords_onto_parallel_coords failed when ps_partition=["field_name"]. It originally occurred to me when I processed SD data partitioned by antenna. But it turned out that it's not sd-specific issue. I reprocuded the issue using "gg084b.ps.zarr", the data for fringefit tutorial. Procedure to reproduce the issue is as follows:

from xradio.measurement_set.open_processing_set import open_processing_set
from astroviper.calibration.fringefit import fringefit_ps, apply_cal_ps
from xradio.measurement_set import convert_msv2_to_processing_set
import toolviper
from graphviper.graph_tools.coordinate_utils import interpolate_data_coords_onto_parallel_coords

import xarray as xa
import astropy.time
import numpy as np
import matplotlib.pyplot as plt

msv2_name = "global_vlbi_gg084b_reduced.ms"
toolviper.utils.data.download(file=msv2_name)

convert_out = "gg084b"
partition_scheme = ["FIELD_ID", "SCAN_NUMBER", "SPW"]

convert_msv2_to_processing_set(
    in_file=msv2_name,
    out_file=convert_out,
    partition_scheme=partition_scheme,
    overwrite=True,
    parallel_mode="partition",
)

ps = open_processing_set("gg084b.ps.zarr")

parallel_coords = {}
node_task_data_mapping = interpolate_data_coords_onto_parallel_coords(
    parallel_coords, ps, ps_partition=["field_name"],
)

It will cause the following error:

ValueError                                Traceback (most recent call last)
Cell In[5], [line 9](vscode-notebook-cell:?execution_count=5&line=9)
      1 from graphviper.graph_tools.coordinate_utils import (
      2     interpolate_data_coords_onto_parallel_coords,
      3     make_parallel_coord,
      4     make_time_coord,
      5     make_frequency_coord,
      6 )
      8 parallel_coords = {}
----> [9](vscode-notebook-cell:?execution_count=5&line=9) node_task_data_mapping = interpolate_data_coords_onto_parallel_coords(
     10     parallel_coords, ps, ps_partition=["field_name"],
     11 )

File /work/dev/nakazato/git/casangi/graphviper/src/graphviper/graph_tools/coordinate_utils.py:376, in interpolate_data_coords_onto_parallel_coords(parallel_coords, input_data, interpolation_method, assume_sorted, ps_partition)
    373     raise ValueError("Cannot split by both spw and frequency")
    375 if len(ps_partition) > 0:
--> [376](https://vscode-remote+ssh-002dremote-002bana01-002ealma-002enao-002eac-002ejp.vscode-resource.vscode-cdn.net/work/dev/nakazato/git/casangi/graphviper/src/graphviper/graph_tools/coordinate_utils.py:376)     partition_map = _partition_ps_by_non_dimensions(input_data, ps_partition)
    377 else:
    378     # By default we iterate over everything
    379     partition_map = {0: [xds_name for xds_name in input_data]}

File /work/dev/nakazato/git/casangi/graphviper/src/graphviper/graph_tools/coordinate_utils.py:588, in _partition_ps_by_non_dimensions(ps, ps_partition_keys)
    585         # OK I think I can punt: the key should probably be an integer but that doesn't feel very Pythonic
    586         # But I *can* reasonably demand it is hashable
...
    590             )
    591         ps_split_map.setdefault(key, {}).setdefault(val_for_xds, []).append(name)
    592 d = {}

ValueError: Can't split by field_name; value [np.str_('J1311-2329_0')] is not suitable for splitting

[taktsutsumi]

Fixed. The problem was that partition info getting from an xds for field_name is a list of strings while a partition mapping function inside interpolate_data_coords_onto_parallel_coords was expecting a value. Now the code goes into the list to extract each element. test_graph_tool.test_partition was also modified to test both 'spectral_window_name' and 'field_name'. @tnakazato, if you can please check with the SD data using 45-bug-interpolate_data_coords_onto_parallel_coords-failed-when-ps_partitionfield_name branch.

[tnakazato]

Thanks @taktsutsumi for the fix. I confirmed that the fix is working for both IF and SD datasets!