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

Closes #34 (closed)

lisainstrument/indexing.py

+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Definitions and routines to handle indexing of measurements.
+
+Authors:
+    Jean-Baptiste Bayle <j2b.bayle@gmail.com>
+"""
+
+import numpy as np
+
+
+SC1, SC2, SC3 = 0, 1, 2
+"""Spacecraft (SC) indices"""
+
+MOSAL, MOSAR = 0, 1
+"""Left and right MOSA indices"""
+
+MOSA12, MOSA13 = (SC1, MOSAL), (SC1, MOSAR)
+MOSA23, MOSA21 = (SC2, MOSAL), (SC2, MOSAR)
+MOSA31, MOSA32 = (SC3, MOSAL), (SC3, MOSAR)
+"""MOSA indices, as combination of SC and MOSA"""
+
+ISC, REF, TM = 0, 1, 2
+"""Interferometer (IFO) indices"""
+
+CARRIER, USB = 0, 1
+"""Modulated (MOD) beam indices"""
+
+OFFSETS, FLUCTUATIONS = 0, 1
+"""Two-variable (VAR) decomposition indices"""
+
+ADJACENT_MOSAS = (slice(None), [MOSAR, MOSAL])
+"""Indices for adjacent MOSAs, as combinations of SC and MOSA
+Note that this uses complex slicing, and therefore copy the data."""
+
+DISTANT_MOSAS = (
+    [[SC2, SC3], [SC3, SC1], [SC1, SC2]],
+    [[MOSAR, MOSAL], [MOSAR, MOSAL], [MOSAR, MOSAL]])
+"""Indices for distant MOSAs, as combinations of SC and MOSA
+Note that this uses complex slicing, and therefore copy the data."""
+
+
+def dict2sc(x, axis=0):
+    """Return a new array whose `axis` is for spacecraft with content of `x`.
+
+    The keys of `x` should be [1, 2, 3] (or convertible to those), and values array-like.
+
+    Note that all values of `x` should have the same shape, or can be automatically broadcasted.
+    If they have shape (N), the resulting array will be of shape (3, N) by default (first axis).
+
+    Args:
+        x: input dictionary
+    """
+    keys = [1, 2, 3]
+    x_int = {int(key): x[key] for key in x}
+    if set(x_int.keys()) != set(keys):
+        raise IndexError(f"invalid spacecraft indices '{set(x.keys())}'")
+
+    return np.stack(tuple(x_int[sc] for sc in keys), axis=axis)
+
+
+def dict2mosa(x, axis=0):
+    """Return a new array whose `axis` is for spacecraft and `axis + 1` for MOSAs, with content of `x`.
+
+    The keys of `x` should be [12, 23, 31, 13, 32, 21] (or convertible to those), and values array-like.
+
+    Note that all values of `x` should have the same shape, or can be automatically broadcasted.
+    If they have shape (N), the resulting array will be of shape (3, 2, N) by default (first axis).
+
+    Args:
+        x: input dictionary
+    """
+    mosas = [12, 23, 31, 13, 32, 21]
+    x_int = {int(key): x[key] for key in x}
+    if set(x_int.keys()) != set(mosas):
+        raise IndexError(f"invalid MOSA indices '{set(x.keys())}'")
+
+    sc1 = np.stack((x_int[12], x_int[13]), axis=axis)
+    sc2 = np.stack((x_int[23], x_int[21]), axis=axis)
+    sc3 = np.stack((x_int[31], x_int[32]), axis=axis)
+    return np.stack((sc1, sc2, sc3), axis=axis)
+
+
+def sc2dict(x, axis=0):
+    """Return a dictionary whose keys are spacecraft indices, from array `x`.
+
+    The array `x` must have `axis` of dimension 3, ordered as [SC1, SC2, SC3].
+
+    The output dictionary containes the keys [1, 2, 3].
+
+    Args:
+        x: array-like input
+    """
+    if x.shape[axis] != 3:
+        raise IndexError(f"invalid size '{x.shape[axis]}' for spacecraft dimension '{axis}'")
+
+    return {
+        1: np.take(x, SC1, axis=axis),
+        2: np.take(x, SC2, axis=axis),
+        3: np.take(x, SC3, axis=axis),
+    }
+
+def mosa2dict(x, axis=0):
+    """Return a dictionary whose keys are MOSA indices, from array `x`.
+
+    The array `x` must have `axis` of dimensions 3, ordered as [SC1, SC2, SC3], and dimension
+    `axis + 1` of dimension 2, ordered as [MOSAL, MOSAR].
+
+    The output dictionary containes the keys [12, 23, 31, 13, 32, 21].
+
+    Args:
+        x: array-like input
+    """
+    if x.shape[axis] != 3:
+        raise IndexError(f"invalid size '{x.shape[axis]}' for spacecraft dimension '{axis}'")
+
+    if x.shape[axis + 1] != 2:
+        raise IndexError(f"invalid size '{x.shape[axis + 1]}' for MOSA dimension '{axis + 1}'")
+
+    take = lambda x, mosa: np.take(np.take(x, mosa[0], axis=axis), mosa[1], axis=axis)
+    return {
+        12: take(x, MOSA12),
+        23: take(x, MOSA23),
+        31: take(x, MOSA31),
+        13: take(x, MOSA13),
+        32: take(x, MOSA32),
+        21: take(x, MOSA21),
+    }