Merge pull request #349 from European-XFEL/reorder_axes_to_shape

Add reorder_axes_to_shape function

Author: takluyver

docs/changelog.md

@@ -30,6 +30,7 @@ Added:
   [DataArray][xarray.DataArray]s properly (!333).
 - Added [hyperslicer2()][extra.utils.hyperslicer2] to make plotting image arrays
   easier (!348).
+- [reorder_axes_to_shape][extra.utils.reorder_axes_to_shape] utility function (!349).
 
 Changed:
 - [Timepix3.spatial_bins()] is now a static method.

docs/utilities.md

@@ -16,6 +16,8 @@ else.
 
 ::: extra.utils.find_nearest_value
 
+::: extra.utils.reorder_axes_to_shape
+
 ## Plotting functions
 
 ::: extra.utils.imshow2

src/extra/components/utils.py

@@ -1,4 +1,4 @@
-import sys
+from ..utils.misc import _isinstance_no_import
 
 # Source prefixes in use at each SASE.
 SASE_TOPICS = {
@@ -63,15 +63,6 @@ def _instrument_to_sase(instrument):
         return 3
 
 
-def _isinstance_no_import(obj, mod: str, cls: str):
-    """Check if isinstance(obj, mod.cls) without loading mod"""
-    m = sys.modules.get(mod)
-    if m is None:
-        return False
-
-    return isinstance(obj, getattr(m, cls))
-
-
 def _select_subcomponent_trains(src, keys, dst=None):
     if dst is None:
         from copy import copy

src/extra/utils/misc.py

@@ -1,9 +1,9 @@
+import sys
+
 import numpy as np
 
 from typing import Any
 
-from ..components.utils import _isinstance_no_import
-
 
 def find_nearest_index(array, value: Any) -> np.int64:
     """Find array index for the nearest value.
@@ -33,6 +33,40 @@ def find_nearest_value(array, value: Any) -> Any:
     return array[find_nearest_index(array, value)]
 
 
+def reorder_axes_to_shape(a, target_shape):
+    """Transpose an array to match the axis order specified by a shape tuple.
+
+    All dimensions must have different sizes. One axis in target_shape may be
+    None, a wildcard for the remainining axis in the array shape.
+    """
+    t = target_shape
+    if len(set(t)) != len(t):
+        raise ValueError(f"Target shape {t} has non-unique axes")
+    if len(t) != len(a.shape):
+        raise ValueError(f"Number of dimensions differs: {a.shape} -> {t}")
+    if None in t:
+        unmatched = set(a.shape) - set(t)
+        if len(unmatched) != 1:
+            raise ValueError(f"Cannot rearrange array shape {a.shape} to {t}")
+        t = list(t)
+        t[t.index(None)] = unmatched.pop()
+
+    if set(t) != set(a.shape):
+        raise ValueError(f"Cannot rearrange array shape {a.shape} to {t}")
+
+    order = tuple([a.shape.index(l) for l in t])
+    return a.transpose(order)
+
+
+def _isinstance_no_import(obj, mod: str, cls: str):
+    """Check if isinstance(obj, mod.cls) without loading mod"""
+    m = sys.modules.get(mod)
+    if m is None:
+        return False
+
+    return isinstance(obj, getattr(m, cls))
+
+
 def imshow2(image, *args, lognorm=False, ax=None, **kwargs):
     """Display an image with reasonable defaults.
 

tests/test_utils.py

@@ -1,7 +1,10 @@
 import numpy as np
+import pytest
 import xarray as xr
 
-from extra.utils import imshow2, hyperslicer2, fit_gaussian, gaussian
+from extra.utils import (
+    imshow2, hyperslicer2, fit_gaussian, gaussian, reorder_axes_to_shape,
+)
 
 
 def test_imshow2():
@@ -43,3 +46,35 @@ def test_fit_gaussian():
     data = gaussian(np.arange(100), *params, norm=False)
     popt = fit_gaussian(data, A_sign=-1)
     assert np.allclose(popt, params)
+
+
+def test_reorder_axes_to_shape():
+    arr = np.zeros((512, 1024, 16), dtype=np.float32)  # E.g. burst mode JUNGFRAU data
+    res = reorder_axes_to_shape(arr, (16, 512, 1024))
+    assert res.shape == (16, 512, 1024)
+    assert res.base is arr
+
+    res = reorder_axes_to_shape(arr, (None, 512, 1024))
+    assert res.shape == (16, 512, 1024)
+    assert res.base is arr
+
+    with pytest.raises(ValueError):
+        reorder_axes_to_shape(arr, (12, 512, 1024))  # Wrong dimension sizes
+
+    with pytest.raises(ValueError):
+        reorder_axes_to_shape(arr, (16, 1, 512, 1024))  # Wrong number of dimensions
+
+    with pytest.raises(ValueError):
+        reorder_axes_to_shape(arr[:, :512], (16, 512, 512))  # Ambiguous order
+
+    with pytest.raises(ValueError):
+        reorder_axes_to_shape(arr, (None, None, 1024))  # Only 1 None allowed
+
+    with pytest.raises(ValueError):
+        reorder_axes_to_shape(arr, (None, 256, 1024))  # Wildcard & wrong number
+
+    # Check we've transposed, not reshaped
+    arr = np.arange(15).reshape(3, 5)
+    res = reorder_axes_to_shape(arr, (5, 3))
+    assert res.shape == (5, 3)
+    np.testing.assert_array_equal(res[0], [0, 5, 10])