From 7df3707c3398a6015e1cb44488b6cfda0d114549 Mon Sep 17 00:00:00 2001
From: Stefan van der Walt <stefan@sun.ac.za>
Date: Fri, 27 Apr 2012 16:02:53 -0700
Subject: [PATCH 1/6] ENH: Implement fast coordinate transformations.

---
 doc/examples/plot_swirl.py              | 83 ++++++++++++++++++++++
 skimage/transform/__init__.py           |  2 +
 skimage/transform/_swirl.py             | 70 +++++++++++++++++++
 skimage/transform/_warp.py              | 91 +++++++++++++++++++++++++
 skimage/transform/project.py            | 10 +--
 skimage/transform/tests/test_project.py |  2 +-
 skimage/transform/tests/test_swirl.py   | 18 +++++
 7 files changed, 268 insertions(+), 8 deletions(-)
 create mode 100644 doc/examples/plot_swirl.py
 create mode 100644 skimage/transform/_swirl.py
 create mode 100644 skimage/transform/_warp.py
 create mode 100644 skimage/transform/tests/test_swirl.py

diff --git a/doc/examples/plot_swirl.py b/doc/examples/plot_swirl.py
new file mode 100644
index 00000000..c49dbabf
--- /dev/null
+++ b/doc/examples/plot_swirl.py
@@ -0,0 +1,83 @@
+r"""
+=====
+Swirl
+=====
+
+Image swirling is a non-linear image deformation that creates a whirlpool
+effect.  This example describes the implementation of this transform in
+``skimage``, as well as the underlying warp mechanism.
+
+Image warping
+`````````````
+When applying a geometric transformation on an image, we typically make use of
+a reverse mapping, i.e., for each pixel in the output image, we compute its
+corresponding position in the input.  The reason is that, if we were to do it
+the other way around (map each input pixel to its new output position), some
+pixels in the output may be left empty.  On the other hand, each output
+coordinate has exactly one corresponding location in (or outside) the input
+image, and even if that position is non-integer, we may use interpolation to
+compute the corresponding image value.
+
+Performing a reverse mapping
+````````````````````````````
+To perform a geometric warp in ``skimage``, you simply need to provide the
+reverse mapping to the ``skimage.transform.warp`` function.  E.g., consider the
+case where we would like to shift an image 50 pixels to the left.  The reverse
+mapping for such a shift would be::
+
+    def shift_left(xy):
+        xy[:, 0] += 50
+        return xy
+
+The corresponding call to warp is::
+
+    from skimage.transform import warp
+    warp(image, shift_left)
+
+The swirl transformation
+````````````````````````
+
+Consider the coordinate :math:`(x, y)` in the output image.  The reverse
+mapping for the swirl transformation first computes, relative to a center
+:math:`(x_0, y_0)`, its polar coordinates,
+
+.. math::
+
+    \theta = \arctan(y/x)
+
+    \rho = \sqrt{(x - x_0)^2 + (y - y_0)^2},
+
+and then transforms them according to
+
+.. math::
+
+    r = \ln(2) \, \mathtt{radius} / 5
+
+    \phi = \mathtt{rotation}
+
+    s = \mathtt{strength}
+
+    \theta' = \phi + s \, e^{-\rho / r + \theta}
+
+where ``strength`` is a parameter for the amount of swirl, ``radius`` indicates
+the extent of the transform in pixels, and ``rotation`` adds a rotation angle.
+The transformation of ``radius`` into :math:`r` is to ensure that the
+transformation decays to :math:`\approx 1/1000^{\mathsf{th}}` within the specified radius.
+"""
+
+from skimage import data
+from skimage.transform import swirl
+
+import matplotlib.pyplot as plt
+
+image = data.checkerboard()
+swirled = swirl(image, rotation=0, strength=10, radius=120, order=2)
+
+f, (ax0, ax1) = plt.subplots(1, 2, figsize=(8, 3))
+
+ax0.imshow(image, cmap=plt.cm.gray, interpolation='none')
+ax0.axis('off')
+ax1.imshow(swirled, cmap=plt.cm.gray, interpolation='none')
+ax1.axis('off')
+
+plt.show()
diff --git a/skimage/transform/__init__.py b/skimage/transform/__init__.py
index 42945fbe..e9450f83 100644
--- a/skimage/transform/__init__.py
+++ b/skimage/transform/__init__.py
@@ -4,3 +4,5 @@ from .finite_radon_transform import *
 from .project import *
 from ._project import homography as fast_homography
 from .integral import *
+from ._warp import warp
+from ._swirl import swirl
diff --git a/skimage/transform/_swirl.py b/skimage/transform/_swirl.py
new file mode 100644
index 00000000..0e144ed1
--- /dev/null
+++ b/skimage/transform/_swirl.py
@@ -0,0 +1,70 @@
+from __future__ import division
+import numpy as np
+
+from ._warp import warp
+
+
+def _swirl_mapping(xy, center, rotation, strength, radius):
+    x, y = xy.T
+    x0, y0 = center
+    radius = radius / 5 * np.log(2)
+
+    rho = np.sqrt((x - x0)**2 + (y - y0)**2)
+    theta = rotation + strength * \
+            np.exp(-rho / radius) + \
+            np.arctan2(y - y0, x - x0)
+
+    xy[..., 0] = x0 + rho * np.cos(theta)
+    xy[..., 1] = y0 + rho * np.sin(theta)
+
+    return xy
+
+def swirl(image, center=None, strength=1, radius=100, rotation=0,
+          output_shape=None, order=1, mode='constant', cval=0):
+    """Perform a swirl transformation.
+
+    Parameters
+    ----------
+    image : ndarray
+        Input image.
+    center : (x,y) tuple or (2,) ndarray
+        Center coordinate of transformation.
+    strength : float
+        The amount of swirling applied.
+    radius : float
+        The extent of the swirling in pixels.  The effect dies out
+        rapidly beyond radius.
+    rotation : float
+        Additional rotation applied to the image.
+
+    Returns
+    -------
+    swirled : ndarray
+        Swirled version of the input.
+
+    Other parameters
+    ----------------
+    output_shape : tuple or ndarray
+        Size of the generated output image.
+    order : int
+        Order of splines used in interpolation, passed as-is to ndimage.
+    mode : string
+        How to handle values outside the image borders, passed as-is
+        to ndimage.
+    cval : string
+        Used in conjunction with mode 'constant', the value outside
+        the image boundaries.
+
+    """
+
+    if center is None:
+        center = np.array(image.shape)[:2] / 2
+
+    warp_args = {'center': center,
+                 'rotation': rotation,
+                 'strength': strength,
+                 'radius': radius}
+
+    return warp(image, _swirl_mapping, tf_args=warp_args,
+                output_shape=output_shape,
+                order=order, mode=mode, cval=cval)
diff --git a/skimage/transform/_warp.py b/skimage/transform/_warp.py
new file mode 100644
index 00000000..00541fef
--- /dev/null
+++ b/skimage/transform/_warp.py
@@ -0,0 +1,91 @@
+__all__ = ['warp']
+
+import numpy as np
+from scipy import ndimage
+from skimage.util import img_as_float
+
+eps = np.finfo(float).eps
+
+def _stackcopy(a, b):
+    """a[:,:,0] = a[:,:,1] = ... = b"""
+    if a.ndim == 3:
+        a.transpose().swapaxes(1, 2)[:] = b
+    else:
+        a[:] = b
+
+def warp(image, coord_tf, tf_args={},
+         output_shape=None, order=1, mode='constant', cval=0.):
+    """Warp an image according to a given coordinate transformation.
+
+    Parameters
+    ----------
+    image : 2-D array
+        Input image.
+    coord_tf : callable xy = f(xy, **kwargs)
+        Function that transforms an Nx2 array of ``(x, y)`` coordinates
+        in the *output image* into their corresponding coordinates in the
+        *source image*.  Note that this is a reverse mapping (also
+        see examples below).
+    tf_args : dict, optional
+        Keyword arguments passed to `coord_tf`.
+    output_shape : tuple (rows, cols)
+        Shape of the output image generated.
+    order : int
+        Order of splines used in interpolation.
+    mode : string
+        How to handle values outside the image borders.  Passed as-is
+        to ndimage.
+    cval : string
+        Used in conjunction with mode 'constant', the value outside
+        the image boundaries.
+
+    Examples
+    --------
+    Shift an image to the right:
+
+    >>> from skimage import data
+    >>> image = data.camera()
+    >>>
+    >>> def shift_right(xy):
+    ...     xy[:, 0] -= 10
+    ...     return xy
+    >>>
+    >>> warp(image, shift_right)
+
+    """
+    if image.ndim < 2:
+        raise ValueError("Input must have more than 1 dimension.")
+
+    image = np.atleast_3d(img_as_float(image))
+    ishape = np.array(image.shape)
+    bands = ishape[2]
+
+    if output_shape is None:
+        output_shape = ishape
+
+    coords = np.empty(np.r_[3, output_shape], dtype=float)
+
+    # Construct transformed coordinates
+    rows, cols = output_shape[:2]
+    tf_coords = np.indices((cols, rows), dtype=float).reshape(2, -1).T
+
+    tf_coords = coord_tf(tf_coords, **tf_args)
+    tf_coords = tf_coords.T.reshape((-1, cols, rows)).swapaxes(1, 2)
+
+    # y-coordinate mapping
+    _stackcopy(coords[1, ...], tf_coords[0, ...])
+
+    # x-coordinate mapping
+    _stackcopy(coords[0, ...], tf_coords[1, ...])
+
+    # colour-coordinate mapping
+    coords[2, ...] = range(bands)
+
+    # Prefilter not necessary for order 1 interpolation
+    prefilter = order > 1
+    mapped = ndimage.map_coordinates(image, coords, prefilter=prefilter,
+                                     mode=mode, order=order, cval=cval)
+
+    # The spline filters sometimes return results outside [0, 1],
+    # so clip to ensure valid data
+    return np.clip(mapped.squeeze(), 0, 1)
diff --git a/skimage/transform/project.py b/skimage/transform/project.py
index 4c8c436e..b786c4a7 100644
--- a/skimage/transform/project.py
+++ b/skimage/transform/project.py
@@ -4,18 +4,12 @@
 
 import numpy as np
 from scipy.ndimage import interpolation as ndii
+from .warp import _stackcopy
 
 __all__ = ['homography']
 
 eps = np.finfo(float).eps
 
-def _stackcopy(a, b):
-    """a[:,:,0] = a[:,:,1] = ... = b"""
-    if a.ndim == 3:
-        a.transpose().swapaxes(1, 2)[:] = b
-    else:
-        a[:] = b
-
 def homography(image, H, output_shape=None, order=1,
                mode='constant', cval=0.):
     """Perform a projective transformation (homography) on an image.
@@ -106,6 +100,8 @@ def homography(image, H, output_shape=None, order=1,
 
     coords = np.empty(np.r_[3, output_shape], dtype=float)
 
+    # TODO: Refactor this method to use transform.warp instead.
+
     # Construct transformed coordinates
     rows, cols = output_shape[:2]
     rows, cols = np.mgrid[:rows, :cols]
diff --git a/skimage/transform/tests/test_project.py b/skimage/transform/tests/test_project.py
index 2482aae0..3446c9a5 100644
--- a/skimage/transform/tests/test_project.py
+++ b/skimage/transform/tests/test_project.py
@@ -1,7 +1,7 @@
 import numpy as np
 from numpy.testing import assert_array_almost_equal
 
-from skimage.transform.project import _stackcopy
+from skimage.transform._warp import _stackcopy
 from skimage.transform import homography, fast_homography
 from skimage import data
 from skimage.color import rgb2gray
diff --git a/skimage/transform/tests/test_swirl.py b/skimage/transform/tests/test_swirl.py
new file mode 100644
index 00000000..e3fcc02e
--- /dev/null
+++ b/skimage/transform/tests/test_swirl.py
@@ -0,0 +1,18 @@
+import numpy as np
+from numpy.testing import assert_array_almost_equal
+
+from skimage import transform as tf, data, img_as_float
+
+
+def test_roundtrip():
+    image = img_as_float(data.checkerboard())
+    swirl_params = {'radius': 80, 'rotation': 0, 'order': 2, 'mode': 'reflect'}
+    unswirled = tf.swirl(
+        tf.swirl(image, strength=10, **swirl_params),
+        strength=-10, **swirl_params
+        )
+
+    assert np.mean(np.abs(image - unswirled)) < 0.01
+
+if __name__ == "__main__":
+    np.testing.run_module_suite()

From e13cc2541fca5542f86bf6e1b5ee87f4359161c0 Mon Sep 17 00:00:00 2001
From: Stefan van der Walt <stefan@sun.ac.za>
Date: Wed, 2 May 2012 21:34:39 -0700
Subject: [PATCH 2/6] BUG: Fix incorrect import.

---
 skimage/transform/project.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/skimage/transform/project.py b/skimage/transform/project.py
index b786c4a7..030e86a9 100644
--- a/skimage/transform/project.py
+++ b/skimage/transform/project.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 from scipy.ndimage import interpolation as ndii
-from .warp import _stackcopy
+from ._warp import _stackcopy
 
 __all__ = ['homography']
 

From a5d8593408ce18ea4138b694a8eda1efdd6b9ee5 Mon Sep 17 00:00:00 2001
From: Stefan van der Walt <stefan@sun.ac.za>
Date: Fri, 4 May 2012 11:50:13 -0700
Subject: [PATCH 3/6] STY: Cleanups after Tony's review.

---
 doc/examples/plot_swirl.py            |  8 ++---
 skimage/transform/_swirl.py           | 18 ++++++----
 skimage/transform/_warp.py            | 47 ++++++++++++++++++---------
 skimage/transform/tests/test_swirl.py |  7 ++--
 4 files changed, 49 insertions(+), 31 deletions(-)

diff --git a/doc/examples/plot_swirl.py b/doc/examples/plot_swirl.py
index c49dbabf..18947dcb 100644
--- a/doc/examples/plot_swirl.py
+++ b/doc/examples/plot_swirl.py
@@ -36,7 +36,6 @@ The corresponding call to warp is::
 
 The swirl transformation
 ````````````````````````
-
 Consider the coordinate :math:`(x, y)` in the output image.  The reverse
 mapping for the swirl transformation first computes, relative to a center
 :math:`(x_0, y_0)`, its polar coordinates,
@@ -60,9 +59,10 @@ and then transforms them according to
     \theta' = \phi + s \, e^{-\rho / r + \theta}
 
 where ``strength`` is a parameter for the amount of swirl, ``radius`` indicates
-the extent of the transform in pixels, and ``rotation`` adds a rotation angle.
-The transformation of ``radius`` into :math:`r` is to ensure that the
-transformation decays to :math:`\approx 1/1000^{\mathsf{th}}` within the specified radius.
+the swirl extent in pixels, and ``rotation`` adds a rotation angle.  The
+transformation of ``radius`` into :math:`r` is to ensure that the
+transformation decays to :math:`\approx 1/1000^{\mathsf{th}}` within the
+specified radius.
 """
 
 from skimage import data
diff --git a/skimage/transform/_swirl.py b/skimage/transform/_swirl.py
index 0e144ed1..6fd92c44 100644
--- a/skimage/transform/_swirl.py
+++ b/skimage/transform/_swirl.py
@@ -7,9 +7,12 @@ from ._warp import warp
 def _swirl_mapping(xy, center, rotation, strength, radius):
     x, y = xy.T
     x0, y0 = center
+    rho = np.sqrt((x - x0)**2 + (y - y0)**2)
+
+    # Ensure that the transformation decays to approximately 1/1000-th
+    # within the specified radius.
     radius = radius / 5 * np.log(2)
 
-    rho = np.sqrt((x - x0)**2 + (y - y0)**2)
     theta = rotation + strength * \
             np.exp(-rho / radius) + \
             np.arctan2(y - y0, x - x0)
@@ -32,8 +35,8 @@ def swirl(image, center=None, strength=1, radius=100, rotation=0,
     strength : float
         The amount of swirling applied.
     radius : float
-        The extent of the swirling in pixels.  The effect dies out
-        rapidly beyond radius.
+        The extent of the swirl in pixels.  The effect dies out
+        rapidly beyond `radius`.
     rotation : float
         Additional rotation applied to the image.
 
@@ -47,10 +50,11 @@ def swirl(image, center=None, strength=1, radius=100, rotation=0,
     output_shape : tuple or ndarray
         Size of the generated output image.
     order : int
-        Order of splines used in interpolation, passed as-is to ndimage.
+        Order of splines used in interpolation.  See
+        `scipy.ndimage.map_coordinates` for detail.
     mode : string
-        How to handle values outside the image borders, passed as-is
-        to ndimage.
+        How to handle values outside the image borders.  See
+        `scipy.ndimage.map_coordinates` for detail.
     cval : string
         Used in conjunction with mode 'constant', the value outside
         the image boundaries.
@@ -65,6 +69,6 @@ def swirl(image, center=None, strength=1, radius=100, rotation=0,
                  'strength': strength,
                  'radius': radius}
 
-    return warp(image, _swirl_mapping, tf_args=warp_args,
+    return warp(image, _swirl_mapping, map_args=warp_args,
                 output_shape=output_shape,
                 order=order, mode=mode, cval=cval)
diff --git a/skimage/transform/_warp.py b/skimage/transform/_warp.py
index 00541fef..76f4fb5d 100644
--- a/skimage/transform/_warp.py
+++ b/skimage/transform/_warp.py
@@ -7,13 +7,21 @@ from skimage.util import img_as_float
 eps = np.finfo(float).eps
 
 def _stackcopy(a, b):
-    """a[:,:,0] = a[:,:,1] = ... = b"""
+    """Copy b into each color layer of a, such that::
+
+      a[:,:,0] = a[:,:,1] = ... = b
+
+    Notes
+    -----
+    Color images are stored as an ``MxNx3`` or ``MxNx4`` arrays.
+
+    """
     if a.ndim == 3:
         a.transpose().swapaxes(1, 2)[:] = b
     else:
         a[:] = b
 
-def warp(image, coord_tf, tf_args={},
+def warp(image, reverse_map, map_args={},
          output_shape=None, order=1, mode='constant', cval=0.):
     """Warp an image according to a given coordinate transformation.
 
@@ -21,20 +29,20 @@ def warp(image, coord_tf, tf_args={},
     ----------
     image : 2-D array
         Input image.
-    coord_tf : callable xy = f(xy, **kwargs)
-        Function that transforms an Nx2 array of ``(x, y)`` coordinates
-        in the *output image* into their corresponding coordinates in the
-        *source image*.  Note that this is a reverse mapping (also
-        see examples below).
-    tf_args : dict, optional
-        Keyword arguments passed to `coord_tf`.
+    reverse_map : callable xy = f(xy, **kwargs)
+        Reverse coordinate map.  A function that transforms an Nx2 array of
+        ``(x, y)`` coordinates in the *output image* into their corresponding
+        coordinates in the *source image*.  Also see examples below.
+    map_args : dict, optional
+        Keyword arguments passed to `reverse_map`.
     output_shape : tuple (rows, cols)
         Shape of the output image generated.
     order : int
-        Order of splines used in interpolation.
+        Order of splines used in interpolation. See
+        `scipy.ndimage.map_coordinates` for detail.
     mode : string
-        How to handle values outside the image borders.  Passed as-is
-        to ndimage.
+        How to handle values outside the image borders.  See
+        `scipy.ndimage.map_coordinates` for detail.
     cval : string
         Used in conjunction with mode 'constant', the value outside
         the image boundaries.
@@ -65,17 +73,24 @@ def warp(image, coord_tf, tf_args={},
 
     coords = np.empty(np.r_[3, output_shape], dtype=float)
 
-    # Construct transformed coordinates
+    ## Construct transformed coordinates
+
     rows, cols = output_shape[:2]
+
+    # Reshape grid coordinates into a (P, 2) array of (x, y) pairs
     tf_coords = np.indices((cols, rows), dtype=float).reshape(2, -1).T
 
-    tf_coords = coord_tf(tf_coords, **tf_args)
+    # Map each (x, y) pair to the source image according to
+    # the user-provided mapping
+    tf_coords = reverse_map(tf_coords, **map_args)
+
+    # Reshape back to a (2, M, N) coordinate grid
     tf_coords = tf_coords.T.reshape((-1, cols, rows)).swapaxes(1, 2)
 
-    # y-coordinate mapping
+    # Place the y-coordinate mapping
     _stackcopy(coords[1, ...], tf_coords[0, ...])
 
-    # x-coordinate mapping
+    # Place the x-coordinate mapping
     _stackcopy(coords[0, ...], tf_coords[1, ...])
 
     # colour-coordinate mapping
diff --git a/skimage/transform/tests/test_swirl.py b/skimage/transform/tests/test_swirl.py
index e3fcc02e..d71f8231 100644
--- a/skimage/transform/tests/test_swirl.py
+++ b/skimage/transform/tests/test_swirl.py
@@ -6,11 +6,10 @@ from skimage import transform as tf, data, img_as_float
 
 def test_roundtrip():
     image = img_as_float(data.checkerboard())
+
     swirl_params = {'radius': 80, 'rotation': 0, 'order': 2, 'mode': 'reflect'}
-    unswirled = tf.swirl(
-        tf.swirl(image, strength=10, **swirl_params),
-        strength=-10, **swirl_params
-        )
+    swirled = tf.swirl(image, strength=10, **swirl_params)
+    unswirled = tf.swirl(swirled, strength=-10, **swirl_params)
 
     assert np.mean(np.abs(image - unswirled)) < 0.01
 

From f1ac4f9a09981be99a089b11ecd705098b201a5a Mon Sep 17 00:00:00 2001
From: Stefan van der Walt <stefan@sun.ac.za>
Date: Fri, 4 May 2012 12:19:15 -0700
Subject: [PATCH 4/6] STY: Move swirl to _warp_zoo so we can add more
 transforms later.

---
 skimage/transform/__init__.py                 | 2 +-
 skimage/transform/{_swirl.py => _warp_zoo.py} | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename skimage/transform/{_swirl.py => _warp_zoo.py} (100%)

diff --git a/skimage/transform/__init__.py b/skimage/transform/__init__.py
index e9450f83..fa4059ee 100644
--- a/skimage/transform/__init__.py
+++ b/skimage/transform/__init__.py
@@ -5,4 +5,4 @@ from .project import *
 from ._project import homography as fast_homography
 from .integral import *
 from ._warp import warp
-from ._swirl import swirl
+from ._warp_zoo import swirl
diff --git a/skimage/transform/_swirl.py b/skimage/transform/_warp_zoo.py
similarity index 100%
rename from skimage/transform/_swirl.py
rename to skimage/transform/_warp_zoo.py

From 54dc1bb59c3eac7234272f8cc8797a39f122defc Mon Sep 17 00:00:00 2001
From: Stefan van der Walt <stefan@sun.ac.za>
Date: Tue, 8 May 2012 12:26:26 -0700
Subject: [PATCH 5/6] STY: Minor cleanups.

---
 skimage/transform/_warp.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/skimage/transform/_warp.py b/skimage/transform/_warp.py
index 76f4fb5d..15c3d2ee 100644
--- a/skimage/transform/_warp.py
+++ b/skimage/transform/_warp.py
@@ -4,7 +4,6 @@ import numpy as np
 from scipy import ndimage
 from skimage.util import img_as_float
 
-eps = np.finfo(float).eps
 
 def _stackcopy(a, b):
     """Copy b into each color layer of a, such that::
@@ -30,7 +29,7 @@ def warp(image, reverse_map, map_args={},
     image : 2-D array
         Input image.
     reverse_map : callable xy = f(xy, **kwargs)
-        Reverse coordinate map.  A function that transforms an Nx2 array of
+        Reverse coordinate map.  A function that transforms a Px2 array of
         ``(x, y)`` coordinates in the *output image* into their corresponding
         coordinates in the *source image*.  Also see examples below.
     map_args : dict, optional

From 4524f994e830cdf5bfd4989185471f80d562b869 Mon Sep 17 00:00:00 2001
From: Stefan van der Walt <stefan@sun.ac.za>
Date: Tue, 8 May 2012 16:21:53 -0700
Subject: [PATCH 6/6] ENH: Better formulation of _stackcopy.

---
 skimage/transform/_warp.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/skimage/transform/_warp.py b/skimage/transform/_warp.py
index 15c3d2ee..29a8531e 100644
--- a/skimage/transform/_warp.py
+++ b/skimage/transform/_warp.py
@@ -10,15 +10,19 @@ def _stackcopy(a, b):
 
       a[:,:,0] = a[:,:,1] = ... = b
 
+    Parameters
+    ----------
+    a : (M, N) or (M, N, P) ndarray
+        Target array.
+    b : (M, N)
+        Source array.
+
     Notes
     -----
     Color images are stored as an ``MxNx3`` or ``MxNx4`` arrays.
 
     """
-    if a.ndim == 3:
-        a.transpose().swapaxes(1, 2)[:] = b
-    else:
-        a[:] = b
+    a[:] = b[:, :, np.newaxis]
 
 def warp(image, reverse_map, map_args={},
          output_shape=None, order=1, mode='constant', cval=0.):