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

Closes #34 (closed)

lisainstrument/indexing.py

@@ -239,3 +239,49 @@ def sc2bothmosas(x):
         raise TypeError("invalid SC array")
     sc_mosas = [sc2index(mosa2sc(mosa)) for mosa in MOSAS]
     return x[sc_mosas]
+
+
+DISTANT_MOSAS = [f'{mosa[1]}{mosa[0]}' for mosa in MOSAS]
+ADJACENT_MOSAS = [f'{mosa[0]}{reflect(mosa[1], axis=mosa[0])}' for mosa in MOSAS]
+
+
+def mosa2distant(x):
+    """Transform MOSA index or array into distant MOSA.
+
+    Args:
+        x: MOSA index (as a string or integer), or MOSA array
+    """
+    # Transform MOSA array
+    if isinstance(x, np.ndarray) and x.shape[0] == 6:
+        transformed_mosas = list(map(mosa2index, DISTANT_MOSAS))
+        return x[transformed_mosas]
+    # Raise an error if we have an array of invalid shape
+    if isinstance(x, np.ndarray):
+        raise TypeError(f"invalid MOSA array shape '{x.shape}'")
+    # Check that index is correct
+    index = str(x)
+    if index not in MOSAS:
+        raise IndexError(f"invalid MOSA index '{index}'")
+    # Transform index
+    return DISTANT_MOSAS[mosa2index(x)]
+
+
+def mosa2adjacent(x):
+    """Transform MOSA index or array into adjacent MOSA.
+
+    Args:
+        x: MOSA index (as a string or integer), or MOSA array
+    """
+    # Transform MOSA array
+    if isinstance(x, np.ndarray) and x.shape[0] == 6:
+        transformed_mosas = list(map(mosa2index, ADJACENT_MOSAS))
+        return x[transformed_mosas]
+    # Raise an error if we have an array of invalid shape
+    if isinstance(x, np.ndarray):
+        raise TypeError(f"invalid MOSA array shape '{x.shape}'")
+    # Check that index is correct
+    index = str(x)
+    if index not in MOSAS:
+        raise IndexError(f"invalid MOSA index '{index}'")
+    # Transform index
+    return ADJACENT_MOSAS[mosa2index(x)]