mirror of
https://github.com/wassname/geopandas.git
synced 2026-08-31 12:01:45 +08:00
124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
from __future__ import absolute_import
|
|
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
|
|
import matplotlib
|
|
matplotlib.use('Agg', warn=False)
|
|
from matplotlib.pyplot import Artist, savefig, clf
|
|
from matplotlib.testing.noseclasses import ImageComparisonFailure
|
|
from matplotlib.testing.compare import compare_images
|
|
from shapely.geometry import Polygon, LineString, Point, MultiPolygon, MultiLineString
|
|
from six.moves import xrange
|
|
|
|
from geopandas import GeoSeries, GeoDataFrame
|
|
|
|
|
|
# If set to True, generate images rather than perform tests (all tests will pass!)
|
|
GENERATE_BASELINE = False
|
|
|
|
BASELINE_DIR = os.path.join(os.path.dirname(__file__), 'baseline_images', 'test_plotting')
|
|
|
|
|
|
class PlotTests(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.tempdir = tempfile.mkdtemp()
|
|
return
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.tempdir)
|
|
return
|
|
|
|
def _compare_images(self, ax, filename, tol=10):
|
|
""" Helper method to do the comparisons """
|
|
assert isinstance(ax, Artist)
|
|
if GENERATE_BASELINE:
|
|
savefig(os.path.join(BASELINE_DIR, filename))
|
|
savefig(os.path.join(self.tempdir, filename))
|
|
err = compare_images(os.path.join(BASELINE_DIR, filename),
|
|
os.path.join(self.tempdir, filename),
|
|
tol, in_decorator=True)
|
|
if err:
|
|
raise ImageComparisonFailure('images not close: %(actual)s '
|
|
'vs. %(expected)s '
|
|
'(RMS %(rms).3f)' % err)
|
|
|
|
def test_poly_plot(self):
|
|
""" Test plotting a simple series of polygons """
|
|
clf()
|
|
filename = 'poly_plot.png'
|
|
t1 = Polygon([(0, 0), (1, 0), (1, 1)])
|
|
t2 = Polygon([(1, 0), (2, 0), (2, 1)])
|
|
polys = GeoSeries([t1, t2])
|
|
ax = polys.plot()
|
|
self._compare_images(ax=ax, filename=filename)
|
|
|
|
def test_multipoly_plot(self):
|
|
""" Test plotting a simple series of multipolygons """
|
|
clf()
|
|
filename = 'multipoly_plot.png'
|
|
t1 = Polygon([(0, 0), (1, 0), (1, 1)])
|
|
t2 = Polygon([(1, 0), (2, 0), (2, 1)])
|
|
t3 = Polygon([(2, 0), (3, 0), (3, 1)], [[(2.2, 0.1), (2.8, 0.1), (2.8, 0.7)]])
|
|
t12 = MultiPolygon([t1, t2, t3])
|
|
mpolys = GeoSeries([t12])
|
|
ax = mpolys.plot()
|
|
self._compare_images(ax=ax, filename=filename)
|
|
|
|
def test_point_plot(self):
|
|
""" Test plotting a simple series of points """
|
|
clf()
|
|
filename = 'points_plot.png'
|
|
N = 10
|
|
points = GeoSeries(Point(i, i) for i in xrange(N))
|
|
ax = points.plot()
|
|
self._compare_images(ax=ax, filename=filename)
|
|
|
|
def test_line_plot(self):
|
|
""" Test plotting a simple series of lines """
|
|
clf()
|
|
filename = 'lines_plot.png'
|
|
N = 10
|
|
lines = GeoSeries([LineString([(0, i), (9, i)]) for i in xrange(N)])
|
|
ax = lines.plot()
|
|
self._compare_images(ax=ax, filename=filename)
|
|
|
|
def test_multiline_plot(self):
|
|
""" Test plotting a multilinestring """
|
|
clf()
|
|
filename = 'multilines_plot.png'
|
|
N = 10
|
|
mlines = GeoSeries(MultiLineString(
|
|
[LineString([(0, i), (9, i)]) for i in xrange(N)]))
|
|
ax = mlines.plot()
|
|
self._compare_images(ax=ax, filename=filename)
|
|
|
|
def test_dataframe_plot(self):
|
|
""" Test plotting a geodataframe"""
|
|
clf()
|
|
df = GeoDataFrame([
|
|
{'geometry': Polygon([(0, 0), (1, 0), (1, 1)]), 'value1': 1, 'cat1': 'low'},
|
|
{'geometry': Polygon([(1, 0), (2, 0), (2, 1)]), 'value1': 2, 'cat1': 'med'},
|
|
{'geometry': Polygon([(2, 0), (3, 0), (3, 1)]), 'value1': 3, 'cat1': 'high'},
|
|
{'geometry': Polygon([(0, 1), (1, 1), (1, 2)]), 'value1': 4, 'cat1': 'low'},
|
|
{'geometry': Polygon([(1, 1), (2, 1), (2, 2)]), 'value1': 5, 'cat1': 'med'},
|
|
{'geometry': Polygon([(2, 1), (3, 1), (3, 2)]), 'value1': 6, 'cat1': 'high'},
|
|
{'geometry': Polygon([(0, 2), (1, 2), (1, 3)]), 'value1': 7, 'cat1': 'low'},
|
|
{'geometry': Polygon([(1, 2), (2, 2), (2, 3)]), 'value1': 8, 'cat1': 'med'},
|
|
{'geometry': Polygon([(2, 2), (3, 2), (3, 3)]), 'value1': 9, 'cat1': 'high'},
|
|
])
|
|
|
|
filename = 'poly_df_plot.png'
|
|
ax = df.plot(column='value1', scheme='QUANTILES', k=3, colormap='OrRd')
|
|
self._compare_images(ax=ax, filename=filename)
|
|
|
|
filename = 'poly_df_category_plot.png'
|
|
ax = df.plot(column='cat1', legend=True, categorical=True, alpha=1.0)
|
|
self._compare_images(ax=ax, filename=filename)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|