如何细分/细化 xarray 数据集中的维度？

Question

摘要：我有一个数据集，其收集方式使得维度最初不可用。我想获取本质上是一大块无差别的数据，并为其添加维度，以便可以对其进行查询、子集化等。这是以下问题的核心。

这是我拥有的 xarray 数据集：

<xarray.Dataset>
Dimensions:  (chain: 1, draw: 2000, rows: 24000)
Coordinates:
  * chain    (chain) int64 0
  * draw     (draw) int64 0 1 2 3 4 5 6 7 ... 1993 1994 1995 1996 1997 1998 1999
  * rows     (rows) int64 0 1 2 3 4 5 6 ... 23994 23995 23996 23997 23998 23999
Data variables:
    obs      (chain, draw, rows) float64 4.304 3.985 4.612 ... 6.343 5.538 6.475
Attributes:
    created_at:                 2019-12-27T17:16:13.847972
    inference_library:          pymc3
    inference_library_version:  3.8

这里的维度rows对应了我需要恢复到数据中的多个子维度。特别是，24,000 行对应于 240 个条件的 100 个样本（这 100 个样本位于连续的块中）。gate这些条件是、input、growth medium和的组合od。

我想最终得到这样的结果：

<xarray.Dataset>
Dimensions:  (chain: 1, draw: 2000, gate: 1, input: 4, growth_medium: 3, sample: 100, rows: 24000)
Coordinates:
  * chain    (chain) int64 0
  * draw     (draw) int64 0 1 2 3 4 5 6 7 ... 1993 1994 1995 1996 1997 1998 1999
  * rows     *MultiIndex*
  * gate     (gate) int64 'AND'
  * input    (input) int64 '00', '01', '10', '11'
  * growth_medium (growth_medium) 'standard', 'rich', 'slow'
  * sample   (sample) int64 0 1 2 3 4 5 6 7 ... 95 96 97 98 99
Data variables:
    obs      (chain, draw, gate, input, growth_medium, samples) float64 4.304 3.985 4.612 ... 6.343 5.538 6.475
Attributes:
    created_at:                 2019-12-27T17:16:13.847972
    inference_library:          pymc3
    inference_library_version:  3.8

我有一个 pandas 数据框，它指定门、输入和生长介质的值 - 每行给出一组门、输入和生长介质的值，以及一个索引，指定相应的rows100 组的位置（在 中）出现样品。目的是该数据框是标记数据集的指南。

我查看了有关“重塑和重组数据”的 xarray 文档，但我不知道如何组合这些操作来完成我需要的操作。我怀疑我需要将它们与 结合起来GroupBy，但我不明白如何做。谢谢！

后来：我有一个解决这个问题的方法，但是它太恶心了，我希望有人能解释我错得有多离谱，以及有什么更优雅的方法是可能的。

因此，首先，我将原始数据中的所有数据提取Dataset为原始 numpy 形式：

foo = qm.idata.posterior_predictive['obs'].squeeze('chain').values.T
foo.shape # (24000, 2000)

然后我根据需要重新塑造它：

bar = np.reshape(foo, (240, 100, 2000))

这大致给出了我想要的形状：有 240 种不同的实验条件，每种条件有 100 个变体，对于每个变体，我的数据集中都有 2000 个蒙特卡洛样本。

现在，我从 Pandas 中提取有关 240 个实验条件的信息DataFrame：

import pandas as pd
# qdf is the original dataframe with the experimental conditions and some
# extraneous information in other columns
new_df = qdf[['gate', 'input', 'output', 'media', 'od_lb', 'od_ub', 'temperature']]
idx = pd.MultiIndex.from_frame(new_df)

最后，我DataArray从 numpy 数组和 pandas重新组装了MultiIndexa ：

<xarray.Dataset>
Dimensions:  (chain: 1, draw: 2000, rows: 24000)
Coordinates:
  * chain    (chain) int64 0
  * draw     (draw) int64 0 1 2 3 4 5 6 7 ... 1993 1994 1995 1996 1997 1998 1999
  * rows     (rows) int64 0 1 2 3 4 5 6 ... 23994 23995 23996 23997 23998 23999
Data variables:
    obs      (chain, draw, rows) float64 4.304 3.985 4.612 ... 6.343 5.538 6.475
Attributes:
    created_at:                 2019-12-27T17:16:13.847972
    inference_library:          pymc3
    inference_library_version:  3.8

结果DataArray有这些坐标，如我所愿：

Coordinates:
  * regions      (regions) MultiIndex
  - gate         (regions) object 'AND' 'AND' 'AND' 'AND' ... 'AND' 'AND' 'AND'
  - input        (regions) object '00' '10' '10' '10' ... '01' '01' '11' '11'
  - output       (regions) object '0' '0' '0' '0' '0' ... '0' '0' '0' '1' '1'
  - media        (regions) object 'standard_media' ... 'high_osm_media_five_percent'
  - od_lb        (regions) float64 0.0 0.001 0.001 ... 0.0001 0.0051 0.0051
  - od_ub        (regions) float64 0.0001 0.0051 0.0051 2.0 ... 0.0003 2.0 2.0
  - temperature  (regions) int64 30 30 37 30 37 30 37 ... 37 30 37 30 37 30 37
  * conditions   (conditions) int64 0 1 2 3 4 5 6 7 ... 92 93 94 95 96 97 98 99
  * draws        (draws) int64 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999

但这是相当可怕的，而且我必须突破所有漂亮的xarray抽象层才能到达这一点，这似乎是错误的。特别是因为这似乎并不是科学工作流程中一个不寻常的部分：获取相对原始的数据集以及需要与数据结合的元数据电子表格。那么我做错了什么？更优雅的解决方案是什么？

Answer 1

给定起始数据集，类似于：

<xarray.Dataset>
Dimensions:  (draw: 2, row: 24)
Coordinates:
  * draw     (draw) int32 0 1
  * row      (row) int32 0 1 2 3 4 5 6 7 8 9 ... 14 15 16 17 18 19 20 21 22 23
Data variables:
    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47

您可以连接多个纯 xarray 命令来细分维度（获取相同形状但使用多重索引的数据），甚至重塑数据集。要细分维度，可以使用以下代码：

multiindex_ds = ds.assign_coords(
    dim_0=["a", "b", "c"], dim_1=[0,1], dim_2=range(4)
).stack(
    dim=("dim_0", "dim_1", "dim_2")
).reset_index(
    "row", drop=True
).rename(
    row="dim"
)
multiindex_ds

其输出是：

<xarray.Dataset>
Dimensions:  (dim: 24, draw: 2)
Coordinates:
  * draw     (draw) int32 0 1
  * dim      (dim) MultiIndex
  - dim_0    (dim) object 'a' 'a' 'a' 'a' 'a' 'a' ... 'c' 'c' 'c' 'c' 'c' 'c'
  - dim_1    (dim) int64 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1
  - dim_2    (dim) int64 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3
Data variables:
    obs      (draw, dim) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47

此外，多重索引可以被拆开，有效地重塑数据集：

reshaped_ds = multiindex_ds.unstack("dim")
reshaped_ds

与输出：

<xarray.Dataset>
Dimensions:  (dim_0: 3, dim_1: 2, dim_2: 4, draw: 2)
Coordinates:
  * draw     (draw) int32 0 1
  * dim_0    (dim_0) object 'a' 'b' 'c'
  * dim_1    (dim_1) int64 0 1
  * dim_2    (dim_2) int64 0 1 2 3
Data variables:
    obs      (draw, dim_0, dim_1, dim_2) int32 0 1 2 3 4 5 ... 42 43 44 45 46 47

我认为仅此并不能完全满足您的需求，因为您想将一个维度转换为二维，其中之一是多重索引。不过，所有的构建模块都在这里。

例如，您可以使用 和 按照此步骤（包括取消堆叠）regions，conditions然后按照此步骤（现在不取消堆叠）转换regions为多索引。另一种选择是从一开始就使用所有维度，取消堆叠它们，然后再次堆叠它们，留下conditions最终的多重索引之外。

详细解答

答案结合了几个完全不相关的命令，并且要了解每个命令在做什么可能很棘手。

assign_coords

第一步是创建新的尺寸和坐标并将其添加到数据集中。这是必要的，因为接下来的方法需要数据集中已经存在的尺寸和坐标。

立即停止会assign_coords产生以下数据集：

<xarray.Dataset>
Dimensions:  (dim_0: 3, dim_1: 2, dim_2: 4, draw: 2, row: 24)
Coordinates:
  * draw     (draw) int32 0 1
  * row      (row) int32 0 1 2 3 4 5 6 7 8 9 ... 14 15 16 17 18 19 20 21 22 23
  * dim_0    (dim_0) <U1 'a' 'b' 'c'
  * dim_1    (dim_1) int32 0 1
  * dim_2    (dim_2) int32 0 1 2 3
Data variables:
    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47

stack

数据集现在包含 3 个维度，最多可添加 24 个元素，但是，由于数据当前相对于这 24 个元素是平坦的，因此我们必须将它们堆叠到单个 24 元素多重索引中，以使它们的形状兼容。

我发现assign_coords接下来是stack最自然的解决方案，但是，另一种可能性是生成一个多重索引，类似于上面的操作方式，并直接使用assign_coords多重索引进行调用，从而不需要堆栈。

此步骤将所有 3 个新维度合并为一个维度：

<xarray.Dataset>
Dimensions:  (dim: 24, draw: 2, row: 24)
Coordinates:
  * draw     (draw) int32 0 1
  * row      (row) int32 0 1 2 3 4 5 6 7 8 9 ... 14 15 16 17 18 19 20 21 22 23
  * dim      (dim) MultiIndex
  - dim_0    (dim) object 'a' 'a' 'a' 'a' 'a' 'a' ... 'c' 'c' 'c' 'c' 'c' 'c'
  - dim_1    (dim) int64 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1
  - dim_2    (dim) int64 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3
Data variables:
    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47

请注意，根据需要，现在我们有 2 个尺寸，尺寸为 24。

reset_index

现在我们的最终维度作为坐标出现在数据集中，我们希望这个新坐标成为用于索引变量的坐标obs。似乎是正确的选择，但是，我们的每个坐标都对自身进行索引（与文档中同时对坐标和坐标进行索引set_index的示例不同），这意味着在这种特殊情况下不能使用它。使用的方法是去掉坐标而不去掉尺寸。set_indexxxaset_indexreset_indexrowrow

在下面的输出中可以看出现在row是一个没有坐标的维度：

<xarray.Dataset>
Dimensions:  (dim: 24, draw: 2, row: 24)
Coordinates:
  * draw     (draw) int32 0 1
  * dim      (dim) MultiIndex
  - dim_0    (dim) object 'a' 'a' 'a' 'a' 'a' 'a' ... 'c' 'c' 'c' 'c' 'c' 'c'
  - dim_1    (dim) int64 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1
  - dim_2    (dim) int64 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3
Dimensions without coordinates: row
Data variables:
    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47

rename

当前的数据集几乎是最后一个，唯一的问题是变量obs仍然具有row维度而不是所需的维度：dim。它看起来并不是真正的预期用法rename，但它可以用于dim吸收，产生所需的最终结果（上面称为） 。rowmultiindex_ds

再次，set_index似乎是要选择的方法，但是，如果使用rename(row="dim"),代替set_index(row="dim")，则多重索引将折叠为由元组组成的索引：

<xarray.Dataset>
Dimensions:  (draw: 2, row: 24)
Coordinates:
  * draw     (draw) int32 0 1
  * row      (row) object ('a', 0, 0) ('a', 0, 1) ... ('c', 1, 2) ('c', 1, 3)
Data variables:
    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47