Skip to content

Feat transpose multiindex #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
OriolAbril merged 9 commits into from
Jun 15, 2026
Merged

Feat transpose multiindex #90

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
edfcc34
Handle transpose via rename not swap_dims
krokosik Apr 15, 2026
1c8d188
Make matrix_transpose dims optional as documented
krokosik Apr 15, 2026
03d4895
Add missing pinv to accessor
krokosik Apr 15, 2026
ed13bf4
Add a multiindex test to show there are no errors
krokosik Apr 15, 2026
428a111
Update docstring
krokosik Apr 15, 2026
a3af36e
Apply tox fixes
krokosik Apr 15, 2026
e2fd7e8
Fix 1D case ;)
krokosik Apr 15, 2026
f95726f
Add missing matmul to accessor
krokosik Apr 15, 2026
ec0a803
add supermatrix use-case test
krokosik Jun 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion src/xarray_einstats/accessors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,12 @@
eigvals,
eigvalsh,
inv,
matmul,
matrix_power,
matrix_rank,
matrix_transpose,
norm,
pinv,
qr,
slogdet,
solve,
Expand All @@ -31,14 +33,18 @@ class LinAlgAccessor:
def __init__(self, xarray_obj):
self._obj = xarray_obj

def matrix_transpose(self, dims):
def matrix_transpose(self, dims=None):
"""Call :func:`xarray_einstats.linalg.matrix_transpose` on this DataArray."""
return matrix_transpose(self._obj, dims=dims)

def matrix_power(self, n, dims=None, **kwargs):
"""Call :func:`xarray_einstats.linalg.matrix_power` on this DataArray."""
return matrix_power(self._obj, n, dims=dims, **kwargs)

def matmul(self, other, dims=None, **kwargs):
"""Call :func:`xarray_einstats.linalg.matmul` with this DataArray as ``a/da``."""
return matmul(self._obj, other, dims=dims, **kwargs)

def cholesky(self, dims=None, **kwargs):
"""Call :func:`xarray_einstats.linalg.cholesky` on this DataArray."""
return cholesky(self._obj, dims=dims, **kwargs)
Expand Down Expand Up @@ -120,6 +126,10 @@ def inv(self, dims=None, **kwargs):
"""Call :func:`xarray_einstats.linalg.inv` on this DataArray."""
return inv(self._obj, dims=dims, **kwargs)

def pinv(self, dims=None, **kwargs):
"""Call :func:`xarray_einstats.linalg.pinv` on this DataArray."""
return pinv(self._obj, dims=dims, **kwargs)


@xr.register_dataarray_accessor("einops")
class EinopsAccessor:
Expand Down
2 changes: 2 additions & 0 deletions src/xarray_einstats/accessors.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ from .linalg import (
eigvals,
eigvalsh,
inv,
matmul,
matrix_power,
matrix_rank,
matrix_transpose,
Expand All @@ -28,6 +29,7 @@ class LinAlgAccessor:
def __init__(self, xarray_obj: Incomplete) -> None: ...
def matrix_transpose(self, dims: Incomplete) -> None: ...
def matrix_power(self, n: Incomplete, dims: Incomplete = ..., **kwargs: Incomplete) -> None: ...
def matmul(self, other: Incomplete, dims: Incomplete = ..., **kwargs: Incomplete) -> None: ...
def cholesky(self, dims: Incomplete = ..., **kwargs: Incomplete) -> None: ...
def qr(
self,
Expand Down
23 changes: 19 additions & 4 deletions src/xarray_einstats/linalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,17 +434,17 @@ def matmul(da, db, dims=None, *, out_append="2", **kwargs):
return matmul_aux


def matrix_transpose(da, dims):
def matrix_transpose(da, dims=None):
"""Transpose the underlying matrix without modifying the dimensions.

This convenience function uses :meth:`~xarray.DataArray.swap_dims` followed
This convenience function uses :meth:`~xarray.DataArray.rename` followed
by :meth:`~xarray.DataArray.transpose` to get the equivalent of a matrix transposition.

Parameters
----------
da : DataArray
Input DataArray
dims : list of str
dims : list of str, optional
Matrix dimensions

Returns
Expand All @@ -455,7 +455,22 @@ def matrix_transpose(da, dims):
if dims is None:
dims = _attempt_default_dims("matrix_transpose", da.dims)
dim1, dim2 = dims
return da.swap_dims({dim1: dim2, dim2: dim1}).transpose(..., *dims)

@OriolAbril OriolAbril Jun 6, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 krokosik Jun 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@krokosik krokosik Jun 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As for non-square matrices, I would say the same linear algebra argument applies (moving between vector spaces). However, I did not deal with such cases much so I won't insist on it due to lack of familiarity.

@krokosik krokosik Jun 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another case for not dropping coordinates is that doing so is fairly straightforward in xarray and can be easily done explicitly (just chain a drop_vars), while restoring them after an implicit drop is more involved/cumbersome, since you need to extract them from the prior object and then assign again, so 2 steps vs 1.

@krokosik krokosik Jun 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As for the final question, I do feel like swapping just the MultiIndex name would be confusing, since you would rename it but get an actually different beast, which breaks semantics.

My algebraic intuition, as outlined, is that we swap two spaces, which are composite spaces:
A=U ⊗ V and B = W ⊗ Y. Should we skip swapping the subspaces, you would swap the definitions of your spaces and get B = U ⊗ V and A = W ⊗ Y, which to my mind breaks the semantics. Even outside of algebra, a stacked coordinate would change its meaning and I don't think this function should meddle with semantics imposed by the user who defined the stacked coordinate and attributed some meaning to a particular way of stacking.

rename_dict = {dim1: dim2, dim2: dim1}

if (
dim1 in da.indexes
and dim2 in da.indexes
and len(da.indexes[dim1].names) == len(da.indexes[dim2].names)
and len(da.indexes[dim1].names) > 1
):
for sub_dim1, sub_dim2 in zip(da.indexes[dim1].names, da.indexes[dim2].names):
rename_dict[sub_dim1] = sub_dim2
rename_dict[sub_dim2] = sub_dim1

da_transposed = da.rename(rename_dict).transpose(..., *dims)

# Purely cosmetic change to preserve order of coordinates in the output
return da_transposed.assign_coords({k: da_transposed.coords[k] for k in da.coords})


def matrix_power(da, n, dims=None, **kwargs):
Expand Down
2 changes: 1 addition & 1 deletion src/xarray_einstats/linalg.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ def matmul(
out_append: str = ...,
**kwargs: Incomplete,
) -> xarray.DataArray: ...
def matrix_transpose(da: xarray.DataArray, dims: list[str]) -> xarray.DataArray: ...
def matrix_transpose(da: xarray.DataArray, dims: list[str] | None = ...) -> xarray.DataArray: ...
def matrix_power(
da: xarray.DataArray, n: int, dims: Sequence[Hashable] | None = ..., **kwargs: Incomplete
) -> xarray.DataArray: ...
Expand Down
36 changes: 36 additions & 0 deletions tests/test_linalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,42 @@ def test_pinv_dataarray_tol(self, matrices, kind):
def test_transpose(self, hermitian):
assert_equal(hermitian, matrix_transpose(hermitian, dims=("dim", "dim2")))

def test_transpose_multiindex(self, matrices):
stacked = matrices.stack(batch_experiment=("batch", "experiment"), dim_dim2=("dim", "dim2"))
out = matrix_transpose(stacked, dims=("batch_experiment", "dim_dim2"))
assert out.sizes["batch_experiment"] == stacked.sizes["dim_dim2"]
assert stacked.sizes["batch_experiment"] == out.sizes["dim_dim2"]

def test_supermatrix_multiindex_transpose(self):
"""Test matrix_transpose with multiindex in a supermatrix use case.

Forms S = A ⊗ B^T via outer product + stack, then verifies
S^T and S^T S against pure numpy.
"""
rng = np.random.default_rng(42)
a_mat_np = rng.normal(size=(2, 2))
b_mat_np = rng.normal(size=(2, 2))

a_mat = xr.DataArray(a_mat_np, dims=["i", "j"])
b_mat = xr.DataArray(b_mat_np, dims=["k", "l"])

# S = A ⊗ B^T using outer product then stacking into multiindex dims
s_outer = a_mat * b_mat # dims: (i, j, k, l)
s_da = s_outer.stack(row=("i", "l"), col=("j", "k"))

# Verify S matches numpy Kronecker product A ⊗ B^T
s_np = np.kron(a_mat_np, b_mat_np.T)
np.testing.assert_allclose(s_da.values, s_np)

# Compute S^T via matrix_transpose and verify against numpy transpose
s_t = matrix_transpose(s_da, dims=("row", "col"))
np.testing.assert_allclose(s_t.values, s_np.T)

# Compute S^T S in numpy
expected = s_np.T @ s_np
result = s_t.values @ s_da.values
np.testing.assert_allclose(result, expected)

def test_matrix_power(self, matrices):
out = matrix_power(matrices, 2, dims=("dim", "dim2"))
assert out.shape == matrices.shape
Expand Down