FEAT: implemented Session.align(other_session) (closes #501)

in the process had to implement Axis.align(*axes), AxisCollection.align(*axis_collections) and Array.align(*arrays)

also deprecated passing non (arrays or scalar) to Array.align and check that each aligned axis is either anonymous or have the same name than others

Author: gdementen

doc/source/changes/version_0_35.rst.inc

@@ -11,6 +11,11 @@ Syntax changes
   (:py:obj:`Array.plot.area()`, :py:obj:`Array.plot.bar()`,
   :py:obj:`Array.plot.barh()`, and :py:obj:`Array.plot.line()`).
 
+* all align() methods (:py:obj:`Axis.align()`, :py:obj:`AxisCollection.align()`
+  and :py:obj:`Array.align()`) only take options (``join``, ``axes`` and/or
+  ``fill_value``) as keywords arguments. Extra positional arguments will be
+  considered as more objects to align (see below).
+
 
 Backward incompatible changes
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -46,6 +51,9 @@ New features
 
     >>> arr.plot.heatmap()
 
+* implemented :py:obj:`Session.align()` to align all the arrays in several
+  sessions at once. Closes :issue:`501`.
+
 * added a feature (see the :ref:`miscellaneous section <misc>` for details). It works on :ref:`api-axis` and
   :ref:`api-group` objects.
 
@@ -89,6 +97,12 @@ Miscellaneous improvements
 
 * made :py:obj:`ipfp()` slightly faster when display_progress is False.
 
+* all align() methods (:py:obj:`Axis.align()`, :py:obj:`AxisCollection.align()`
+  and :py:obj:`Array.align()`) now support aligning more than two objects at
+  once by passing them as positional arguments. For example:
+
+    >>> array1.align(array2, array3, join='outer')
+
 
 Fixes
 ^^^^^

larray/core/array.py

@@ -54,6 +54,7 @@
 from larray.core.group import (Group, IGroup, LGroup, _to_key, _to_keys,
                                _translate_sheet_name, _translate_group_key_hdf)
 from larray.core.axis import Axis, AxisReference, AxisCollection, X, _make_axis         # noqa: F401
+from larray.core.axis import align_axis_collections
 from larray.core.plot import PlotObject
 from larray.util.misc import (table2str, size2str, ReprString,
                               float_error_handler_factory, light_product, common_dtype,
@@ -853,6 +854,34 @@ def np_array_to_pd_index(array, name=None, tupleize_cols=True):
     return pd.Index(array, dtype=dtype, name=name, tupleize_cols=tupleize_cols)
 
 
+def align_arrays(values, join='outer', fill_value=nan, axes=None):
+    bad_values = [value for value in values
+                 if not isinstance(value, Array) and not np.isscalar(value)]
+    if bad_values:
+        bad_types = set(type(v) for v in bad_values)
+        bad_type_names = sorted(t.__name__ for t in bad_types)
+        raise TypeError("align only supports Arrays and scalars but got:"
+                        f"{', '.join(bad_type_names)}")
+    axis_collections = [
+        value.axes if isinstance(value, Array) else AxisCollection()
+        for value in values
+    ]
+    # fail early because reindex does not currently support anonymous axes
+    if any(any(name is None for name in axis_col.names)
+           for axis_col in axis_collections):
+        raise ValueError("arrays with anonymous axes are currently not "
+                         "supported by Array.align")
+    try:
+       aligned_axis_collections = align_axis_collections(axis_collections,
+                                                         join=join, axes=axes)
+    except ValueError as e:
+        raise ValueError(f"Arrays are not aligned because {e}")
+    return tuple(value.reindex(aligned_axes, fill_value=fill_value)
+                    if isinstance(value, Array)
+                    else value
+                 for value, aligned_axes in zip(values, aligned_axis_collections))
+
+
 class Array(ABCArray):
     r"""
     An Array object represents a multidimensional, homogeneous array of fixed-size items with labeled axes.
@@ -1817,14 +1846,14 @@ def get_group(res_axes, self_axis):
         else:
             return res
 
-    def align(self, other, join='outer', fill_value=nan, axes=None) -> Tuple['Array', 'Array']:
-        r"""Align two arrays on their axes with the specified join method.
+    def align(self, *other, join='outer', fill_value=nan, axes=None) -> Tuple['Array', 'Array']:
+        r"""Align array with other(s) on their axes with the specified join method.
 
         In other words, it ensure all common axes are compatible. Those arrays can then be used in binary operations.
 
         Parameters
         ----------
-        other : Array-like
+        *other : Array-like
         join : {'outer', 'inner', 'left', 'right', 'exact'}, optional
             Join method. For each axis common to both arrays:
               - outer: will use a label if it is in either arrays axis (ordered like the first array).
@@ -1837,13 +1866,13 @@ def align(self, other, join='outer', fill_value=nan, axes=None) -> Tuple['Array'
             Value used to fill cells corresponding to label combinations which are not common to both arrays.
             Defaults to NaN.
         axes : AxisReference or sequence of them, optional
-            Axes to align. Need to be valid in both arrays. Defaults to None (all common axes). This must be specified
+            Axes to align. Need to be valid in all arrays. Defaults to None (all common axes). This must be specified
             when mixing anonymous and non-anonymous axes.
 
         Returns
         -------
-        (left, right) : (Array, Array)
-            Aligned objects
+        arrays : tuple of Array
+            Aligned arrays
 
         Notes
         -----
@@ -1989,18 +2018,11 @@ def align(self, other, join='outer', fill_value=nan, axes=None) -> Tuple['Array'
         >>> arr1.align(arr2, join='exact')   # doctest: +NORMALIZE_WHITESPACE
         Traceback (most recent call last):
         ...
-        ValueError: Both arrays are not aligned because align method with join='exact'
+        ValueError: Arrays are not aligned because align method with join='exact'
         expected Axis(['a0', 'a1'], 'a') to be equal to Axis(['a0', 'a1', 'a2'], 'a')
         """
-        other = asarray(other)
-        # reindex does not currently support anonymous axes
-        if any(name is None for name in self.axes.names) or any(name is None for name in other.axes.names):
-            raise ValueError("arrays with anonymous axes are currently not supported by Array.align")
-        try:
-            left_axes, right_axes = self.axes.align(other.axes, join=join, axes=axes)
-        except ValueError as e:
-            raise ValueError(f"Both arrays are not aligned because {e}")
-        return self.reindex(left_axes, fill_value=fill_value), other.reindex(right_axes, fill_value=fill_value)
+        return align_arrays((self, *other),
+                            join=join, fill_value=fill_value, axes=axes)
 
     @deprecate_kwarg('reverse', 'ascending', {True: False, False: True})
     def sort_values(self, key=None, axis=None, ascending=True) -> 'Array':

larray/core/axis.py

@@ -17,9 +17,9 @@
 from larray.util.misc import (duplicates, array_lookup2, ReprString, index_by_id, renamed_to, LHDFStore,
                               lazy_attribute, _isnoneslice, unique_list, unique_multi, Product, argsort, has_duplicates,
                               exactly_one, concatenate_ndarrays)
+from larray.util.misc import first
 from larray.util.types import Scalar
 
-
 np_frompyfunc = np.frompyfunc
 
 
@@ -1330,12 +1330,12 @@ def difference(self, other) -> 'Axis':
         to_drop = set(other)
         return Axis([label for label in self.labels if label not in to_drop], self.name)
 
-    def align(self, other, join='outer') -> 'Axis':
+    def align(self, *other, join='outer') -> 'Axis':
         r"""Align axis with other object using specified join method.
 
         Parameters
         ----------
-        other : Axis or label sequence
+        *other : Axis or label sequence
         join : {'outer', 'inner', 'left', 'right', 'exact'}, optional
             Defaults to 'outer'.
 
@@ -1366,22 +1366,16 @@ def align(self, other, join='outer') -> 'Axis':
         ValueError: align method with join='exact' expected
         Axis(['a0', 'a1', 'a2'], 'a') to be equal to Axis(['a1', 'a2', 'a3'], 'a')
         """
-        assert join in {'outer', 'inner', 'left', 'right', 'exact'}
-        if join == 'outer':
-            return self.union(other)
-        elif join == 'inner':
-            return self.intersection(other)
-        elif join == 'left':
-            return self
-        elif join == 'right':
-            if not isinstance(other, Axis):
-                other = Axis(other)
-            return other
-        elif join == 'exact':
-            if not self.equals(other):
-                raise ValueError(f"align method with join='exact' expected {self!r} to be equal to {other!r}")
-            else:
-                return self
+        bad_objs = [obj for obj in other if not isinstance(obj, Axis)]
+        if bad_objs:
+            for obj in bad_objs:
+                obj_type = type(obj).__name__
+                warnings.warn(f"aligning an Axis to a non-Axis object "
+                              f"({obj_type}) is deprecated. Please convert to "
+                              f"an Axis first.", FutureWarning, stacklevel=2)
+            other = [obj if isinstance(obj, Axis) else Axis(obj)
+                     for obj in other]
+        return align_axes((self, *other), join=join)
 
     def to_hdf(self, filepath, key=None) -> None:
         r"""
@@ -1462,6 +1456,50 @@ def ignore_labels(self) -> 'Axis':
         return Axis(len(self), self.name)
 
 
+def align_axes(axes: Sequence[Axis], join: str = 'outer') -> Axis:
+    if not all(isinstance(axis, Axis) for axis in axes):
+        raise TypeError("all objects to align must be Axis objects")
+
+    if join not in {'outer', 'inner', 'left', 'right', 'exact'}:
+        raise ValueError(f"join must be one of 'outer', 'inner', 'left', "
+                         f"'right' or 'exact', got {join!r}")
+
+    names = [axis.name for axis in axes]
+    first_name = first((name for name in names if name is not None),
+                       default=None)
+    if first_name is not None:
+        if not all(name is None or name == first_name for name in names):
+            raise ValueError("In align, all axes must be anonymous or "
+                             "have the same name: "
+                             f"{', '.join(repr(name) for name in names)}")
+
+    def join_left(axis1, axis2):
+        return axis1
+    def join_right(axis1, axis2):
+        return axis2
+    def join_exact(axis1, axis2):
+        if not axis1.equals(axis2):
+            raise ValueError(f"align method with join='exact' expected "
+                             f"{axis1!r} to be equal to {axis2!r}")
+        else:
+            return axis1
+    if join == 'outer':
+        join_labels_func = Axis.union
+    elif join == 'inner':
+        join_labels_func = Axis.intersection
+    elif join == 'left':
+        join_labels_func = join_left
+    elif join == 'right':
+        join_labels_func = join_right
+    else:
+        assert join == 'exact'
+        join_labels_func = join_exact
+    aligned_axis = axes[0]
+    for axis in axes[1:]:
+        aligned_axis = join_labels_func(aligned_axis, axis)
+    return aligned_axis
+
+
 def _make_axis(obj) -> Axis:
     if isinstance(obj, Axis):
         return obj
@@ -3552,23 +3590,25 @@ def _prepare_split_axes(self, axes, names, sep) -> dict:
 
     split_axis = renamed_to(split_axes, 'split_axis', raise_error=True)
 
-    def align(self, other, join='outer', axes=None) -> Tuple['AxisCollection', 'AxisCollection']:
-        r"""Align this axis collection with another.
+    def align(self, *other, join='outer', axes=None) -> Tuple['AxisCollection']:
+        r"""Align this AxisCollection with (an)other AxisCollection(s).
 
         This ensures all common axes are compatible.
 
         Parameters
         ----------
-        other : AxisCollection
+        *other : AxisCollection
+            AxisCollection(s) to align with this one.
         join : {'outer', 'inner', 'left', 'right', 'exact'}, optional
             Defaults to 'outer'.
         axes : AxisReference or sequence of them, optional
-            Axes to align. Need to be valid in both arrays. Defaults to None (all common axes). This must be specified
+            Axes to align. Need to be valid in all axis collections.
+            Defaults to None (all common axes). This must be specified
             when mixing anonymous and non-anonymous axes.
 
         Returns
         -------
-        (left, right) : (AxisCollection, AxisCollection)
+        tuple of AxisCollection
             Aligned collections
 
         See Also
@@ -3631,31 +3671,20 @@ def align(self, other, join='outer', axes=None) -> Tuple['AxisCollection', 'Axis
             Axis(['c0'], None)
         ])
         """
-        if join not in {'outer', 'inner', 'left', 'right', 'exact'}:
-            raise ValueError("join should be one of 'outer', 'inner', 'left', 'right' or 'exact'")
-        other = other if isinstance(other, AxisCollection) else AxisCollection(other)
-
-        # if axes not specified
-        if axes is None:
-            # and we have only anonymous axes on both sides
-            if all(name is None for name in self.names) and all(name is None for name in other.names):
-                # use N first axes by position
-                join_axes = list(range(min(len(self), len(other))))
-            elif any(name is None for name in self.names) or any(name is None for name in other.names):
-                raise ValueError("axes collections with mixed anonymous/non anonymous axes are not supported by align"
-                                 "without specifying axes explicitly")
-            else:
-                assert all(name is not None for name in self.names) and all(name is not None for name in other.names)
-                # use all common axes
-                join_axes = list(OrderedSet(self.names) & OrderedSet(other.names))
-        else:
-            if isinstance(axes, (int, str, Axis)):
-                axes = [axes]
-            join_axes = axes
-        new_axes = [self_axis.align(other_axis, join=join)
-                    for self_axis, other_axis in zip(self[join_axes], other[join_axes])]
-        axes_changes = list(zip(join_axes, new_axes))
-        return self.replace(axes_changes), other.replace(axes_changes)
+        # For backward compatibility with older code using align with a
+        # non-AxisCollection second argument, we only support aligning more
+        # than two collection when other contains actual AxisCollection objects
+        bad_objs = [obj for obj in other if not isinstance(obj, AxisCollection)]
+        if bad_objs:
+            for obj in bad_objs:
+                obj_type = type(obj).__name__
+                warnings.warn(f"aligning an AxisCollection to a "
+                              f"non-AxisCollection object ({obj_type}) is "
+                              f"deprecated. Please convert to an AxisCollection "
+                              f"first.", FutureWarning, stacklevel=2)
+            other = [AxisCollection(obj) for obj in other]
+
+        return align_axis_collections((self, *other), join=join, axes=axes)
 
     # XXX: make this into a public method/property? AxisCollection.flat_labels[flat_indices]?
     def _flat_lookup(self, flat_indices):
@@ -3802,6 +3831,57 @@ def _adv_keys_to_combined_axes(self, key, wildcard=False, sep='_'):
         return AxisCollection(combined_axis)
 
 
+def align_axis_collections(axis_collections, join='outer', axes=None):
+    if join not in {'outer', 'inner', 'left', 'right', 'exact'}:
+        raise ValueError("join should be one of 'outer', 'inner', 'left', "
+                         "'right' or 'exact'")
+
+    # if axes not specified
+    if axes is None:
+        # and we have only anonymous axes
+        if all(name is None for col in axis_collections
+               for name in col.names):
+            # use all axes by position
+            max_length = max(len(col) for col in axis_collections)
+            join_axes_refs = list(range(max_length))
+        elif any(name is None for col in axis_collections
+                 for name in col.names):
+            raise ValueError(
+                "axes collections with mixed anonymous/non anonymous axes "
+                "are not supported by align without specifying axes "
+                "explicitly")
+        else:
+            assert all(name is not None for col in axis_collections
+                       for name in col.names)
+            # use all axes by name
+            join_axes_refs = OrderedSet(axis_collections[0].names)
+            for col in axis_collections[1:]:
+                join_axes_refs |= OrderedSet(col.names)
+    else:
+        if isinstance(axes, (int, str, Axis)):
+            axes = [axes]
+        join_axes_refs = axes
+
+    # first compute all aligned axes for all collections
+    axes_changes = {
+        axis_ref: align_axes([axis_col[axis_ref]
+                              for axis_col in axis_collections
+                              if axis_ref in axis_col],
+                             join=join)
+        for axis_ref in join_axes_refs
+    }
+
+    # then apply the changed axes for the collections where the axis exists
+    return tuple(
+        axis_col.replace({
+            axis_ref: aligned_axis
+            for axis_ref, aligned_axis in axes_changes.items()
+            if axis_ref in axis_col
+        })
+        for axis_col in axis_collections
+    )
+
+
 class AxisReference(ABCAxisReference, ExprNode, Axis):
     def __init__(self, name):
         self.name = name