Merge branch 'main' of https://github.com/geo-engine/geoengine-python into pydantic-2

Author: ChristianBeilschmidt

examples/copernicus_dataspace.ipynb …s/data-spaces/copernicus_dataspace.ipynbexamples/copernicus_dataspace.ipynb renamed to examples/data-spaces/copernicus_dataspace.ipynb examples/copernicus_dataspace.ipynb renamed to examples/data-spaces/copernicus_dataspace.ipynb

@@ -29,16 +29,33 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Connect to the Virtual Data Trustee"
+    "## Note\n",
+    "\n",
+    "1. It is important that the Geo Engine is running with\n",
+    "```toml\n",
+    "[quota]\n",
+    "mode = \"track\"\n",
+    "```\n",
+    "\n",
+    "2. You have to do some work beforehand to get the data from the log.\n",
+    "\n",
+    "This is a simple example to show how to access the log."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Connect to Geo Engine"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "ge.initialize(\"http://localhost:3030/api\", (\"admin@localhost\", \"adminadmin\")) # TODO: load from .env for demo TODO: set remote url"
+    "ge.initialize(\"http://localhost:3030/api\", (\"admin@localhost\", \"adminadmin\"))"
    ]
   },
   {
@@ -296,7 +313,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -312,6 +329,9 @@
    ],
    "source": [
     "df = ge.data_usage_summary(ge.UsageSummaryGranularity.MINUTES)\n",
+    "if df.empty:\n",
+    "    print('No data usage found')\n",
+    "    exit()\n",
     "df['timestamp'] = pd.to_datetime(df['timestamp']).dt.tz_localize(None)\n",
     "\n",
     "pivot_df = df.pivot(index='timestamp', columns='data', values='count').fillna(0)\n",
@@ -324,13 +344,6 @@
     "plt.legend(title='Data')\n",
     "plt.show()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

examples/data_usage.ipynb

@@ -10,11 +10,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "from datetime import datetime\n",
+    "from IPython.display import display\n",
     "\n",
     "import geoengine as ge"
    ]

examples/plots.ipynb

@@ -213,6 +213,12 @@ def spatial_partition(self) -> gety.SpatialPartition2D:
     def spatial_resolution(self) -> gety.SpatialResolution:
         return self.geo_transform.spatial_resolution()
 
+    def is_empty(self) -> bool:
+        ''' Returns true if the tile is empty'''
+        num_pixels = self.size_x * self.size_y
+        num_nulls = self.data.null_count
+        return num_pixels == num_nulls
+
     @staticmethod
     def from_ge_record_batch(record_batch: pa.RecordBatch) -> RasterTile2D:
         '''Create a RasterTile2D from an Arrow record batch recieved from the Geo Engine'''

geoengine/raster.py

@@ -168,20 +168,37 @@ def __create_new_dataset(self, query: QueryRectangle):
             **self.rio_kwargs
         )
 
+        for i, b in enumerate(self.bands, start=1):
+            b_n = b.name
+            b_m = str(b.measurement)
+            rio_dataset.update_tags(i, band_name=b_n, band_measurement=b_m)
+
         self.current_dataset = rio_dataset
 
-    async def query_and_write(self, query: QueryRectangle):
-        ''' Query the raster workflow and write the tiles to the dataset.'''
+    async def query_and_write(self, query: QueryRectangle, skip_empty_times=True):
+        '''
+        Query the raster workflow and write the resulting tiles to a GDAL dataset per timeslice.
+
+        :param query: The QueryRectangle to write to GDAL dataset(s)
+        :param skip_empty_times: Skip timeslices where all pixels are empty/nodata
+        '''
 
         self.create_tiling_geo_transform_width_height(query)
 
-        assert self.workflow is not None, "The workflow must be set"
+        assert self.bands is not None, "The bands must be set"
+        bands = list(range(0, len(self.bands)))
 
+        assert self.workflow is not None, "The workflow must be set"
         try:
-            async for tile in self.workflow.raster_stream(query):
+            async for tile in self.workflow.raster_stream(query, bands=bands):
                 if self.current_time != tile.time:
                     self.close_current_dataset()
                     self.current_time = tile.time
+
+                if tile.is_empty() and skip_empty_times:
+                    continue
+
+                if self.current_dataset is None:
                     self.__create_new_dataset(query)
 
                 assert self.current_time == tile.time, "The time of the current dataset does not match the tile"

geoengine/raster_workflow_rio_writer.py

@@ -7,13 +7,6 @@
 import shutil
 import sys
 
-# TODO: fix the blacklisted notebooks
-BLACKLIST = [
-    'copernicus_dataspace.ipynb',
-    'data_usage.ipynb',
-    'plots.ipynb',
-]
-
 
 def eprint(*args, **kwargs):
     '''Print to stderr.'''
@@ -57,9 +50,6 @@ def main() -> int:
             if not file.endswith('.ipynb'):
                 eprint(f"Skipping non-notebook file {file}")
                 continue
-            if file in BLACKLIST:
-                eprint(f"Skipping blacklisted notebook {file}")
-                continue
             notebook_path = os.path.join(root, file)
             if not run_test_notebook(notebook_path):
                 return -1