Feat transpose multiindex

[krokosik]

Thank you for the amazing package! I was happy to see the linalg operations magically working xarray style with minimal effort, even with stacked dimensions! Unfortunately that was not the case for matrix_transpose, so here is a fix that makes it not error out when the dimensions being swapped both use a MultiIndex.

Additionally, the method was the only one, where dims were acutally not optional, so I fixed this.

Also fixed a small annoyance where pinv and matmul were not members of the accessor.

[read-the-docs-community]

Documentation build overview

📚 xarray-einstats | 🛠️ Build #33056318 | 📁 Comparing ec0a803 against latest (2dcbd7e)

  🔍 Preview build  

81 files changed · ± 81 modified

± Modified

• changelog.html
• getting_started.html
• api/linalg.html
• api/stats.html
• background/einops.html
• background/index.html
• background/linalg.html
• background/stats.html
• tutorials/dask_support.html
• tutorials/einops-basics-port.html
• and 71 more...

[OriolAbril]

Thanks for the PRs! I don't think I'll be able to review them this week but hopefully the next one. I am also happy to have users contributing as it has been a while since I used it regularly so I am maintaining and making releases but it is a bit frozen.

[krokosik]

No problem, I am currently integrating this package into my project and had to make a fork with several adjustments to make it work for my use case, but I'd prefer to be in sync with upsteam :)

There is no hurry for me, just thought it would be nice to contribute.

[OriolAbril]

Sorry for the extreme slowness. I added a comment on dim->coord pairing behaviour and suggested an extension to the test but all the other fixes should definitely go in without extra changes.

[OriolAbril]

I do think that the function should work even when the dimensions to transpose have a multiindex, but I am not sure about the handling of coords. I would be very interested in having your feedback and usecases. I had only considered the possibility of using this function for square matrices where there is basically a repeated dimension so I had dim+dim2 or xyz+xyz_bis with same length and same coordinate values.

The function does already work on non multiindex but different length dimensions, and had I realized before, my intuition would have been to remove the coordinate values from the output in such cases. This is what sort does.

What do/would you use this changing of dim->coords pairings for?

For example, if I add coordiate values to one of the example datasets:

matrices = tutorial.generate_matrices_dataarray()
matrices = matrices.assign_coords({name: [f"{name}{i}" for i in range(matrices.sizes[name])] for name in matrices.dims})

<xarray.DataArray (batch: 10, experiment: 3, dim: 4, dim2: 4)> Size: 4kB
0.4345 2.96 0.2627 0.4519 0.5507 0.1642 ... 0.2144 0.9432 0.5874 0.5385 2.987
Coordinates:
  * batch       (batch) <U6 240B 'batch0' 'batch1' ... 'batch8' 'batch9'
  * experiment  (experiment) <U11 132B 'experiment0' 'experiment1' 'experiment2'
  * dim         (dim) <U4 64B 'dim0' 'dim1' 'dim2' 'dim3'
  * dim2        (dim2) <U5 80B 'dim20' 'dim21' 'dim22' 'dim23'

and after matrix_transpose (no multiindex, I don't think it is relevant in this particular case):

matrix_transpose(matrices, dims=["batch", "experiment"])

<xarray.DataArray (dim: 4, dim2: 4, batch: 3, experiment: 10)> Size: 4kB
0.4345 0.4027 3.489 0.3505 0.2236 0.744 ... 0.1844 3.124 0.4926 0.255 2.987
Coordinates:
  * dim         (dim) <U4 64B 'dim0' 'dim1' 'dim2' 'dim3'
  * dim2        (dim2) <U5 80B 'dim20' 'dim21' 'dim22' 'dim23'
  * batch       (batch) <U11 132B 'experiment0' 'experiment1' 'experiment2'
  * experiment  (experiment) <U6 240B 'batch0' 'batch1' ... 'batch8' 'batch9'

Having the batch dimension now hold the experiment coordinates seems strange, independently of them having a multiindex. I was also curious about the constraint of both multiindexes having the same levels, I guess this is so the levels can also be renamed, any thoughts on swapping only the top multiindex name and not the levels?

[krokosik]

Thanks for the review, and no worries on the timing! I've added a new test_supermatrix_multiindex_transpose that demonstrates the use case I had in mind.

The context where I ran into this is when constructing a supermatrix from smaller matrices. For example, given:

• A with dims (i, j)
• B with dims (k, l)
  We can form S = A ⊗ B^T via an outer product A * B followed by stack(row=("i", "l"), col=("j", "k")). This gives us a single large square matrix with two multi-index dimensions row and col. Once we have S, we naturally want to compute things like S^T or S^T S — which is where matrix_transpose comes in.

Later, I want to decompose the result of some computation back into the original (i, j) dims. Tbh, I have not considered what happens in non-square cases, just followed what the existing code seemed to handle. To my mind, the whole reason behind using this library is to keep using pure xarray syntax and keeping the clear declarative code in linalg scripts.

Manipulating dimensions like this is kind of wonky in xarray, and depending what a person wants to achieve, requires different sequences of renames/dropping vars etc. What I would expect of an xarray related linalg.transpose is to offer a transposition in the linear algebra sense, meaning that a vector from one space (with its base/coordinates) is swapped with another space (so essentially, their coordinates are swapped).

From ML standpoint, the batch/experiment swapping seems weird, but the transpose is for linear algebra and this choice of keeping coordinates (vector space basis) seems appropriate to me.