fix #782 : implemented AxisCollection.set_labels()

+ moved _guess_axis() from LArray to AxisCollection

Author: alixdamman

doc/source/api.rst

@@ -202,6 +202,7 @@ Modifying/Selecting
    AxisCollection.insert
    AxisCollection.rename
    AxisCollection.replace
+   AxisCollection.set_labels
    AxisCollection.without
    AxisCollection.combine_axes
    AxisCollection.split_axes

doc/source/changes/version_0_30.rst.inc

@@ -243,6 +243,8 @@ Miscellaneous improvements
   the overall performance improvement is negligible over most of the real-world models using larray that we tested these
   changes on.
 
+* implemented :py:obj:`AxisCollection.set_labels()` (closes :issue:`782`).
+
 
 Fixes
 ^^^^^

larray/core/array.py

@@ -2082,37 +2082,6 @@ def _translate_axis_key(self, axis_key):
         else:
             return self._translate_axis_key_chunk(axis_key)
 
-    def _guess_axis(self, axis_key):
-        if isinstance(axis_key, Group):
-            group_axis = axis_key.axis
-            if group_axis is not None:
-                # we have axis information but not necessarily an Axis object from self.axes
-                real_axis = self.axes[group_axis]
-                if group_axis is not real_axis:
-                    axis_key = axis_key.with_axis(real_axis)
-                return axis_key
-
-        # TODO: instead of checking all axes, we should have a big mapping
-        # (in AxisCollection or LArray):
-        # label -> (axis, index)
-        # or possibly (for ambiguous labels)
-        # label -> {axis: index}
-        # but for Pandas, this wouldn't work, we'd need label -> axis
-        valid_axes = []
-        for axis in self.axes:
-            try:
-                axis.index(axis_key)
-                valid_axes.append(axis)
-            except KeyError:
-                continue
-        if not valid_axes:
-            raise ValueError("%s is not a valid label for any axis" % axis_key)
-        elif len(valid_axes) > 1:
-            valid_axes = ', '.join(a.name if a.name is not None else '{{{}}}'.format(self.axes.index(a))
-                                   for a in valid_axes)
-            raise ValueError('%s is ambiguous (valid in %s)' % (axis_key, valid_axes))
-        return valid_axes[0][axis_key]
-
     def __getitem__(self, key, collapse_slices=False, translate_key=True):
         data = self.data
         # FIXME: I have a huge problem with boolean axis labels + non points
@@ -2788,13 +2757,13 @@ def to_labelgroup(key, stack_depth=1):
                 # new axis, and igroups are not the same that LGroups in this regard (I wonder if ideally it shouldn't
                 # be the same???)
                 # groups = tuple(self._translate_axis_key(k) for k in key)
-                groups = tuple(self._guess_axis(_to_key(k, stack_depth + 1)) for k in key)
+                groups = tuple(self.axes._guess_axis(_to_key(k, stack_depth + 1)) for k in key)
                 axis = groups[0].axis
                 if not all(g.axis.equals(axis) for g in groups[1:]):
                     raise ValueError("group with different axes: %s" % str(key))
                 return groups
             if isinstance(key, (Group, int, basestring, list, slice)):
-                return self._guess_axis(key)
+                return self.axes._guess_axis(key)
             else:
                 raise NotImplementedError("%s has invalid type (%s) for a group aggregate key"
                                           % (key, type(key).__name__))
@@ -7146,7 +7115,7 @@ def __array__(self, dtype=None):
 
     # TODO: this should be a thin wrapper around a method in AxisCollection
     def set_labels(self, axis=None, labels=None, inplace=False, **kwargs):
-        r"""Replaces the labels of an axis of array.
+        r"""Replaces the labels of one or several axes of the array.
 
         Parameters
         ----------
@@ -7170,6 +7139,10 @@ def set_labels(self, axis=None, labels=None, inplace=False, **kwargs):
         LArray
             Array with modified labels.
 
+        See Also
+        --------
+        AxisCollection.set_labels
+
         Examples
         --------
         >>> a = ndtest('nat=BE,FO;sex=M,F')
@@ -7234,33 +7207,7 @@ def set_labels(self, axis=None, labels=None, inplace=False, **kwargs):
         Belgian    0  1
              FO    2  3
         """
-        if axis is None:
-            changes = {}
-        elif isinstance(axis, dict):
-            changes = axis
-        elif isinstance(axis, (basestring, Axis, int)):
-            changes = {axis: labels}
-        else:
-            raise ValueError("Expected None or a string/int/Axis/dict instance for axis argument")
-        changes.update(kwargs)
-        # TODO: we should implement the non-dict behavior in Axis.replace, so that we can simplify this code to:
-        # new_axes = [self.axes[old_axis].replace(axis_changes) for old_axis, axis_changes in changes.items()]
-        new_axes = []
-        for old_axis, axis_changes in changes.items():
-            try:
-                real_axis = self.axes[old_axis]
-            except KeyError:
-                axis_changes = {old_axis: axis_changes}
-                real_axis = self._guess_axis(old_axis).axis
-            if isinstance(axis_changes, dict):
-                new_axis = real_axis.replace(axis_changes)
-            elif callable(axis_changes):
-                new_axis = real_axis.apply(axis_changes)
-            else:
-                new_axis = Axis(axis_changes, real_axis.name)
-            new_axes.append((real_axis, new_axis))
-        axes = self.axes.replace(new_axes)
-
+        axes = self.axes.set_labels(axis, labels, **kwargs)
         if inplace:
             self.axes = axes
             return self

larray/core/axis.py

@@ -2324,6 +2324,163 @@ def replace(self, axes_to_replace=None, new_axis=None, inplace=False, **kwargs):
         else:
             return AxisCollection(axes)
 
+    def _guess_axis(self, axis_key):
+        if isinstance(axis_key, Group):
+            group_axis = axis_key.axis
+            if group_axis is not None:
+                # we have axis information but not necessarily an Axis object from self.axes
+                real_axis = self[group_axis]
+                if group_axis is not real_axis:
+                    axis_key = axis_key.with_axis(real_axis)
+                return axis_key
+
+        # TODO: instead of checking all axes, we should have a big mapping
+        # (in AxisCollection or LArray):
+        # label -> (axis, index)
+        # or possibly (for ambiguous labels)
+        # label -> {axis: index}
+        # but for Pandas, this wouldn't work, we'd need label -> axis
+        valid_axes = []
+        for axis in self:
+            try:
+                axis.index(axis_key)
+                valid_axes.append(axis)
+            except KeyError:
+                continue
+        if not valid_axes:
+            raise ValueError("%s is not a valid label for any axis" % axis_key)
+        elif len(valid_axes) > 1:
+            valid_axes = ', '.join(a.name if a.name is not None else '{{{}}}'.format(self.axes.index(a))
+                                   for a in valid_axes)
+            raise ValueError('%s is ambiguous (valid in %s)' % (axis_key, valid_axes))
+        return valid_axes[0][axis_key]
+
+    def set_labels(self, axis=None, labels=None, inplace=False, **kwargs):
+        r"""Replaces the labels of one or several axes.
+
+        Parameters
+        ----------
+        axis : string or Axis or dict
+            Axis for which we want to replace labels, or mapping {axis: changes} where changes can either be the
+            complete list of labels, a mapping {old_label: new_label} or a function to transform labels.
+            If there is no ambiguity (two or more axes have the same labels), `axis` can be a direct mapping
+            {old_label: new_label}.
+        labels : int, str, iterable or mapping or function, optional
+            Integer or list of values usable as the collection of labels for an Axis. If this is mapping, it must be
+            {old_label: new_label}. If it is a function, it must be a function accepting a single argument (a
+            label) and returning a single value. This argument must not be used if axis is a mapping.
+        inplace : bool, optional
+            Whether or not to modify the original object or return a new AxisCollection and leave the original intact.
+            Defaults to False.
+        **kwargs :
+            `axis`=`labels` for each axis you want to set labels.
+
+        Returns
+        -------
+        AxisCollection
+            AxisCollection with modified labels.
+
+        Examples
+        --------
+        >>> from larray import ndtest
+        >>> axes = AxisCollection('nat=BE,FO;sex=M,F')
+        >>> axes
+        AxisCollection([
+            Axis(['BE', 'FO'], 'nat'),
+            Axis(['M', 'F'], 'sex')
+        ])
+        >>> axes.set_labels('sex', ['Men', 'Women'])
+        AxisCollection([
+            Axis(['BE', 'FO'], 'nat'),
+            Axis(['Men', 'Women'], 'sex')
+        ])
+
+        when passing a single string as labels, it will be interpreted to create the list of labels, so that one can
+        use the same syntax than during axis creation.
+
+        >>> axes.set_labels('sex', 'Men,Women')
+        AxisCollection([
+            Axis(['BE', 'FO'], 'nat'),
+            Axis(['Men', 'Women'], 'sex')
+        ])
+
+        to replace only some labels, one must give a mapping giving the new label for each label to replace
+
+        >>> axes.set_labels('sex', {'M': 'Men'})
+        AxisCollection([
+            Axis(['BE', 'FO'], 'nat'),
+            Axis(['Men', 'F'], 'sex')
+        ])
+
+        to transform labels by a function, use any function accepting and returning a single argument:
+
+        >>> axes.set_labels('nat', str.lower)
+        AxisCollection([
+            Axis(['be', 'fo'], 'nat'),
+            Axis(['M', 'F'], 'sex')
+        ])
+
+        to replace labels for several axes at the same time, one should give a mapping giving the new labels for each
+        changed axis
+
+        >>> axes.set_labels({'sex': 'Men,Women', 'nat': 'Belgian,Foreigner'})
+        AxisCollection([
+            Axis(['Belgian', 'Foreigner'], 'nat'),
+            Axis(['Men', 'Women'], 'sex')
+        ])
+
+        or use keyword arguments
+
+        >>> axes.set_labels(sex='Men,Women', nat='Belgian,Foreigner')
+        AxisCollection([
+            Axis(['Belgian', 'Foreigner'], 'nat'),
+            Axis(['Men', 'Women'], 'sex')
+        ])
+
+        one can also replace some labels in several axes by giving a mapping of mappings
+
+        >>> axes.set_labels({'sex': {'M': 'Men'}, 'nat': {'BE': 'Belgian'}})
+        AxisCollection([
+            Axis(['Belgian', 'FO'], 'nat'),
+            Axis(['Men', 'F'], 'sex')
+        ])
+
+        when there is no ambiguity (two or more axes have the same labels), it is possible to give a mapping
+        between old and new labels
+
+        >>> axes.set_labels({'M': 'Men', 'BE': 'Belgian'})
+        AxisCollection([
+            Axis(['Belgian', 'FO'], 'nat'),
+            Axis(['Men', 'F'], 'sex')
+        ])
+        """
+        if axis is None:
+            changes = {}
+        elif isinstance(axis, dict):
+            changes = axis
+        elif isinstance(axis, (basestring, Axis, int)):
+            changes = {axis: labels}
+        else:
+            raise ValueError("Expected None or a string/int/Axis/dict instance for axis argument")
+        changes.update(kwargs)
+        # TODO: we should implement the non-dict behavior in Axis.replace, so that we can simplify this code to:
+        # new_axes = [self[old_axis].replace(axis_changes) for old_axis, axis_changes in changes.items()]
+        new_axes = []
+        for old_axis, axis_changes in changes.items():
+            try:
+                real_axis = self[old_axis]
+            except KeyError:
+                axis_changes = {old_axis: axis_changes}
+                real_axis = self._guess_axis(old_axis).axis
+            if isinstance(axis_changes, dict):
+                new_axis = real_axis.replace(axis_changes)
+            elif callable(axis_changes):
+                new_axis = real_axis.apply(axis_changes)
+            else:
+                new_axis = Axis(axis_changes, real_axis.name)
+            new_axes.append((real_axis, new_axis))
+        return self.replace(new_axes, inplace=inplace)
+
     # TODO: deprecate method (should use __sub__ instead)
     def without(self, axes):
         """