cecil 0.0.28__tar.gz → 0.0.31__tar.gz

{cecil-0.0.28 → cecil-0.0.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cecil
-Version: 0.0.28
+Version: 0.0.31
 Summary: Python SDK for Cecil Earth
 License-Expression: MIT
 License-File: LICENSE.txt

{cecil-0.0.28 → cecil-0.0.31}/src/cecil/client.py

@@ -1,16 +1,15 @@
 import os
+from typing import Dict, List, Optional
+from warnings import warn
 
 import pandas as pd
 import requests
 import snowflake.connector
-import xarray
-
+from cryptography.hazmat.primitives import serialization
 from pydantic import BaseModel
 from requests import auth
-from cryptography.hazmat.primitives import serialization
-from typing import Dict, List, Optional
-from warnings import warn
 
+import xarray
 from .errors import (
     Error,
     _handle_bad_request,
@@ -35,11 +34,15 @@ from .models import (
     TransformationCreate,
     User,
     UserCreate,
-    DataRequestMetadata,
-    DataRequestParquetFiles,
+    SubscriptionMetadata,
+    SubscriptionParquetFiles,
+    SubscriptionListFiles,
+    Subscription,
+    SubscriptionCreate,
 )
 from .version import __version__
 from .xarray import load_xarray
+from .xarray import load_xarray_v2
 
 
 class Client:
@@ -69,6 +72,11 @@ class Client:
     def create_data_request(
         self, aoi_id: str, dataset_id: str, external_ref: Optional[str] = None
     ) -> DataRequest:
+        warn(
+            "create_data_request() is deprecated, use create_subscription() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         res = self._post(
             url="/v0/data-requests",
             model=DataRequestCreate(
@@ -78,22 +86,120 @@ class Client:
         return DataRequest(**res)
 
     def get_data_request(self, id: str) -> DataRequest:
+        warn(
+            "get_data_request() is deprecated, use get_subscription() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         res = self._get(url=f"/v0/data-requests/{id}")
         return DataRequest(**res)
 
     def list_data_requests(self) -> List[DataRequest]:
+        warn(
+            "list_data_requests() is deprecated, use list_subscriptions() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         res = self._get(url="/v0/data-requests")
         return [DataRequest(**record) for record in res["records"]]
 
-    def load_xarray(self, data_request_id: str) -> xarray.Dataset:
-        res = self._get(url=f"/v0/data-requests/{data_request_id}/metadata")
-        metadata = DataRequestMetadata(**res)
-        return load_xarray(metadata)
+    def list_subscriptions(self) -> List[Subscription]:
+        res = self._get(url="/v0/data-requests")
+        return [Subscription(**record) for record in res["records"]]
+
+    def create_subscription(
+        self, aoi_id: str, dataset_id: str, external_ref: Optional[str] = None
+    ) -> Subscription:
+        res = self._post(
+            url="/v0/data-requests",
+            model=SubscriptionCreate(
+                aoi_id=aoi_id, dataset_id=dataset_id, external_ref=external_ref
+            ),
+        )
+
+        return Subscription(**res)
+
+    def get_subscription(self, id: str) -> Subscription:
+        res = self._get(url=f"/v0/data-requests/{id}")
+        return Subscription(**res)
+
+    def load_xarray(
+        self,
+        subscription_id: Optional[str] = None,
+        data_request_id: Optional[str] = None,
+    ) -> xarray.Dataset:
+        if subscription_id is None and data_request_id is None:
+            raise TypeError("load_xarray() missing argument: 'subscription_id'")
+
+        if subscription_id is not None and data_request_id is not None:
+            raise ValueError(
+                "load_xarray() only accepts one argument but two were provided"
+            )
+
+        if data_request_id:
+            warn(
+                "data_request_id is deprecated, use subscription_id instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            subscription_id = data_request_id
+
+        res = SubscriptionMetadata(
+            **self._get(url=f"/v0/data-requests/{subscription_id}/metadata")
+        )
+        return load_xarray(res)
+
+    def _load_xarray_v2(
+        self,
+        subscription_id: Optional[str] = None,
+        data_request_id: Optional[str] = None,
+    ) -> xarray.Dataset:
+        if subscription_id is None and data_request_id is None:
+            raise TypeError("load_xarray_v2() missing argument: 'subscription_id'")
+
+        if subscription_id is not None and data_request_id is not None:
+            raise ValueError(
+                "load_xarray_v2() only accepts one argument but two were provided"
+            )
+
+        if data_request_id:
+            warn(
+                "data_request_id is deprecated, use subscription_id instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            subscription_id = data_request_id
+
+        res = SubscriptionListFiles(
+            **self._get(url=f"/v0/data-requests/{subscription_id}/files/tiff")
+        )
+        return load_xarray_v2(res)
 
-    def load_dataframe(self, data_request_id: str) -> pd.DataFrame:
-        res = self._get(url=f"/v0/data-requests/{data_request_id}/parquet-files")
-        metadata = DataRequestParquetFiles(**res)
-        df = pd.concat((pd.read_parquet(f) for f in metadata.files))
+    def load_dataframe(
+        self,
+        subscription_id: Optional[str] = None,
+        data_request_id: Optional[str] = None,
+    ) -> pd.DataFrame:
+        if subscription_id is None and data_request_id is None:
+            raise TypeError("load_dataframe missing argument: 'subscription_id'")
+
+        if subscription_id is not None and data_request_id is not None:
+            raise ValueError(
+                "load_dataframe only accepts one argument but two were provided"
+            )
+
+        if data_request_id:
+            warn(
+                "data_request_id is deprecated, use subscription_id instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            subscription_id = data_request_id
+
+        res = SubscriptionParquetFiles(
+            **self._get(url=f"/v0/data-requests/{subscription_id}/parquet-files")
+        )
+        df = pd.concat((pd.read_parquet(f) for f in res.files))
         return df[
             [col for col in df.columns if col not in ("organisation_id", "created_at")]
         ]

{cecil-0.0.28 → cecil-0.0.31}/src/cecil/models.py

@@ -126,7 +126,7 @@ class File(BaseModel):
     bands: List[Band]
 
 
-class DataRequestMetadata(BaseModel):
+class SubscriptionMetadata(BaseModel):
     model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
     provider_name: str
     dataset_id: str
@@ -137,6 +137,55 @@ class DataRequestMetadata(BaseModel):
     files: List[File]
 
 
-class DataRequestParquetFiles(BaseModel):
+class Bucket(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+    name: str
+    prefix: str
+
+
+class BucketCredentials(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+    access_key_id: str
+    secret_access_key: str
+    session_token: str
+    expiration: datetime.datetime
+
+
+class FileMapping(BaseModel):
+    type: str
+    bands: List
+
+
+class SubscriptionListFiles(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+    provider_name: str
+    dataset_id: str
+    dataset_name: str
+    aoi_id: str
+    data_request_id: str
+    bucket: Bucket
+    credentials: BucketCredentials
+    allowed_actions: List
+    file_mapping: Dict[str, FileMapping]
+
+
+class SubscriptionParquetFiles(BaseModel):
     model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
     files: List[str]
+
+
+class Subscription(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+    id: str
+    aoi_id: str
+    dataset_id: str
+    external_ref: Optional[str]
+    created_at: datetime.datetime
+    created_by: str
+
+
+class SubscriptionCreate(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)
+    aoi_id: str
+    dataset_id: str
+    external_ref: Optional[str]

cecil-0.0.31/src/cecil/version.py

@@ -0,0 +1 @@
+__version__ = "0.0.31"

cecil-0.0.31/src/cecil/xarray.py

@@ -0,0 +1,415 @@
+import re
+import time
+from datetime import datetime
+
+import boto3
+import dask
+import rasterio
+import rasterio.session
+import rioxarray
+import xarray
+
+from .errors import Error
+from .models import SubscriptionMetadata, SubscriptionListFiles
+
+# v1
+
+
+def load_xarray(metadata: SubscriptionMetadata) -> xarray.Dataset:
+    data_vars = {}
+
+    for f in metadata.files:
+        try:
+            dataset = _retry_with_exponential_backoff(_load_file, 5, 1, 2, f.url)
+        except Exception as e:
+            raise ValueError(f"failed to load file: {e}")
+
+        for b in f.bands:
+            band = dataset.sel(band=b.number, drop=True)
+
+            if b.time and b.time_pattern:
+                t = datetime.strptime(b.time, b.time_pattern)
+                band = band.expand_dims("time")
+                band = band.assign_coords(time=[t])
+
+            band.name = b.variable_name
+
+            if b.variable_name not in data_vars:
+                data_vars[b.variable_name] = []
+
+            data_vars[b.variable_name].append(band)
+
+    for variable_name, time_series in data_vars.items():
+        if "time" in time_series[0].dims:
+            data_vars[variable_name] = xarray.concat(
+                time_series, dim="time", join="exact"
+            )
+        else:
+            data_vars[variable_name] = time_series[0]
+
+    return xarray.Dataset(
+        data_vars=data_vars,
+        attrs={
+            "provider_name": metadata.provider_name,
+            "dataset_name": metadata.dataset_name,
+            "dataset_id": metadata.dataset_id,
+            "aoi_id": metadata.aoi_id,
+            "subscription_id": metadata.data_request_id,
+        },
+    )
+
+
+def _retry_with_exponential_backoff(
+    func, retries, start_delay, multiplier, *args, **kwargs
+):
+    delay = start_delay
+    for attempt in range(1, retries + 1):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            if attempt == retries:
+                raise e
+            time.sleep(delay)
+            delay *= multiplier
+    return None
+
+
+def _load_file(url: str):
+    return rioxarray.open_rasterio(
+        url,
+        chunks={"x": 2000, "y": 2000},
+    )
+
+
+# v2
+
+
+def load_xarray_v2(res: SubscriptionListFiles) -> xarray.Dataset:
+    session = boto3.session.Session(
+        aws_access_key_id=res.credentials.access_key_id,
+        aws_secret_access_key=res.credentials.secret_access_key,
+        aws_session_token=res.credentials.session_token,
+    )
+
+    keys = _list_keys_v2(session, res.bucket.name, res.bucket.prefix)
+
+    if not keys:
+        return xarray.Dataset()
+
+    timestamp_pattern = re.compile(r"\d{4}/\d{2}/\d{2}/\d{2}/\d{2}/\d{2}")
+    data_vars = {}
+
+    for key in keys:
+        try:
+            file_da = _retry_with_exponential_backoff(
+                _load_file_v2,
+                5,
+                1,
+                2,
+                session,
+                f"s3://{res.bucket.name}/{key}",
+            )
+        except Exception as e:
+            raise ValueError(f"failed to load file: {e}")
+
+        filename = key.split("/")[-1]
+
+        file_info = res.file_mapping.get(filename)
+        if not file_info:
+            continue
+
+        timestamp_str = timestamp_pattern.search(key).group()
+
+        for band_num, var_name in enumerate(file_info.bands, start=1):
+            band_da = file_da.sel(band=band_num, drop=True)
+            band_da.name = var_name
+
+            # Dataset with time dimension
+            if timestamp_str != "0000/00/00/00/00/00":
+                t = datetime.strptime(timestamp_str, "%Y/%m/%d/%H/%M/%S")
+                band_da = band_da.expand_dims("time")
+                band_da = band_da.assign_coords(time=[t])
+
+            if var_name not in data_vars:
+                data_vars[var_name] = []
+
+            data_vars[var_name].append(band_da)
+
+    for var_name, time_series in data_vars.items():
+        if "time" in time_series[0].dims:
+            data_vars[var_name] = xarray.concat(time_series, dim="time", join="exact")
+        else:
+            data_vars[var_name] = time_series[0]
+
+    return xarray.Dataset(
+        data_vars=data_vars,
+        attrs={
+            "provider_name": res.provider_name,
+            "dataset_name": res.dataset_name,
+            "dataset_id": res.dataset_id,
+            "aoi_id": res.aoi_id,
+            "subscription_id": res.data_request_id,
+        },
+    )
+
+
+def _list_keys_v2(session: boto3.session.Session, bucket_name, prefix) -> list[str]:
+    s3_client = session.client("s3")
+    paginator = s3_client.get_paginator("list_objects_v2")
+    page_iterator = paginator.paginate(
+        Bucket=bucket_name,
+        Prefix=prefix,
+    )
+
+    keys = []
+    for page in page_iterator:
+        for obj in page.get("Contents", []):
+            keys.append(obj["Key"])
+
+    return keys
+
+
+def _load_file_v2(aws_session: boto3.session.Session, url: str):
+    with rasterio.env.Env(
+        session=rasterio.session.AWSSession(aws_session),
+        GDAL_DISABLE_READDIR_ON_OPEN=True,
+    ):
+        return rioxarray.open_rasterio(
+            url,
+            chunks={"x": 2000, "y": 2000},
+        )
+
+
+# v3
+
+
+def load_xarray_v3(res: SubscriptionListFiles) -> xarray.Dataset:
+    session = boto3.session.Session(
+        aws_access_key_id=res.credentials.access_key_id,
+        aws_secret_access_key=res.credentials.secret_access_key,
+        aws_session_token=res.credentials.session_token,
+    )
+
+    keys = _list_keys_v3(session, res.bucket.name, res.bucket.prefix)
+
+    if not keys:
+        return xarray.Dataset()
+
+    timestamp_pattern = re.compile(r"\d{4}/\d{2}/\d{2}/\d{2}/\d{2}/\d{2}")
+    data_vars = {}
+
+    with rasterio.env.Env(
+        session=rasterio.session.AWSSession(session),
+        GDAL_DISABLE_READDIR_ON_OPEN=True,
+    ):
+        first_file = rioxarray.open_rasterio(
+            f"s3://{res.bucket.name}/{keys[0]}", chunks="auto"
+        )
+
+    for key in keys:
+        filename = key.split("/")[-1]
+
+        file_info = res.file_mapping.get(filename)
+        if not file_info:
+            continue
+
+        lazy_array = dask.array.from_delayed(
+            dask.delayed(_load_file_v3)(session, f"s3://{res.bucket.name}/{key}"),
+            shape=first_file.shape,
+            dtype=file_info.type,
+        )
+        lazy_da = xarray.DataArray(
+            lazy_array,
+            dims=first_file.dims,
+            coords=dict(first_file.coords),
+            # attrs=first_file.attrs.copy() # TODO: not the same for all files
+        )
+        # lazy_da.encoding = first_file.encoding.copy()
+        # lazy_da.rio.write_crs(first_file.rio.crs, inplace=True)
+        # lazy_da.rio.write_transform(first_file.rio.transform(), inplace=True)
+
+        timestamp_str = timestamp_pattern.search(key).group()
+
+        for band_num, var_name in enumerate(file_info.bands, start=1):
+            band_da = lazy_da.sel(band=band_num, drop=True)
+            band_da.name = var_name
+
+            # Dataset with time dimension
+            if timestamp_str != "0000/00/00/00/00/00":
+                t = datetime.strptime(timestamp_str, "%Y/%m/%d/%H/%M/%S")
+                band_da = band_da.expand_dims("time")
+                band_da = band_da.assign_coords(time=[t])
+
+            if var_name not in data_vars:
+                data_vars[var_name] = []
+
+            data_vars[var_name].append(band_da)
+
+    for var_name, time_series in data_vars.items():
+        if "time" in time_series[0].dims:
+            data_vars[var_name] = xarray.concat(time_series, dim="time", join="exact")
+        else:
+            data_vars[var_name] = time_series[0]
+
+    return xarray.Dataset(
+        data_vars=data_vars,
+        attrs={
+            "provider_name": res.provider_name,
+            "dataset_name": res.dataset_name,
+            "dataset_id": res.dataset_id,
+            "aoi_id": res.aoi_id,
+            "subscription_id": res.data_request_id,
+        },
+    )
+
+
+def _load_file_v3(aws_session: boto3.session.Session, url: str):
+    with rasterio.env.Env(
+        session=rasterio.session.AWSSession(aws_session),
+        GDAL_DISABLE_READDIR_ON_OPEN=True,
+    ):
+        return rioxarray.open_rasterio(
+            url,
+            chunks="auto",
+        ).values
+        # ).sel(band=num_band, drop=True)
+        # ).sel(band=num_band, drop=True).values
+        # ).isel(band=num_band-1).values
+
+
+def _list_keys_v3(session: boto3.session.Session, bucket_name, prefix) -> list[str]:
+    s3_client = session.client("s3")
+    paginator = s3_client.get_paginator("list_objects_v2")
+    page_iterator = paginator.paginate(
+        Bucket=bucket_name,
+        Prefix=prefix,
+    )
+
+    keys = []
+    for page in page_iterator:
+        for obj in page.get("Contents", []):
+            keys.append(obj["Key"])
+
+    return keys
+
+
+# v4
+
+
+def load_xarray_v4(res: SubscriptionListFiles) -> xarray.Dataset:
+
+    session = boto3.session.Session(
+        aws_access_key_id=res.credentials.access_key_id,
+        aws_secret_access_key=res.credentials.secret_access_key,
+        aws_session_token=res.credentials.session_token,
+    )
+
+    keys = _list_keys_v2(session, res.bucket.name, res.bucket.prefix)
+
+    if not keys:
+        return xarray.Dataset()
+
+    first_file_metadata = _get_file_metadata_v4(session, res.bucket.name, keys[0])
+
+    timestamp_pattern = re.compile(r"\d{4}/\d{2}/\d{2}/\d{2}/\d{2}/\d{2}")
+
+    data_vars = {}
+    for key in keys:
+        filename = key.split("/")[-1].rsplit(".", 1)[0]
+
+        file_info = res.file_mapping.get(filename)
+        if not file_info:
+            continue
+
+        timestamp_str = timestamp_pattern.search(key).group()
+
+        for band_num, band_name in enumerate(file_info.bands, start=1):
+            array = _create_dask_array_v4(
+                session,
+                f"s3://{res.bucket.name}/{key}",
+                band_num,
+                first_file_metadata["height"],
+                first_file_metadata["width"],
+                file_info.type,
+            )
+            da = xarray.DataArray(
+                array,
+                dims=("y", "x"),
+            )
+            da.name = band_name
+
+            # Dataset with time dimension
+            if timestamp_str != "0000/00/00/00/00/00":
+                time = datetime.strptime(timestamp_str, "%Y/%m/%d/%H/%M/%S")
+                da = da.expand_dims("time")
+                da = da.assign_coords(time=[time])
+
+            if band_name not in data_vars:
+                data_vars[band_name] = []
+
+            data_vars[band_name].append(da)
+
+    for variable_name, time_series in data_vars.items():
+        if "time" in time_series[0].dims:
+            data_vars[variable_name] = xarray.concat(
+                time_series,
+                dim="time",
+                join="exact",
+            )
+        else:
+            data_vars[variable_name] = time_series[0]
+
+    ds = xarray.Dataset(
+        data_vars=data_vars,
+        coords={
+            "y": first_file_metadata["y"],
+            "x": first_file_metadata["x"],
+        },
+        attrs={
+            "provider_name": res.provider_name,
+            "dataset_name": res.dataset_name,
+            "dataset_id": res.dataset_id,
+            "aoi_id": res.aoi_id,
+            "subscription_id": res.data_request_id,
+        },
+    )
+    ds = ds.rio.write_crs(first_file_metadata["crs"])
+
+    return ds
+
+
+def _get_file_metadata_v4(session, bucket: str, path: str):
+    with rasterio.env.Env(
+        rasterio.session.AWSSession(session), GDAL_DISABLE_READDIR_ON_OPEN=True
+    ):
+        da = xarray.open_dataarray(f"s3://{bucket}/{path}", engine="rasterio")
+
+    return {
+        "crs": da.rio.crs,
+        "height": da.rio.height,
+        "width": da.rio.width,
+        "x": da.x.values,
+        "y": da.y.values,
+    }
+
+
+def _create_dask_array_v4(
+    session: boto3.session.Session,
+    file_path: str,
+    band_num: int,
+    height: int,
+    width: int,
+    dtype: str,
+):
+    rasterio_session = rasterio.session.AWSSession(session)
+
+    def read_chunk():
+        with rasterio.env.Env(
+            session=rasterio_session, GDAL_DISABLE_READDIR_ON_OPEN=True
+        ):
+            with rasterio.open(file_path) as src:
+                return src.read(band_num)
+
+    return dask.array.from_delayed(
+        dask.delayed(read_chunk)(), shape=(height, width), dtype=dtype
+    )

cecil-0.0.28/src/cecil/version.py

@@ -1 +0,0 @@
-__version__ = "0.0.28"

cecil-0.0.28/src/cecil/xarray.py

@@ -1,74 +0,0 @@
-import os
-import rioxarray
-import xarray
-
-from datetime import datetime
-
-from .errors import Error
-from .models import DataRequestMetadata
-
-os.environ["GDAL_NUM_THREADS"] = "1"
-os.environ["GDAL_DISABLE_READDIR_ON_OPEN"] = "FALSE"
-
-
-def align_pixel_grids(time_series):
-    # Use the first timestep as reference
-    reference_da = time_series[0]
-    aligned_series = [reference_da]
-
-    # Align all other timesteps to the reference grid
-    for i, da in enumerate(time_series[1:], 1):
-        try:
-            aligned_da = da.rio.reproject_match(reference_da)
-            aligned_series.append(aligned_da)
-        except Exception as e:
-            raise Error
-            continue
-
-    return aligned_series
-
-
-def load_xarray(metadata: DataRequestMetadata) -> xarray.Dataset:
-    data_vars = {}
-
-    for f in metadata.files:
-        dataset = rioxarray.open_rasterio(
-            f.url,
-            chunks={"x": 2000, "y": 2000},
-        )
-
-        for b in f.bands:
-            band = dataset.sel(band=b.number, drop=True)
-
-            if b.time and b.time_pattern:
-                time = datetime.strptime(b.time, b.time_pattern)
-                band = band.expand_dims("time")
-                band = band.assign_coords(time=[time])
-
-            band.name = b.variable_name
-
-            if b.variable_name not in data_vars:
-                data_vars[b.variable_name] = []
-
-            data_vars[b.variable_name].append(band)
-
-    for variable_name, time_series in data_vars.items():
-        if "time" in time_series[0].dims:
-            # time_series = align_pixel_grids(time_series)
-            data_vars[variable_name] = xarray.concat(
-                time_series, dim="time", join="exact"
-            )
-        else:
-            data_vars[variable_name] = time_series[0]
-
-    return xarray.Dataset(
-        data_vars=data_vars,
-        attrs={
-            "provider_name": metadata.provider_name,
-            "dataset_id": metadata.dataset_id,
-            "dataset_name": metadata.dataset_name,
-            "dataset_crs": metadata.dataset_crs,
-            "aoi_id": metadata.aoi_id,
-            "data_request_id": metadata.data_request_id,
-        },
-    )