1from __future__ import absolute_import, division, print_function
2
3import io
4
5import numpy as np
6from numpy.testing import assert_array_almost_equal
7import pytest
8
9from matplotlib import (
10    collections, path, pyplot as plt, transforms as mtransforms, rcParams)
11from matplotlib.image import imread
12from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
13from matplotlib.figure import Figure
14from matplotlib.testing.decorators import image_comparison
15
16
17def test_repeated_save_with_alpha():
18    # We want an image which has a background color of bluish green, with an
19    # alpha of 0.25.
20
21    fig = Figure([1, 0.4])
22    canvas = FigureCanvas(fig)
23    fig.set_facecolor((0, 1, 0.4))
24    fig.patch.set_alpha(0.25)
25
26    # The target color is fig.patch.get_facecolor()
27
28    buf = io.BytesIO()
29
30    fig.savefig(buf,
31                facecolor=fig.get_facecolor(),
32                edgecolor='none')
33
34    # Save the figure again to check that the
35    # colors don't bleed from the previous renderer.
36    buf.seek(0)
37    fig.savefig(buf,
38                facecolor=fig.get_facecolor(),
39                edgecolor='none')
40
41    # Check the first pixel has the desired color & alpha
42    # (approx: 0, 1.0, 0.4, 0.25)
43    buf.seek(0)
44    assert_array_almost_equal(tuple(imread(buf)[0, 0]),
45                              (0.0, 1.0, 0.4, 0.250),
46                              decimal=3)
47
48
49def test_large_single_path_collection():
50    buff = io.BytesIO()
51
52    # Generates a too-large single path in a path collection that
53    # would cause a segfault if the draw_markers optimization is
54    # applied.
55    f, ax = plt.subplots()
56    collection = collections.PathCollection(
57        [path.Path([[-10, 5], [10, 5], [10, -5], [-10, -5], [-10, 5]])])
58    ax.add_artist(collection)
59    ax.set_xlim(10**-3, 1)
60    plt.savefig(buff)
61
62
63def test_marker_with_nan():
64    # This creates a marker with nans in it, which was segfaulting the
65    # Agg backend (see #3722)
66    fig, ax = plt.subplots(1)
67    steps = 1000
68    data = np.arange(steps)
69    ax.semilogx(data)
70    ax.fill_between(data, data*0.8, data*1.2)
71    buf = io.BytesIO()
72    fig.savefig(buf, format='png')
73
74
75def test_long_path():
76    buff = io.BytesIO()
77
78    fig, ax = plt.subplots()
79    np.random.seed(0)
80    points = np.random.rand(70000)
81    ax.plot(points)
82    fig.savefig(buff, format='png')
83
84
85@image_comparison(baseline_images=['agg_filter'],
86                  extensions=['png'], remove_text=True)
87def test_agg_filter():
88    def smooth1d(x, window_len):
89        s = np.r_[2*x[0] - x[window_len:1:-1],
90                  x,
91                  2*x[-1] - x[-1:-window_len:-1]]
92        w = np.hanning(window_len)
93        y = np.convolve(w/w.sum(), s, mode='same')
94        return y[window_len-1:-window_len+1]
95
96    def smooth2d(A, sigma=3):
97        window_len = max(int(sigma), 3)*2 + 1
98        A1 = np.array([smooth1d(x, window_len) for x in np.asarray(A)])
99        A2 = np.transpose(A1)
100        A3 = np.array([smooth1d(x, window_len) for x in A2])
101        A4 = np.transpose(A3)
102
103        return A4
104
105    class BaseFilter(object):
106        def prepare_image(self, src_image, dpi, pad):
107            ny, nx, depth = src_image.shape
108            padded_src = np.zeros([pad*2 + ny, pad*2 + nx, depth], dtype="d")
109            padded_src[pad:-pad, pad:-pad, :] = src_image[:, :, :]
110
111            return padded_src  # , tgt_image
112
113        def get_pad(self, dpi):
114            return 0
115
116        def __call__(self, im, dpi):
117            pad = self.get_pad(dpi)
118            padded_src = self.prepare_image(im, dpi, pad)
119            tgt_image = self.process_image(padded_src, dpi)
120            return tgt_image, -pad, -pad
121
122    class OffsetFilter(BaseFilter):
123        def __init__(self, offsets=None):
124            if offsets is None:
125                self.offsets = (0, 0)
126            else:
127                self.offsets = offsets
128
129        def get_pad(self, dpi):
130            return int(max(*self.offsets)/72.*dpi)
131
132        def process_image(self, padded_src, dpi):
133            ox, oy = self.offsets
134            a1 = np.roll(padded_src, int(ox/72.*dpi), axis=1)
135            a2 = np.roll(a1, -int(oy/72.*dpi), axis=0)
136            return a2
137
138    class GaussianFilter(BaseFilter):
139        "simple gauss filter"
140
141        def __init__(self, sigma, alpha=0.5, color=None):
142            self.sigma = sigma
143            self.alpha = alpha
144            if color is None:
145                self.color = (0, 0, 0)
146            else:
147                self.color = color
148
149        def get_pad(self, dpi):
150            return int(self.sigma*3/72.*dpi)
151
152        def process_image(self, padded_src, dpi):
153            tgt_image = np.zeros_like(padded_src)
154            aa = smooth2d(padded_src[:, :, -1]*self.alpha,
155                          self.sigma/72.*dpi)
156            tgt_image[:, :, -1] = aa
157            tgt_image[:, :, :-1] = self.color
158            return tgt_image
159
160    class DropShadowFilter(BaseFilter):
161        def __init__(self, sigma, alpha=0.3, color=None, offsets=None):
162            self.gauss_filter = GaussianFilter(sigma, alpha, color)
163            self.offset_filter = OffsetFilter(offsets)
164
165        def get_pad(self, dpi):
166            return max(self.gauss_filter.get_pad(dpi),
167                       self.offset_filter.get_pad(dpi))
168
169        def process_image(self, padded_src, dpi):
170            t1 = self.gauss_filter.process_image(padded_src, dpi)
171            t2 = self.offset_filter.process_image(t1, dpi)
172            return t2
173
174    fig = plt.figure()
175    ax = fig.add_subplot(111)
176
177    # draw lines
178    l1, = ax.plot([0.1, 0.5, 0.9], [0.1, 0.9, 0.5], "bo-",
179                  mec="b", mfc="w", lw=5, mew=3, ms=10, label="Line 1")
180    l2, = ax.plot([0.1, 0.5, 0.9], [0.5, 0.2, 0.7], "ro-",
181                  mec="r", mfc="w", lw=5, mew=3, ms=10, label="Line 1")
182
183    gauss = DropShadowFilter(4)
184
185    for l in [l1, l2]:
186
187        # draw shadows with same lines with slight offset.
188
189        xx = l.get_xdata()
190        yy = l.get_ydata()
191        shadow, = ax.plot(xx, yy)
192        shadow.update_from(l)
193
194        # offset transform
195        ot = mtransforms.offset_copy(l.get_transform(), ax.figure,
196                                     x=4.0, y=-6.0, units='points')
197
198        shadow.set_transform(ot)
199
200        # adjust zorder of the shadow lines so that it is drawn below the
201        # original lines
202        shadow.set_zorder(l.get_zorder() - 0.5)
203        shadow.set_agg_filter(gauss)
204        shadow.set_rasterized(True)  # to support mixed-mode renderers
205
206    ax.set_xlim(0., 1.)
207    ax.set_ylim(0., 1.)
208
209    ax.xaxis.set_visible(False)
210    ax.yaxis.set_visible(False)
211
212
213def test_too_large_image():
214    fig = plt.figure(figsize=(300, 1000))
215    buff = io.BytesIO()
216    with pytest.raises(ValueError):
217        fig.savefig(buff)
218
219
220def test_chunksize():
221    x = range(200)
222
223    # Test without chunksize
224    fig, ax = plt.subplots()
225    ax.plot(x, np.sin(x))
226    fig.canvas.draw()
227
228    # Test with chunksize
229    fig, ax = plt.subplots()
230    rcParams['agg.path.chunksize'] = 105
231    ax.plot(x, np.sin(x))
232    fig.canvas.draw()
233
234
235@pytest.mark.backend('Agg')
236def test_jpeg_dpi():
237    Image = pytest.importorskip("PIL.Image")
238    # Check that dpi is set correctly in jpg files.
239    plt.plot([0, 1, 2], [0, 1, 0])
240    buf = io.BytesIO()
241    plt.savefig(buf, format="jpg", dpi=200)
242    im = Image.open(buf)
243    assert im.info['dpi'] == (200, 200)
244