Merge pull request #671 from jpolchlo/feature/layoutdefinition

Add interface for key/layout computations

docs/docs/geopyspark.geotrellis.key_conversion.rst

@@ -0,0 +1,6 @@
+geopyspark.geotrellis.key_conversion module
+=====================================
+
+.. automodule:: geopyspark.geotrellis.key_conversion
+   :members:
+   :inherited-members:

docs/docs/geopyspark.geotrellis.rst

@@ -73,6 +73,7 @@ geopyspark.geotrellis package
   geopyspark.geotrellis.geotiff
   geopyspark.geotrellis.hillshade
   geopyspark.geotrellis.histogram
+  geopyspark.geotrellis.key_conversion
   geopyspark.geotrellis.layer
   geopyspark.geotrellis.neighborhood
   geopyspark.geotrellis.protobufcodecs

geopyspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/util/GeometryUtils.scala

@@ -0,0 +1,10 @@
+package geopyspark.geotrellis.util
+
+import geotrellis.vector.Geometry
+import geotrellis.vector.io.wkb.WKB
+
+object GeometryUtil {
+
+  def wkbToScalaGeometry(wkb: Array[Byte]): Geometry = WKB.read(wkb)
+
+}

geopyspark/geotrellis/__init__.py

@@ -801,6 +801,7 @@ def __str__(self):
 from . import neighborhood
 from . import s3
 from . import tms
+from . import key_conversion
 
 from .catalog import *
 from .color import *
@@ -817,6 +818,7 @@ def __str__(self):
 from .tms import *
 from .union import *
 from .combine_bands import *
+from .key_conversion import *
 
 __all__ += catalog.__all__
 __all__ += color.__all__
@@ -834,3 +836,4 @@ def __str__(self):
 __all__ += tms.__all__
 __all__ += ['union']
 __all__ += ['combine_bands']
+__all__ += key_conversion.__all__

geopyspark/geotrellis/color.py

@@ -43,7 +43,7 @@ def get_colors_from_matplotlib(ramp_name, num_colors=1<<8):
         import colortools
         import matplotlib.cm as mpc
     except:
-       raise Exception('matplotlib>=2.0.0 and colortools>=0.1.2 required')
+       raise ImportError('matplotlib>=2.0.0 and colortools>=0.1.2 required')
 
     ramp = mpc.get_cmap(ramp_name)
     return  [struct.unpack('>L', bytes(map(lambda x: int(x*255), ramp(x / (num_colors - 1)))))[0] for x in range(0, num_colors)]

geopyspark/geotrellis/key_conversion.py

@@ -0,0 +1,158 @@
+'''
+A module to provide facilites for converting between layout keys and spatial objects.  These facilities aim to bridge between the various layouts, ``SpatialKey``\s, and geometry.
+'''
+
+from math import ceil
+
+import geopyspark as gps
+from . import Extent
+
+__all__ = ["KeyTransform"]
+
+"""The WebMercator (EPSG=3857) world extent"""
+WEB_MERCATOR = Extent(-20037508.342789244, -20037508.342789244, 20037508.342789244, 20037508.342789244)
+
+"""The LatLng (EPSG=4326) world extent"""
+LATLNG = Extent(-180.0, -89.99999, 179.99999, 89.99999)
+
+class KeyTransform(object):
+    """Provides functions to move from keys to geometry and vice-versa.
+
+    Tile Layers have an underlying RDD which is keyed by either :class:`~geopyspark.geotrellis.SpatialKey` or
+    :class:`~geopyspark.geotrellis.SpaceTimeKey`.  Each key represents a region in space, depending on a choice of layout.
+    In order to enable the conversion of keys to regions, and of geometry to keys, the ``KeyTransform`` class is provided.
+    This class is constructed with a layout, which is either ``GlobalLayout``, ``LocalLayout``, or a ``LayoutDefinition``.
+    Global layouts use power-of-two pyramids over the world extent, while local layouts operate over a defined extent and
+    cellsize.
+
+    NOTE: LocalLayouts will encompass the requested extent, but the final layout may include ``SpatialKey``s which only
+    partially cover the requested extent.  The upper-left corner of the resulting layout will match the requested extent,
+    but the right and bottom edges may be beyond the boundaries of the requested extent.
+
+    NOTE: GlobalLayouts require pyproj to be installed.
+
+    Args:
+        layout(:class:`~geopyspark.geotrellis.GlobalLayout` or :class:`~geopyspark.geotrellis.LocalLayout`
+            or :class:`~geopyspark.geotrellis.LayoutDefinition`): a definition of the layout scheme defining the key structure.
+        crs (str or int): Used only when `layout` is :class:`~geopyspark.geotrellis.GlobalLayout`.  Target CRS of reprojection.
+            Either EPSG code, well-known name, or a PROJ.4 string
+        extent (:class:`~geopyspark.geotrellis.Extent`): Used only for ``LocalLayout``s.  The area of interest.
+        cellsize (tup of (float, float)): Used only for ``LocalLayout``s.  The (width, height) in extent units of a pixel.
+            Cannot be specified simultaneously with ``dimensions``.
+        dimensions (tup of (int, int)): Used only for ``LocalLayout``s.  The number of (columns, rows) of pixels over the
+            entire extent.  Cannot be specified simultaneously with ``cellsize``.
+    """
+
+    def __init__(self, layout, crs=None, extent=None, cellsize=None, dimensions=None):
+        self.__jvm = gps.get_spark_context()._gateway.jvm
+
+        if isinstance(layout, gps.LocalLayout):
+            if not extent:
+                raise ValueError("Must specify an extent when using LocalLayout")
+
+            if dimensions and not cellsize:
+                cellsize = ((extent.xmax - extent.xmin)/dimensions[0], (extent.ymax - extent.ymin)/dimensions[1])
+                dimensions = None
+
+            if cellsize and not dimensions:
+                tilewidth = layout.tile_cols * cellsize[0]
+                tileheight = layout.tile_rows * cellsize[1]
+                rows = ceil((extent.xmax - extent.xmin) / tilewidth)
+                cols = ceil((extent.ymax - extent.ymin) / tileheight)
+                extent = gps.Extent(extent.xmin, extent.ymax - rows * tileheight, extent.xmin + cols * tilewidth, extent.ymax)
+                tl = gps.TileLayout(cols, rows, layout.tile_cols, layout.tile_rows)
+            else:
+                raise ValueError("For LocalLayout, must specify exactly one: cellsize or dimension")
+        elif isinstance(layout, gps.GlobalLayout):
+            try:
+                from pyproj import Proj, transform
+            except:
+                raise ImportError('pyproj is required for GlobalLayout')
+
+            if not layout.zoom:
+                raise ValueError("Must specify a zoom level when using GlobalLayout")
+
+            if not crs:
+                raise ValueError("Must specify a crs when using GlobalLayout")
+
+            if isinstance(crs, int):
+                crs = "{}".format(crs)
+
+            gtcrs = self.__jvm.geopyspark.geotrellis.TileLayer.getCRS(crs).get()
+
+            if gtcrs.epsgCode().isDefined() and gtcrs.epsgCode().get() == 3857:
+                extent = WEB_MERCATOR
+            elif gtcrs.epsgCode().isDefined() and gtcrs.epsgCode().get() == 4326:
+                extent = LATLNG
+            else:
+                llex = LATLNG
+                proj4str = gtcrs.toProj4String()
+                target = Proj(proj4str)
+                xmin, ymin = target(llex.xmin, llex.ymin)
+                xmax, ymax = target(llex.xmax, llex.ymax)
+                extent = gps.Extent(xmin, ymin, xmax, ymax)
+
+            layout_rows_cols = int(pow(2, layout.zoom))
+            tl = gps.TileLayout(layout_rows_cols, layout_rows_cols, layout.tile_size, layout.tile_size)
+        elif isinstance(layout, gps.LayoutDefinition):
+            extent = layout.extent
+            tl = layout.tileLayout
+
+        ex = self.__jvm.geotrellis.vector.Extent(float(extent.xmin), float(extent.ymin), float(extent.xmax), float(extent.ymax))
+        tilelayout = self.__jvm.geotrellis.raster.TileLayout(int(tl[0]), int(tl[1]), int(tl[2]), int(tl[3]))
+        self.layout = gps.LayoutDefinition(extent, tl)
+        self.__layout = self.__jvm.geotrellis.spark.tiling.LayoutDefinition(ex, tilelayout)
+
+    def key_to_extent(self, key, *args):
+        """Returns the Extent corresponding to a given key.
+
+        Args:
+            key (:class:`~geopyspark.geotrellis.SpatialKey` or :class:`~geopyspark.geotrellis.SpaceTimeKey` or int): The
+                key to find the extent for.  If of type int, then this parameter is the column of the key, and the call
+                must provide a single additional int value in the args parameter to serve as the row of the key.
+
+        Returns:
+            :class:`~geopyspark.geotrellis.Extent`
+        """
+        if isinstance(key, (gps.SpatialKey, gps.SpaceTimeKey)):
+            skey = self.__jvm.geotrellis.spark.SpatialKey(key.col, key.row)
+        elif isinstance(key, tuple):
+            skey = self.__jvm.geotrellis.spark.SpatialKey(key[0], key[1])
+        elif isinstance(key, int) and len(args) == 1 and isinstance(args[0], int):
+            skey = self.__jvm.geotrellis.spark.SpatialKey(key, args[0])
+        else:
+            raise ValueError("Please supply either gps.SpatialKey, gps.SpaceTimeKey, (int, int), or two ints")
+        ex = self.__layout.mapTransform().apply(skey)
+        return gps.Extent(ex.xmin(), ex.ymin(), ex.xmax(), ex.ymax())
+
+    def extent_to_keys(self, extent):
+        """Returns the keys in the layout intersecting/covered by a given extent.
+
+        Args:
+            extent (:class:`~geopyspark.geotrellis.Extent`): The extent to find the matching keys for.
+
+        Returns:
+            [:class:`~geopyspark.geotrellis.SpatialKey`]
+        """
+        ex = self.__jvm.geotrellis.vector.Extent(float(extent.xmin), float(extent.ymin), float(extent.xmax), float(extent.ymax))
+        gridbnd = self.__layout.mapTransform().apply(ex)
+        cmin = gridbnd.colMin()
+        cmax = gridbnd.colMax()
+        rmin = gridbnd.rowMin()
+        rmax = gridbnd.rowMax()
+        return (gps.SpatialKey(c, r) for c in range(cmin, cmax + 1) for r in range(rmin, rmax + 1))
+
+    def geometry_to_keys(self, geom):
+        """Returns the keys corresponding to grid cells that intersect/are covered by a given Shapely geometry.
+
+        Args:
+            geom (:class:`~shapely.geometry.Geometry`): The geometry to find the matching keys for.
+
+        Returns:
+            [:class:`~geopyspark.geotrellis.SpatialKey`]
+        """
+        from shapely.wkb import dumps
+        jts_geom = self.__jvm.geopyspark.geotrellis.util.GeometryUtil.wkbToScalaGeometry(dumps(geom))
+        scala_key_set = self.__layout.mapTransform().keysForGeometry(jts_geom)
+        key_set = self.__jvm.scala.collection.JavaConverters.setAsJavaSetConverter(scala_key_set).asJava()
+        return [gps.SpatialKey(key.col(), key.row()) for key in key_set]

geopyspark/tests/geotrellis/key_conversion_test.py

@@ -0,0 +1,33 @@
+import pytest
+import unittest
+
+from shapely.geometry import Point
+import geopyspark as gps
+from geopyspark.tests.base_test_class import BaseTestClass
+
+class KeyTransformTest(BaseTestClass):
+
+    layout = gps.LayoutDefinition(gps.Extent(0,0,1,1), gps.TileLayout(5,5,2,2))
+
+    def test_key_to_extent(self):
+        kt_layout = gps.KeyTransform(self.layout)
+        self.assertEqual(gps.Extent(0.0, 0.8, 0.2, 1.0), kt_layout.key_to_extent(gps.SpatialKey(0,0)))
+
+        kt_local = gps.KeyTransform(gps.LocalLayout(2), extent=gps.Extent(0,0,1,1), dimensions=(10,10))
+        self.assertEqual(gps.Extent(0.0, 0.8, 0.2, 1.0), kt_layout.key_to_extent(gps.SpatialKey(0,0)))
+
+        kt_global = gps.KeyTransform(gps.GlobalLayout(zoom=1), crs=4326)
+        nw_global_extent = kt_global.key_to_extent(gps.SpatialKey(0,0))
+        self.assertTrue(abs(nw_global_extent.xmin + 180.0) <= 1e-4 and
+                        abs(nw_global_extent.xmax) <= 1e-4 and
+                        abs(nw_global_extent.ymin) <= 1e-4 and
+                        abs(nw_global_extent.ymax -90) <= 1e-4
+        )
+
+    def test_extent_to_key(self):
+        kt = gps.KeyTransform(self.layout)
+        self.assertTrue(set(kt.extent_to_keys(gps.Extent(0,0,0.4,0.4))) == set([gps.SpatialKey(x,y) for x in [0,1] for y in [3,4]]))
+
+    def test_geom_to_key(self):
+        kt = gps.KeyTransform(self.layout)
+        self.assertTrue(kt.geometry_to_keys(Point(0.1,0.1)) == [gps.SpatialKey(0,4)])