Merge pull request #94 from mayrajeo/zarr-downloads

Add zarr to supported download formats

Author: avsm

geotessera/cli.py

@@ -12,7 +12,12 @@
 import numpy as np
 import geopandas as gpd
 
-from .registry import Registry, EMBEDDINGS_DIR_NAME, tile_to_geotiff_path
+from .registry import (
+    Registry,
+    EMBEDDINGS_DIR_NAME,
+    tile_to_geotiff_path,
+    tile_to_zarr_path,
+)
 
 try:
     import importlib.metadata
@@ -63,7 +68,7 @@ class GeoTessera:
     Core functionality:
     - Download tiles to local embeddings_dir
     - Sample embeddings at point locations from local tiles
-    - Export individual tiles as GeoTIFF files with correct metadata
+    - Export individual tiles as GeoTIFF files or zarr archives with correct metadata
     - Manage registry and data access
 
     Typical workflows:
@@ -1838,6 +1843,271 @@ def merge_geotiffs_to_mosaic(
 
         return str(output_path)
 
+    def export_embedding_zarr(
+        self,
+        lon: float,
+        lat: float,
+        output_path: Union[str, Path],
+        year: int = 2024,
+        bands: Optional[List[int]] = None,
+    ) -> str:
+        """Export a single embedding tile as a zarr archive with native UTM projection.
+
+        Args:
+            lon: Tile center longitude
+            lat: Tile center latitude
+            output_path: Output path for zarr file
+            year: Year of embeddings to export
+            bands: List of band indices to export (None = all 128 bands)
+            compress: Compression method for GeoTIFF
+
+        Returns:
+            Path to created zarr file
+
+        Raises:
+            ImportError: If xarray, rioxarray, zarr or dask is not available
+            RuntimeError: If landmask tile or embedding data cannot be fetched
+            FileNotFoundError: If registry files are missing
+        """
+        try:
+            import xarray as xr
+            import rioxarray as rxr  # noqa: F401 - needed for .rio accessor
+            import zarr  # noqa: F401 - needed for .to_zarr()
+            import dask  # noqa: F401 - needed for chunking
+
+            import warnings
+
+            warnings.filterwarnings("ignore", category=UserWarning)
+        except ImportError:
+            raise ImportError(
+                "saving to zarr requires xarray, rioxarray, zarr and dask"
+            )
+
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Fetch single tile with CRS info
+        embedding, crs, transform = self.fetch_embedding(lon, lat, year)
+
+        # Select bands
+        if bands is not None:
+            data = embedding[:, :, bands].copy()
+            band_count = len(bands)
+        else:
+            data = embedding.copy()
+            band_count = 128
+
+        # Get dimensions for GeoTIFF
+        height, width = data.shape[:2]
+
+        ds = xr.Dataset(
+            {"embedding": (("y", "x", "band"), data)},
+            coords={
+                "y": np.arange(height),
+                "x": np.arange(width),
+                "band": np.arange(band_count),
+            },
+            attrs={
+                "TESSERA_DATASET_VERSION": self.dataset_version,
+                "TESSERA_YEAR": year,
+                "TESSERA_TILE_LAT": f"{lat:.2f}",
+                "TESSERA_TILE_LON": f"{lon:.2f}",
+                "TESSERA_DESCRIPTION": "GeoTessera satellite embedding tile",
+                "GEOTESSERA_VERSION": __version__,
+            },
+        )
+
+        x_coords = [transform.c + (i + 0.5) * transform.a for i in range(width)]
+        y_coords = [transform.f + (j + 0.5) * transform.e for j in range(height)]
+
+        # Add band descriptions
+        if bands is not None:
+            output_bands = [f"Tessera_Band_{band_idx}" for band_idx in bands]
+        else:
+            output_bands = [f"Tessera_Band_{j}" for j in range(128)]
+
+        ds = ds.assign_coords(
+            x=("x", x_coords),
+            y=("y", y_coords),
+            band=("band", np.array(output_bands, dtype=np.dtypes.StringDType())),
+        )
+        ds = (
+            ds.rio.write_crs(crs)
+            .rio.set_spatial_dims(x_dim="x", y_dim="y")
+            .rio.write_coordinate_system()
+        )
+        ds = ds.rio.write_transform(transform)
+        ds.to_zarr(output_path, zarr_format=3)
+        return str(output_path)
+
+    def export_embedding_zarrs(
+        self,
+        tiles_to_fetch: Iterable[Tuple[int, float, float]],
+        output_dir: Union[str, Path],
+        bands: Optional[List[int]] = None,
+        progress_callback: Optional[callable] = None,
+    ) -> List[str]:
+        """Export all embedding tiles in bounding box as individual zarr files with native UTM projections.
+        The list of tiles to fetch can be obtained by registry.load_blocks_for_region().
+
+        Args:
+            tiles_to_fetch: List of tiles to fetch as (year, tile_lon, tile_lat) tuples
+            output_dir: Directory to save GeoTIFF files
+            bands: List of band indices to export (None = all 128 bands)
+            compress: Compression method for GeoTIFF
+            progress_callback: Optional callback function(current, total) for progress tracking
+
+        Returns:
+            List of paths to created zarr files
+
+        Raises:
+            ImportError: If rasterio is not available
+            RuntimeError: If landmask tiles or embedding data cannot be fetched
+            FileNotFoundError: If registry files are missing
+        """
+        try:
+            import xarray as xr
+            import rioxarray as rxr  # noqa: F401 - needed for .rio accessor
+            import zarr  # noqa: F401 - needed for .to_zarr()
+            import dask  # noqa: F401 - needed for chunking
+
+            import warnings
+
+            warnings.filterwarnings("ignore", category=UserWarning)
+        except ImportError:
+            raise ImportError(
+                "saving to zarr requires xarray, rioxarray, zarr and dask"
+            )
+
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Create a wrapper callback to handle two-phase progress
+        def fetch_progress_callback(current: int, total: int, status: str = None):
+            # Phase 1: Fetching tiles (0-50% of total progress)
+            overall_progress = int((current / total) * 50)
+            display_status = status or f"Fetching tile {current}/{total}"
+            progress_callback(overall_progress, 100, display_status)
+
+        # Fetch tiles with progress tracking
+        if progress_callback:
+            progress_callback(0, 100, "Loading registry blocks...")
+
+        tiles = list(
+            self.fetch_embeddings(
+                tiles_to_fetch, fetch_progress_callback if progress_callback else None
+            )
+        )
+        if progress_callback:
+            total_tiles = len(tiles_to_fetch)
+
+        if not tiles:
+            self.logger.warning("No tiles found in bounding box")
+            return []
+
+        if progress_callback:
+            progress_callback(
+                50, 100, f"Fetched {total_tiles} tiles, starting GeoTIFF export..."
+            )
+
+        created_files = []
+
+        # Sequential zarr writing
+        for i, (year, tile_lon, tile_lat, embedding, crs, transform) in enumerate(
+            tiles
+        ):
+            # Use centralized path construction from registry
+            zarr_rel_path = tile_to_zarr_path(tile_lon, tile_lat, year)
+            output_path = output_dir / EMBEDDINGS_DIR_NAME / zarr_rel_path
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+
+            # Update progress to show we're starting this file
+            if progress_callback:
+                export_progress = int(50 + (i / total_tiles) * 50)
+                progress_callback(
+                    export_progress, 100, f"Creating {output_path.name}..."
+                )
+
+            # Select bands
+            if bands is not None:
+                data = embedding[:, :, bands].copy()
+                band_count = len(bands)
+            else:
+                data = embedding.copy()
+                band_count = 128
+
+            # Get dimensions for GeoTIFF
+            height, width = data.shape[:2]
+
+            # Select bands
+            if bands is not None:
+                data = embedding[:, :, bands].copy()
+                band_count = len(bands)
+            else:
+                data = embedding.copy()
+                band_count = 128
+
+            # Get dimensions for GeoTIFF
+            height, width = data.shape[:2]
+            ds = xr.Dataset(
+                {"embedding": (("y", "x", "band"), data)},
+                coords={
+                    "y": np.arange(height),
+                    "x": np.arange(width),
+                    "band": np.arange(band_count),
+                },
+                attrs={
+                    "TESSERA_DATASET_VERSION": self.dataset_version,
+                    "TESSERA_YEAR": year,
+                    "TESSERA_TILE_LAT": f"{tile_lat:.2f}",
+                    "TESSERA_TILE_LON": f"{tile_lon:.2f}",
+                    "TESSERA_DESCRIPTION": "GeoTessera satellite embedding tile",
+                    "GEOTESSERA_VERSION": __version__,
+                },
+            )
+
+            x_coords = [transform.c + (i + 0.5) * transform.a for i in range(width)]
+            y_coords = [transform.f + (j + 0.5) * transform.e for j in range(height)]
+
+            # Add band descriptions
+            if bands is not None:
+                output_bands = [f"Tessera_Band_{band_idx}" for band_idx in bands]
+            else:
+                output_bands = [f"Tessera_Band_{j}" for j in range(128)]
+
+            ds = ds.assign_coords(
+                x=("x", x_coords),
+                y=("y", y_coords),
+                band=("band", np.array(output_bands, dtype=np.dtypes.StringDType())),
+            )
+            ds = (
+                ds.rio.write_crs(crs)
+                .rio.set_spatial_dims(x_dim="x", y_dim="y")
+                .rio.write_coordinate_system()
+            )
+            ds = ds.rio.write_transform(transform)
+
+            ds.to_zarr(output_path, zarr_format=3)
+            created_files.append(str(output_path))
+
+            # Update progress for zarr export phase
+            if progress_callback:
+                # Phase 2: Exporting zarr (50-100% of total progress)
+                export_progress = int(50 + ((i + 1) / total_tiles) * 50)
+                progress_callback(
+                    export_progress,
+                    100,
+                    f"Exported {output_path.name} ({i + 1}/{total_tiles})",
+                )
+
+        if progress_callback:
+            progress_callback(
+                100, 100, f"Completed! Exported {len(created_files)} zarr files"
+            )
+
+        self.logger.info(f"Exported {len(created_files)} zarr files to {output_dir}")
+        return created_files
+
     def apply_pca_to_embeddings(
         self,
         embeddings: List[Tuple[int, float, float, np.ndarray, object, object]],

geotessera/core.py

@@ -276,6 +276,21 @@ def tile_to_geotiff_path(lon: float, lat: float, year: int) -> str:
     return f"{year}/{grid_name}/{grid_name}_{year}.tiff"
 
 
+def tile_to_zarr_path(lon: float, lat: float, year: int) -> str:
+    """Generate GeoTIFF file path for a tile.
+
+    Args:
+        lon: Tile center longitude
+        lat: Tile center latitude
+        year: Year of embeddings
+
+    Returns:
+        str: Relative path like "{year}/grid_{lon}_{lat}/grid_{lon}_{lat}_{year}.zarr"
+    """
+    grid_name = tile_to_grid_name(lon, lat)
+    return f"{year}/{grid_name}/{grid_name}_{year}.zarr"
+
+
 def tile_to_landmask_filename(lon: float, lat: float) -> str:
     """Generate landmask filename for a tile.