Merge pull request #146 from jwass/explode_collect

Explode collect

geopandas/base.py

@@ -6,7 +6,7 @@ from shapely.ops import cascaded_union, unary_union
 import shapely.affinity as affinity
 
 import numpy as np
-from pandas import Series, DataFrame
+from pandas import Series, DataFrame, MultiIndex
 
 import geopandas as gpd
 
@@ -411,6 +411,47 @@ class GeoPandasBase(object):
             use_radians=use_radians) for s in self.geometry],
             index=self.index, crs=self.crs)
 
+    def explode(self):
+        """
+        Explode multi-part geometries into multiple single geometries.
+
+        Single rows can become multiple rows.
+        This is analogous to PostGIS's ST_Dump(). The 'path' index is the
+        second level of the returned MultiIndex
+
+        Returns
+        ------
+        A GeoSeries with a MultiIndex. The levels of the MultiIndex are the
+        original index and an integer.
+
+        Example
+        -------
+        >>> gdf  # gdf is GeoSeries of MultiPoints
+        0                 (POINT (0 0), POINT (1 1))
+        1    (POINT (2 2), POINT (3 3), POINT (4 4))
+
+        >>> gdf.explode()
+        0  0    POINT (0 0)
+           1    POINT (1 1)
+        1  0    POINT (2 2)
+           1    POINT (3 3)
+           2    POINT (4 4)
+        dtype: object
+
+        """
+        index = []
+        geometries = []
+        for idx, s in self.geometry.iteritems():
+            if s.type.startswith('Multi') or s.type == 'GeometryCollection':
+                geoms = s.geoms
+                idxs = [(idx, i) for i in range(len(geoms))]
+            else:
+                geoms = [s]
+                idxs = [(idx, 0)]
+            index.extend(idxs)
+            geometries.extend(geoms)
+        return gpd.GeoSeries(geometries,
+            index=MultiIndex.from_tuples(index)).__finalize__(self)
 
 def _array_input(arr):
     if isinstance(arr, (MultiPoint, MultiLineString, MultiPolygon)):

geopandas/tools/__init__.py

@@ -3,11 +3,12 @@ from __future__ import absolute_import
 from .geocoding import geocode, reverse_geocode
 from .overlay import overlay
 from .sjoin import sjoin
+from .util import collect
 
 __all__ = [
     'overlay',
     'sjoin',
     'geocode',
-    'reverse_geocode'
+    'reverse_geocode',
+    'collect',
 ]
-

geopandas/tools/util.py

@@ -0,0 +1,52 @@
+import pandas as pd
+import geopandas as gpd
+from shapely.geometry import (
+    Point,
+    LineString,
+    Polygon,
+    MultiPoint,
+    MultiLineString,
+    MultiPolygon
+)
+from shapely.geometry.base import BaseGeometry
+
+_multi_type_map = {
+    'Point': MultiPoint,
+    'LineString': MultiLineString,
+    'Polygon': MultiPolygon
+}
+
+def collect(x, multi=False):
+    """
+    Collect single part geometries into their Multi* counterpart
+
+    Parameters
+    ----------
+    x : an iterable or Series of Shapely geometries, a GeoSeries, or
+        a single Shapely geometry        
+    multi : boolean, default False
+        if True, force returned geometries to be Multi* even if they
+        only have one component.
+
+    """
+    if isinstance(x, BaseGeometry):
+        x = [x]
+    elif isinstance(x, pd.Series):
+        x = list(x)
+
+    # We cannot create GeometryCollection here so all types
+    # must be the same. If there is more than one element,
+    # they cannot be Multi*, i.e., can't pass in combination of
+    # Point and MultiPoint... or even just MultiPoint
+    t = x[0].type
+    if not all(g.type == t for g in x):
+        raise ValueError('Geometry type must be homogenous')
+    if len(x) > 1 and t.startswith('Multi'):
+        raise ValueError(
+            'Cannot collect {0}. Must have single geometries'.format(t))
+
+    if len(x) == 1 and (t.startswith('Multi') or not multi):
+        # If there's only one single part geom and we're not forcing to
+        # multi, then just return it
+        return x[0]
+    return _multi_type_map[t](x)

tests/test_geom_methods.py

@@ -5,8 +5,10 @@ import string
 import numpy as np
 from numpy.testing import assert_array_equal
 from pandas.util.testing import assert_series_equal, assert_frame_equal
-from pandas import Series, DataFrame
-from shapely.geometry import Point, LinearRing, LineString, Polygon
+from pandas import Series, DataFrame, MultiIndex
+from shapely.geometry import (
+    Point, LinearRing, LineString, Polygon, MultiPoint
+)
 from shapely.geometry.collection import GeometryCollection
 
 from geopandas import GeoSeries, GeoDataFrame
@@ -389,6 +391,19 @@ class TestGeomMethods(unittest.TestCase):
                            'col1': range(len(self.landmarks))})
         self.assert_(df.total_bounds, bbox)
 
+    def test_explode(self):
+        s = GeoSeries([MultiPoint([(0,0), (1,1)]),
+                      MultiPoint([(2,2), (3,3), (4,4)])])
+
+        index = [(0, 0), (0, 1), (1, 0), (1, 1), (1, 2)]
+        expected = GeoSeries([Point(0,0), Point(1,1), Point(2,2), Point(3,3),
+                              Point(4,4)], index=MultiIndex.from_tuples(index))
+
+        assert_geoseries_equal(expected, s.explode())
+
+        df = self.gdf1[:2].set_geometry(s)
+        assert_geoseries_equal(expected, df.explode())
+
     #
     # Test '&', '|', '^', and '-'
     # The left can only be a GeoSeries. The right hand side can be a

tests/test_tools.py

@@ -0,0 +1,49 @@
+from __future__ import absolute_import
+from shapely.geometry import Point, MultiPoint, LineString
+from geopandas import GeoSeries
+from geopandas.tools import collect
+from .util import unittest
+
+class TestTools(unittest.TestCase):
+    def setUp(self):
+        self.p1 = Point(0,0)
+        self.p2 = Point(1,1)
+        self.p3 = Point(2,2)
+        self.mpc = MultiPoint([self.p1, self.p2, self.p3])
+
+        self.mp1 = MultiPoint([self.p1, self.p2])
+        self.line1 = LineString([(3,3), (4,4)])
+
+    def test_collect_single(self):
+        result = collect(self.p1)
+        self.assert_(self.p1.equals(result))
+
+    def test_collect_single_force_multi(self):
+        result = collect(self.p1, multi=True)
+        expected = MultiPoint([self.p1])
+        self.assert_(expected.equals(result))
+
+    def test_collect_multi(self):
+        result = collect(self.mp1)
+        self.assert_(self.mp1.equals(result))
+
+    def test_collect_multi_force_multi(self):
+        result = collect(self.mp1)
+        self.assert_(self.mp1.equals(result))
+
+    def test_collect_list(self):
+        result = collect([self.p1, self.p2, self.p3])
+        self.assert_(self.mpc.equals(result))
+
+    def test_collect_GeoSeries(self):
+        s = GeoSeries([self.p1, self.p2, self.p3])
+        result = collect(s)
+        self.assert_(self.mpc.equals(result))
+
+    def test_collect_mixed_types(self):
+        with self.assertRaises(ValueError):
+            collect([self.p1, self.line1])
+
+    def test_collect_mixed_multi(self):
+        with self.assertRaises(ValueError):
+            collect([self.mpc, self.mp1])