From bd3b47ba0844cb3da1d9a719e540e38ec2cf1ac4 Mon Sep 17 00:00:00 2001
From: Andrew Davison <andrew.davison@cnrs.fr>
Date: Tue, 9 Nov 2021 12:26:14 +0100
Subject: [PATCH 1/2] Partial fix for #17 - handles the case where 2nd argument
 is a scalar

---
 lazyarray.py       | 37 ++++++++++++++++++++++++++++---------
 setup.py           |  2 +-
 test/test_ufunc.py |  8 +++++++-
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/lazyarray.py b/lazyarray.py
index a202f15..307e536 100644
--- a/lazyarray.py
+++ b/lazyarray.py
@@ -554,6 +554,15 @@ def __call__(self, arg):
     __abs__ = lazy_unary_operation('abs')
 
 
+class VectorizedIterable(object):
+    """
+    Base class for any class which has a method `next(n)`, i.e., where you
+    can choose how many values to return rather than just returning one at a
+    time.
+    """
+    pass
+
+
 def _build_ufunc(func):
     """Return a ufunc that works with lazy arrays"""
     def larray_compatible_ufunc(x):
@@ -566,18 +575,28 @@ def larray_compatible_ufunc(x):
     return larray_compatible_ufunc
 
 
-class VectorizedIterable(object):
-    """
-    Base class for any class which has a method `next(n)`, i.e., where you
-    can choose how many values to return rather than just returning one at a
-    time.
-    """
-    pass
+def _build_ufunc_2nd_arg(func):
+    """Return a ufunc taking a second, non-array argument, that works with lazy arrays"""
+    def larray_compatible_ufunc2(x1, x2):
+        if not isinstance(x2, numbers.Number):
+            raise TypeError("lazyarry ufuncs do not accept an array as the second argument")
+        if isinstance(x1, larray):
+            def partial(x):
+                return func(x, x2)
+            y = deepcopy(x1)
+            y.apply(partial)
+            return y
+        else:
+            return func(x1, x2)
+    return larray_compatible_ufunc2
 
 
 # build lazy-array compatible versions of NumPy ufuncs
 namespace = globals()
 for name in dir(np):
     obj = getattr(np, name)
-    if isinstance(obj, np.ufunc):
-        namespace[name] = _build_ufunc(obj)
+    if isinstance(obj, np.ufunc) and name not in namespace:
+        if name in ("power", "fmod", "arctan2, hypot, ldexp, maximum, minimum"):
+            namespace[name] = _build_ufunc_2nd_arg(obj)
+        else:
+            namespace[name] = _build_ufunc(obj)
diff --git a/setup.py b/setup.py
index 9d3c725..27f1e6f 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@
 
 setup(
     name='lazyarray',
-    version='0.4.0',
+    version='0.5.0.dev',
     py_modules=['lazyarray'],
     license='Modified BSD',
     author="Andrew P. Davison",
diff --git a/test/test_ufunc.py b/test/test_ufunc.py
index d7e3ead..eb069ad 100644
--- a/test/test_ufunc.py
+++ b/test/test_ufunc.py
@@ -4,7 +4,7 @@
 Copyright Andrew P. Davison, 2012-2017
 """
 
-from lazyarray import larray, sqrt, cos
+from lazyarray import larray, sqrt, cos, power
 import numpy as np
 from numpy.testing import assert_array_equal, assert_array_almost_equal
 
@@ -42,3 +42,9 @@ def clock():
                               np.array([[1.0, 0.0],
                                            [-1.0, 0.0]]),
                               decimal=15)
+
+
+def testpower_from_array():
+    A = larray(np.array([1, 4, 9, 16, 25]))
+    assert_array_equal(power(A, 0.5).evaluate(),
+                       np.arange(1, 6))

From 238de5ae7faa4052b2c6575b64b43ad5a0660678 Mon Sep 17 00:00:00 2001
From: Andrew Davison <andrew.davison@cnrs.fr>
Date: Tue, 9 Nov 2021 13:20:48 +0100
Subject: [PATCH 2/2] test some edge cases

---
 test/test_ufunc.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/test/test_ufunc.py b/test/test_ufunc.py
index eb069ad..6be5c36 100644
--- a/test/test_ufunc.py
+++ b/test/test_ufunc.py
@@ -4,9 +4,10 @@
 Copyright Andrew P. Davison, 2012-2017
 """
 
-from lazyarray import larray, sqrt, cos, power
+from lazyarray import larray, sqrt, cos, power, fmod
 import numpy as np
 from numpy.testing import assert_array_equal, assert_array_almost_equal
+from nose.tools import assert_raises
 
 
 def test_sqrt_from_array():
@@ -44,7 +45,18 @@ def clock():
                               decimal=15)
 
 
-def testpower_from_array():
+def test_power_from_array():
     A = larray(np.array([1, 4, 9, 16, 25]))
     assert_array_equal(power(A, 0.5).evaluate(),
                        np.arange(1, 6))
+
+
+def test_power_with_plain_array():
+    A = np.array([1, 4, 9, 16, 25])
+    assert_array_equal(power(A, 0.5), np.arange(1, 6))
+
+
+def test_fmod_with_array_as_2nd_arg():
+    A = larray(np.array([1, 4, 9, 16, 25]))
+    B = larray(np.array([1, 4, 9, 16, 25]))
+    assert_raises(TypeError, fmod, A, B)