From 99efa1b771bb5679737e93093dee9aaea42b6899 Mon Sep 17 00:00:00 2001
From: Nicolas Pinto <pinto@alum.mit.edu>
Date: Mon, 13 Feb 2012 21:20:31 -0500
Subject: [PATCH] ENH: add montage2d with tests

---
 skimage/util/montage.py            | 87 ++++++++++++++++++++++++++++++
 skimage/util/tests/test_montage.py | 59 ++++++++++++++++++++
 2 files changed, 146 insertions(+)
 create mode 100644 skimage/util/montage.py
 create mode 100644 skimage/util/tests/test_montage.py

diff --git a/skimage/util/montage.py b/skimage/util/montage.py
new file mode 100644
index 00000000..15646440
--- /dev/null
+++ b/skimage/util/montage.py
@@ -0,0 +1,87 @@
+__all__ = ['montage2d']
+
+import numpy as np
+
+
+def montage2d(arr_in, fill='mean'):
+    """Create a 2-dimensional 'montage' from a 3-dimensional input array
+    representing an ensemble of equally shaped 2-dimensional images.
+
+    For example, montage2d(arr_in, fill) with the following `arr_in`
+
+    +---+---+---+
+    | 1 | 2 | 3 |
+    +---+---+---+
+
+    will return:
+
+    +---+---+
+    | 1 | 2 |
+    +---+---+
+    | 3 | * |
+    +---+---+
+
+    Where the '*' patch will be determined by the `fill` parameter.
+
+    Parameters
+    ----------
+    arr_in: ndarray, shape=[n_images, height, width]
+        3-dimensional input array representing an ensemble of n_images
+        of equal shape (i.e. [height, width]).
+
+    fill: float or 'mean'
+        How to fill the 2-dimensional output array when sqrt(n_images)
+        is not an integer. If 'mean' is chosen, then fill = arr_in.mean().
+
+    Returns
+    -------
+    arr_out: ndarray, shape=[alpha * height, alpha * width]
+        Output array where 'alpha' has been determined automatically to
+        fit (at least) the `n_images` in `arr_in`.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> from skimage.util.montage import montage2d
+    >>> arr_in = np.arange(3 * 2 * 2).reshape(3, 2, 2)
+    >>> print arr_in
+    [[[ 0  1]
+      [ 2  3]]
+    <BLANKLINE>
+     [[ 4  5]
+      [ 6  7]]
+    <BLANKLINE>
+     [[ 8  9]
+      [10 11]]]
+    >>> arr_out = montage2d(arr_in)
+    >>> print arr_out.shape
+    (4, 4)
+    >>> print arr_out
+    [[  0.    1.    4.    5. ]
+     [  2.    3.    6.    7. ]
+     [  8.    9.    5.5   5.5]
+     [ 10.   11.    5.5   5.5]]
+    >>> print arr_in.mean()
+    5.5
+    """
+    assert arr_in.ndim == 3
+
+    n_images, height, width = arr_in.shape
+
+    # -- determine alpha
+    alpha = int(np.ceil(np.sqrt(n_images)))
+
+    # -- fill missing patches
+    if fill == 'mean':
+        fill = arr_in.mean()
+
+    n_missing = int((alpha ** 2.) - n_images)
+    missing = np.ones((n_missing, height, width), dtype=arr_in.dtype) * fill
+    arr_out = np.vstack((arr_in, missing))
+
+    # -- reshape to 2d montage, step by step
+    arr_out = arr_out.reshape(alpha, alpha, height, width)
+    arr_out = arr_out.swapaxes(1, 2)
+    arr_out = arr_out.reshape(alpha * height, alpha * width)
+
+    return arr_out
diff --git a/skimage/util/tests/test_montage.py b/skimage/util/tests/test_montage.py
new file mode 100644
index 00000000..cd7151ab
--- /dev/null
+++ b/skimage/util/tests/test_montage.py
@@ -0,0 +1,59 @@
+from nose.tools import assert_equal, raises
+from numpy.testing import assert_array_equal
+
+import numpy as np
+from skimage.util.montage import montage2d
+
+
+def test_simple():
+    n_images = 3
+    height, width = 2, 3,
+    arr_in = np.arange(n_images * height * width)
+    arr_in = arr_in.reshape(n_images, height, width)
+
+    arr_out = montage2d(arr_in)
+
+    gt = np.array(
+        [[  0. ,   1. ,   2. ,   6. ,   7. ,   8. ],
+         [  3. ,   4. ,   5. ,   9. ,  10. ,  11. ],
+         [ 12. ,  13. ,  14. ,   8.5,   8.5,   8.5],
+         [ 15. ,  16. ,  17. ,   8.5,   8.5,   8.5]]
+    )
+
+    assert_array_equal(arr_out, gt)
+
+
+def test_fill():
+    n_images = 3
+    height, width = 2, 3,
+    arr_in = np.arange(n_images * height * width)
+    arr_in = arr_in.reshape(n_images, height, width)
+
+    arr_out = montage2d(arr_in, fill=0)
+
+    gt = np.array(
+        [[  0. ,   1. ,   2. ,   6. ,   7. ,   8. ],
+         [  3. ,   4. ,   5. ,   9. ,  10. ,  11. ],
+         [ 12. ,  13. ,  14. ,   0. ,   0. ,   0. ],
+         [ 15. ,  16. ,  17. ,   0. ,   0. ,   0. ]]
+    )
+
+    assert_array_equal(arr_out, gt)
+
+
+def test_shape():
+    n_images = 15
+    height, width = 11, 7
+    arr_in = np.arange(n_images * height * width)
+    arr_in = arr_in.reshape(n_images, height, width)
+
+    alpha = int(np.ceil(np.sqrt(n_images)))
+
+    arr_out = montage2d(arr_in)
+    assert_equal(arr_out.shape, (alpha * height, alpha * width))
+
+
+@raises(AssertionError)
+def test_error_ndim():
+    arr_error = np.random.randn(1, 2, 3, 4)
+    montage2d(arr_error)