mlx-geospatial

GPU-accelerated spatial operations for Apple silicon, built on MLX: haversine distances, point-in-polygon, and grid-indexed spatial joins.

The pitch beyond speed: mlx-geospatial is array-native in unified memory. Re-center a point table once with PolygonSet.to_local, and zone assignment, distance-to-POI features and nearest-neighbor lookups run entirely on the GPU and come back as MLX arrays: no NumPy round trip, no copies, no serialization, straight into an MLX training loop or embedding model (see examples/mlx_feature_pipeline.py). GeoPandas users keep NumPy: the predicates return the container you pass in.

Install

pip install mlx-geospatial

Requires macOS on Apple silicon (Metal), Python ≥ 3.10, MLX ≥ 0.32.0, numpy, shapely ≥ 2.0. geopandas is optional (pip install "mlx-geospatial[geo]").

Quickstart

The first block runs as written (synthetic data stands in for a real zone layer); the second continues from it and elides the model.

GeoPandas / NumPy users: pass NumPy, shapely or GeoSeries, get NumPy back.

import numpy as np, shapely
import mlx_geospatial as ms

rng = np.random.default_rng(0)
pickups = np.column_stack([rng.uniform(-74.03, -73.9, 200_000), rng.uniform(40.7, 40.82, 200_000)])
stations = np.column_stack([rng.uniform(-74.03, -73.9, 500), rng.uniform(40.7, 40.82, 500)])
gx, gy = np.linspace(-74.03, -73.9, 21), np.linspace(40.7, 40.82, 21)
zones = [shapely.box(gx[i], gy[j], gx[i + 1], gy[j + 1]) for j in range(20) for i in range(20)]
# real data: zones = gpd.read_file("taxi_zones.geojson").geometry

# --- haversine distances (coords are (lon, lat) degrees; meters out) ------
d = ms.haversine(pickups[:-1], pickups[1:])              # elementwise, mx.array
dist, idx = ms.haversine_knn(pickups, stations, k=5)     # (N, 5) each, mx.array

# --- point-in-polygon / spatial join --------------------------------------
zone_idx = ms.point_in_polygon(pickups, zones)          # (N,) numpy int32, -1 = none
polyset = ms.PolygonSet(zones)                           # ingest + index once, reuse
labels = ms.point_in_polygon(pickups, polyset)

# GeoPandas-flavored join (pandas DataFrames with a geometry column work too)
# joined = ms.sjoin(trips_gdf, zones_gdf, how="inner")   # adds index_right + zone cols

MLX users: re-center once, slice per batch, everything stays on the GPU.

import mlx.core as mx, mlx.nn as nn

polyset = ms.PolygonSet(zones)
index = ms.GridIndex(polyset)
local = polyset.to_local(pickups)         # float64 -> fp32 local frame, once; LocalPoints
pickups_mx = mx.array(np.radians(pickups).astype(np.float32))   # raw lon/lat for distances
pois_mx = mx.array(np.radians(stations[:32]).astype(np.float32))
n_zones = len(polyset)
embed = nn.Embedding(n_zones + 1, 16)     # extra row for "no zone"

for start in range(0, len(local), 65_536):
    batch = local[start : start + 65_536]                  # slicing keeps the frame
    zone = index.query(batch)                              # lazy mx.array int32, GPU only
    zone = mx.where(zone < 0, n_zones, zone)               # -1 would read the last row
    dist = ms.haversine_cross(pickups_mx[start : start + 65_536], pois_mx, degrees=False)
    feats = mx.concatenate([embed(zone), dist / 1e4], axis=1)
    mx.eval(feats)                                         # your mlx.nn step goes here

LocalPoints are re-centered coordinates, so the distance functions reject them by design; keep the raw lon/lat around for haversine*.

API

Output containers: the predicates (point_in_polygon, GridIndex.query, contains_matrix) return the container you pass in — NumPy for NumPy, lists, shapely and GeoSeries; MLX arrays for mx.array and LocalPoints, computed without leaving unified memory (float32 inputs and LocalPoints entirely on the GPU; float64 re-centering on the MLX CPU stream). The distance functions always return MLX arrays; sjoin returns pandas.

GridIndex cells default to a quarter of the median polygon bbox's longer side (about three candidates per point on typical zone layers), backing off to a half or the median itself when a few large polygons among many small ones would inflate the index, and coarsening further to fit the cell/entry budgets. Construction is adaptive: low-coverage bbox expansions use stable integer sorting, more expanded grids use lower-memory packed-key sorting, and material high-occupancy grids use a sort-free direct CSR fill. Every route retains ascending polygon IDs within each cell. Pass a prebuilt PolygonSet to point_in_polygon/sjoin, or a GridIndex via point_in_polygon(points, index=grid), to reuse ingestion and the index across calls.

Conventions: coordinates are (x, y) = (lon, lat), degrees; planar math for containment (project first if you need metric correctness near poles); spherical math for haversine. No reprojection — that stays pyproj's job. When both sides carry a CRS (GeoSeries/GeoDataFrame), sjoin, point_in_polygon, contains_matrix and PolygonSet.to_local raise on a mismatch; a CRS-less side is never checked.

Precision, honestly

Metal GPUs are fp32-only, and raw lon/lat in fp32 carries ~1 m rounding — enough to flip containment at polygon edges. mlx-geospatial handles this structurally rather than hopefully:

• All GPU predicates run in locally re-centered coordinates: a float64 origin (dataset bbox center) is subtracted before casting to fp32, restoring ~millimeter resolution at city/state extents.
• Boundary tolerance policy: points within PolygonSet.boundary_tolerance (a few fp32 ulps of the local extent — centimeters for a metro area) of a polygon edge may classify either way. Outside that band, results match shapely/GEOS; the test suite enforces exactly this contract. predicate ("intersects", the default, or "within") only matters for points exactly on a boundary, inside that band, so by default both give the same answer; pass exact_boundary=True to re-check the near-boundary points with GEOS and get the predicate honored exactly. Finding them is a host pass over all points (an STRtree over tolerance-widened edge boxes plus a float64 distance test: about 0.3 s per million points on 48-vertex zones, 0.5 s on 10k zones), cheaper than a shapely join but far from GPU speed, so reserve it for boundary-sensitive work; only the band is re-checked, so interior points of invalid geometry keep the even-odd answer. The band scales with the dataset extent: one stray geometry far from the rest (a placeholder at (0, 0), a wrong-hemisphere typo) widens it for everyone, so clean outliers before ingesting.
• Cell assignment runs on the GPU in fp32. Polygon cell ranges are planned in float64 and padded by 16 ulps of the local extent — more than the kernel's bbox widening plus the fp32 assignment error — so whichever cell a point near a cell line lands in lists every polygon whose kernel bbox accepts it: rounding can only add candidate pairs, never lose a true hit, and host and device inputs give bit-identical answers. Tilings whose edges fall on cell lines pay for this with each tile listed on both sides of those lines.
• Device inputs: host float64 coordinates and mx.array float64 inputs (build them with mx.array(x, dtype=mx.float64); a plain mx.array of a float64 NumPy array is silently float32) are re-centered in float64 and keep this contract. A float32 (or 16-bit) MLX array already carries ~0.5 ulp of raw lon/lat rounding (about 0.4 m at longitude 74), wider than the band at city or state extents, so the join re-centers it on the device and emits mlx_geospatial.PrecisionWarning once per call site; silence it with warnings.filterwarnings("ignore", category=ms.PrecisionWarning) or make it strict with warnings.simplefilter("error", ms.PrecisionWarning).
• One containment contract: containment is the even-odd rule over all rings on every backend. The CPU fallback assigns grid cells with the device graph's fp32 arithmetic and evaluates the same fp32 predicate in NumPy with the kernels' operation order, so backend="cpu" and backend="gpu" are bit-identical for every input and auto dispatch can never change an answer (the one escape is fp32 subnormals, below ~1e-38, which Metal flushes to zero and NumPy keeps: a coordinate, an edge delta or an intermediate of the crossing test that small can differ, which needs local coordinates themselves near 1e-38 and never occurs for real-world units). Even-odd equals GEOS fill semantics for valid polygons; on invalid input (self-intersections, overlapping multipolygon parts) it differs from GEOS, which you meet only through exact_boundary=True, and only for the near-boundary points it re-checks. Pass PolygonSet(geoms, validate=True) to reject invalid geometry up front.
• Long polygons are adaptive: Metal containment chooses scalar traversal, one 32-lane SIMDgroup, or one 128-thread group per predicate. Lanes evaluate disjoint edges and XOR integer crossing parity, so the cooperative paths preserve the scalar fp32 even-odd result exactly. They are enabled only for rings with at least 1,024 edges and bounded work sizes; indexed first-match stays scalar unless every candidate in every used cell has a long ring.
• Haversine: absolute error ~1-2 m from fp32 differencing (relative ~1e-4 at 10 km), far below the ~0.5% sphere-vs-ellipsoid modeling error any haversine implementation carries. One exception: within ~2° of the exact antipode, arcsin amplification grows absolute error to 10²-10⁴ m (relative error stays < 2e-4).

Small host-resident workloads automatically fall back to the NumPy evaluation of the same predicate: auto dispatch compares its cost (per point plus per candidate polygon edge, estimated from the index) with the fixed GPU cost of the route, which puts the crossover at a few hundred points for typical zone layers, a couple of thousand for sparse tilings and under ten for detailed polygons, later for all-match queries whose compaction passes cost more (measured in the crossover table below; python bench/bench.py --op crossover re-derives the constants for another machine). The fallback exists for consistency, not speed: it is about 1.5-2x slower than GEOS per point on 48-vertex zones (up to ~7x on 2048-vertex zones, faster on sparse tilings) and carries ~50 µs of fixed setup, so a loop of tiny queries pays roughly 0.05-0.2 ms per call. Device-resident inputs (mx.array, LocalPoints) run on the GPU whenever Metal is available, because the fallback would need a host copy. Force a path with backend="gpu" | "cpu" or MLX_GEOSPATIAL_BACKEND; backend="cpu" on device inputs copies the fp32 local coordinates to the host and runs the same predicate there (a parity/debug path).

Benchmarks

Apple M4 Max (16 cores: 12 performance + 4 efficiency), macOS 26.2, mlx 0.31.2, Python 3.13. These are pre-release reference measurements; v0.1.0 requires MLX ≥ 0.32.0 for an upstream custom-kernel cache correctness fix, so reproduce them on your supported runtime before making capacity decisions. End-to-end wall time from host-resident numpy/shapely inputs — ingestion, host→device transfer, and device result evaluation/synchronization included; host-input join rows return NumPy labels as a zero-copy view of the MLX result; best of 3 (exceptions in bench/README.md). CPU baselines: shapely/GEOS, GeoPandas and scikit-learn run on one core, DuckDB on its default 16 threads; the GPU rows use the whole GPU, with clocks warmed before each warm-index timing (steady-state burst numbers, like the crossover tables). Reproduce with python bench/bench.py --all (methodology in bench/README.md).

Point-in-polygon join. The baselines evaluate every (point, polygon) pair (shapely and GeoPandas return them, DuckDB is timed as count(*) over the join), so their speedups (in parentheses) are against our all_matches=True column, which answers the same question. The first-match columns are the cheaper contract most joins want: NumPy in → NumPy out, and the array-native pattern (LocalPoints re-centered once, MLX array out, mx.eval inside the timing). GEOS gets its best configuration: the polygons are the prepared side (STRtree(points).query(polygons), polygons on the left of sjoin), which is 8-60x faster on detailed polygons than querying a polygon tree with points and returns the same pairs. The synthetic zones are random star-shaped polygons that overlap and leave gaps (not a tiling, not convex): about two thirds of the points fall in a zone, in 1.6 zones on average, so the all-pairs outputs hold ~1.1 pairs per point. The coastline row is one 150,000-vertex polygon covering 22% of the extent. Cold runs with polygon ingestion and index construction included: 1.8 ms at 100k × 260, 4.0 ms at 1M × 260, 22 ms at 1M × 10k, 45 ms on the 2,048-vertex zones, 0.55 s on the coastline.

Three things the table shows that a headline would hide. Against single-core GEOS in its best configuration the advantage is 26-131x on 48-vertex zones and shrinks with geometric detail: 22x on 2,048-vertex zones, 7.7x on the 150,000-vertex polygon. Prepared GEOS polygons carry an indexed point-in-area locator that costs roughly log(edges) per test, while our kernels walk every edge of each candidate ring; GPU parallelism absorbs that up to a few thousand edges per ring and no further (indexing long rings is on the backlog). The unprepared orientation, STRtree(polygons).query(points), is what this harness used before and what produced detailed-geometry speedups in the hundreds; those numbers measured GEOS's worst configuration, not ours. DuckDB, which has no prepared geometries but uses 16 threads, stays 60-107x behind throughout. On the 48-vertex zones per-point overhead dominates our time (cell assignment and the count/scan/compaction passes, plus ~1.1 ms of fixed launch cost for the all-pairs route, the floor the crossover table shows: over half of the 100k row, about 14% at 1M), not edge tests, and the baselines are within about 2.5x of each other (1.4-1.6x from 1M points up). And on the single huge polygon first match is slower than all pairs: the first-match kernel gives every point a thread, and the 61% of points outside the polygon's bbox (the grid covers only that bbox, so they have no candidates) exit at once and leave their SIMD lanes idle while their neighbours in the same 32-lane group walk 150,000 edges, whereas the all-pairs route compacts candidate pairs first; reordering the same points so the candidate-bearing 39% are contiguous makes first match as fast as all pairs. Compacting first-match candidates is a known follow-up.

Crossover (python bench/bench.py --op crossover; 260 zones, warm index, host-resident NumPy points, best of 7). The CPU column is the host fallback the auto dispatch uses: the same fp32 even-odd predicate evaluated in NumPy, so it returns exactly what the GPU returns. GPU cells are steady-state numbers, measured in a back-to-back burst with the GPU clocks warmed first; the note below the tables says what an isolated query costs. First-match queries:

An axis-aligned square tiling of 256 tiles has only ~9 candidate edges per point, so the fallback stays cheap for longer (first match / all matches):

All-match queries (all_matches=True) carry a larger fixed GPU cost for the count, scan and compaction passes, so they switch later:

auto picked the faster side in every row except three near the thresholds (first match at 200 simple points and at 2,000 tiles, all matches at 1,000 simple points), each within about 0.2 ms. The fixed GPU cost depends on what the GPU was doing just before: about 0.2 ms per first-match query and 1.1 ms per all-match query inside a burst, but roughly 0.8-1.1 ms and 2.5 ms for a query issued after tens of milliseconds of host-only work have let the GPU idle (bench.py --op crossover prints both). The dispatch constants (GPU_LAUNCH_US = 350, GPU_ALL_MATCHES_LAUNCH_US = 1500) sit between the two regimes: picking the CPU while the GPU was warm loses at most ~0.2 ms (~0.45 ms for all matches), picking the GPU after it has idled loses up to ~0.6 ms (~1 ms for all matches) just above the threshold, so a loop of small host queries interleaved with heavy host work may prefer backend="cpu" for a few hundred points more. The GPU path is a near-constant launch cost until kernel work shows (tens of thousands of points for simple polygons, a few thousand for detailed ones) while the fallback grows per point and per candidate edge, so the model switches first-match queries at about 250 points for 48-vertex zones, about 2,100 for the square tiling and under ten for 2048-vertex zones, and all-match queries roughly four times later. Device-resident inputs skip the comparison and run on the GPU whenever Metal is available.

Haversine (mx.compile-fused distances; kNN via adaptive scalar/SIMDgroup selection kernels). Our kNN is exact brute force with nothing to prebuild, so the BallTree baseline is shown both with the tree built inside the timing (a one-off query) and query-only on a prebuilt tree (repeated queries against a fixed reference set, where the tree amortizes):

Honesty notes: haversine_knn is exact brute force. For k <= 32 it runs a shape-adaptive fused Metal selector: 32 lanes cooperate on each query when query parallelism is scarce, while large query batches retain the faster one-thread-per-query kernel. Neither materializes an (N, M) matrix or sorts the full reference set. Brute force costs queries × references, a BallTree queries × log(references), so the reference-set size decides the winner, not the query count: against 200k references brute force is cheaper per query (about 3 µs vs 6 µs here) and the GPU wins 2x, against 10k references 30x; against 1M references a prebuilt tree is cheaper per query (about 10 µs vs 20-25 µs) at any query count, and only the tree-building cost puts us ahead (1k × 1M above; at 100k queries even build + query beats brute force). With k > 32 (or without an active Metal device) kNN falls back to a chunked distance-matrix + argpartition path that is sort-bound; at large reference sets a BallTree can win there too. Below a few hundred host points (fewer for detailed polygons) the join auto-dispatches to the host fallback, because the fixed GPU launch overhead dominates there; the crossover table shows both sides.

Limitations (v0.1)

• Points-in-polygons only: no polygon-polygon predicates, overlays, or distance-to-polygon (backlog).
• Uniform grid index; heavily skewed polygon sizes may want the quadtree planned for v0.2.
• all_matches=True uses bounded adaptive GPU count/scan/ordered compaction without NumPy candidate-pair expansion. A long-ring candidate or at least 8,192 candidate edges in one used grid cell selects the pair-parallel route; its returned arrays still require memory proportional to the number of hits.
• Containment walks every edge of a candidate ring, so on rings with tens of thousands of vertices single-core GEOS with prepared polygons (an indexed point-in-area locator) comes within about 8x of the GPU. Indexing long rings is on the backlog.
• First match gives every point a GPU thread, so on a layer dominated by one huge polygon (the 150k-vertex coastline row) points outside its bbox idle their SIMD lanes while neighbours walk the ring, and all_matches=True, which compacts candidate pairs first, is faster. Compacting first-match candidates the same way is on the backlog.
• sjoin is a first-match (lowest polygon index) point-in-polygon join with the GeoPandas default predicate intersects; predicate="within" only differs on boundary points and needs exact_boundary=True to be honored. Use all_matches=True on point_in_polygon for multi-match joins.

License

MIT