terrakio-core 0.5.3__tar.gz → 0.5.4__tar.gz

{terrakio_core-0.5.3 → terrakio_core-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: terrakio-core
-Version: 0.5.3
+Version: 0.5.4
 Summary: Core package for the terrakio-python-api
 Requires-Python: >=3.11
 Requires-Dist: aiofiles>=24.1.0

{terrakio_core-0.5.3 → terrakio_core-0.5.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "terrakio-core"
-version = "0.5.3"
+version = "0.5.4"
 description = "Core package for the terrakio-python-api"
 readme = "README.md"
 requires-python = ">=3.11"

{terrakio_core-0.5.3 → terrakio_core-0.5.4}/terrakio_core/__init__.py

@@ -8,7 +8,7 @@ Core components for Terrakio API clients.
 # Suppress ONNX Runtime GPU device discovery warnings - MUST BE FIRST!
 import os
 os.environ['ORT_LOGGING_LEVEL'] = '3'
-__version__ = "0.5.3"
+__version__ = "0.5.4"
 
 from .async_client import AsyncClient
 from .sync_client import SyncClient as Client

{terrakio_core-0.5.3 → terrakio_core-0.5.4}/terrakio_core/async_client.py

@@ -185,7 +185,7 @@ class AsyncClient(BaseClient):
         result, status_code = await self._terrakio_request("POST", "geoquery", json=payload)
 
         if status_code != 200:
-            raise GeoQueryError("Something went wrong with the geoquery request. Please try again later.", status_code=status_code)
+            raise GeoQueryError(result['detail'], status_code=status_code)
 
         return result
 

{terrakio_core-0.5.3 → terrakio_core-0.5.4}/terrakio_core/endpoints/mass_stats.py

@@ -1,13 +1,23 @@
 import json
 import os
 import time
-from typing import Dict, Any, Optional, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
+import dateutil.parser
+import geopandas as gpd
+import pyproj
+import shapely.geometry
 import typer
 from dateutil import parser
 from rich.console import Console
-from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeElapsedColumn
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
 
 from ..exceptions import (
     CancelAllTasksError,
@@ -26,15 +36,58 @@ from ..exceptions import (
     TaskNotFoundError,
     UploadArtifactsError,
     UploadRequestsError,
+    QuotaInsufficientError,
 )
 from ..helper.decorators import require_api_key
 
+import enum
+
+class OutputTypes(enum.Enum):
+    geotiff = 'geotiff'
+    png = 'png'
+    netcdf = 'netcdf'
+    json = 'json'
+    json_v2 = 'json_v2'
+    csv = 'csv'
+
+class Region(str, enum.Enum):
+    aus = "aus"
+    eu = "eu"
+    us = "us"
+
+regions = {
+    Region.aus : {
+        "name" : "australia-southeast1", 
+        "url" : "https://terrakio-server-candidate-573248941006.australia-southeast1.run.app", 
+        "bucket" : "terrakio-mass-requests"
+    },
+    
+    Region.eu : {
+        "name" : "europe-west4", 
+        "url" : "https://terrakio-server-candidate-573248941006.europe-west4.run.app", 
+        "bucket" : "terrakio-mass-requests-eu"
+    },
+    
+    Region.us : {
+        "name" : "us-central1", 
+        "url" : "https://terrakio-server-candidate-573248941006.us-central1.run.app", 
+        "bucket" : "terrakio-mass-requests-us"
+    },
+
+}
+
+class Dataset_Dtype(enum.Enum):
+    uint8 = 'uint8'
+    float32 = 'float32'
 
 class MassStats:
     def __init__(self, client):
         self._client = client
         self.console = Console()
-
+        self.OutputTypes = OutputTypes
+        self.Region = Region
+        self.regions = regions
+        self.Dataset_Dtype = Dataset_Dtype
     async def track_progress(self, task_id):
         task_info = await self.get_task(task_id=task_id)
         number_of_jobs = task_info["task"]["total"]
@@ -156,7 +209,6 @@ class MassStats:
         
         self.console.print(f"[bold green]All {number_of_jobs} jobs finished![/bold green]")
 
-    # below are functions related to collection
     @require_api_key
     async def create_collection(
         self, 
@@ -193,9 +245,8 @@ class MassStats:
             payload["location"] = location
         
         response, status = await self._client._terrakio_request("POST", f"collections/{collection}", json=payload)
-
         if status != 200:
-            if status == 400:
+            if status == 400 or status == 409:
                 raise CollectionAlreadyExistsError(f"Collection {collection} already exists", status_code=status)
             if status == 422:
                 raise InvalidCollectionTypeError(f"Invalid collection type: {collection_type}", status_code=status)
@@ -289,13 +340,15 @@ class MassStats:
             CollectionNotFoundError: If the collection is not found
             DeleteCollectionError: If the API request fails due to unknown reasons
         """
-        payload = {
-            "full": full,
-            "outputs": outputs,
-            "data": data
+        params = {
+            "full": str(full).lower(),
+            "data": str(data).lower()
         }
         
-        response, status = await self._client._terrakio_request("DELETE", f"collections/{collection}", json=payload)
+        if outputs:
+            params["outputs"] = outputs
+        
+        response, status = await self._client._terrakio_request("DELETE", f"collections/{collection}", params=params)
 
         if status != 200:
             if status == 404:
@@ -564,7 +617,7 @@ class MassStats:
         
         return response
 
-    async def _upload_requests(
+    async def _get_upload_url(
         self,
         collection: str
     ) -> Dict[str, Any]:
@@ -641,6 +694,8 @@ class MassStats:
         response = await self._client._regular_request("PUT", url, data=body, headers=headers)
         return response
 
+
+
     @require_api_key
     async def generate_data(
         self,
@@ -672,7 +727,7 @@ class MassStats:
         
         await self.get_collection(collection = collection)
 
-        upload_urls = await self._upload_requests(
+        upload_urls = await self._get_upload_url(
             collection = collection
         )
         
@@ -696,6 +751,7 @@ class MassStats:
         
         return response
 
+
     @require_api_key
     async def post_processing(
         self,
@@ -1017,7 +1073,7 @@ class MassStats:
         """
         await self.get_collection(collection = collection)
 
-        upload_urls = await self._upload_requests(collection=collection)
+        upload_urls = await self._get_upload_url(collection=collection)
         url = upload_urls['url']
         
         # Handle requests_file - either file path (str) or file object
@@ -1124,4 +1180,337 @@ class MassStats:
         task_id = response["task_id"]
         await self.track_progress(task_id)
 
-        return {"task_id": task_id}
+        return {"task_id": task_id}
+
+    def tile_generator(self, x_min, y_min, x_max, y_max, aoi, crs, res, tile_size, expression, output = OutputTypes.netcdf, fully_cover=True):
+        
+        i_max = int((x_max-x_min)/(tile_size*res))
+        j_max = int((y_max-y_min)/(tile_size*res))
+        if fully_cover:
+            i_max += 1
+            j_max += 1
+        for j in range(0, int(j_max)):
+            for i in range(0, int(i_max)):
+                x = x_min + i*(tile_size*res)
+                y = y_max - j*(tile_size*res)
+                bbox = shapely.geometry.box(x, y-(tile_size*res), x + (tile_size*res), y)
+                if not aoi.geometry[0].intersects(bbox):
+                    continue
+                feat  = {"type": "Feature", "geometry": bbox.__geo_interface__}
+                data = {
+                    "feature": feat,
+                    "in_crs": crs,
+                    "out_crs": crs,
+                    "resolution": res,
+                    "expr" : expression,
+                    "output" : output.value,
+                }
+                yield data, i , j
+
+
+    def get_bounds(self, aoi, crs, to_crs = None):
+        gdf : gpd.GeoDataFrame = gpd.read_file(aoi)
+        gdf = gdf.set_crs(crs, allow_override=True)
+        if to_crs:
+            gdf = gdf.to_crs(to_crs)
+        bounds = gdf.geometry[0].bounds
+        return *bounds, gdf
+
+    def validate_date(self, date: str) -> str:
+        try:
+            date = dateutil.parser.parse(date)
+            return date
+        except ValueError:
+            print(f"Invalid date: {date}")
+            raise typer.BadParameter(f"Invalid date: {date}")
+
+
+    @require_api_key
+    async def dataset(
+        self,
+        products: List[str],
+        name: str,
+        bucket: str = "terrakio",
+        location: str = "testing/MSWXsmall",
+        aoi: Optional[str] = None,
+        expression: Optional[str] = None,
+        date: Optional[str] = "2021-01-01",
+        tile_size: float = 100,
+        crs: str = "epsg:4326",
+        res: float = 10,
+        out_res: float = 10,
+        no_data: float = -9999,
+        dtype: str = "float32",
+        create_doc: bool = False,
+        skip_test: bool = False,
+        force_res: bool = False,
+        to_crs: Optional[str] = None,
+        fully_cover: bool = True,
+        skip_existing: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Generate a dataset with the specified parameters.
+
+        Args:
+            products: List of product names
+            name: Name of the dataset
+            bucket: Storage bucket
+            location: Storage location
+            aoi: Path to GeoJSON file containing area of interest
+            expression: Expression for data processing
+            date: Date in YYYY-MM-DD format
+            tile_size: Size of tiles (default: 100)
+            crs: Coordinate reference system (default: "epsg:4326")
+            res: Resolution (default: 10)
+            out_res: Output resolution (default: 10)
+            no_data: No data value (default: -9999)
+            dtype: Data type (default: "float32")
+            create_doc: Add dataset to the DB (default: False)
+            skip_test: Skip testing the expression (default: False)
+            force_res: Force resolution in case requests are too large (default: False)
+            to_crs: Target coordinate reference system
+            fully_cover: Fully cover the area (default: True)
+            skip_existing: Skip existing data (default: False)
+
+        Returns:
+            Response containing task_id and collection name
+
+        Raises:
+            CollectionNotFoundError: If the collection is not found
+            GetTaskError: If the API request fails
+        """
+        await self.create_collection(collection = name, bucket = bucket, location = location)
+        date = self.validate_date(date)
+        sample = None
+        reqs = []
+        x_min, y_min, x_max, y_max, aoi = self.get_bounds(aoi, crs, to_crs)
+        if to_crs is None:
+            to_crs = crs
+        c=0
+        for tile_req, i, j in self.tile_generator(x_min, y_min, x_max, y_max, aoi, to_crs, res, tile_size, expression, fully_cover = fully_cover):
+            c+=1
+            if force_res:
+                tile_req["force_res"] = True
+            req_names = []
+            for product in products:
+                req_names.append(f"{product}_{date.strftime('%Y%m%d')}000000_{i:03d}_{j:03d}_00")
+            reqs.append({"group": name, "file": req_names, "request": tile_req})
+            if sample is None:
+                sample = tile_req["expr"]
+        i_max = int((x_max-x_min)/(tile_size*res))
+        j_max = int((y_max-y_min)/(tile_size*res))
+        geot = [x_min, out_res, 0, y_max, 0, -out_res]
+        if not skip_test:
+            result = await self._client.geoquery(**reqs[0]["request"], debug = "requests")
+            request_count = result.get('request_count', 0)
+
+        user_quota = await self._client.auth.get_user_quota()
+        user_quota = user_quota.get('quota', -9999)
+        
+        if user_quota !=-9999 and user_quota < len(reqs) * request_count:
+            raise QuotaInsufficientError(f"User quota is insufficient. Please contact support to increase your quota.")
+
+        upload_urls = await self._get_upload_url(collection=name)
+        url = upload_urls['url']
+        await self._upload_json_data(reqs, url, use_gzip=True)
+
+        payload = {"output": "snp", "skip_existing": skip_existing}
+        
+        task_id, status = await self._client._terrakio_request("POST", f"collections/{name}/generate_data", json=payload)
+        task_id = task_id["task_id"]
+        if dtype == self.Dataset_Dtype.uint8.value:
+            no_data = int(no_data)
+        if create_doc:
+            await self._client.datasets.create_dataset(
+                name=name,
+                products=products,
+                dates_iso8601=[date.isoformat()],
+                proj4=pyproj.CRS.to_proj4(aoi.crs),
+                i_max=i_max,
+                j_max=j_max,
+                x_size=int((res*tile_size)/out_res),
+                y_size=int((res*tile_size)/out_res),
+                geotransform=geot,
+                no_data=no_data,
+                data_type=dtype,
+                bucket=bucket,
+                path=f"{location}/%s_%s_%03d_%03d_%02d.snp",
+                max_zoom=0,
+            )
+
+        await self.track_progress(task_id)
+
+
+    @require_api_key
+    async def tiles(
+        self,
+        collection: str,
+        name: str = "irrigation_2019",
+        aoi: Optional[str] = None,
+        expression: str = "NSWIrrigation.landuse@(year=2019)",
+        output: OutputTypes = OutputTypes.netcdf,
+        tile_size: float = 10000,
+        crs: str = "epsg:3577",
+        res: float = 10,
+        skip_test: bool = False,
+        force_res: bool = False,
+        to_crs: Optional[str] = None,
+        fully_cover: bool = True,
+        skip_existing: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Generate tiles with the specified parameters.
+
+        Args:
+            collection: Name of the collection
+            name: Name of the dataset (default: "irrigation_2019")
+            aoi: Path to GeoJSON file containing area of interest
+            expression: Expression for data processing (default: "NSWIrrigation.landuse@(year=2019)")
+            output: Output format (default: "netcdf")
+            tile_size: Size of tiles (default: 10000)
+            crs: Coordinate reference system (default: "epsg:3577")
+            res: Resolution (default: 10)
+            skip_test: Skip testing the expression (default: False)
+            force_res: Force resolution in case requests are too large (default: False)
+            to_crs: Target coordinate reference system
+            fully_cover: Fully cover the area (default: True)
+            skip_existing: Skip existing data (default: False)
+
+        Returns:
+            Response containing task_id
+
+        Raises:
+            CollectionNotFoundError: If the collection is not found
+            GetTaskError: If the API request fails
+        """
+        
+
+        await self.get_collection(collection=collection)
+
+        reqs = []
+        sample = None
+        x_min, y_min, x_max, y_max, aoi = self.get_bounds(aoi, crs, to_crs)
+        if to_crs is None:
+            to_crs = crs
+        for tile_req, i, j in self.tile_generator(x_min, y_min, x_max, y_max, aoi, to_crs, res, tile_size, expression, output, fully_cover):
+            if force_res:  
+                tile_req["force_res"] = True
+            req_name = f"{name}_{i:02d}_{j:02d}"
+            reqs.append({"group": "tiles", "file": req_name, "request": tile_req})
+            if sample is None:
+                sample = tile_req["expr"]
+
+        if not skip_test:
+            print("the reqs are ", reqs[0])
+            result = await self._client.geoquery(**reqs[0]["request"], debug = "requests")
+            request_count = result.get('request_count', 0)
+
+        user_quota = await self._client.auth.get_user_quota()
+        user_quota = user_quota.get('quota', -9999)
+        
+        if user_quota !=-9999 and user_quota < len(reqs) * request_count:
+            raise QuotaInsufficientError(f"User quota is insufficient. Please contact support to increase your quota.")
+
+        count = len(reqs)
+        groups = list(set(dic["group"] for dic in reqs))
+        print(f"[green]{count}[/green] requests with [blue]{len(groups)}[/blue] groups identified.")
+        upload_urls = await self._get_upload_url(collection=collection)
+        url = upload_urls['url']
+        await self._upload_json_data(reqs, url, use_gzip=True)
+
+        payload = {"output": output.value, "skip_existing": skip_existing}
+        
+        task_id, status = await self._client._terrakio_request("POST", f"collections/{collection}/generate_data", json=payload)
+        task_id = task_id["task_id"]
+
+        await self.track_progress(task_id)
+
+    @require_api_key
+    async def polygons(
+        self,
+        collection: str,
+        aoi: str,
+        expression: str = "mean:space(MSWX.air_temperature@(year=2022))",
+        output: OutputTypes = OutputTypes.netcdf,
+        id_field: str = "GID_0",
+        crs: str = "epsg:4326",
+        res: float = -1,
+        skip_test: bool = False,
+        skip_existing: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Generate mass-stats for polygons in a GeoJSON file using the same expression.
+
+        Args:
+            collection: Name of the collection
+            name: Name of the dataset
+            aoi: Path to GeoJSON file containing area of interest
+            expression: Expression for data processing (default: "mean:space(MSWX.air_temperature@(year=2022))")
+            output: Output format (default: "netcdf")
+            id_field: Field name to use as identifier (default: "GID_0")
+            crs: Coordinate reference system (default: "epsg:4326")
+            res: Resolution (default: -1)
+            skip_test: Skip testing the expression (default: False)
+            skip_existing: Skip existing data (default: False)
+
+        Returns:
+            Response containing task_id
+
+        Raises:
+            CollectionNotFoundError: If the collection is not found
+            GetTaskError: If the API request fails
+            ValueError: If id_field not found in feature properties
+        """
+        await self.get_collection(collection=collection)
+
+        gdf = gpd.read_file(aoi)
+        sample = None
+        features = gdf.__geo_interface__
+        reqs = []
+        for feature in features["features"]:
+            feat = {
+                "type": "Feature",
+                "properties": {},
+                "geometry": feature["geometry"],
+            }
+            request = {
+                "feature": feat,
+                "expr": expression,
+                "output": output.value,
+                "in_crs": crs, 
+                "out_crs": crs,
+                "resolution": res,
+            }
+            if id_field not in feature["properties"]:
+                raise ValueError(f"ID field {id_field} not found in feature properties.")
+            reqs.append({"group": "polygons", "file": feature["properties"][id_field], "request": request, "metadata": feature["properties"]})
+            if sample is None:
+                sample = request["expr"]
+
+        # Test first request to ensure expression is valid
+        if not skip_test:
+            result = await self._client.geoquery(**reqs[0]["request"], debug="requests")
+            request_count = result.get('request_count', 0)
+
+        user_quota = await self._client.auth.get_user_quota()
+        user_quota = user_quota.get('quota', -9999)
+        
+        if user_quota != -9999 and user_quota < len(reqs) * request_count:
+            raise QuotaInsufficientError(f"User quota is insufficient. Please contact support to increase your quota.")
+        
+        count = len(reqs)
+        groups = list(set(dic["group"] for dic in reqs))
+        print(f"[green]{count}[/green] requests with [blue]{len(groups)}[/blue] groups identified.")
+        
+        upload_urls = await self._get_upload_url(collection=collection)
+        url = upload_urls['url']
+        await self._upload_json_data(reqs, url, use_gzip=True)
+
+        payload = {"output": output.value, "skip_existing": skip_existing}
+        
+        task_id, status = await self._client._terrakio_request("POST", f"collections/{collection}/generate_data", json=payload)
+        task_id = task_id["task_id"]
+
+        await self.track_progress(task_id)
+
+        return {"task_id": task_id, "collection": collection}

{terrakio_core-0.5.3 → terrakio_core-0.5.4}/terrakio_core/exceptions.py

@@ -391,6 +391,12 @@ class CancelAllTasksError(Exception):
 
 class GeoQueryError(Exception):
    """Exception raised for geo query errors."""
+   def __init__(self, message: str, status_code: int = None):
+       super().__init__(message)
+       self.status_code = status_code
+    
+class QuotaInsufficientError(Exception):
+   """Exception raised for quota insufficient errors."""
    def __init__(self, message: str, status_code: int = None):
        super().__init__(message)
        self.status_code = status_code