From cf861ed9cdacf2256751c2ac5cd85e576fd17c93 Mon Sep 17 00:00:00 2001
From: atmorling <atmorling@gmail.com>
Date: Fri, 30 Aug 2024 21:16:19 +0200
Subject: [PATCH] Allow `io.grid_to_raster` to return a `rio.MemoryFile` (#257)

---
 .../04. EcoMap & EcoPlot/EcoMap.ipynb         |  2 +-
 ecoscope/io/raster.py                         | 27 ++++++++++++++-----
 ecoscope/mapping/map.py                       |  9 ++++---
 tests/test_ecomap.py                          | 23 ++++++++++++++++
 4 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/doc/source/notebooks/04. EcoMap & EcoPlot/EcoMap.ipynb b/doc/source/notebooks/04. EcoMap & EcoPlot/EcoMap.ipynb
index 98fc3e60..b6d7a4b2 100644
--- a/doc/source/notebooks/04. EcoMap & EcoPlot/EcoMap.ipynb	
+++ b/doc/source/notebooks/04. EcoMap & EcoPlot/EcoMap.ipynb	
@@ -332,7 +332,7 @@
     "\n",
     "m = EcoMap(width=800, height=600)\n",
     "m.add_layer(EcoMap.get_named_tile_layer(\"OpenStreetMap\"))\n",
-    "m.add_geotiff(path=os.path.join(output_dir, \"mara_dem.tif\"), zoom=True, cmap=\"jet\")\n",
+    "m.add_geotiff(tiff=os.path.join(output_dir, \"mara_dem.tif\"), zoom=True, cmap=\"jet\")\n",
     "m"
    ]
   },
diff --git a/ecoscope/io/raster.py b/ecoscope/io/raster.py
index 0f5a0722..4a0ff064 100644
--- a/ecoscope/io/raster.py
+++ b/ecoscope/io/raster.py
@@ -208,7 +208,7 @@ def raster_to_gdf(raster_path):
         )
 
 
-def grid_to_raster(grid=None, val_column="", out_dir="", raster_name="grid.tif", xlen=5000, ylen=5000, grid_crs=4326):
+def grid_to_raster(grid=None, val_column="", out_dir="", raster_name=None, xlen=5000, ylen=5000):
     """
     Save a GeoDataFrame grid to a raster.
     """
@@ -223,7 +223,7 @@ def grid_to_raster(grid=None, val_column="", out_dir="", raster_name="grid.tif",
 
     raster_profile = ecoscope.io.raster.RasterProfile(
         pixel_size=xlen,
-        crs=grid_crs,
+        crs=grid.crs.to_epsg(),
         nodata_value=np.nan,
         band_count=1,
     )
@@ -231,8 +231,21 @@ def grid_to_raster(grid=None, val_column="", out_dir="", raster_name="grid.tif",
     raster_profile.raster_extent = ecoscope.io.raster.RasterExtent(
         x_min=bounds[0], x_max=bounds[2], y_min=bounds[1], y_max=bounds[3]
     )
-    ecoscope.io.raster.RasterPy.write(
-        ndarray=vals,
-        fp=os.path.join(out_dir, raster_name),
-        **raster_profile,
-    )
+
+    if raster_name:
+        ecoscope.io.raster.RasterPy.write(
+            ndarray=vals,
+            fp=os.path.join(out_dir, raster_name),
+            **raster_profile,
+        )
+    else:
+        memfile = rio.MemoryFile()
+        memfile.open(
+            driver="GTiff",
+            width=raster_profile.pop("columns"),
+            height=raster_profile.pop("rows"),
+            count=raster_profile.pop("band_count"),
+            **raster_profile,
+        ).write(vals, 1)
+
+        return memfile
diff --git a/ecoscope/mapping/map.py b/ecoscope/mapping/map.py
index 469ee3f9..54ab2321 100644
--- a/ecoscope/mapping/map.py
+++ b/ecoscope/mapping/map.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 import pandas as pd
+import rasterio as rio
 from io import BytesIO
 from typing import Dict, IO, List, Optional, TextIO, Union
 from pathlib import Path
@@ -368,7 +369,7 @@ def zoom_to_bounds(self, feat: Union[BaseLayer, List[BaseLayer], gpd.GeoDataFram
 
     def add_geotiff(
         self,
-        path: str,
+        tiff: str | rio.MemoryFile,
         zoom: bool = False,
         cmap: Union[str, mpl.colors.Colormap] = None,
         opacity: float = 0.7,
@@ -380,8 +381,8 @@ def add_geotiff(
 
         Parameters
         ----------
-        path : str
-            The path to the local tiff
+        tiff : str | rio.MemoryFile
+            The string path to a tiff on disk or a rio.MemoryFile
         zoom : bool
             Whether to zoom the map to the bounds of the tiff
         cmap: str or matplotlib.colors.Colormap
@@ -389,7 +390,7 @@ def add_geotiff(
         opacity: float
             The opacity of the overlay
         """
-        with rasterio.open(path) as src:
+        with rasterio.open(tiff) as src:
             transform, width, height = rasterio.warp.calculate_default_transform(
                 src.crs, "EPSG:4326", src.width, src.height, *src.bounds
             )
diff --git a/tests/test_ecomap.py b/tests/test_ecomap.py
index e38ca5b6..e30e1ecc 100644
--- a/tests/test_ecomap.py
+++ b/tests/test_ecomap.py
@@ -1,3 +1,5 @@
+import os
+import random
 import ee
 import geopandas as gpd
 import pandas as pd
@@ -215,6 +217,27 @@ def test_add_geotiff_with_cmap():
     assert isinstance(m.layers[1], BitmapLayer)
 
 
+def test_add_geotiff_in_mem_with_cmap():
+    AOI = gpd.read_file(os.path.join("tests/sample_data/vector", "maec_4zones_UTM36S.gpkg"))
+
+    grid = gpd.GeoDataFrame(
+        geometry=ecoscope.base.utils.create_meshgrid(
+            AOI.unary_union, in_crs=AOI.crs, out_crs=AOI.crs, xlen=5000, ylen=5000, return_intersecting_only=False
+        )
+    )
+
+    grid["fake_density"] = grid.apply(lambda _: random.randint(1, 50), axis=1)
+    raster = ecoscope.io.raster.grid_to_raster(
+        grid,
+        val_column="fake_density",
+    )
+
+    m = EcoMap()
+    m.add_geotiff(raster, cmap="jet")
+    assert len(m.layers) == 2
+    assert isinstance(m.layers[1], BitmapLayer)
+
+
 def test_add_datashader_gdf(point_gdf):
     m = EcoMap()
     img, bounds = datashade_gdf(point_gdf, "point")