diff --git a/README.md b/README.md index c04b640f..a99f224b 100644 --- a/README.md +++ b/README.md @@ -145,6 +145,7 @@ Native GeoTIFF and Cloud Optimized GeoTIFF reader/writer. No GDAL required. |:-----|:------------|:-----:|:----:|:--------:|:-------------:|:-----:| | [open_geotiff](xrspatial/geotiff/__init__.py) | Read GeoTIFF / COG / VRT | ✅️ | ✅️ | ✅️ | ✅️ | ✅️ | | [to_geotiff](xrspatial/geotiff/__init__.py) | Write DataArray as GeoTIFF / COG | ✅️ | ✅️ | ✅️ | ✅️ | ✅️ | +| [write_geotiff_gpu](xrspatial/geotiff/__init__.py) | GPU-accelerated GeoTIFF / COG write | | | ✅️ | | | | [write_vrt](xrspatial/geotiff/__init__.py) | Generate VRT mosaic from GeoTIFFs | ✅️ | | | | | `open_geotiff` and `to_geotiff` auto-dispatch to the correct backend: @@ -163,6 +164,10 @@ open_geotiff('mosaic.vrt') # VRT mosaic (auto-detected to_geotiff(cupy_array, 'out.tif') # auto-detects GPU to_geotiff(data, 'out.tif', gpu=True) # force GPU compress to_geotiff(data, 'ortho.tif', compression='jpeg') # JPEG for orthophotos +to_geotiff(data, 'cog.tif', cog=True) # COG with auto overviews +to_geotiff(data, 'cog.tif', cog=True, # COG with explicit levels + overview_levels=[2, 4, 8], + overview_resampling='nearest') write_vrt('mosaic.vrt', ['tile1.tif', 'tile2.tif']) # generate VRT open_geotiff('dem.tif', dtype='float32') # half memory @@ -189,7 +194,7 @@ ds.xrs.open_geotiff('large_dem.tif') # read windowed to Dataset - Cloud storage: S3 (`s3://`), GCS (`gs://`), Azure (`az://`) via fsspec - GPUDirect Storage: SSD→GPU direct DMA via KvikIO (optional) - Thread-safe mmap reads, atomic writes, HTTP connection reuse (urllib3) -- Overview generation: mean, nearest, min, max, median, mode, cubic +- Overview generation (CPU and GPU): mean, nearest, min, max, median, mode, cubic - Planar config, big-endian byte swap, PixelIsArea/PixelIsPoint **Read performance** (real-world files, A6000 GPU): diff --git a/docs/source/reference/geotiff.rst b/docs/source/reference/geotiff.rst new file mode 100644 index 00000000..134ae61c --- /dev/null +++ b/docs/source/reference/geotiff.rst @@ -0,0 +1,22 @@ +.. _reference.geotiff: + +*************** +GeoTIFF / COG +*************** + +Reading +======= +.. autosummary:: + :toctree: _autosummary + + xrspatial.geotiff.open_geotiff + xrspatial.geotiff.read_vrt + +Writing +======= +.. autosummary:: + :toctree: _autosummary + + xrspatial.geotiff.to_geotiff + xrspatial.geotiff.write_geotiff_gpu + xrspatial.geotiff.write_vrt diff --git a/docs/source/reference/index.rst b/docs/source/reference/index.rst index a268b5cd..bcdcb1f4 100644 --- a/docs/source/reference/index.rst +++ b/docs/source/reference/index.rst @@ -14,6 +14,7 @@ Reference fire flood focal + geotiff hydrology interpolation kde diff --git a/examples/user_guide/50_COG_Overview_Generation.ipynb b/examples/user_guide/50_COG_Overview_Generation.ipynb new file mode 100644 index 00000000..4043d4a9 --- /dev/null +++ b/examples/user_guide/50_COG_Overview_Generation.ipynb @@ -0,0 +1,230 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Cloud Optimized GeoTIFF: Overview (Pyramid) Generation\n", + "\n", + "Cloud Optimized GeoTIFFs (COGs) store internal overview images at reduced\n", + "resolution, so map viewers and tiling servers can fetch the right zoom\n", + "level without downloading the full file.\n", + "\n", + "xarray-spatial generates these overviews natively during `to_geotiff()`\n", + "with `cog=True`. No GDAL or `gdaladdo` post-processing required.\n", + "\n", + "This notebook covers:\n", + "- Writing a COG with automatic overviews\n", + "- Choosing resampling methods\n", + "- Specifying explicit overview levels\n", + "- Verifying the result" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import xarray as xr\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from xrspatial.geotiff import open_geotiff, to_geotiff\n", + "from xrspatial.geotiff._header import parse_header, parse_all_ifds" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Create synthetic terrain data\n", + "\n", + "A 256x256 elevation surface with some peaks, to give the overviews\n", + "something visually meaningful to downsample." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rng = np.random.RandomState(42)\n", + "y = np.linspace(45.0, 44.0, 256)\n", + "x = np.linspace(-120.0, -119.0, 256)\n", + "\n", + "# Smooth terrain from Gaussian peaks\n", + "xx, yy = np.meshgrid(x, y)\n", + "terrain = (\n", + " 800 * np.exp(-((xx + 119.5)**2 + (yy - 44.5)**2) / 0.02)\n", + " + 600 * np.exp(-((xx + 119.3)**2 + (yy - 44.7)**2) / 0.05)\n", + " + 100 * rng.rand(256, 256)\n", + ").astype(np.float32)\n", + "\n", + "da = xr.DataArray(\n", + " terrain, dims=['y', 'x'],\n", + " coords={'y': y, 'x': x},\n", + " attrs={'crs': 4326},\n", + " name='elevation',\n", + ")\n", + "\n", + "fig, ax = plt.subplots(figsize=(6, 5))\n", + "da.plot(ax=ax, cmap='terrain')\n", + "ax.set_title('Synthetic elevation')\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Write a COG with automatic overviews\n", + "\n", + "Pass `cog=True` and xarray-spatial handles the rest: it halves the\n", + "dimensions until the smallest overview fits within a single tile,\n", + "then writes all levels into the file with IFDs at the start (the COG\n", + "layout requirement for efficient HTTP range requests)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import tempfile, os\n", + "\n", + "tmpdir = tempfile.mkdtemp(prefix='cog_demo_')\n", + "cog_path = os.path.join(tmpdir, 'auto_overviews.tif')\n", + "\n", + "to_geotiff(da, cog_path, cog=True, compression='deflate')\n", + "\n", + "# Check how many IFDs (resolution levels) were written\n", + "with open(cog_path, 'rb') as f:\n", + " raw = f.read()\n", + "\n", + "header = parse_header(raw)\n", + "ifds = parse_all_ifds(raw, header)\n", + "for i, ifd in enumerate(ifds):\n", + " label = 'Full resolution' if i == 0 else f'Overview {i}'\n", + " print(f'{label}: {ifd.width} x {ifd.height}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Resampling methods\n", + "\n", + "Different data types need different resampling:\n", + "- **mean** (default): continuous data like elevation or temperature\n", + "- **nearest**: categorical or index data (land cover classes)\n", + "- **mode**: majority-class downsampling for classified rasters\n", + "- **min** / **max** / **median**: when you need conservative bounds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "methods = ['mean', 'nearest', 'min', 'max']\n", + "fig, axes = plt.subplots(1, len(methods), figsize=(14, 3))\n", + "\n", + "for ax, method in zip(axes, methods):\n", + " path = os.path.join(tmpdir, f'cog_{method}.tif')\n", + " to_geotiff(da, path, cog=True, compression='deflate',\n", + " overview_resampling=method)\n", + "\n", + " # Read back the first overview level\n", + " result = open_geotiff(path, overview_level=1)\n", + " result.plot(ax=ax, cmap='terrain', add_colorbar=False)\n", + " ax.set_title(f'{method} ({result.shape[0]}x{result.shape[1]})')\n", + " ax.set_xlabel('')\n", + " ax.set_ylabel('')\n", + "\n", + "plt.suptitle('Overview level 1 by resampling method', y=1.02)\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Explicit overview levels\n", + "\n", + "For fine-grained control, pass `overview_levels` as a list of\n", + "decimation factors. Each factor halves the previous level once." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "explicit_path = os.path.join(tmpdir, 'explicit_levels.tif')\n", + "to_geotiff(da, explicit_path, cog=True, compression='deflate',\n", + " overview_levels=[2, 4, 8])\n", + "\n", + "with open(explicit_path, 'rb') as f:\n", + " raw = f.read()\n", + "\n", + "header = parse_header(raw)\n", + "ifds = parse_all_ifds(raw, header)\n", + "for i, ifd in enumerate(ifds):\n", + " label = 'Full resolution' if i == 0 else f'Overview {i}'\n", + " print(f'{label}: {ifd.width} x {ifd.height}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify round-trip\n", + "\n", + "Full-resolution pixel values are preserved through the COG write." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "result = open_geotiff(cog_path)\n", + "max_diff = float(np.max(np.abs(result.values - terrain)))\n", + "print(f'Max pixel difference: {max_diff}')\n", + "assert max_diff < 1e-5, 'Values should round-trip exactly'" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Clean up\n", + "import shutil\n", + "shutil.rmtree(tmpdir)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/xrspatial/geotiff/__init__.py b/xrspatial/geotiff/__init__.py index 2fe46115..1a300afc 100644 --- a/xrspatial/geotiff/__init__.py +++ b/xrspatial/geotiff/__init__.py @@ -465,7 +465,10 @@ def to_geotiff(data: xr.DataArray | np.ndarray, path: str, *, compression=compression, compression_level=compression_level, tile_size=tile_size, - predictor=predictor) + predictor=predictor, + cog=cog, + overview_levels=overview_levels, + overview_resampling=overview_resampling) return except (ImportError, Exception): pass # fall through to CPU path @@ -1154,7 +1157,10 @@ def write_geotiff_gpu(data, path: str, *, compression: str = 'zstd', compression_level: int | None = None, tile_size: int = 256, - predictor: bool = False) -> None: + predictor: bool = False, + cog: bool = False, + overview_levels: list[int] | None = None, + overview_resampling: str = 'mean') -> None: """Write a CuPy-backed DataArray as a GeoTIFF with GPU compression. Tiles are extracted and compressed on the GPU via nvCOMP, then @@ -1162,6 +1168,10 @@ def write_geotiff_gpu(data, path: str, *, throughout compression -- only the compressed bytes transfer to CPU for file writing. + When ``cog=True``, generates overview pyramids on GPU and writes a + Cloud Optimized GeoTIFF with all IFDs at the file start for + efficient range-request access. + Falls back to CPU compression if nvCOMP is not available. Parameters @@ -1184,13 +1194,22 @@ def write_geotiff_gpu(data, path: str, *, Tile size in pixels (default 256). predictor : bool Apply horizontal differencing predictor. + cog : bool + Write as Cloud Optimized GeoTIFF with overviews. + overview_levels : list[int] or None + Overview decimation factors (e.g. [2, 4, 8]). Only used when + cog=True. If None and cog=True, auto-generates levels by + halving until the smallest overview fits in a single tile. + overview_resampling : str + Resampling method for overviews: 'mean' (default), 'nearest', + 'min', 'max', 'median', or 'mode'. """ try: import cupy except ImportError: raise ImportError("cupy is required for GPU writes") - from ._gpu_decode import gpu_compress_tiles + from ._gpu_decode import gpu_compress_tiles, make_overview_gpu from ._writer import ( _compression_tag, _assemble_tiff, _write_bytes, GeoTransform as _GT, @@ -1245,28 +1264,45 @@ def write_geotiff_gpu(data, path: str, *, comp_tag = _compression_tag(compression) pred_val = 2 if predictor else 1 - # GPU compress - compressed_tiles = gpu_compress_tiles( - arr, tile_size, tile_size, width, height, - comp_tag, pred_val, np_dtype, samples) - - # Build offset/bytecount lists - rel_offsets = [] - byte_counts = [] - offset = 0 - for tile in compressed_tiles: - rel_offsets.append(offset) - byte_counts.append(len(tile)) - offset += len(tile) - - # Assemble TIFF on CPU (only metadata + compressed bytes) - # _assemble_tiff needs an array in parts[0] to detect samples_per_pixel - shape_stub = np.empty((1, 1, samples) if samples > 1 else (1, 1), dtype=np_dtype) - parts = [(shape_stub, width, height, rel_offsets, byte_counts, compressed_tiles)] + def _gpu_compress_to_part(gpu_arr, w, h, spp): + """Compress a GPU array into a (stub, w, h, offsets, counts, tiles) tuple.""" + compressed = gpu_compress_tiles( + gpu_arr, tile_size, tile_size, w, h, + comp_tag, pred_val, np_dtype, spp) + rel_off = [] + bc = [] + off = 0 + for tile in compressed: + rel_off.append(off) + bc.append(len(tile)) + off += len(tile) + stub = np.empty((1, 1, spp) if spp > 1 else (1, 1), dtype=np_dtype) + return (stub, w, h, rel_off, bc, compressed) + + # Full resolution + parts = [_gpu_compress_to_part(arr, width, height, samples)] + + # Overview generation + if cog: + if overview_levels is None: + overview_levels = [] + oh, ow = height, width + while oh > tile_size and ow > tile_size: + oh //= 2 + ow //= 2 + if oh > 0 and ow > 0: + overview_levels.append(len(overview_levels) + 1) + + current = arr + for _ in overview_levels: + current = make_overview_gpu(current, method=overview_resampling) + oh, ow = current.shape[:2] + parts.append(_gpu_compress_to_part(current, ow, oh, samples)) file_bytes = _assemble_tiff( width, height, np_dtype, comp_tag, predictor, True, tile_size, - parts, geo_transform, epsg, nodata, is_cog=False, + parts, geo_transform, epsg, nodata, + is_cog=(cog and len(parts) > 1), raster_type=raster_type) _write_bytes(file_bytes, path) diff --git a/xrspatial/geotiff/_gpu_decode.py b/xrspatial/geotiff/_gpu_decode.py index 2b596ef0..480d0e48 100644 --- a/xrspatial/geotiff/_gpu_decode.py +++ b/xrspatial/geotiff/_gpu_decode.py @@ -2317,3 +2317,80 @@ def gpu_compress_tiles(d_image, tile_width, tile_height, result.append(cpu_compress(tile_data, compression)) return result + + +# --------------------------------------------------------------------------- +# GPU overview (pyramid) generation +# --------------------------------------------------------------------------- + +GPU_OVERVIEW_METHODS = ('mean', 'nearest', 'min', 'max', 'median', 'mode') + + +def _block_reduce_2d_gpu(arr2d, method): + """2x block-reduce a single 2D CuPy plane using *method*.""" + import cupy + + h, w = arr2d.shape + h2 = (h // 2) * 2 + w2 = (w // 2) * 2 + cropped = arr2d[:h2, :w2] + oh, ow = h2 // 2, w2 // 2 + + if method == 'nearest': + return cropped[::2, ::2].copy() + + if method == 'mode': + # Mode is expensive on GPU; fall back to CPU + cpu_arr = arr2d.get() + from ._writer import _block_reduce_2d + cpu_result = _block_reduce_2d(cpu_arr, 'mode') + return cupy.asarray(cpu_result) + + # Block reshape for mean/min/max/median + if arr2d.dtype.kind == 'f': + blocks = cropped.reshape(oh, 2, ow, 2) + else: + blocks = cropped.astype(cupy.float64).reshape(oh, 2, ow, 2) + + if method == 'mean': + result = cupy.nanmean(blocks, axis=(1, 3)) + elif method == 'min': + result = cupy.nanmin(blocks, axis=(1, 3)) + elif method == 'max': + result = cupy.nanmax(blocks, axis=(1, 3)) + elif method == 'median': + flat = blocks.transpose(0, 2, 1, 3).reshape(oh, ow, 4) + result = cupy.nanmedian(flat, axis=2) + else: + raise ValueError( + f"Unknown GPU overview resampling method: {method!r}. " + f"Use one of: {GPU_OVERVIEW_METHODS}") + + if arr2d.dtype.kind != 'f': + return cupy.around(result).astype(arr2d.dtype) + return result.astype(arr2d.dtype) + + +def make_overview_gpu(arr, method='mean'): + """Generate a 2x decimated overview on GPU. + + Parameters + ---------- + arr : cupy.ndarray + 2D or 3D (height, width, bands) array on GPU. + method : str + Resampling method: 'mean', 'nearest', 'min', 'max', 'median', + or 'mode'. + + Returns + ------- + cupy.ndarray + Half-resolution array on GPU. + """ + import cupy + + if arr.ndim == 3: + bands = [_block_reduce_2d_gpu(arr[:, :, b], method) + for b in range(arr.shape[2])] + return cupy.stack(bands, axis=2) + return _block_reduce_2d_gpu(arr, method) diff --git a/xrspatial/geotiff/_header.py b/xrspatial/geotiff/_header.py index 1f0751f7..8c6ddce8 100644 --- a/xrspatial/geotiff/_header.py +++ b/xrspatial/geotiff/_header.py @@ -15,6 +15,7 @@ ) # Well-known TIFF tag IDs +TAG_NEW_SUBFILE_TYPE = 254 TAG_IMAGE_WIDTH = 256 TAG_IMAGE_LENGTH = 257 TAG_BITS_PER_SAMPLE = 258 diff --git a/xrspatial/geotiff/_writer.py b/xrspatial/geotiff/_writer.py index 0b430b50..980fa95e 100644 --- a/xrspatial/geotiff/_writer.py +++ b/xrspatial/geotiff/_writer.py @@ -39,6 +39,7 @@ TAG_MODEL_TIEPOINT, ) from ._header import ( + TAG_NEW_SUBFILE_TYPE, TAG_IMAGE_WIDTH, TAG_IMAGE_LENGTH, TAG_BITS_PER_SAMPLE, @@ -593,6 +594,11 @@ def _assemble_tiff(width: int, height: int, dtype: np.dtype, for level_idx, (arr, lw, lh, rel_offsets, byte_counts, comp_data) in enumerate(pixel_data_parts): tags = [] + # Mark overview IFDs as reduced-resolution images (TIFF tag 254). + # GDAL/rasterio use this tag to identify overview sub-IFDs. + if level_idx > 0: + tags.append((TAG_NEW_SUBFILE_TYPE, LONG, 1, 1)) + tags.append((TAG_IMAGE_WIDTH, LONG, 1, lw)) tags.append((TAG_IMAGE_LENGTH, LONG, 1, lh)) if samples_per_pixel > 1: diff --git a/xrspatial/geotiff/tests/test_cog.py b/xrspatial/geotiff/tests/test_cog.py index 4fdea88f..5c255409 100644 --- a/xrspatial/geotiff/tests/test_cog.py +++ b/xrspatial/geotiff/tests/test_cog.py @@ -131,6 +131,248 @@ def test_write_rejects_4d(self, tmp_path): to_geotiff(arr, str(tmp_path / 'bad.tif')) +class TestCOGOverviewResampling: + """Test overview resampling methods produce correct results.""" + + def test_overview_mean(self, tmp_path): + arr = np.array([[1, 3, 5, 7], + [2, 4, 6, 8], + [9, 11, 13, 15], + [10, 12, 14, 16]], dtype=np.float32) + path = str(tmp_path / 'cog_1150_mean.tif') + write(arr, path, compression='deflate', tiled=True, tile_size=4, + cog=True, overview_levels=[1], overview_resampling='mean') + + with open(path, 'rb') as f: + data = f.read() + header = parse_header(data) + ifds = parse_all_ifds(data, header) + assert len(ifds) == 2 + # Overview should be 2x2 + ov_ifd = ifds[1] + assert ov_ifd.width == 2 + assert ov_ifd.height == 2 + + def test_overview_nearest(self, tmp_path): + arr = np.arange(64, dtype=np.float32).reshape(8, 8) + path = str(tmp_path / 'cog_1150_nearest.tif') + write(arr, path, compression='deflate', tiled=True, tile_size=4, + cog=True, overview_levels=[1], overview_resampling='nearest') + + result, _ = read_to_array_local(path) + np.testing.assert_array_equal(result, arr) + + def test_overview_mode(self, tmp_path): + # Categorical data: mode should pick the most common value + arr = np.array([[1, 1, 2, 2], + [1, 1, 2, 2], + [3, 3, 4, 4], + [3, 3, 4, 4]], dtype=np.int32) + path = str(tmp_path / 'cog_1150_mode.tif') + write(arr, path, compression='deflate', tiled=True, tile_size=4, + cog=True, overview_levels=[1], overview_resampling='mode') + + with open(path, 'rb') as f: + data = f.read() + header = parse_header(data) + ifds = parse_all_ifds(data, header) + assert len(ifds) == 2 + + @pytest.mark.parametrize('method', ['min', 'max', 'median']) + def test_overview_other_methods(self, tmp_path, method): + arr = np.arange(256, dtype=np.float32).reshape(16, 16) + path = str(tmp_path / f'cog_1150_{method}.tif') + write(arr, path, compression='deflate', tiled=True, tile_size=8, + cog=True, overview_levels=[1], overview_resampling=method) + + with open(path, 'rb') as f: + data = f.read() + header = parse_header(data) + ifds = parse_all_ifds(data, header) + assert len(ifds) >= 2 + + +class TestCOGMultipleOverviews: + def test_multiple_overview_levels(self, tmp_path): + """Multiple explicit overview levels produce correct number of IFDs.""" + arr = np.arange(4096, dtype=np.float32).reshape(64, 64) + path = str(tmp_path / 'cog_1150_multi.tif') + write(arr, path, compression='deflate', tiled=True, tile_size=8, + cog=True, overview_levels=[1, 2, 3]) + + with open(path, 'rb') as f: + data = f.read() + header = parse_header(data) + ifds = parse_all_ifds(data, header) + # Full res + 3 overviews + assert len(ifds) == 4 + + def test_auto_overviews_large_raster(self, tmp_path): + """Auto-generation on a larger raster produces multiple levels.""" + arr = np.random.RandomState(42).rand(512, 512).astype(np.float32) + path = str(tmp_path / 'cog_1150_auto_large.tif') + write(arr, path, compression='deflate', tiled=True, tile_size=64, + cog=True) + + with open(path, 'rb') as f: + data = f.read() + header = parse_header(data) + ifds = parse_all_ifds(data, header) + # 512 -> 256 -> 128 -> 64: should stop, so 3 overview levels + full = 4 + assert len(ifds) >= 3 + + def test_cog_overview_round_trip_values(self, tmp_path): + """Full-res values are preserved through COG write with overviews.""" + arr = np.random.RandomState(99).rand(32, 32).astype(np.float32) + gt = GeoTransform(-120.0, 45.0, 0.001, -0.001) + path = str(tmp_path / 'cog_1150_rt_values.tif') + write(arr, path, geo_transform=gt, crs_epsg=4326, + compression='deflate', tiled=True, tile_size=16, + cog=True, overview_levels=[1, 2]) + + result, geo = read_to_array_local(path) + np.testing.assert_array_equal(result, arr) + assert geo.crs_epsg == 4326 + + +class TestCOGPublicAPIOverviews: + def test_to_geotiff_cog_with_overviews(self, tmp_path): + """Public to_geotiff() with cog=True writes overviews.""" + y = np.linspace(45.0, 44.0, 32) + x = np.linspace(-120.0, -119.0, 32) + data = np.random.RandomState(42).rand(32, 32).astype(np.float32) + + da = xr.DataArray( + data, dims=['y', 'x'], + coords={'y': y, 'x': x}, + attrs={'crs': 4326}, + ) + + path = str(tmp_path / 'cog_1150_api.tif') + to_geotiff(da, path, compression='deflate', cog=True, + tile_size=16, overview_levels=[1]) + + result = open_geotiff(path) + np.testing.assert_array_almost_equal(result.values, data, decimal=5) + + # Verify COG structure + with open(path, 'rb') as f: + raw = f.read() + header = parse_header(raw) + ifds = parse_all_ifds(raw, header) + assert len(ifds) >= 2 + + def test_to_geotiff_cog_auto_overviews(self, tmp_path): + """Public API auto-generates overviews when only cog=True.""" + data = np.random.RandomState(7).rand(64, 64).astype(np.float32) + da = xr.DataArray(data, dims=['y', 'x']) + + path = str(tmp_path / 'cog_1150_api_auto.tif') + to_geotiff(da, path, compression='deflate', cog=True, tile_size=16) + + with open(path, 'rb') as f: + raw = f.read() + header = parse_header(raw) + ifds = parse_all_ifds(raw, header) + assert len(ifds) >= 2 + + +try: + import cupy + _HAS_CUPY = True +except ImportError: + _HAS_CUPY = False + + +@pytest.mark.skipif(not _HAS_CUPY, reason="CuPy not installed") +class TestGPUCOGOverviews: + """GPU-specific COG overview tests (require CuPy + CUDA).""" + + def test_gpu_cog_round_trip(self, tmp_path): + import cupy + arr = np.random.RandomState(42).rand(32, 32).astype(np.float32) + gpu_arr = cupy.asarray(arr) + + path = str(tmp_path / 'cog_1150_gpu_rt.tif') + from xrspatial.geotiff import write_geotiff_gpu + write_geotiff_gpu(gpu_arr, path, crs=4326, compression='deflate', + cog=True, overview_levels=[1]) + + result = open_geotiff(path) + np.testing.assert_array_almost_equal(result.values, arr, decimal=5) + + with open(path, 'rb') as f: + raw = f.read() + header = parse_header(raw) + ifds = parse_all_ifds(raw, header) + assert len(ifds) >= 2 + + def test_gpu_cog_auto_overviews(self, tmp_path): + import cupy + arr = np.random.RandomState(7).rand(64, 64).astype(np.float32) + gpu_arr = cupy.asarray(arr) + + path = str(tmp_path / 'cog_1150_gpu_auto.tif') + from xrspatial.geotiff import write_geotiff_gpu + write_geotiff_gpu(gpu_arr, path, compression='deflate', + cog=True, tile_size=16) + + with open(path, 'rb') as f: + raw = f.read() + header = parse_header(raw) + ifds = parse_all_ifds(raw, header) + assert len(ifds) >= 2 + + def test_gpu_overview_resampling_nearest(self, tmp_path): + import cupy + arr = np.arange(64, dtype=np.float32).reshape(8, 8) + gpu_arr = cupy.asarray(arr) + + path = str(tmp_path / 'cog_1150_gpu_nearest.tif') + from xrspatial.geotiff import write_geotiff_gpu + write_geotiff_gpu(gpu_arr, path, compression='deflate', + cog=True, overview_levels=[1], + overview_resampling='nearest') + + result = open_geotiff(path) + np.testing.assert_array_equal(result.values, arr) + + def test_gpu_make_overview_values(self): + """GPU overview block-reduce matches CPU for simple case.""" + import cupy + from xrspatial.geotiff._gpu_decode import make_overview_gpu + from xrspatial.geotiff._writer import _make_overview + + arr = np.random.RandomState(42).rand(16, 16).astype(np.float32) + gpu_arr = cupy.asarray(arr) + + for method in ('mean', 'nearest', 'min', 'max'): + cpu_ov = _make_overview(arr, method=method) + gpu_ov = make_overview_gpu(gpu_arr, method=method).get() + np.testing.assert_allclose(gpu_ov, cpu_ov, rtol=1e-5, + err_msg=f"Mismatch for method={method}") + + def test_gpu_to_geotiff_dispatches_with_overviews(self, tmp_path): + """to_geotiff auto-dispatches CuPy data with overview params.""" + import cupy + arr = np.random.RandomState(11).rand(32, 32).astype(np.float32) + da = xr.DataArray(cupy.asarray(arr), dims=['y', 'x'], + attrs={'crs': 4326}) + + path = str(tmp_path / 'cog_1150_gpu_dispatch.tif') + to_geotiff(da, path, compression='deflate', cog=True, + overview_levels=[1]) + + result = open_geotiff(path) + np.testing.assert_array_almost_equal(result.values, arr, decimal=5) + + with open(path, 'rb') as f: + raw = f.read() + header = parse_header(raw) + ifds = parse_all_ifds(raw, header) + assert len(ifds) >= 2 + + def read_to_array_local(path): """Helper to call read_to_array for local files.""" from xrspatial.geotiff._reader import read_to_array