Adding testcases for visualization.image_writer

yt/visualization/image_writer.py

@@ -28,6 +28,7 @@
 import yt.utilities.lib.image_utilities as au
 import yt.utilities.png_writer as pw
 from yt.extern.six.moves import builtins
+from yt.extern.six import string_types
 
 
 def scale_image(image, mi=None, ma=None):
@@ -148,8 +149,10 @@ def write_bitmap(bitmap_array, filename, max_val = None, transpose=False):
         that the first element resides in the upper-left corner.  If True, the
         first element will be placed in the lower-left corner.
     """
-    if len(bitmap_array.shape) != 3 or bitmap_array.shape[-1] not in (3,4):
-        raise RuntimeError
+    if len(bitmap_array.shape) != 3 or bitmap_array.shape[-1] not in (3, 4):
+        raise RuntimeError("Expecting image array of shape (N,M,3) or "
+                           "(N,M,4), received %s" % str(bitmap_array.shape))
+
     if bitmap_array.dtype != np.uint8:
         s1, s2 = bitmap_array.shape[:2]
         if bitmap_array.shape[-1] == 3:
@@ -294,9 +297,11 @@ def strip_colormap_data(fn = "color_map_data.py",
     f.write("### Auto-generated colormap tables, taken from Matplotlib ###\n\n")
     f.write("from numpy import array\n")
     f.write("color_map_luts = {}\n\n\n")
-    if cmaps is None: cmaps = rcm.ColorMaps
+    if cmaps is None:
+        cmaps = rcm.ColorMaps
+    if isinstance(cmaps, string_types):
+        cmaps = [cmaps]
     for cmap_name in sorted(cmaps):
-        print("Stripping", cmap_name)
         vals = rcm._extract_lookup_table(cmap_name)
         f.write("### %s ###\n\n" % (cmap_name))
         f.write("color_map_luts['%s'] = \\\n" % (cmap_name))
@@ -413,14 +418,11 @@ def write_projection(data, filename, colorbar=True, colorbar_label=None,
         suffix = '.png'
         filename = "%s%s" % (filename, suffix)
     mylog.info("Saving plot %s", filename)
-    if suffix == ".png":
-        canvas = FigureCanvasAgg(fig)
-    elif suffix == ".pdf":
+    if suffix == ".pdf":
         canvas = FigureCanvasPdf(fig)
     elif suffix in (".eps", ".ps"):
         canvas = FigureCanvasPS(fig)
     else:
-        mylog.warning("Unknown suffix %s, defaulting to Agg", suffix)
         canvas = FigureCanvasAgg(fig)
 
     fig.tight_layout()

yt/visualization/tests/test_image_writer.py

@@ -0,0 +1,78 @@
+"""
+Tests for visualization.image_writer
+"""
+# -----------------------------------------------------------------------------
+# Copyright (c) 2018, yt Development Team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file COPYING.txt, distributed with this software.
+# -----------------------------------------------------------------------------
+import os
+import shutil
+import tempfile
+import unittest
+
+import numpy as np
+from nose.tools import assert_raises
+
+from yt.testing import fake_random_ds, assert_equal
+from yt.visualization.image_writer import multi_image_composite, splat_points, \
+    write_bitmap, apply_colormap, strip_colormap_data
+
+
+class TestImageWriter(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.tmpdir = tempfile.mkdtemp()
+        cls.curdir = os.getcwd()
+        os.chdir(cls.tmpdir)
+
+    @classmethod
+    def tearDownClass(cls):
+        os.chdir(cls.curdir)
+        shutil.rmtree(cls.tmpdir)
+
+    def test_multi_image_composite(self):
+        ds = fake_random_ds(64, nprocs=4, particles=16**3)
+        center = [0.5, 0.5, 0.5]
+        normal = [1, 1, 1]
+        cut = ds.cutting(normal, center)
+        frb = cut.to_frb((0.75, 'unitary'), 64)
+        multi_image_composite("multi_channel1.png", frb["x"], frb["y"])
+
+        # Test multi_image_composite with user specified scaling values
+        mi = ds.quan(0.1, 'code_length')
+        ma = ds.quan(0.9, 'code_length')
+        multi_image_composite("multi_channel2.png", (frb["x"], mi, ma),
+                              [frb["y"], mi, None], green_channel=frb["z"],
+                              alpha_channel=frb["density"])
+
+        # Test with numpy integer array
+        x = np.array(np.random.randint(0, 256, size=(10, 10)), dtype='uint8')
+        y = np.array(np.random.randint(0, 256, size=(10, 10)), dtype='uint8')
+        multi_image_composite("multi_channel3.png", x, y)
+
+    def test_write_bitmap(self):
+        image = np.zeros([16, 16, 4], dtype='uint8')
+        xs = np.random.rand(100)
+        ys = np.random.rand(100)
+        image = splat_points(image, xs, ys)
+        png_str = write_bitmap(image.copy(), None)
+
+        image_trans = image.swapaxes(0, 1).copy(order="C")
+        png_str_trans = write_bitmap(image_trans, None, transpose=True)
+        assert_equal(png_str, png_str_trans)
+
+        with assert_raises(RuntimeError) as ex:
+            write_bitmap(np.ones([16, 16]), None)
+        desired = ("Expecting image array of shape (N,M,3) "
+                   "or (N,M,4), received (16, 16)")
+        assert_equal(str(ex.exception)[:50], desired[:50])
+
+    def test_strip_colormap_data(self):
+        strip_colormap_data("_stripped_cmap.py", ("arbre"))
+
+    def test_apply_colormap(self):
+        x = np.array(np.random.randint(0, 256, size=(10, 10)), dtype='uint8')
+        apply_colormap(x, color_bounds=None, cmap_name=None, func=lambda x:x**2)

yt/visualization/tests/test_offaxisprojection.py

@@ -11,77 +11,80 @@
 # Distributed under the terms of the Modified BSD License.
 #
 # The full license is in the file COPYING.txt, distributed with this software.
-#-----------------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 import os
-import os.path
-import tempfile
 import shutil
+import tempfile
+import unittest
+
+from yt.visualization.image_writer import write_projection
 from yt.testing import \
-    fake_random_ds, assert_equal, expand_keywords
-from yt.mods import write_projection
+    fake_random_ds, expand_keywords, assert_fname
 from yt.visualization.volume_rendering.api import off_axis_projection
 
 
-def setup():
-    """Test specific setup."""
-    from yt.config import ytcfg
-    ytcfg["yt", "__withintesting"] = "True"
-
-
-def test_oap(tmpdir=True):
-    """Tests functionality of off_axis_projection and write_projection."""
-    # Perform I/O in safe place instead of yt main dir
-    if tmpdir:
-        tmpdir = tempfile.mkdtemp()
-        curdir = os.getcwd()
-        os.chdir(tmpdir)
-
-    # args for off_axis_projection
-    test_ds = fake_random_ds(64)
-    c = test_ds.domain_center 
-    norm = [0.5, 0.5, 0.5]
-    W = test_ds.arr([0.5,0.5,1.0], 'unitary')
-    N = 256
-    field = ("gas","density")
-    oap_args = [test_ds, c, norm, W, N, field]
-
-    # kwargs for off_axis_projection
-    oap_kwargs = {}
-    oap_kwargs['weight'] = (None, 'cell_mass')
-    oap_kwargs['no_ghost'] = (True, False)
-    oap_kwargs['interpolated'] = (False,)
-    oap_kwargs['north_vector'] = ((1,0,0), (0,0.5,1.0))
-    oap_kwargs_list = expand_keywords(oap_kwargs)
-
-    # args or write_projection
-    fn = "test_%d.png"
-
-    # kwargs for write_projection
-    wp_kwargs = {}
-    wp_kwargs['colorbar'] = (True, False)
-    wp_kwargs['colorbar_label'] = ('test')
-    wp_kwargs['title'] = ('test')
-    wp_kwargs['limits'] = (None, (1e3, 1e5))
-    wp_kwargs['take_log'] = (True, False)
-    wp_kwargs['figsize'] = ((8,6), [1,1])
-    wp_kwargs['dpi'] = (100, 50)
-    wp_kwargs['cmap_name'] = ('arbre', 'kelp')
-    wp_kwargs_list = expand_keywords(wp_kwargs)
-
-    # test all off_axis_projection kwargs and write_projection kwargs
-    # make sure they are able to be projected, then remove and try next
-    # iteration
-    for i, oap_kwargs in enumerate(oap_kwargs_list):
-        image = off_axis_projection(*oap_args, **oap_kwargs)
-        for wp_kwargs in wp_kwargs_list:
-            write_projection(image, fn % i, **wp_kwargs)
-            assert_equal(os.path.exists(fn % i), True)
-
-    if tmpdir:
-        os.chdir(curdir)
-        # clean up
-        shutil.rmtree(tmpdir)
-
-if __name__ == "__main__":
-    for test in test_oap(tmpdir=False):
-        pass
+class TestOffAxisProjection(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.tmpdir = tempfile.mkdtemp()
+        cls.curdir = os.getcwd()
+        os.chdir(cls.tmpdir)
+
+    @classmethod
+    def tearDownClass(cls):
+        os.chdir(cls.curdir)
+        shutil.rmtree(cls.tmpdir)
+
+    def test_oap(self):
+        """Tests functionality of off_axis_projection and write_projection."""
+
+        # args for off_axis_projection
+        test_ds = fake_random_ds(64)
+        c = test_ds.domain_center
+        norm = [0.5, 0.5, 0.5]
+        W = test_ds.arr([0.5,0.5,1.0], 'unitary')
+        N = 256
+        field = ("gas","density")
+        oap_args = [test_ds, c, norm, W, N, field]
+
+        # kwargs for off_axis_projection
+        oap_kwargs = {}
+        oap_kwargs['weight'] = (None, 'cell_mass')
+        oap_kwargs['no_ghost'] = (True, False)
+        oap_kwargs['interpolated'] = (False,)
+        oap_kwargs['north_vector'] = ((1, 0, 0), (0, 0.5, 1.0))
+        oap_kwargs_list = expand_keywords(oap_kwargs)
+
+        # args or write_projection
+        fn = "test_%d.png"
+
+        # kwargs for write_projection
+        wp_kwargs = {}
+        wp_kwargs['colorbar'] = (True, False)
+        wp_kwargs['colorbar_label'] = ('test')
+        wp_kwargs['title'] = ('test')
+        wp_kwargs['limits'] = (None, (1e3, 1e5))
+        wp_kwargs['take_log'] = (True, False)
+        wp_kwargs['figsize'] = ((8,6), [1,1])
+        wp_kwargs['dpi'] = (100, 50)
+        wp_kwargs['cmap_name'] = ('arbre', 'kelp')
+        wp_kwargs_list = expand_keywords(wp_kwargs)
+
+        # test all off_axis_projection kwargs and write_projection kwargs
+        # make sure they are able to be projected, then remove and try next
+        # iteration
+        for i, oap_kwargs in enumerate(oap_kwargs_list):
+            image = off_axis_projection(*oap_args, **oap_kwargs)
+            for wp_kwargs in wp_kwargs_list:
+                write_projection(image, fn % i, **wp_kwargs)
+                assert_fname(fn % i)
+
+        # Test remaining parameters of write_projection
+        write_projection(image, "test_2", xlabel="x-axis", ylabel="y-axis")
+        assert_fname("test_2.png")
+
+        write_projection(image, "test_3.pdf", xlabel="x-axis", ylabel="y-axis")
+        assert_fname("test_3.pdf")
+
+        write_projection(image, "test_4.eps", xlabel="x-axis", ylabel="y-axis")
+        assert_fname("test_4.eps")