FEAT: evaluate X.axis_name expressions in make_numpy_broadcastable

gdementen · gdementen · commit 22ac89a6eb23 · 2023-10-23T16:02:41.000+02:00
(used in all wrapped numpy functions, most notably where())
diff --git a/doc/source/changes/version_0_34_2.rst.inc b/doc/source/changes/version_0_34_2.rst.inc
@@ -1,55 +1,21 @@
 ﻿.. py:currentmodule:: larray
 
 
-Syntax changes
-^^^^^^^^^^^^^^
-
-* renamed ``Array.old_method_name()`` to :py:obj:`Array.new_method_name()` (closes :issue:`1`).
-
-* renamed ``old_argument_name`` argument of :py:obj:`Array.method_name()` to ``new_argument_name``.
-
-
-Backward incompatible changes
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-* other backward incompatible changes
-
-
 New features
 ^^^^^^^^^^^^
 
-* added a feature (see the :ref:`miscellaneous section <misc>` for details). It works on :ref:`api-axis` and
-  :ref:`api-group` objects.
-
-  Here is an example of the new feature:
-
-    >>> arr = ndtest((2, 3))
-    >>> arr
-    a\b  b0  b1  b2
-     a0   0   1   2
-     a1   3   4   5
-
-  And it can also be used like this:
+* added support for evaluating expressions using X.axis_name when calling some
+  built-in functions, most notably `where()`. For example, the following code
+  now works (previously it seemed to work but produced the wrong result -- see the
+  fixes section below): ::
 
-    >>> arr = ndtest("a=a0..a2")
+    >>> arr = ndtest("age=0..3")
     >>> arr
-    a  a0  a1  a2
-        0   1   2
-
-* added another feature in the editor (closes :editor_issue:`1`).
-
-  .. note::
-
-     - It works for foo bar !
-     - It does not work for foo baz !
-
-
-.. _misc:
-
-Miscellaneous improvements
-^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-* improved something.
+    age  0  1  2  3
+         0  1  2  3
+    >>> where(X.age == 2, 42, arr)
+    age  0  1   2  3
+         0  1  42  3
 
 
 Fixes
@@ -68,7 +34,7 @@ Fixes
   there is no array to extract the axis labels from) instead of always
   evaluating to True. This was especially dangerous in the context of a
   where() function, which always evaluated to its left side
-  (e.g. where(X.age > 0, arr, 0) evaluated to 'arr' for all ages).
+  (e.g. `where(X.age > 0, arr, 0)` evaluated to `arr` for all ages).
   Closes :issue:`1083`.
 
 * expressions using `X.axis_name` and an Array now evaluate correctly when
diff --git a/larray/core/array.py b/larray/core/array.py
@@ -9949,13 +9949,19 @@ def make_numpy_broadcastable(values, min_axes=None) -> Tuple[List[Array], AxisCo
     --------
     Axis.iscompatible : tests if axes are compatible between them.
     """
-    all_axes = AxisCollection.union(*[get_axes(v) for v in values])
+    axes_union = AxisCollection.union(*[get_axes(v) for v in values])
     if min_axes is not None:
         if not isinstance(min_axes, AxisCollection):
             min_axes = AxisCollection(min_axes)
-        all_axes = min_axes | all_axes
-    return [v.broadcast_with(all_axes) if isinstance(v, Array) else v
-            for v in values], all_axes
+        axes_union = min_axes | axes_union
+    def broadcasted_value(value):
+        if isinstance(value, Array):
+            return value.broadcast_with(axes_union)
+        elif isinstance(value, ExprNode):
+            return value.evaluate(axes_union)
+        else:
+            return value
+    return [broadcasted_value(value) for value in values], axes_union
 
 
 def raw_broadcastable(values, min_axes=None) -> Tuple[Tuple[Any, ...], AxisCollection]:
diff --git a/larray/tests/test_array.py b/larray/tests/test_array.py
@@ -5011,23 +5011,45 @@ def test_ufuncs(small_array):
     assert_nparray_equal(clip(small_array, low, high).data,
                          np.clip(raw, raw_low, raw_high))
 
+    # round
+    small_float = small_array + 0.6
+    rounded = round(small_float)
+    assert_nparray_equal(rounded.data, np.round(raw + 0.6))
+
+
+def test_where():
+    arr = ndtest((2, 3))
+    # a\b  b0  b1  b2
+    #  a0   0   1   2
+    #  a1   3   4   5
+
+    expected = from_string(r"""a\b  b0  b1  b2
+                                a0  -1  -1  -1
+                                a1  -1   4   5""")
+
     # where (no broadcasting)
-    assert_nparray_equal(where(small_array < 5, -5, small_array).data,
-                         np.where(raw < 5, -5, raw))
+    res = where(arr < 4, -1, arr)
+    assert_larray_equal(res, expected)
 
     # where (transposed no broadcasting)
-    assert_nparray_equal(where(small_array < 5, -5, small_array.T).data,
-                         np.where(raw < 5, -5, raw))
+    res = where(arr < 4, -1, arr.T)
+    assert_larray_equal(res, expected)
 
     # where (with broadcasting)
-    result = where(small_array['d1'] < 5, -5, small_array)
-    assert result.axes.names == ['c', 'd']
-    assert_nparray_equal(result.data, np.where(raw[:, [0]] < 5, -5, raw))
+    res = where(arr['b1'] < 4, -1, arr)
+    assert_larray_equal(res, from_string(r"""a\b  b0  b1  b2
+                                              a0  -1  -1  -1
+                                              a1   3   4   5"""))
 
-    # round
-    small_float = small_array + 0.6
-    rounded = round(small_float)
-    assert_nparray_equal(rounded.data, np.round(raw + 0.6))
+    # with expressions (issue #1083)
+    arr = ndtest("age=0..5")
+    res = where(X.age == 3, 42, arr)
+    assert_larray_equal(res, from_string("""age  0  1  2   3  4  5
+                                             \t  0  1  2  42  4  5"""))
+
+    res = where(X.age == 3, arr, 42)
+    assert_larray_equal(res, from_string("""age   0   1   2  3   4   5
+                                             \t  42  42  42  3  42  42"""))
 
 
 def test_eye():
