Merge pull request #18 from apdavison/two-arg-ufunc

Partial fix for #17 - handles the case where 2nd argument is a scalar

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)

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",

test/test_ufunc.py

@@ -4,9 +4,10 @@
 Copyright Andrew P. Davison, 2012-2017
 """
 
-from lazyarray import larray, sqrt, cos
+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():
@@ -42,3 +43,20 @@ def clock():
                               np.array([[1.0, 0.0],
                                            [-1.0, 0.0]]),
                               decimal=15)
+
+
+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)