1"""
2Testing that skewed axes properly work
3"""
4from __future__ import absolute_import, division, print_function
5
6import itertools
7
8import matplotlib.pyplot as plt
9from matplotlib.testing.decorators import image_comparison
10
11from matplotlib.axes import Axes
12import matplotlib.transforms as transforms
13import matplotlib.axis as maxis
14import matplotlib.spines as mspines
15import matplotlib.patches as mpatch
16from matplotlib.projections import register_projection
17
18
19# The sole purpose of this class is to look at the upper, lower, or total
20# interval as appropriate and see what parts of the tick to draw, if any.
21class SkewXTick(maxis.XTick):
22    def update_position(self, loc):
23        # This ensures that the new value of the location is set before
24        # any other updates take place
25        self._loc = loc
26        super(SkewXTick, self).update_position(loc)
27
28    def _has_default_loc(self):
29        return self.get_loc() is None
30
31    def _need_lower(self):
32        return (self._has_default_loc() or
33                transforms.interval_contains(self.axes.lower_xlim,
34                                             self.get_loc()))
35
36    def _need_upper(self):
37        return (self._has_default_loc() or
38                transforms.interval_contains(self.axes.upper_xlim,
39                                             self.get_loc()))
40
41    @property
42    def gridOn(self):
43        return (self._gridOn and (self._has_default_loc() or
44                transforms.interval_contains(self.get_view_interval(),
45                                             self.get_loc())))
46
47    @gridOn.setter
48    def gridOn(self, value):
49        self._gridOn = value
50
51    @property
52    def tick1On(self):
53        return self._tick1On and self._need_lower()
54
55    @tick1On.setter
56    def tick1On(self, value):
57        self._tick1On = value
58
59    @property
60    def label1On(self):
61        return self._label1On and self._need_lower()
62
63    @label1On.setter
64    def label1On(self, value):
65        self._label1On = value
66
67    @property
68    def tick2On(self):
69        return self._tick2On and self._need_upper()
70
71    @tick2On.setter
72    def tick2On(self, value):
73        self._tick2On = value
74
75    @property
76    def label2On(self):
77        return self._label2On and self._need_upper()
78
79    @label2On.setter
80    def label2On(self, value):
81        self._label2On = value
82
83    def get_view_interval(self):
84        return self.axes.xaxis.get_view_interval()
85
86
87# This class exists to provide two separate sets of intervals to the tick,
88# as well as create instances of the custom tick
89class SkewXAxis(maxis.XAxis):
90    def _get_tick(self, major):
91        return SkewXTick(self.axes, None, '', major=major)
92
93    def get_view_interval(self):
94        return self.axes.upper_xlim[0], self.axes.lower_xlim[1]
95
96
97# This class exists to calculate the separate data range of the
98# upper X-axis and draw the spine there. It also provides this range
99# to the X-axis artist for ticking and gridlines
100class SkewSpine(mspines.Spine):
101    def _adjust_location(self):
102        pts = self._path.vertices
103        if self.spine_type == 'top':
104            pts[:, 0] = self.axes.upper_xlim
105        else:
106            pts[:, 0] = self.axes.lower_xlim
107
108
109# This class handles registration of the skew-xaxes as a projection as well
110# as setting up the appropriate transformations. It also overrides standard
111# spines and axes instances as appropriate.
112class SkewXAxes(Axes):
113    # The projection must specify a name.  This will be used be the
114    # user to select the projection, i.e. ``subplot(111,
115    # projection='skewx')``.
116    name = 'skewx'
117
118    def _init_axis(self):
119        # Taken from Axes and modified to use our modified X-axis
120        self.xaxis = SkewXAxis(self)
121        self.spines['top'].register_axis(self.xaxis)
122        self.spines['bottom'].register_axis(self.xaxis)
123        self.yaxis = maxis.YAxis(self)
124        self.spines['left'].register_axis(self.yaxis)
125        self.spines['right'].register_axis(self.yaxis)
126
127    def _gen_axes_spines(self):
128        spines = {'top': SkewSpine.linear_spine(self, 'top'),
129                  'bottom': mspines.Spine.linear_spine(self, 'bottom'),
130                  'left': mspines.Spine.linear_spine(self, 'left'),
131                  'right': mspines.Spine.linear_spine(self, 'right')}
132        return spines
133
134    def _set_lim_and_transforms(self):
135        """
136        This is called once when the plot is created to set up all the
137        transforms for the data, text and grids.
138        """
139        rot = 30
140
141        # Get the standard transform setup from the Axes base class
142        Axes._set_lim_and_transforms(self)
143
144        # Need to put the skew in the middle, after the scale and limits,
145        # but before the transAxes. This way, the skew is done in Axes
146        # coordinates thus performing the transform around the proper origin
147        # We keep the pre-transAxes transform around for other users, like the
148        # spines for finding bounds
149        self.transDataToAxes = (self.transScale +
150                                (self.transLimits +
151                                 transforms.Affine2D().skew_deg(rot, 0)))
152
153        # Create the full transform from Data to Pixels
154        self.transData = self.transDataToAxes + self.transAxes
155
156        # Blended transforms like this need to have the skewing applied using
157        # both axes, in axes coords like before.
158        self._xaxis_transform = (transforms.blended_transform_factory(
159            self.transScale + self.transLimits,
160            transforms.IdentityTransform()) +
161            transforms.Affine2D().skew_deg(rot, 0)) + self.transAxes
162
163    @property
164    def lower_xlim(self):
165        return self.axes.viewLim.intervalx
166
167    @property
168    def upper_xlim(self):
169        pts = [[0., 1.], [1., 1.]]
170        return self.transDataToAxes.inverted().transform(pts)[:, 0]
171
172
173# Now register the projection with matplotlib so the user can select
174# it.
175register_projection(SkewXAxes)
176
177
178@image_comparison(baseline_images=['skew_axes'], remove_text=True)
179def test_set_line_coll_dash_image():
180    fig = plt.figure()
181    ax = fig.add_subplot(1, 1, 1, projection='skewx')
182    ax.set_xlim(-50, 50)
183    ax.set_ylim(50, -50)
184    ax.grid(True)
185
186    # An example of a slanted line at constant X
187    ax.axvline(0, color='b')
188
189
190@image_comparison(baseline_images=['skew_rects'], remove_text=True)
191def test_skew_rectangle():
192
193    fix, axes = plt.subplots(5, 5, sharex=True, sharey=True, figsize=(8, 8))
194    axes = axes.flat
195
196    rotations = list(itertools.product([-3, -1, 0, 1, 3], repeat=2))
197
198    axes[0].set_xlim([-3, 3])
199    axes[0].set_ylim([-3, 3])
200    axes[0].set_aspect('equal', share=True)
201
202    for ax, (xrots, yrots) in zip(axes, rotations):
203        xdeg, ydeg = 45 * xrots, 45 * yrots
204        t = transforms.Affine2D().skew_deg(xdeg, ydeg)
205
206        ax.set_title('Skew of {0} in X and {1} in Y'.format(xdeg, ydeg))
207        ax.add_patch(mpatch.Rectangle([-1, -1], 2, 2,
208                                      transform=t + ax.transData,
209                                      alpha=0.5, facecolor='coral'))
210
211    plt.subplots_adjust(wspace=0, left=0.01, right=0.99, bottom=0.01, top=0.99)
212