Skip to content

S2Pandas module

S2Pandas module for S2 cell operations on pandas DataFrames and GeoDataFrames.

S2Pandas

Source code in vgridpandas/s2pandas.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
@pd.api.extensions.register_dataframe_accessor("s2")
class S2Pandas:
    def __init__(self, df: DataFrame):
        self._df = df

    def latlon2s2(
        self,
        resolution: int,
        lat_col: str = "lat",
        lon_col: str = "lon",
        set_index: bool = False,
    ) -> AnyDataFrame:
        """Adds S2 token to (Geo)DataFrame.

        pd.DataFrame: uses `lat_col` and `lon_col` (default `lat` and `lon`)
        gpd.GeoDataFrame: uses `geometry`

        Assumes coordinates in epsg=4326.

        Parameters
        ----------
        resolution : int
            S2 resolution
        lat_col : str
            Name of the latitude column (if used), default 'lat'
        lon_col : str
            Name of the longitude column (if used), default 'lon'
        set_index : bool
            If True, the column with S2 token is set as index, default False

        Returns
        -------
        (Geo)DataFrame with S2 token and resolution columns added

        """

        if isinstance(self._df, gpd.GeoDataFrame):
            lons = self._df.geometry.x
            lats = self._df.geometry.y
        else:
            lons = self._df[lon_col]
            lats = self._df[lat_col]

        s2_tokens = [latlon_to_s2(lat, lon, resolution) for lat, lon in zip(lats, lons)]

        s2_col = S2_COL
        df = self._df.assign(**{s2_col: s2_tokens, f"{s2_col}_res": resolution})
        if set_index:
            return df.set_index(s2_col)
        return df

    def s22geo(
        self, s2_col: str = None, fix_antimeridian: Optional[str] = None
    ) -> GeoDataFrame:
        """Add geometry with S2 geometry to the DataFrame."""
        if s2_col is not None:
            if s2_col not in self._df.columns:
                raise ValueError(f"Column '{s2_col}' not found in DataFrame")
            ids = self._df[s2_col]
        else:
            if S2_COL not in self._df.columns:
                raise ValueError(f"Column '{S2_COL}' not found in DataFrame")
            ids = self._df[S2_COL]
        return dggs_ids_to_geodataframe(
            self._df, ids, s2_to_geo, fix_antimeridian=fix_antimeridian
        )

    def polyfill(
        self,
        resolution: int,
        predicate: str = None,
        compact: bool = False,
        explode: bool = False,
        fix_antimeridian: Optional[str] = None,
    ) -> AnyDataFrame:
        """
        Parameters
        ----------
        resolution : int
            S2 resolution
        predicate : str, optional
            Spatial predicate to apply ('intersect', 'within', 'centroid_within', 'largest_overlap')
        compact : bool, optional
            Whether to compact the S2 tokens
        explode : bool
            If True, will explode the resulting list vertically.
            All other columns' values are copied.
            Default: False
        fix_antimeridian : str, optional
            Antimeridian fix: 'shift', 'shift_balanced', 'shift_west', 'shift_east', or 'split'

        Returns
        -------
        (Geo)DataFrame with S2 tokens in column 's2', exploded to one row per cell if explode=True.
        """
        result = self._df.geometry.apply(
            lambda geom: polyfill_row(
                geom, resolution, predicate, compact, fix_antimeridian
            )
        )

        if not explode:
            return self._df.assign(**{S2_COL: result})

        result = result.explode().to_frame(S2_COL)

        return self._df.join(result)

    def s2bin(
        self,
        resolution: int,
        agg: str = "count",
        numeric_col: str = None,
        category_col: str = None,
        lat_col: str = "lat",
        lon_col: str = "lon",
        fix_antimeridian: Optional[str] = None,
    ) -> GeoDataFrame:
        """Bin points into S2 cells and compute statistics."""
        s2_col = S2_COL
        df = self.latlon2s2(resolution, lat_col, lon_col)
        result = aggregate_bin(df, s2_col, agg, numeric_col, category_col)
        return result.s2.s22geo(s2_col=s2_col, fix_antimeridian=fix_antimeridian)

    def linetrace(
        self,
        resolution: int,
        explode: bool = False,
        fix_antimeridian: Optional[str] = None,
    ) -> AnyDataFrame:
        """S2 cell representation of a (Multi)LineString traced along its vertices.

        Uses the same neighbor walk as ``vgrid.conversion.vector2dggs.vector2s2.polyline2s2``.
        """
        result = self._df.apply(
            lambda row: list(
                linetrace(
                    row.geometry, resolution, fix_antimeridian=fix_antimeridian
                )
            ),
            axis=1,
        )
        if not explode:
            return self._df.assign(**{S2_COL: result})
        result = result.explode().to_frame(S2_COL)
        return self._df.join(result)

latlon2s2(resolution, lat_col='lat', lon_col='lon', set_index=False)

Adds S2 token to (Geo)DataFrame.

pd.DataFrame: uses lat_col and lon_col (default lat and lon) gpd.GeoDataFrame: uses geometry

Assumes coordinates in epsg=4326.

Parameters

resolution : int S2 resolution lat_col : str Name of the latitude column (if used), default 'lat' lon_col : str Name of the longitude column (if used), default 'lon' set_index : bool If True, the column with S2 token is set as index, default False

Returns

(Geo)DataFrame with S2 token and resolution columns added

Source code in vgridpandas/s2pandas.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def latlon2s2(
    self,
    resolution: int,
    lat_col: str = "lat",
    lon_col: str = "lon",
    set_index: bool = False,
) -> AnyDataFrame:
    """Adds S2 token to (Geo)DataFrame.

    pd.DataFrame: uses `lat_col` and `lon_col` (default `lat` and `lon`)
    gpd.GeoDataFrame: uses `geometry`

    Assumes coordinates in epsg=4326.

    Parameters
    ----------
    resolution : int
        S2 resolution
    lat_col : str
        Name of the latitude column (if used), default 'lat'
    lon_col : str
        Name of the longitude column (if used), default 'lon'
    set_index : bool
        If True, the column with S2 token is set as index, default False

    Returns
    -------
    (Geo)DataFrame with S2 token and resolution columns added

    """

    if isinstance(self._df, gpd.GeoDataFrame):
        lons = self._df.geometry.x
        lats = self._df.geometry.y
    else:
        lons = self._df[lon_col]
        lats = self._df[lat_col]

    s2_tokens = [latlon_to_s2(lat, lon, resolution) for lat, lon in zip(lats, lons)]

    s2_col = S2_COL
    df = self._df.assign(**{s2_col: s2_tokens, f"{s2_col}_res": resolution})
    if set_index:
        return df.set_index(s2_col)
    return df

linetrace(resolution, explode=False, fix_antimeridian=None)

S2 cell representation of a (Multi)LineString traced along its vertices.

Uses the same neighbor walk as vgrid.conversion.vector2dggs.vector2s2.polyline2s2.

Source code in vgridpandas/s2pandas.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
def linetrace(
    self,
    resolution: int,
    explode: bool = False,
    fix_antimeridian: Optional[str] = None,
) -> AnyDataFrame:
    """S2 cell representation of a (Multi)LineString traced along its vertices.

    Uses the same neighbor walk as ``vgrid.conversion.vector2dggs.vector2s2.polyline2s2``.
    """
    result = self._df.apply(
        lambda row: list(
            linetrace(
                row.geometry, resolution, fix_antimeridian=fix_antimeridian
            )
        ),
        axis=1,
    )
    if not explode:
        return self._df.assign(**{S2_COL: result})
    result = result.explode().to_frame(S2_COL)
    return self._df.join(result)

polyfill(resolution, predicate=None, compact=False, explode=False, fix_antimeridian=None)

Parameters

resolution : int S2 resolution predicate : str, optional Spatial predicate to apply ('intersect', 'within', 'centroid_within', 'largest_overlap') compact : bool, optional Whether to compact the S2 tokens explode : bool If True, will explode the resulting list vertically. All other columns' values are copied. Default: False fix_antimeridian : str, optional Antimeridian fix: 'shift', 'shift_balanced', 'shift_west', 'shift_east', or 'split'

Returns

(Geo)DataFrame with S2 tokens in column 's2', exploded to one row per cell if explode=True.

Source code in vgridpandas/s2pandas.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def polyfill(
    self,
    resolution: int,
    predicate: str = None,
    compact: bool = False,
    explode: bool = False,
    fix_antimeridian: Optional[str] = None,
) -> AnyDataFrame:
    """
    Parameters
    ----------
    resolution : int
        S2 resolution
    predicate : str, optional
        Spatial predicate to apply ('intersect', 'within', 'centroid_within', 'largest_overlap')
    compact : bool, optional
        Whether to compact the S2 tokens
    explode : bool
        If True, will explode the resulting list vertically.
        All other columns' values are copied.
        Default: False
    fix_antimeridian : str, optional
        Antimeridian fix: 'shift', 'shift_balanced', 'shift_west', 'shift_east', or 'split'

    Returns
    -------
    (Geo)DataFrame with S2 tokens in column 's2', exploded to one row per cell if explode=True.
    """
    result = self._df.geometry.apply(
        lambda geom: polyfill_row(
            geom, resolution, predicate, compact, fix_antimeridian
        )
    )

    if not explode:
        return self._df.assign(**{S2_COL: result})

    result = result.explode().to_frame(S2_COL)

    return self._df.join(result)

s22geo(s2_col=None, fix_antimeridian=None)

Add geometry with S2 geometry to the DataFrame.

Source code in vgridpandas/s2pandas.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def s22geo(
    self, s2_col: str = None, fix_antimeridian: Optional[str] = None
) -> GeoDataFrame:
    """Add geometry with S2 geometry to the DataFrame."""
    if s2_col is not None:
        if s2_col not in self._df.columns:
            raise ValueError(f"Column '{s2_col}' not found in DataFrame")
        ids = self._df[s2_col]
    else:
        if S2_COL not in self._df.columns:
            raise ValueError(f"Column '{S2_COL}' not found in DataFrame")
        ids = self._df[S2_COL]
    return dggs_ids_to_geodataframe(
        self._df, ids, s2_to_geo, fix_antimeridian=fix_antimeridian
    )

s2bin(resolution, agg='count', numeric_col=None, category_col=None, lat_col='lat', lon_col='lon', fix_antimeridian=None)

Bin points into S2 cells and compute statistics.

Source code in vgridpandas/s2pandas.py 
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
def s2bin(
    self,
    resolution: int,
    agg: str = "count",
    numeric_col: str = None,
    category_col: str = None,
    lat_col: str = "lat",
    lon_col: str = "lon",
    fix_antimeridian: Optional[str] = None,
) -> GeoDataFrame:
    """Bin points into S2 cells and compute statistics."""
    s2_col = S2_COL
    df = self.latlon2s2(resolution, lat_col, lon_col)
    result = aggregate_bin(df, s2_col, agg, numeric_col, category_col)
    return result.s2.s22geo(s2_col=s2_col, fix_antimeridian=fix_antimeridian)

linetrace(geometry, resolution, fix_antimeridian=None)

Trace a (Multi)LineString with S2 cells (same walk as polyline2s2).

Source code in vgridpandas/s2pandas.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def linetrace(
    geometry: MultiLineOrLine,
    resolution: int,
    fix_antimeridian: Optional[str] = None,
) -> Iterator[str]:
    """Trace a (Multi)LineString with S2 cells (same walk as ``polyline2s2``)."""
    resolution = validate_s2_resolution(resolution)

    def segment_cells(start_xy, end_xy):
        return _s2_segment_cells(
            resolution, start_xy, end_xy, fix_antimeridian=fix_antimeridian
        )

    yield from linetrace_polyline(geometry, segment_cells)

poly2s2(geometry, resolution, predicate=None, compact=False, fix_antimeridian=None)

Convert polygon geometries to S2 grid cell tokens.

Source code in vgridpandas/s2pandas.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def poly2s2(geometry, resolution, predicate=None, compact=False, fix_antimeridian=None):
    """Convert polygon geometries to S2 grid cell tokens."""
    resolution = validate_s2_resolution(resolution)
    s2_tokens = []
    if isinstance(geometry, Polygon):
        polys = [geometry]
    elif isinstance(geometry, MultiPolygon):
        polys = list(geometry.geoms)
    else:
        return []

    for poly in polys:
        min_lon, min_lat, max_lon, max_lat = poly.bounds
        level = resolution
        coverer = s2.RegionCoverer()
        coverer.min_level = level
        coverer.max_level = level
        region = s2.LatLngRect(
            s2.LatLng.from_degrees(min_lat, min_lon),
            s2.LatLng.from_degrees(max_lat, max_lon),
        )
        covering = coverer.get_covering(region)
        cell_ids = covering
        if compact:
            covering = s2.CellUnion(covering)
            covering.normalize()
            cell_ids = covering.cell_ids()

        for cell_id in cell_ids:
            cell_token = s2.CellId.to_token(cell_id)
            cell_polygon = s2_to_geo(cell_token, fix_antimeridian=fix_antimeridian)
            if not check_predicate(cell_polygon, poly, predicate):
                continue
            s2_tokens.append(cell_token)

    return s2_tokens

polyfill_row(geometry, resolution, predicate=None, compact=False, fix_antimeridian=None)

Return S2 tokens covering a single row geometry.

Source code in vgridpandas/s2pandas.py 
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def polyfill_row(
    geometry,
    resolution: int,
    predicate: str = None,
    compact: bool = False,
    fix_antimeridian: Optional[str] = None,
) -> list:
    """Return S2 tokens covering a single row geometry."""
    if isinstance(geometry, (Polygon, MultiPolygon)):
        tokens = set(
            poly2s2(geometry, resolution, predicate, compact, fix_antimeridian)
        )
    elif isinstance(geometry, (LineString, MultiLineString)):
        tokens = set(linetrace(geometry, resolution, fix_antimeridian))
    else:
        raise TypeError(f"Unknown type {type(geometry)}")
    return list(tokens)