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

Closes #34 (closed)

lisainstrument/indexing.py

+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Utility functions for indexing MOSAs and SC.
+
+Authors:
+    Jean-Baptiste Bayle <j2b.bayle@gmail.com>
+"""
+
+import numpy as np
+
+
+MOSAS = ['12', '23', '31', '13', '32', '21']
+SC = ['1', '2', '3']
+
+
+def index2mosa(index):
+    """Return MOSA index associated with a given array index.
+
+    Args:
+        index: array index
+
+    Returns:
+        MOSA index as a two-figure string, e.g., '12'.
+    """
+    try:
+        return MOSAS[index]
+    except IndexError as error:
+        raise IndexError(f"invalid array index '{index}' for MOSA") from error
+
+
+def mosa2index(mosa):
+    """Return array index associated with a given MOSA index.
+
+    Args:
+        mosa: MOSA index as a two-figure string or integer
+
+    Returns:
+        Array index as an integer.
+    """
+    try:
+        return MOSAS.index(str(mosa))
+    except ValueError as error:
+        raise ValueError(f"invalid MOSA index '{mosa}'") from error
+
+
+MOSA12 = mosa2index('12')
+MOSA23 = mosa2index('23')
+MOSA31 = mosa2index('31')
+MOSA13 = mosa2index('13')
+MOSA32 = mosa2index('32')
+MOSA21 = mosa2index('21')
+
+
+def index2sc(index):
+    """Return SC index associated with a given array index.
+
+    Args:
+        index: array index
+
+    Returns:
+        SC index as a one-figure string, e.g., '1'.
+    """
+    try:
+        return SC[index]
+    except IndexError as error:
+        raise IndexError(f"invalid array index '{index}' for SC") from error
+
+
+def sc2index(mosa):
+    """Return array index associated with a given SC index.
+
+    Args:
+        mosa: SC index as a one-figure string, or integer
+
+    Returns:
+        Array index as an integer.
+    """
+    try:
+        return SC.index(str(mosa))
+    except ValueError as error:
+        raise ValueError(f"invalid SC index '{mosa}'") from error
+
+
+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
+    """
+    axis = str(axis)
+    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)
+
+
+def mosa2sc(mosa):
+    """Return SC index from MOSA index.
+
+    Args:
+        mosa: MOSA index, as a two-figure string or integer
+    """
+    if str(mosa) not in MOSAS:
+        raise ValueError(f"invalid MOSA index '{mosa}'")
+    return str(mosa)[0]
+
+
+def sc2leftdistant(sc):
+    """Return index of left distant SC.
+
+    Note that this is equivalent to applying a 120-degree rotation clockwise.
+
+    Args:
+        sc: SC index, as one-figure string or integer
+    """
+    if str(sc) not in SC:
+        raise ValueError(f"invalid SC index '{sc}'")
+    return rotate(sc, nrot=1)
+
+
+def sc2rightdistant(sc):
+    """Return index of right distant SC.
+
+    Note that this is equivalent to applying a 120-degree rotation counterclockwise.
+
+    Args:
+        sc: SC index, as one-figure string or integer
+    """
+    if str(sc) not in SC:
+        raise ValueError(f"invalid SC index '{sc}'")
+    return rotate(sc, nrot=-1)
+
+
+def sc2leftmosa(sc):
+    """Return left MOSA index.
+
+    Args:
+        sc: SC index, as one-figure string or integer
+    """
+    if str(sc) not in SC:
+        raise ValueError(f"invalid SC index '{sc}'")
+    return f'{sc}{sc2leftdistant(sc)}'
+
+
+def sc2rightmosa(sc):
+    """Return right MOSA index.
+
+    Args:
+        sc: SC index, as one-figure string or integer
+    """
+    if str(sc) not in SC:
+        raise ValueError(f"invalid SC index '{sc}'")
+    return f'{sc}{sc2rightdistant(sc)}'
+
+
+def sc2bothmosas(x):
+    """Broadcast a SC array to a MOSA array.
+
+    We broadcast by assuming the same values for both MOSAs of each SC.
+
+    Args:
+        sc: SC array
+
+    Returns:
+        A MOSA array.
+    """
+    if not isinstance(x, np.ndarray) or x.shape[0] != 3:
+        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)]
+
+
+def empty_mosa(shape, *args, **kwargs):
+    """Return a new MOSA array of given shape and type, without initializing entries.
+
+    Returns:
+        MOSA array of shape (6, shape).
+    """
+    if isinstance(shape, int):
+        shape = (6, shape)
+    else:
+        shape = (6, *shape)
+    return np.empty(shape, *args, **kwargs)
+
+
+def empty_sc(shape, *args, **kwargs):
+    """Return a new SC array of given shape and type, without initializing entries.
+
+    Returns:
+        SC array of shape (3, shape).
+    """
+    if isinstance(shape, int):
+        shape = (3, shape)
+    else:
+        shape = (3, *shape)
+    return np.empty(shape, *args, **kwargs)
+
+
+def ones_mosa(shape, *args, **kwargs):
+    """Return a new MOSA array of given shape and type, filled with ones.
+
+    Returns:
+        MOSA array of shape (6, shape).
+    """
+    if isinstance(shape, int):
+        shape = (6, shape)
+    else:
+        shape = (6, *shape)
+    return np.ones(shape, *args, **kwargs)
+
+
+def ones_sc(shape, *args, **kwargs):
+    """Return a new SC array of given shape and type, filled with ones.
+
+    Returns:
+        SC array of shape (3, shape).
+    """
+    if isinstance(shape, int):
+        shape = (3, shape)
+    else:
+        shape = (3, *shape)
+    return np.ones(shape, *args, **kwargs)
+
+
+def zeros_mosa(shape, *args, **kwargs):
+    """Return a new MOSA array of given shape and type, filled with zeros.
+
+    Returns:
+        MOSA array of shape (6, shape).
+    """
+    if isinstance(shape, int):
+        shape = (6, shape)
+    else:
+        shape = (6, *shape)
+    return np.zeros(shape, *args, **kwargs)
+
+
+def zeros_sc(shape, *args, **kwargs):
+    """Return a new SC array of given shape and type, filled with zeros.
+
+    Returns:
+        SC array of shape (3, shape).
+    """
+    if isinstance(shape, int):
+        shape = (3, shape)
+    else:
+        shape = (3, *shape)
+    return np.zeros(shape, *args, **kwargs)
+
+
+def full_mosa(shape, fill_value, *args, **kwargs):
+    """Return a new MOSA array of given shape and type, filled with `fill_value`.
+
+    Returns:
+        MOSA array of shape (6, shape).
+    """
+    if isinstance(shape, int):
+        shape = (6, shape)
+    else:
+        shape = (6, *shape)
+    return np.full(shape, fill_value, *args, **kwargs)
+
+
+def full_sc(shape, fill_value, *args, **kwargs):
+    """Return a new SC array of given shape and type, filled with `fill_value`.
+
+    Returns:
+        SC array of shape (3, shape).
+    """
+    if isinstance(shape, int):
+        shape = (3, shape)
+    else:
+        shape = (3, *shape)
+    return np.full(shape, fill_value, *args, **kwargs)
+
+
+def dict2mosa(dictionary):
+    """Return a new MOSA array with values given in `dictionary`.
+
+    Args:
+        dictionary: dictionary with MOSA indices as keys
+    """
+    dict_str = {str(key): dictionary[key] for key in dictionary}
+    try:
+        values = [dict_str[mosa] for mosa in MOSAS]
+    except KeyError as error:
+        raise ValueError(f"incomplete dictionary '{dictionary}', must contain '{MOSAS}' as keys") from error
+    return np.stack(values, axis=0)
+
+
+def dict2sc(dictionary):
+    """Return a new SC array with values given in `dictionary`.
+
+    Args:
+        dictionary: dictionary with SC indices as keys
+    """
+    dict_str = {str(key): dictionary[key] for key in dictionary}
+    try:
+        values = [dict_str[sc] for sc in SC]
+    except KeyError as error:
+        raise ValueError(f"incomplete dictionary '{dictionary}', must contain '{SC}' as keys") from error
+    return np.stack(values, axis=0)