diff --git a/src/cal_tools/tools.py b/src/cal_tools/tools.py index de53ba77e02ca2879d2502299124c8988d91eb3e..314e2084da62ca1355ca24c98332244038c4ff56 100644 --- a/src/cal_tools/tools.py +++ b/src/cal_tools/tools.py @@ -1030,3 +1030,17 @@ def write_compressed_frames( dataset.id.write_direct_chunk(chunk_start, compressed) return dataset + + +def reorder_axes(a, from_order, to_order): + """Rearrange axes of array a from from_order to to_order + + This does the same as np.transpose(), but making the before & after axes + more explicit. from_order is a sequence of strings labelling the axes of a, + and to_order is a similar sequence for the axes of the result. + """ + assert len(from_order) == a.ndim + assert sorted(from_order) == sorted(to_order) + from_order = list(from_order) + order = tuple([from_order.index(lbl) for lbl in to_order]) + return a.transpose(order) diff --git a/tests/test_cal_tools.py b/tests/test_cal_tools.py index 36d9d324021b3ce650434d3d651ca1834207b972..5929944c59b1e32836ba7f56387882b328de816f 100644 --- a/tests/test_cal_tools.py +++ b/tests/test_cal_tools.py @@ -22,6 +22,7 @@ from cal_tools.tools import ( recursive_update, send_to_db, write_constants_fragment, + reorder_axes, ) # AGIPD operating conditions. @@ -614,3 +615,13 @@ def test_write_constants_fragment(tmp_path: Path): }, } } + + +def test_reorder_axes(): + a = np.zeros((10, 32, 256, 3)) + from_order = ('cells', 'slow_scan', 'fast_scan', 'gain') + to_order = ('slow_scan', 'fast_scan', 'cells', 'gain') + assert reorder_axes(a, from_order, to_order).shape == (32, 256, 10, 3) + + to_order = ('gain', 'fast_scan', 'slow_scan', 'cells') + assert reorder_axes(a, from_order, to_order).shape == (3, 256, 32, 10)