How to ask for/query data in dataset with arrays as parameters.

[miguelcarcamov]

Hi a have the following xarray dataset with the antenna-baseline relationships. Basically two different antenna IDs make a baseline. As you see: 0 & 1 = Baseline 0, 0&2 = Baseline 1... etc....

<xarray.Dataset>
Dimensions:          (row: 7021)
Coordinates:
    ANTENNA1         (row) int64 0 0 0 0 0 0 0 0 ... 114 115 115 115 116 116 117
    ANTENNA2         (row) int64 1 2 3 4 5 6 7 8 ... 118 116 117 118 117 118 118
    ROWID            (row) int64 0 1 2 3 4 5 6 ... 7015 7016 7017 7018 7019 7020
Dimensions without coordinates: row
Data variables:
    BASELINE_LENGTH  (row) float64 dask.array<chunksize=(7021,), meta=Quantity>
Attributes:
    description:  Baseline-Antenna relationship

And I have the following xarray DataArrays with antennas which IDs are in the dataset, but they might not have all the IDs in the dataset.

<xarray.DataArray 'ANTENNA1' (row: 48959)>
array([ 0,  0,  0, ..., 25, 25, 25], dtype=int32)  --------> IDs
Coordinates:
    ROWID    (row) int64 0 1 2 3 4 5 6 ... 97169 97170 97171 97172 97173 97174
Dimensions without coordinates: row

<xarray.DataArray 'ANTENNA2' (row: 48959)>
array([ 1,  2,  3, ..., 30, 31, 32], dtype=int32) ------------> IDs
Coordinates:
    ROWID    (row) int64 0 1 2 3 4 5 6 ... 97169 97170 97171 97172 97173 97174
Dimensions without coordinates: row

What I would like is to query the dataset with the values of my arrays to get the baseline ID (rowid in the dataset) and get something like this as result:

ANTENNA1           ANTENNA2           BASELINE_ID
0                      1                  0
0                      2                  1
0                      3                  2
....
25                   30                2654
25                   31                2655
25                   32                2656

So first columns are basically the latter two arrays, and the last column is the ROWID got from the query on the dataset.

I have tried to do this using

baseline.dataset.where((baseline.dataset.ANTENNA1==25)&(baseline.dataset.ANTENNA2==32), drop=True).ROWID but instead of scalars using numpy arrays and I didn't get much luck.

How can I do this using xarray?

[keewis]

Could you post the text reprs of your data (use xr.set_options(display_style="text") for this)? Otherwise it's pretty difficult to understand what you're trying to do.

[miguelcarcamov]

@keewis Done :)

[miguelcarcamov]

I have solve this myself. I will attach the piece of code just in case anyone need it:

baseline_id = da.argmax((baseline_obj.dataset.ANTENNA1.data[None, :] == antenna1_ids[:, None]) & (baseline_obj.dataset.ANTENNA2.data[None, :] == antenna2_ids[:, None]), axis=1).astype(np.int32)