Merge pull request #364 from creare-com/feature/geotiff-export

Feature/geotiff export

podpac/core/coordinates/coordinates.py

@@ -20,6 +20,7 @@
 import xarray.core.coordinates
 from six import string_types
 import pyproj
+import logging
 
 import podpac
 from podpac.core.settings import settings
@@ -32,6 +33,14 @@
 from podpac.core.coordinates.dependent_coordinates import DependentCoordinates
 from podpac.core.coordinates.rotated_coordinates import RotatedCoordinates
 
+# Optional dependencies
+from lazy_import import lazy_module, lazy_class
+
+rasterio = lazy_module("rasterio")
+
+# Set up logging
+_logger = logging.getLogger(__name__)
+
 
 class Coordinates(tl.HasTraits):
     """
@@ -442,6 +451,42 @@ def from_url(cls, url):
 
         return cls.from_definition(coords)
 
+    @classmethod
+    def from_geotransform(cls, geotransform, shape, crs=None):
+        """ Creates Coordinates from GDAL Geotransform. 
+        
+        """
+        # Handle the case of rotated coordinates
+        try:
+            rcoords = RotatedCoordinates.from_geotransform(geotransform)
+            coords = Coordinates([rcoords], dims=["lat,lon"], crs=crs)
+        except:
+            rcoords = None
+            _logger.debug("Rasterio source dataset does not have Rotated Coordinates")
+
+        if rcoords is not None and rcoords.theta != 0:
+            # These are Rotated coordinates and we can return
+            return coords
+
+        # Handle the case of uniform coordinates (not rotated, but N-S E-W aligned)
+        affine = rasterio.Affine.from_gdal(*geotransform)
+        if affine.e == affine.a == 0:
+            order = -1
+            step = np.array([affine.d, affine.b])
+        else:
+            order = 1
+            step = np.array([affine.e, affine.a])
+
+        origin = affine.f + step[0] / 2, affine.c + step[1] / 2
+        end = origin[0] + step[0] * (shape[::order][0] - 1), origin[1] + step[1] * (shape[::order][1] - 1)
+        coords = Coordinates(
+            [
+                podpac.clinspace(origin[0], end[0], shape[::order][0], "lat"),
+                podpac.clinspace(origin[1], end[1], shape[::order][1], "lon"),
+            ][::order]
+        )
+        return coords
+
     @classmethod
     def from_definition(cls, d):
         """
@@ -692,6 +737,10 @@ def shape(self):
 
         return tuple(size for c in self._coords.values() for size in c.shape)
 
+    @property
+    def ushape(self):
+        return tuple(self[dim].size for dim in self.udims)
+
     @property
     def ndim(self):
         """:int: Number of dimensions. """
@@ -801,6 +850,53 @@ def hash(self):
         json_d = json.dumps(self.full_definition, separators=(",", ":"), cls=podpac.core.utils.JSONEncoder)
         return hash_alg(json_d.encode("utf-8")).hexdigest()
 
+    @property
+    def geotransform(self):
+        """ :tuple: GDAL geotransform. """
+        # Make sure we only have 1 time and alt dimension
+        if "time" in self.udims and self["time"].size > 1:
+            raise TypeError(
+                'Only 2-D coordinates have a GDAL transform. This array has a "time" dimension of {} > 1'.format(
+                    self["time"].size
+                )
+            )
+        if "alt" in self.udims and self["alt"].size > 1:
+            raise TypeError(
+                'Only 2-D coordinates have a GDAL transform. This array has a "alt" dimension of {} > 1'.format(
+                    self["alt"].size
+                )
+            )
+
+        # Do the uniform coordinates case
+        if (
+            "lat" in self.dims
+            and isinstance(self._coords["lat"], UniformCoordinates1d)
+            and "lon" in self.dims
+            and isinstance(self._coords["lon"], UniformCoordinates1d)
+        ):
+            if self.dims.index("lon") < self.dims.index("lat"):
+                first, second = "lat", "lon"
+            else:
+                first, second = "lon", "lat"  # This case will have the exact correct geotransform
+            transform = rasterio.transform.Affine.translation(
+                self[first].start - self[first].step / 2, self[second].start - self[second].step / 2
+            ) * rasterio.transform.Affine.scale(self[first].step, self[second].step)
+            transform = transform.to_gdal()
+        elif "lat,lon" in self.dims and isinstance(self._coords["lat,lon"], RotatedCoordinates):
+            transform = self._coords["lat,lon"].geotransform
+        elif "lon,lat" in self.dims and isinstance(self._coords["lon,lat"], RotatedCoordinates):
+            transform = self._coords["lon,lat"].geotransform
+        else:
+            raise TypeError(
+                "Only 2-D coordinates that are uniform or rotated have a GDAL transform. These coordinates "
+                "{} do not.".format(self)
+            )
+        if self.udims.index("lon") < self.udims.index("lat"):
+            # transform = (transform[3], transform[5], transform[4], transform[0], transform[2], transform[1])
+            transform = transform[3:] + transform[:3]
+
+        return transform
+
     # ------------------------------------------------------------------------------------------------------------------
     # Methods
     # ------------------------------------------------------------------------------------------------------------------

podpac/core/coordinates/rotated_coordinates.py

@@ -116,10 +116,11 @@ def _validate_step(self, d):
     @classmethod
     def from_geotransform(cls, geotransform, shape, dims=None, ctypes=None, segment_lengths=None):
         affine = rasterio.Affine.from_gdal(*geotransform)
-        origin = affine.c, affine.f
+        origin = affine.f, affine.c
         deg = affine.rotation_angle
         scale = ~affine.rotation(deg) * ~affine.translation(*origin) * affine
-        step = np.array([scale.a, scale.e])
+        step = np.array([scale.e, scale.a])
+        origin = affine.f + step[0] / 2, affine.c + step[1] / 2
         return cls(shape, np.deg2rad(deg), origin, step, dims=dims, ctypes=ctypes, segment_lengths=segment_lengths)
 
     @classmethod
@@ -237,8 +238,15 @@ def corner(self):
 
     @property
     def geotransform(self):
-        """ :tuple: GDAL geotransform. """
-        return self.affine.to_gdal()
+        """ :tuple: GDAL geotransform. 
+        Note: This property may not provide the correct order of lat/lon in the geotransform as this class does not
+        always have knowledge of the dimension order of the specified dataset. As such it always supplies
+        geotransforms assuming that dims = ['lat', 'lon']
+        """
+        t = rasterio.Affine.translation(self.origin[1] - self.step[1] / 2, self.origin[0] - self.step[0] / 2)
+        r = rasterio.Affine.rotation(self.deg)
+        s = rasterio.Affine.scale(*self.step[::-1])
+        return (t * r * s).to_gdal()
 
     @property
     def coordinates(self):

podpac/core/coordinates/test/test_coordinates.py

@@ -1618,3 +1618,115 @@ def test_concat_crs(self):
 
         with pytest.raises(ValueError, match="Cannot concat Coordinates"):
             concat([c1, c2])
+
+
+class TestCoordinatesGeoTransform(object):
+    def uniform_working(self):
+        # order: -lat, lon
+        c = Coordinates([clinspace(1.5, 0.5, 5, "lat"), clinspace(1, 2, 9, "lon")])
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf, np.array([[c["lon"].area_bounds[0], c["lon"].step, 0], [c["lat"].area_bounds[1], 0, c["lat"].step]])
+        )
+        # order: lon, lat
+        c = Coordinates([clinspace(0.5, 1.5, 5, "lon"), clinspace(1, 2, 9, "lat")])
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf, np.array([[c["lon"].area_bounds[0], 0, c["lon"].step], [c["lat"].area_bounds[0], c["lat"].step, 0]])
+        )
+
+        # order: lon, -lat, time
+        c = Coordinates([clinspace(0.5, 1.5, 5, "lon"), clinspace(2, 1, 9, "lat"), crange(10, 11, 2, "time")])
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf, np.array([[c["lon"].area_bounds[0], 0, c["lon"].step], [c["lat"].area_bounds[1], c["lat"].step, 0]])
+        )
+        # order: -lon, -lat, time, alt
+        c = Coordinates(
+            [
+                clinspace(1.5, 0.5, 5, "lon"),
+                clinspace(2, 1, 9, "lat"),
+                crange(10, 11, 2, "time"),
+                crange(10, 11, 2, "alt"),
+            ]
+        )
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf, np.array([[c["lon"].area_bounds[1], 0, c["lon"].step], [c["lat"].area_bounds[1], c["lat"].step, 0]])
+        )
+
+    def error_time_alt_too_big(self):
+        # time
+        c = Coordinates(
+            [
+                clinspace(1.5, 0.5, 5, "lon"),
+                clinspace(2, 1, 9, "lat"),
+                crange(1, 11, 2, "time"),
+                crange(1, 11, 2, "alt"),
+            ]
+        )
+        with pytest.raises(
+            TypeError, match='Only 2-D coordinates have a GDAL transform. This array has a "time" dimension of'
+        ):
+            c.geotransform
+        # alt
+        c = Coordinates([clinspace(1.5, 0.5, 5, "lon"), clinspace(2, 1, 9, "lat"), crange(1, 11, 2, "alt")])
+        with pytest.raises(
+            TypeError, match='Only 2-D coordinates have a GDAL transform. This array has a "alt" dimension of'
+        ):
+            c.geotransform
+
+    def rot_coords_working(self):
+        # order -lat, lon
+        rc = RotatedCoordinates(shape=(4, 3), theta=np.pi / 8, origin=[10, 20], step=[-2.0, 1.0], dims=["lat", "lon"])
+        c = Coordinates([rc], dims=["lat,lon"])
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf,
+            np.array(
+                [
+                    [rc.origin[1] - rc.step[1] / 2, rc.step[1] * np.cos(rc.theta), -rc.step[0] * np.sin(rc.theta)],
+                    [rc.origin[0] - rc.step[0] / 2, rc.step[1] * np.sin(rc.theta), rc.step[0] * np.cos(rc.theta)],
+                ]
+            ),
+        )
+        # order lon, lat
+        rc = RotatedCoordinates(shape=(4, 3), theta=np.pi / 8, origin=[10, 20], step=[2.0, 1.0], dims=["lon", "lat"])
+        c = Coordinates([rc], dims=["lon,lat"])
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf,
+            np.array(
+                [
+                    [rc.origin[0] - rc.step[0] / 2, rc.step[1] * np.sin(rc.theta), rc.step[0] * np.cos(rc.theta)],
+                    [rc.origin[1] - rc.step[1] / 2, rc.step[1] * np.cos(rc.theta), -rc.step[0] * np.sin(rc.theta)],
+                ]
+            ),
+        )
+
+        # order -lon, lat
+        rc = RotatedCoordinates(shape=(4, 3), theta=np.pi / 8, origin=[10, 20], step=[-2.0, 1.0], dims=["lon", "lat"])
+        c = Coordinates([rc], dims=["lon,lat"])
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf,
+            np.array(
+                [
+                    [rc.origin[0] - rc.step[0] / 2, rc.step[1] * np.sin(rc.theta), rc.step[0] * np.cos(rc.theta)],
+                    [rc.origin[1] - rc.step[1] / 2, rc.step[1] * np.cos(rc.theta), -rc.step[0] * np.sin(rc.theta)],
+                ]
+            ),
+        )
+        # order -lat, -lon
+        rc = RotatedCoordinates(shape=(4, 3), theta=np.pi / 8, origin=[10, 20], step=[-2.0, -1.0], dims=["lat", "lon"])
+        c = Coordinates([rc], dims=["lat,lon"])
+        tf = np.array(c.geotransform).reshape(2, 3)
+        np.testing.assert_almost_equal(
+            tf,
+            np.array(
+                [
+                    [rc.origin[1] - rc.step[1] / 2, rc.step[1] * np.cos(rc.theta), -rc.step[0] * np.sin(rc.theta)],
+                    [rc.origin[0] - rc.step[0] / 2, rc.step[1] * np.sin(rc.theta), rc.step[0] * np.cos(rc.theta)],
+                ]
+            ),
+        )

podpac/core/coordinates/test/test_rotated_coordinates.py

@@ -129,11 +129,17 @@ def test_copy(self):
 class TestRotatedCoordinatesGeotransform(object):
     def test_geotransform(self):
         c = RotatedCoordinates(shape=(3, 4), theta=np.pi / 4, origin=[10, 20], step=[1.0, 2.0], dims=["lat", "lon"])
-        assert_allclose(c.geotransform, (10.0, 0.7071068, -1.4142136, 20.0, 0.7071068, 1.4142136))
+        assert_allclose(c.geotransform, (19.0, 1.4142136, -0.7071068, 9.5, 1.4142136, 0.7071068))
 
         c2 = RotatedCoordinates.from_geotransform(c.geotransform, c.shape, dims=["lat", "lon"])
         assert c == c2
 
+        c = RotatedCoordinates(shape=(3, 4), theta=np.pi / 4, origin=[10, 20], step=[1.0, 2.0], dims=["lon", "lat"])
+        assert_allclose(c.geotransform, (19.0, 1.4142136, -0.7071068, 9.5, 1.4142136, 0.7071068))
+
+        c2 = RotatedCoordinates.from_geotransform(c.geotransform, c.shape, dims=["lon", "lat"])
+        assert c == c2
+
 
 class TestRotatedCoordinatesStandardMethods(object):
     def test_eq_type(self):

podpac/core/data/file.py

@@ -12,11 +12,14 @@
 import traitlets as tl
 import pandas as pd
 import xarray as xr
+import pyproj
+import logging
 
 from podpac.core.settings import settings
 from podpac.core.utils import common_doc, trait_is_defined
 from podpac.core.data.datasource import COMMON_DATA_DOC, DataSource
 from podpac.core.coordinates import Coordinates, UniformCoordinates1d, ArrayCoordinates1d, StackedCoordinates
+from podpac.core.coordinates import RotatedCoordinates
 from podpac.core.coordinates.utils import Dimension, VALID_DIMENSION_NAMES
 
 # Optional dependencies
@@ -30,6 +33,9 @@
 zarrGroup = lazy_class("zarr.Group")
 s3fs = lazy_module("s3fs")
 
+# Set up logging
+_logger = logging.getLogger(__name__)
+
 
 @common_doc(COMMON_DATA_DOC)
 class DatasetSource(DataSource):
@@ -676,8 +682,22 @@ class Rasterio(DataSource):
     source = tl.Union([tl.Unicode(), tl.Instance(BytesIO)]).tag(readonly=True)
     dataset = tl.Any().tag(readonly=True)
 
+    @property
+    def nan_vals(self):
+        return list(self.dataset.nodatavals)
+
     # node attrs
-    band = tl.CInt(1).tag(attr=True)
+    band = tl.CInt(allow_none=True).tag(attr=True)
+
+    @tl.default("band")
+    def _band_default(self):
+        if (self.outputs is not None) and (self.output is not None):
+            band = self.outputs.index(self.output)
+        elif self.outputs is None:
+            band = 1
+        else:
+            band = None  # All bands
+        return band
 
     @tl.default("dataset")
     def _open_dataset(self):
@@ -719,21 +739,18 @@ def get_native_coordinates(self):
 
         # check to see if the coordinates are rotated used affine
         affine = self.dataset.transform
-        if affine[1] != 0.0 or affine[3] != 0.0:
-            raise NotImplementedError("Rotated coordinates are not yet supported")
 
-        try:
+        if isinstance(self.dataset.crs, rasterio.crs.CRS):
+            crs = self.dataset.crs.wkt
+        elif isinstance(self.dataset.crs, dict) and "init" in self.dataset.crs:
             crs = self.dataset.crs["init"].upper()
-        except:
-            crs = None
-
-        # get bounds
-        left, bottom, right, top = self.dataset.bounds
+        else:
+            try:
+                crs = pyproj.CRS(self.dataset.crs).to_wkt()
+            except:
+                raise RuntimeError("Unexpected rasterio crs '%s'" % self.dataset.crs)
 
-        # rasterio reads data upside-down from coordinate conventions, so lat goes from top to bottom
-        lat = UniformCoordinates1d(top, bottom, size=self.dataset.height, name="lat")
-        lon = UniformCoordinates1d(left, right, size=self.dataset.width, name="lon")
-        return Coordinates([lat, lon], dims=["lat", "lon"], crs=crs)
+        return Coordinates.from_geotransform(affine.to_gdal(), self.dataset.shape, crs)
 
     @common_doc(COMMON_DATA_DOC)
     def get_data(self, coordinates, coordinates_index):
@@ -744,7 +761,12 @@ def get_data(self, coordinates, coordinates_index):
 
         # read data within coordinates_index window
         window = ((slc[0].start, slc[0].stop), (slc[1].start, slc[1].stop))
-        raster_data = self.dataset.read(self.band, out_shape=tuple(coordinates.shape), window=window)
+
+        if self.outputs is not None:  # read all the bands
+            raster_data = self.dataset.read(out_shape=(len(self.outputs),) + tuple(coordinates.shape), window=window)
+            raster_data = np.moveaxis(raster_data, 0, 2)
+        else:  # read the requested band
+            raster_data = self.dataset.read(self.band, out_shape=tuple(coordinates.shape), window=window)
 
         # set raster data to output array
         data.data.ravel()[:] = raster_data.ravel()