Merge pull request #265 from odlgroup/apply_on_boundary_oop

ENH: made apply_on_boundary use out parameter

odl/discr/lp_discr.py

@@ -249,11 +249,7 @@ def _inner(self, x, y):
             # TODO: implement without copying x
             func_list = _scaling_func_list(bdry_fracs)
 
-            x_arr = x.asarray()
-            if not x_arr.flags.owndata:
-                x_arr = x_arr.copy()
-
-            apply_on_boundary(x_arr, func=func_list, only_once=False)
+            x_arr = apply_on_boundary(x, func=func_list, only_once=False)
             return super()._inner(self.element(x_arr), y)
 
     def _norm(self, x):
@@ -264,13 +260,9 @@ def _norm(self, x):
             return super()._norm(x)
         else:
             # TODO: implement without copying x
-            func_list = _scaling_func_list(bdry_fracs,
-                                           exponent=self.exponent)
-            x_arr = x.asarray()
-            if not x_arr.flags.owndata:
-                x_arr = x_arr.copy()
+            func_list = _scaling_func_list(bdry_fracs, exponent=self.exponent)
 
-            apply_on_boundary(x_arr, func=func_list, only_once=False)
+            x_arr = apply_on_boundary(x, func=func_list, only_once=False)
             return super()._norm(self.element(x_arr))
 
     def _dist(self, x, y):
@@ -281,19 +273,12 @@ def _dist(self, x, y):
             return super()._dist(x, y)
         else:
             # TODO: implement without copying x
-            func_list = _scaling_func_list(bdry_fracs,
-                                           exponent=self.exponent)
-
-            x_arr, y_arr = x.asarray(), y.asarray()
-            if not x_arr.flags.owndata:
-                x_arr = x_arr.copy()
-            if not y_arr.flags.owndata:
-                y_arr = y_arr.copy()
+            func_list = _scaling_func_list(bdry_fracs, exponent=self.exponent)
 
-            for arr in (x_arr, y_arr):
-                apply_on_boundary(arr, func=func_list, only_once=False)
+            arrs = [apply_on_boundary(vec, func=func_list, only_once=False)
+                    for vec in (x, y)]
 
-            return super()._dist(self.element(x_arr), self.element(y_arr))
+            return super()._dist(self.element(arrs[0]), self.element(arrs[1]))
 
     def __repr__(self):
         """Return ``repr(self).``"""

odl/util/__init__.py

@@ -37,4 +37,8 @@
 from .graphics import *
 __all__ += graphics.__all__
 
+from .import numerics
+from .numerics import *
+__all__ += numerics.__all__
+
 from . import ufuncs

odl/util/numerics.py

@@ -24,20 +24,22 @@
 from future import standard_library
 standard_library.install_aliases()
 
-# External module imports
+import numpy as np
 
 
 __all__ = ('apply_on_boundary',)
 
 
 def apply_on_boundary(array, func, only_once=True, which_boundaries=None,
-                      axis_order=None):
+                      axis_order=None, out=None):
     """Apply a function of the boundary of an n-dimensional array.
 
+    All other values are preserved as is.
+
     Parameters
     ----------
-    array : `numpy.ndarray`
-        Modify the boundary of this array in place
+    array : array-like
+        Modify the boundary of this array
     func : `callable` or `sequence`
         If a single function is given, assign
         ``array[slice] = func(array[slice])`` on the boundary slices,
@@ -61,36 +63,50 @@ def apply_on_boundary(array, func, only_once=True, which_boundaries=None,
         the left, the second for the right boundary. The length of the
         sequence must be ``array.ndim``. `None` is interpreted as
         'all boundaries'.
-    axis_order : `sequence` of `int`
+    axis_order : `sequence` of `int`, optional
         Permutation of ``range(array.ndim)`` defining the order in which
         to process the axes. If combined with ``only_once`` and a
         function list, this determines which function is evaluated in
         the points that are potentially processed multiple times.
+    out : `numpy.ndarray`, optional
+        Location in which to store the result, can be the same as ``array``.
+        Default: copy of ``array``
 
     Examples
     --------
     >>> import numpy as np
     >>> arr = np.ones((3, 3))
     >>> apply_on_boundary(arr, lambda x: x / 2)
-    >>> arr
     array([[ 0.5,  0.5,  0.5],
            [ 0.5,  1. ,  0.5],
            [ 0.5,  0.5,  0.5]])
-    >>>
-    >>> arr = np.ones((3, 3))
+
+    If called with ``only_once=False``, applies function repeatedly
+
     >>> apply_on_boundary(arr, lambda x: x / 2, only_once=False)
-    >>> arr
     array([[ 0.25,  0.5 ,  0.25],
            [ 0.5 ,  1.  ,  0.5 ],
            [ 0.25,  0.5 ,  0.25]])
-    >>> arr = np.ones((3, 3))
+
     >>> apply_on_boundary(arr, lambda x: x / 2, only_once=True,
     ...                   which_boundaries=((True, False), True))
-    >>> arr
     array([[ 0.5,  0.5,  0.5],
            [ 0.5,  1. ,  0.5],
            [ 0.5,  1. ,  0.5]])
+
+    Also accepts out parameter
+
+    >>> out = np.empty_like(arr)
+    >>> result = apply_on_boundary(arr, lambda x: x / 2, out=out)
+    >>> result
+    array([[ 0.5,  0.5,  0.5],
+           [ 0.5,  1. ,  0.5],
+           [ 0.5,  0.5,  0.5]])
+    >>> result is out
+    True
     """
+    array = np.asarray(array)
+
     if callable(func):
         func = [func] * array.ndim
     elif len(func) != array.ndim:
@@ -109,6 +125,11 @@ def apply_on_boundary(array, func, only_once=True, which_boundaries=None,
         raise ValueError('axis_order has length {}, expected {}.'
                          ''.format(len(axis_order), array.ndim))
 
+    if out is None:
+        out = array.copy()
+    else:
+        out[:] = array  # Self assignment is free, in case out is array
+
     # The 'only_once' functionality is implemented by storing for each axis
     # if the left and right boundaries have been processed. This information
     # is stored in a list of slices which is reused for the next axis in the
@@ -141,13 +162,13 @@ def apply_on_boundary(array, func, only_once=True, which_boundaries=None,
             mod_left = mod_right = which
 
         if mod_left and func_l is not None:
-            array[slc_l] = func_l(array[slc_l])
+            out[slc_l] = func_l(out[slc_l])
             start = 1
         else:
             start = None
 
         if mod_right and func_r is not None:
-            array[slc_r] = func_r(array[slc_r])
+            out[slc_r] = func_r(out[slc_r])
             end = -1
         else:
             end = None
@@ -156,6 +177,9 @@ def apply_on_boundary(array, func, only_once=True, which_boundaries=None,
         # Start and end include the boundary if it was processed.
         slices[ax] = slice(start, end)
 
+    return out
+
+
 if __name__ == '__main__':
     from doctest import testmod, NORMALIZE_WHITESPACE
     testmod(optionflags=NORMALIZE_WHITESPACE)

test/util/numerics_test.py

@@ -34,73 +34,74 @@ def test_apply_on_boundary_default():
 
     # 1d
     arr = np.ones(5)
-    apply_on_boundary(arr, lambda x: x * 2)
-    assert all_equal(arr, [2, 1, 1, 1, 2])
+    result = apply_on_boundary(arr, lambda x: x * 2)
+    assert all_equal(arr, [1, 1, 1, 1, 1])
+    assert all_equal(result, [2, 1, 1, 1, 2])
 
     # 3d
     arr = np.ones((3, 4, 5))
-    apply_on_boundary(arr, lambda x: x * 2)
+    result = apply_on_boundary(arr, lambda x: x * 2)
     true_arr = 2 * np.ones((3, 4, 5))
     true_arr[1:-1, 1:-1, 1:-1] = 1
-    assert all_equal(arr, true_arr)
+    assert all_equal(result, true_arr)
 
 
 def test_apply_on_boundary_func_sequence_2d():
 
     arr = np.ones((3, 5))
-    apply_on_boundary(arr, [lambda x: x * 2, lambda x: x * 3])
-    assert all_equal(arr, [[2, 2, 2, 2, 2],
-                           [3, 1, 1, 1, 3],
-                           [2, 2, 2, 2, 2]])
+    result = apply_on_boundary(arr, [lambda x: x * 2, lambda x: x * 3])
+    assert all_equal(result, [[2, 2, 2, 2, 2],
+                              [3, 1, 1, 1, 3],
+                              [2, 2, 2, 2, 2]])
 
 
 def test_apply_on_boundary_multiple_times_2d():
 
     arr = np.ones((3, 5))
-    apply_on_boundary(arr, lambda x: x * 2, only_once=False)
-    assert all_equal(arr, [[4, 2, 2, 2, 4],
-                           [2, 1, 1, 1, 2],
-                           [4, 2, 2, 2, 4]])
+    result = apply_on_boundary(arr, lambda x: x * 2, only_once=False)
+    assert all_equal(result, [[4, 2, 2, 2, 4],
+                              [2, 1, 1, 1, 2],
+                              [4, 2, 2, 2, 4]])
 
 
 def test_apply_on_boundary_which_boundaries():
 
     # 1d
     arr = np.ones(5)
     which = ((True, False),)
-    apply_on_boundary(arr, lambda x: x * 2, which_boundaries=which)
-    assert all_equal(arr, [2, 1, 1, 1, 1])
+    result = apply_on_boundary(arr, lambda x: x * 2, which_boundaries=which)
+    assert all_equal(result, [2, 1, 1, 1, 1])
 
     # 2d
     arr = np.ones((3, 5))
     which = ((True, False), True)
-    apply_on_boundary(arr, lambda x: x * 2, which_boundaries=which)
-    assert all_equal(arr, [[2, 2, 2, 2, 2],
-                           [2, 1, 1, 1, 2],
-                           [2, 1, 1, 1, 2]])
+    result = apply_on_boundary(arr, lambda x: x * 2, which_boundaries=which)
+    assert all_equal(result, [[2, 2, 2, 2, 2],
+                              [2, 1, 1, 1, 2],
+                              [2, 1, 1, 1, 2]])
 
 
 def test_apply_on_boundary_which_boundaries_multiple_times_2d():
 
     # 2d
     arr = np.ones((3, 5))
     which = ((True, False), True)
-    apply_on_boundary(arr, lambda x: x * 2, which_boundaries=which,
-                      only_once=False)
-    assert all_equal(arr, [[4, 2, 2, 2, 4],
-                           [2, 1, 1, 1, 2],
-                           [2, 1, 1, 1, 2]])
+    result = apply_on_boundary(arr, lambda x: x * 2, which_boundaries=which,
+                               only_once=False)
+    assert all_equal(result, [[4, 2, 2, 2, 4],
+                              [2, 1, 1, 1, 2],
+                              [2, 1, 1, 1, 2]])
 
 
 def test_apply_on_boundary_axis_order_2d():
 
     arr = np.ones((3, 5))
     axis_order = (-1, -2)
-    apply_on_boundary(arr, [lambda x: x * 3, lambda x: x * 2],
-                      axis_order=axis_order)
-    assert all_equal(arr, [[3, 2, 2, 2, 3],
-                           [3, 1, 1, 1, 3],
-                           [3, 2, 2, 2, 3]])
+    result = apply_on_boundary(arr, [lambda x: x * 3, lambda x: x * 2],
+                               axis_order=axis_order)
+    assert all_equal(result, [[3, 2, 2, 2, 3],
+                              [3, 1, 1, 1, 3],
+                              [3, 2, 2, 2, 3]])
 
 if __name__ == '__main__':
     pytest.main(str(__file__.replace('\\', '/')) + ' -v')