Xarray: Make two DataArrays in the same Dataset use the same coordinate system

问题

I have an ArviZ InferenceData posterior trace which is an XArray Dataset.

In there, posterior traces for two of my random variables, a_mu_org and b_mu_org are DataArrays. Their coordinates are:

a_mu_org: (chain, draws, a_mu_org), with lengths (1, 2000, 15) respectively.
b_mu_org: (chain, draws, b_mu_org), with lengths (1, 2000, 15) respectively.

Semantically, a_mu_org and b_mu_org should really be indexed by a single categorical coordinate system of 15 organisms, rather than be separate indexes.

For a bit more clarity, here is the full dataset string repr:

<xarray.Dataset>
Dimensions:             (L_dim_0: 34281, a_dim_0: 456260, a_prot_shift_dim_0: 34281, b_dim_0: 456260, b_mu_org_dim_0: 15, b_prot_shift_dim_0: 34281, chain: 1, draw: 2000, organism: 15, sigma_dim_0: 34281, t50_org_dim_0: 15, t50_prot_dim_0: 39957)
Coordinates:
  * chain               (chain) int64 0
  * draw                (draw) int64 0 1 2 3 4 5 ... 1995 1996 1997 1998 1999
  * a_prot_shift_dim_0  (a_prot_shift_dim_0) object 'A0A023PXQ4_YMR173W-A' ... 'Z4YNA9_AB124611'
  * b_prot_shift_dim_0  (b_prot_shift_dim_0) object 'A0A023PXQ4_YMR173W-A' ... 'Z4YNA9_AB124611'
  * L_dim_0             (L_dim_0) object 'A0A023PXQ4_YMR173W-A' ... 'Z4YNA9_AB124611'
    a_mu_org_dim_0      (organism) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14
  * a_dim_0             (a_dim_0) object 'ytzI' 'mtlF' ... 'atpG2' 'atpB2'
  * b_mu_org_dim_0      (b_mu_org_dim_0) int64 0 1 2 3 4 5 ... 9 10 11 12 13 14
  * b_dim_0             (b_dim_0) object 'ytzI' 'mtlF' ... 'atpG2' 'atpB2'
  * t50_prot_dim_0      (t50_prot_dim_0) <U65 'Bacillus subtilis_168_lysate_R1-C0H3Q1_ytzI' ... 'Oleispira antarctica_RB-8_lysate_R1-R4YVF0_atpB2'
  * t50_org_dim_0       (t50_org_dim_0) <U43 'Arabidopsis thaliana seedling lysate' ... 'Thermus thermophilus HB27 lysate'
  * sigma_dim_0         (sigma_dim_0) object 'A0A023PXQ4_YMR173W-A' ... 'Z4YNA9_AB124611'
Dimensions without coordinates: organism
Data variables:
    a_org_pop           (chain, draw) float32 519.3236 518.8292 ... 517.84784
    a_prot_shift        (chain, draw, a_prot_shift_dim_0) float32 ...
    b_org_pop           (chain, draw) float32 11.509291 11.445394 ... 11.929538
    b_prot_shift        (chain, draw, b_prot_shift_dim_0) float32 ...
    L_pop               (chain, draw) float32 3.445896 3.4300675 ... 3.3917112
    L                   (chain, draw, L_dim_0) float32 ...
    a_mu_org            (chain, draw, organism) float32 430.56827 ... 813.2518
    a                   (chain, draw, a_dim_0) float32 ...
    b_mu_org            (chain, draw, b_mu_org_dim_0) float32 9.997488 ... 8.389757
    b                   (chain, draw, b_dim_0) float32 ...
    t50_prot            (chain, draw, t50_prot_dim_0) float32 39.249863 ... 52.19809
    t50_org             (chain, draw, t50_org_dim_0) float32 43.067646 ... 96.93388
    sigma               (chain, draw, sigma_dim_0) float32 ...
Attributes:
    created_at:                 2020-04-23T08:54:58.300091
    arviz_version:              0.7.0
    inference_library:          pymc3
    inference_library_version:  3.8

I would like to make a_mu_org and b_mu_org take on dimensions (chain, draw, organism) instead of their separate a_mu_org and b_mu_org. Things I have already tried include:

Adding a coordinate called organism, and then doing trace.posterior.swap_dims({"a_mu_org_dim_0": "organism"}), but I get an error stating that "replacement dimension 'organism' is not a 1D variable along the old dimension 'a_mu_org_dim_0'".
Renaming the dimension a_mu_org_dim_0 to organism, but then I also can't swap b_mu_org_dim_0 to the new organism.

Is what I'm trying to accomplish possible?

回答1:

I am not sure my solution is very good practice, it feels a little too hacky. Also, terminology is quite tricky, I'll try to stick to xarray terminology but may fail in doing so. The trick is to remove the coordinates so that a_dim_0 and b_dim_0 become only dimensions (now dimensions without coordinates). Afterwards, they can be renamed to the same thing and assigned to a new coord. Here is one example:

Starting from the following dataset called ds:

<xarray.Dataset>
Dimensions:  (a_dim_0: 15, b_dim_0: 15, chain: 4, draw: 100)
Coordinates:
  * chain    (chain) int64 0 1 2 3
  * draw     (draw) int64 0 1 2 3 4 5 6 7 8 9 ... 90 91 92 93 94 95 96 97 98 99
  * a_dim_0  (a_dim_0) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14
  * b_dim_0  (b_dim_0) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14
Data variables:
    a        (chain, draw, a_dim_0) float64 0.8152 1.189 ... 1.32 -0.2023
    b        (chain, draw, b_dim_0) float64 0.6447 -0.8059 ... -0.06435 -0.8666

the following 3 commands do the trick (the place of the assign_coord does not seem to affect the output, which makes sense, but it is key to first remove coordinates and then rename):

organism_names = [f"o{i}" for i in range(15)]
ds.reset_index(["a_dim_0", "b_dim_0"], drop=True) \
    .assign_coords(organism=organism_names) \
    .rename({"a_dim_0": "organism", "b_dim_0": "organism"})

Output:

<xarray.Dataset>
Dimensions:   (chain: 4, draw: 100, organism: 15)
Coordinates:
  * chain     (chain) int64 0 1 2 3
  * draw      (draw) int64 0 1 2 3 4 5 6 7 8 9 ... 90 91 92 93 94 95 96 97 98 99
  * organism  (organism) <U3 'o0' 'o1' 'o2' 'o3' ... 'o11' 'o12' 'o13' 'o14'
Data variables:
    a         (chain, draw, organism) float64 0.8152 1.189 ... 1.32 -0.2023
    b         (chain, draw, organism) float64 0.6447 -0.8059 ... -0.8666

来源：https://stackoverflow.com/questions/61827061/xarray-make-two-dataarrays-in-the-same-dataset-use-the-same-coordinate-system

标签

python

python-xarray