1import numpy as np
2import platform
3
4import matplotlib.pyplot as plt
5from matplotlib.path import Path
6from matplotlib.projections import PolarAxes
7from matplotlib.transforms import Affine2D, Transform
8from matplotlib.testing.decorators import image_comparison
9
10from mpl_toolkits.axes_grid1.parasite_axes import ParasiteAxes
11from mpl_toolkits.axisartist import SubplotHost
12from mpl_toolkits.axes_grid1.parasite_axes import host_subplot_class_factory
13from mpl_toolkits.axisartist import angle_helper
14from mpl_toolkits.axisartist.axislines import Axes
15from mpl_toolkits.axisartist.grid_helper_curvelinear import \
16    GridHelperCurveLinear
17
18
19@image_comparison(['custom_transform.png'], style='default',
20                  tol=0.03 if platform.machine() == 'x86_64' else 0.04)
21def test_custom_transform():
22    class MyTransform(Transform):
23        input_dims = output_dims = 2
24
25        def __init__(self, resolution):
26            """
27            Resolution is the number of steps to interpolate between each input
28            line segment to approximate its path in transformed space.
29            """
30            Transform.__init__(self)
31            self._resolution = resolution
32
33        def transform(self, ll):
34            x, y = ll.T
35            return np.column_stack([x, y - x])
36
37        transform_non_affine = transform
38
39        def transform_path(self, path):
40            ipath = path.interpolated(self._resolution)
41            return Path(self.transform(ipath.vertices), ipath.codes)
42
43        transform_path_non_affine = transform_path
44
45        def inverted(self):
46            return MyTransformInv(self._resolution)
47
48    class MyTransformInv(Transform):
49        input_dims = output_dims = 2
50
51        def __init__(self, resolution):
52            Transform.__init__(self)
53            self._resolution = resolution
54
55        def transform(self, ll):
56            x, y = ll.T
57            return np.column_stack([x, y + x])
58
59        def inverted(self):
60            return MyTransform(self._resolution)
61
62    fig = plt.figure()
63
64    SubplotHost = host_subplot_class_factory(Axes)
65
66    tr = MyTransform(1)
67    grid_helper = GridHelperCurveLinear(tr)
68    ax1 = SubplotHost(fig, 1, 1, 1, grid_helper=grid_helper)
69    fig.add_subplot(ax1)
70
71    ax2 = ParasiteAxes(ax1, tr, viewlim_mode="equal")
72    ax1.parasites.append(ax2)
73    ax2.plot([3, 6], [5.0, 10.])
74
75    ax1.set_aspect(1.)
76    ax1.set_xlim(0, 10)
77    ax1.set_ylim(0, 10)
78
79    ax1.grid(True)
80
81
82@image_comparison(['polar_box.png'], style='default',
83                  tol={'aarch64': 0.04}.get(platform.machine(), 0.03))
84def test_polar_box():
85    # Remove this line when this test image is regenerated.
86    plt.rcParams['text.kerning_factor'] = 6
87
88    fig = plt.figure(figsize=(5, 5))
89
90    # PolarAxes.PolarTransform takes radian. However, we want our coordinate
91    # system in degree
92    tr = Affine2D().scale(np.pi / 180., 1.) + PolarAxes.PolarTransform()
93
94    # polar projection, which involves cycle, and also has limits in
95    # its coordinates, needs a special method to find the extremes
96    # (min, max of the coordinate within the view).
97    extreme_finder = angle_helper.ExtremeFinderCycle(20, 20,
98                                                     lon_cycle=360,
99                                                     lat_cycle=None,
100                                                     lon_minmax=None,
101                                                     lat_minmax=(0, np.inf))
102
103    grid_locator1 = angle_helper.LocatorDMS(12)
104    tick_formatter1 = angle_helper.FormatterDMS()
105
106    grid_helper = GridHelperCurveLinear(tr,
107                                        extreme_finder=extreme_finder,
108                                        grid_locator1=grid_locator1,
109                                        tick_formatter1=tick_formatter1)
110
111    ax1 = SubplotHost(fig, 1, 1, 1, grid_helper=grid_helper)
112
113    ax1.axis["right"].major_ticklabels.set_visible(True)
114    ax1.axis["top"].major_ticklabels.set_visible(True)
115
116    # let right axis shows ticklabels for 1st coordinate (angle)
117    ax1.axis["right"].get_helper().nth_coord_ticks = 0
118    # let bottom axis shows ticklabels for 2nd coordinate (radius)
119    ax1.axis["bottom"].get_helper().nth_coord_ticks = 1
120
121    fig.add_subplot(ax1)
122
123    ax1.axis["lat"] = axis = grid_helper.new_floating_axis(0, 45, axes=ax1)
124    axis.label.set_text("Test")
125    axis.label.set_visible(True)
126    axis.get_helper().set_extremes(2, 12)
127
128    ax1.axis["lon"] = axis = grid_helper.new_floating_axis(1, 6, axes=ax1)
129    axis.label.set_text("Test 2")
130    axis.get_helper().set_extremes(-180, 90)
131
132    # A parasite axes with given transform
133    ax2 = ParasiteAxes(ax1, tr, viewlim_mode="equal")
134    assert ax2.transData == tr + ax1.transData
135    # Anything you draw in ax2 will match the ticks and grids of ax1.
136    ax1.parasites.append(ax2)
137    ax2.plot(np.linspace(0, 30, 50), np.linspace(10, 10, 50))
138
139    ax1.set_aspect(1.)
140    ax1.set_xlim(-5, 12)
141    ax1.set_ylim(-5, 10)
142
143    ax1.grid(True)
144
145
146@image_comparison(['axis_direction.png'], style='default', tol=0.03)
147def test_axis_direction():
148    # Remove this line when this test image is regenerated.
149    plt.rcParams['text.kerning_factor'] = 6
150
151    fig = plt.figure(figsize=(5, 5))
152
153    # PolarAxes.PolarTransform takes radian. However, we want our coordinate
154    # system in degree
155    tr = Affine2D().scale(np.pi / 180., 1.) + PolarAxes.PolarTransform()
156
157    # polar projection, which involves cycle, and also has limits in
158    # its coordinates, needs a special method to find the extremes
159    # (min, max of the coordinate within the view).
160
161    # 20, 20 : number of sampling points along x, y direction
162    extreme_finder = angle_helper.ExtremeFinderCycle(20, 20,
163                                                     lon_cycle=360,
164                                                     lat_cycle=None,
165                                                     lon_minmax=None,
166                                                     lat_minmax=(0, np.inf),
167                                                     )
168
169    grid_locator1 = angle_helper.LocatorDMS(12)
170    tick_formatter1 = angle_helper.FormatterDMS()
171
172    grid_helper = GridHelperCurveLinear(tr,
173                                        extreme_finder=extreme_finder,
174                                        grid_locator1=grid_locator1,
175                                        tick_formatter1=tick_formatter1)
176
177    ax1 = SubplotHost(fig, 1, 1, 1, grid_helper=grid_helper)
178
179    for axis in ax1.axis.values():
180        axis.set_visible(False)
181
182    fig.add_subplot(ax1)
183
184    ax1.axis["lat1"] = axis = grid_helper.new_floating_axis(
185        0, 130,
186        axes=ax1, axis_direction="left")
187    axis.label.set_text("Test")
188    axis.label.set_visible(True)
189    axis.get_helper().set_extremes(0.001, 10)
190
191    ax1.axis["lat2"] = axis = grid_helper.new_floating_axis(
192        0, 50,
193        axes=ax1, axis_direction="right")
194    axis.label.set_text("Test")
195    axis.label.set_visible(True)
196    axis.get_helper().set_extremes(0.001, 10)
197
198    ax1.axis["lon"] = axis = grid_helper.new_floating_axis(
199        1, 10,
200        axes=ax1, axis_direction="bottom")
201    axis.label.set_text("Test 2")
202    axis.get_helper().set_extremes(50, 130)
203    axis.major_ticklabels.set_axis_direction("top")
204    axis.label.set_axis_direction("top")
205
206    grid_helper.grid_finder.grid_locator1.set_params(nbins=5)
207    grid_helper.grid_finder.grid_locator2.set_params(nbins=5)
208
209    ax1.set_aspect(1.)
210    ax1.set_xlim(-8, 8)
211    ax1.set_ylim(-4, 12)
212
213    ax1.grid(True)
214