[CIF-165] Write layers directly

city_metrix/layers/layer.py

@@ -1,7 +1,11 @@
+import os
 from abc import abstractmethod
-from typing import Union, Tuple, List
+from typing import Union, Tuple
+from uuid import uuid4
+from osgeo import gdal
 
 import ee
+import boto3
 from dask.diagnostics import ProgressBar
 from ee import ImageCollection
 from geocube.api.core import make_geocube
@@ -52,6 +56,37 @@ def groupby(self, zones, layer=None):
         """
         return LayerGroupBy(self.aggregate, zones, layer, self.masks)
 
+    def write(self, bbox, output_path, tile_degrees=None):
+        """
+        Write the layer to a path. Does not apply masks.
+
+        :param bbox: (min x, min y, max x, max y)
+        :param output_path: local or s3 path to output to
+        :param tile_degrees: optional param to tile the results into multiple files with a VRT.
+            Degrees to tile by. `output_path` should be a folder path to store the tiles.
+        :return:
+        """
+
+        if tile_degrees is not None:
+            tiles = create_fishnet_grid(*bbox, tile_degrees)
+
+            if not os.path.exists(output_path):
+                os.makedirs(output_path)
+
+            file_names = []
+            for tile in tiles["geometry"]:
+                data = self.aggregate.get_data(tile.bounds)
+
+                file_name = f"{output_path}/{uuid4()}.tif"
+                file_names.append(file_name)
+
+                write_layer(file_name, data)
+
+            gdal.BuildVRT(f"{output_path}.vrt", file_names)
+        else:
+            data = self.aggregate.get_data(bbox)
+            write_layer(output_path, data)
+
 
 class LayerGroupBy:
     def __init__(self, aggregate, zones, layer=None, masks=[]):
@@ -240,3 +275,21 @@ def get_image_collection(
 
     return data
 
+
+def write_layer(path, data):
+    if isinstance(data, xr.DataArray):
+        # for rasters, need to write to locally first then copy to cloud storage
+        if path.startswith("s3://"):
+            tmp_path = f"{uuid4()}.tif"
+            data.rio.to_raster(raster_path=tmp_path, driver="COG")
+
+            s3 = boto3.client('s3')
+            s3.upload_file(tmp_path, path.split('/')[2], '/'.join(path.split('/')[3:]))
+
+            os.remove(tmp_path)
+        else:
+            data.rio.to_raster(raster_path=path, driver="COG")
+    elif isinstance(data, gpd.GeoDataFrame):
+        data.to_file(path, driver="GeoJSON")
+    else:
+        raise NotImplementedError("Can only write DataArray or GeoDataFrame")

environment.yml

@@ -17,6 +17,7 @@ dependencies:
   - dask[complete]=2023.11.0
   - matplotlib=3.8.2
   - jupyterlab=4.0.10
+  - s3fs=2024.5.0
   - pip=23.3.1
   - pip:
     - cartoframes==1.2.5