Draft: Resolve "Use Numpy arrays instead of `ForEachObject` instances"

Closes #34 (closed)

lisainstrument/indexing.py

@@ -7,6 +7,7 @@ Authors:
     Jean-Baptiste Bayle <j2b.bayle@gmail.com>
 """
 
+import numpy as np
 import logging
 
 logger = logging.getLogger(__name__)
@@ -91,3 +92,80 @@ def sc2index(mosa):
 SC1 = sc2index(1)
 SC2 = sc2index(2)
 SC3 = sc2index(3)
+
+
+def transform(x, mapping):
+    """Transform indices or arrays with a mapping of SC.
+
+    If `x` is an MOSA or SC index, it is transformed according to the mapping of SC indices.
+    If `x` is a MOSA or SC array, the first axis is re-ordered according to the mapping.
+
+    Args:
+        x: MOSA or SC index (as a string or integer), or MOSA or SC array
+        mapping: dictionary describing the mapping from SC indices to new SC indices
+    """
+    # Check that we have a complete mapping
+    sc_set = set(SC)
+    mapping_str = {str(key): str(mapping[key]) for key in mapping}
+    if set(mapping_str.keys()) != sc_set:
+        raise ValueError(f"incomplete mapping '{mapping}', should have '{sc_set}' as keys")
+    if set(mapping_str.values()) != sc_set:
+        raise ValueError(f"incomplete mapping '{mapping}', should have '{sc_set}' as values")
+    # Transform MOSA array
+    if isinstance(x, np.ndarray) and x.shape[0] == 6:
+        transformed_mosas = [mosa2index(transform(mosa, mapping)) for mosa in MOSAS]
+        return x[transformed_mosas]
+    # Transform SC array
+    if isinstance(x, np.ndarray) and x.shape[0] == 3:
+        transformed_sc = [sc2index(transform(sc, mapping)) for sc in SC]
+        return x[transformed_sc]
+    # Raise an error if we have an array of invalid shape
+    if isinstance(x, np.ndarray):
+        raise TypeError(f"invalid MOSA or SC array shape '{x.shape}'")
+    # Check that index is correct
+    index = str(x)
+    if index not in MOSAS and index not in SC:
+        raise IndexError(f"invalid MOSA or SC index '{index}'")
+    # Transform index
+    mapped_chars = [mapping_str[char] for char in index]
+    return ''.join(mapped_chars)
+
+
+def rotate(x, nrot=1):
+    """Rotate SC indices in clockwise direction.
+
+    A rotation is a circular permutation of indices.
+
+    Args:
+        x: MOSA or SC index (as a string or integer), or MOSA or SC array
+        nrot: number of 120-degree rotations
+    """
+    nrot = nrot % 3
+    if nrot == 0:
+        mapping = {1: 1, 2: 2, 3: 3}
+    elif nrot == 1:
+        mapping = {1: 2, 2: 3, 3: 1}
+    elif nrot == 2:
+        mapping = {1: 3, 2: 1, 3: 2}
+
+    return transform(x, mapping)
+
+
+def reflect(x, axis):
+    """Reflect SC indices around an axis.
+
+    A reflection leaves the axis unchanged, and swaps the others.
+
+    Args:
+        axis: SC index around which the reflection is performed
+    """
+    if axis == 1:
+        mapping = {1: 1, 2: 3, 3: 2}
+    elif axis == 2:
+        mapping = {1: 3, 2: 2, 3: 1}
+    elif axis == 3:
+        mapping = {1: 2, 2: 1, 3: 3}
+    else:
+        raise ValueError(f"invalid SC index '{axis}'")
+
+    return transform(x, mapping)