diff --git a/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py b/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py index 44a07a2..5d8c616 100644 --- a/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py +++ b/darts-acquisition/src/darts_acquisition/arcticdem/datacube.py @@ -298,6 +298,7 @@ def load_arcticdem_tile( resolution: RESOLUTIONS, chunk_size: int = 6000, buffer: int = 0, + persist: bool = True, ) -> xr.Dataset: """Get the corresponding ArcticDEM tile for the given geobox. @@ -308,6 +309,8 @@ def load_arcticdem_tile( chunk_size (int, optional): The chunk size for the datacube. Only relevant for the initial creation. Has no effect otherwise. Defaults to 6000. buffer (int, optional): The buffer around the geobox in pixels. Defaults to 0. + persist (bool, optional): If the data should be persisted in memory. + If not, this will return a Dask backed Dataset. Defaults to True. Returns: xr.Dataset: The ArcticDEM tile, with a buffer applied. @@ -364,10 +367,13 @@ def load_arcticdem_tile( arcticdem_aoi["datamask"] = arcticdem_aoi.datamask.astype("uint8") # The following code would load the data from disk - tick_sload = time.perf_counter() - arcticdem_aoi = arcticdem_aoi.compute() - tick_eload = time.perf_counter() - logger.debug(f"ArcticDEM AOI loaded from disk in {tick_eload - tick_sload:.2f} seconds") + if persist: + tick_sload = time.perf_counter() + arcticdem_aoi = arcticdem_aoi.compute() + tick_eload = time.perf_counter() + logger.debug(f"ArcticDEM AOI loaded from disk in {tick_eload - tick_sload:.2f} seconds") - logger.info(f"ArcticDEM tile loaded in {time.perf_counter() - tick_fstart:.2f} seconds") + logger.info( + f"ArcticDEM tile {'loaded' if persist else 'lazy-opened'} in {time.perf_counter() - tick_fstart:.2f} seconds" + ) return arcticdem_aoi diff --git a/darts-preprocessing/src/darts_preprocessing/engineering/arcticdem.py b/darts-preprocessing/src/darts_preprocessing/engineering/arcticdem.py index ed292ff..a9e278b 100644 --- a/darts-preprocessing/src/darts_preprocessing/engineering/arcticdem.py +++ b/darts-preprocessing/src/darts_preprocessing/engineering/arcticdem.py @@ -34,7 +34,7 @@ def calculate_topographic_position_index( f"{inner_radius}-{outer_radius} ({inner_radius_m}-{outer_radius_m}) cells." ) - tpi = arcticdem_ds.dem - convolution.convolution_2d(arcticdem_ds.dem.values, kernel) / kernel.sum() + tpi = arcticdem_ds.dem - convolution.convolution_2d(arcticdem_ds.dem, kernel) / kernel.sum() tpi.attrs = { "long_name": "Topographic Position Index", "units": "m", diff --git a/darts-preprocessing/src/darts_preprocessing/preprocess.py b/darts-preprocessing/src/darts_preprocessing/preprocess.py index 32e6259..7ffcc27 100644 --- a/darts-preprocessing/src/darts_preprocessing/preprocess.py +++ b/darts-preprocessing/src/darts_preprocessing/preprocess.py @@ -3,8 +3,9 @@ import logging import time -import rasterio +import odc.geo.xr # noqa: F401 import xarray as xr +from xrspatial.utils import has_cuda_and_cupy from darts_preprocessing.engineering.arcticdem import calculate_slope, calculate_topographic_position_index from darts_preprocessing.engineering.indices import calculate_ndvi @@ -12,6 +13,14 @@ logger = logging.getLogger(__name__.replace("darts_", "darts.")) +if has_cuda_and_cupy(): + import cupy_xarray # noqa: F401 + + logger.info("GPU-accelerated xrspatial functions are available.") +else: + logger.info("GPU-accelerated xrspatial functions are not available.") + + def preprocess_legacy( ds_optical: xr.Dataset, ds_arcticdem: xr.Dataset, @@ -49,6 +58,7 @@ def preprocess_legacy_fast( ds_data_masks: xr.Dataset, tpi_outer_radius: int = 30, tpi_inner_radius: int = 25, + use_gpu: bool = True, ) -> xr.Dataset: """Preprocess optical data with legacy (DARTS v1) preprocessing steps, but with new data concepts. @@ -70,6 +80,7 @@ def preprocess_legacy_fast( in number of cells. Defaults to 30. tpi_inner_radius (int, optional): The inner radius of the annulus kernel for the tpi calculation in number of cells. Defaults to 25. + use_gpu (bool, optional): Whether to use GPU-accelerated functions. Defaults to True. Returns: xr.Dataset: The preprocessed dataset. @@ -86,9 +97,20 @@ def preprocess_legacy_fast( # Calculate TPI and slope from ArcticDEM # We need to calculate them before reprojecting, hence we cant merge the data yet + # Move to GPU if available + if use_gpu and has_cuda_and_cupy(): + logger.debug("Moving arcticdem to GPU.") + # Check if dem is dask, if not persist it, since tpi and slope can't be calculated from cupy-dask arrays + if ds_arcticdem.chunks is not None: + ds_arcticdem = ds_arcticdem.persist() + ds_arcticdem = ds_arcticdem.cupy.as_cupy() + ds_arcticdem = calculate_topographic_position_index(ds_arcticdem, tpi_outer_radius, tpi_inner_radius) ds_arcticdem = calculate_slope(ds_arcticdem) - ds_arcticdem = ds_arcticdem.rio.reproject_match(ds_optical, resampling=rasterio.enums.Resampling.cubic) + # Move back to CPU + if use_gpu and has_cuda_and_cupy(): + ds_arcticdem = ds_arcticdem.cupy.as_numpy() + ds_arcticdem = ds_arcticdem.odc.reproject(ds_optical.odc.geobox, resampling="cubic") ds_merged["dem"] = ds_arcticdem.dem ds_merged["relative_elevation"] = ds_arcticdem.tpi diff --git a/darts/src/darts/cli.py b/darts/src/darts/cli.py index 6df963b..1dba927 100644 --- a/darts/src/darts/cli.py +++ b/darts/src/darts/cli.py @@ -13,7 +13,7 @@ from darts.utils.config import ConfigParser from darts.utils.logging import add_logging_handlers, setup_logging -root = Path(__name__).resolve() +root_file = Path(__file__).resolve() logger = logging.getLogger(__name__) console = Console() @@ -76,7 +76,7 @@ def launcher( # noqa: D103 ): command, bound = app.parse_args(tokens) add_logging_handlers(command.__name__, console, log_dir) - logger.debug(f"Running on Python version {sys.version} from {root}") + logger.debug(f"Running on Python version {sys.version} from {__name__} ({root_file})") return command(*bound.args, **bound.kwargs) diff --git a/notebooks/test-arcticdem-datacube.ipynb b/notebooks/test-arcticdem-datacube.ipynb index 0ac42a6..e92cc37 100644 --- a/notebooks/test-arcticdem-datacube.ipynb +++ b/notebooks/test-arcticdem-datacube.ipynb @@ -9,6 +9,7 @@ "import logging\n", "from pathlib import Path\n", "\n", + "import cupy_xarray # noqa: F401\n", "import hvplot.xarray # noqa: F401\n", "import rioxarray # noqa: F401\n", "from darts_acquisition.planet import load_planet_scene\n", @@ -26,7 +27,7 @@ " datefmt=\"[%X]\",\n", " handlers=[RichHandler(rich_tracebacks=True)],\n", ")\n", - "traceback.install(show_locals=False)\n", + "traceback.install(show_locals=True)\n", "client = Client()\n", "configure_rio(cloud_defaults=True, aws={\"aws_unsigned\": True}, client=client)\n", "client" @@ -52,7 +53,7 @@ "outputs": [], "source": [ "# load planet scene\n", - "ds_planet = load_planet_scene(fpath).isel(x=slice(0, 2000), y=slice(6000, 8000))\n", + "ds_planet = load_planet_scene(fpath) # .isel(x=slice(0, 2000), y=slice(6000, 8000))\n", "ds_planet" ] }, @@ -71,7 +72,7 @@ "source": [ "from darts_acquisition.arcticdem import load_arcticdem_tile\n", "\n", - "ds = load_arcticdem_tile(ds_planet, arcticdem_dir, resolution=2, buffer=0)\n", + "ds = load_arcticdem_tile(ds_planet.odc.geobox, arcticdem_dir, resolution=2, buffer=0, persist=True)\n", "ds" ] }, @@ -92,8 +93,12 @@ "outputs": [], "source": [ "crs = ds_planet.rio.crs.to_string()\n", - "dem_plot = ds.dem.rio.reproject_match(ds_planet).hvplot.image(aggregator=\"max\", rasterize=True, cmap=\"terrain\")\n", - "red_plot = ds_planet.red.hvplot.image(x=\"x\", y=\"y\", aggregator=\"mean\", rasterize=True, cmap=\"reds\")\n", + "dem_plot = ds.dem.rio.reproject_match(ds_planet).hvplot.image(\n", + " aggregator=\"max\", rasterize=True, cmap=\"terrain\", data_aspect=1, crs=crs, projection=crs\n", + ")\n", + "red_plot = ds_planet.red.hvplot.image(\n", + " x=\"x\", y=\"y\", aggregator=\"mean\", rasterize=True, cmap=\"reds\", data_aspect=1, crs=crs, projection=crs\n", + ")\n", "dem_plot + red_plot" ] }, @@ -101,7 +106,22 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Calculate Relative Elevation and Slope" + "## Calculate Relative Elevation and Slope\n", + "\n", + "### Numpy\n", + "\n", + "Topographic Position Index calculated in 652.77 seconds.\n", + "Slope calculated in 3.67 seconds.\n", + "\n", + "### Dask (4 worker)\n", + "\n", + "Topographic Position Index calculated in 135.34 seconds.\n", + "Slope calculated in 4.33 seconds.\n", + "\n", + "### Cupy\n", + "\n", + "Topographic Position Index calculated in 12.69 seconds. \n", + "Slope calculated in 0.16 seconds.\n" ] }, { @@ -111,21 +131,18 @@ "outputs": [], "source": [ "from darts_preprocessing.engineering.arcticdem import calculate_slope, calculate_topographic_position_index\n", + "from xrspatial.utils import has_cuda_and_cupy\n", "\n", + "use_cupy = True\n", + "if use_cupy and has_cuda_and_cupy():\n", + " ds = ds.cupy.as_cupy()\n", "ds = calculate_topographic_position_index(ds)\n", "ds = calculate_slope(ds)\n", + "if ds.cupy.is_cupy:\n", + " ds = ds.cupy.as_numpy()\n", "ds" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ds = ds.persist()" - ] - }, { "cell_type": "code", "execution_count": null, @@ -134,13 +151,13 @@ "source": [ "crs = ds.rio.crs.to_string()\n", "dem_plot = ds.dem.hvplot.image(\n", - " x=\"x\", y=\"y\", aggregator=\"max\", rasterize=True, cmap=\"terrain\", crs=crs, projection=crs, title=\"DEM\"\n", + " x=\"x\", y=\"y\", aggregator=\"max\", rasterize=True, cmap=\"terrain\", data_aspect=1, crs=crs, projection=crs, title=\"DEM\"\n", ")\n", "tpi_plot = ds.tpi.hvplot.image(\n", - " x=\"x\", y=\"y\", aggregator=\"max\", rasterize=True, cmap=\"terrain\", crs=crs, projection=crs, title=\"TPI\"\n", + " x=\"x\", y=\"y\", aggregator=\"max\", rasterize=True, cmap=\"terrain\", data_aspect=1, crs=crs, projection=crs, title=\"TPI\"\n", ")\n", "slope_plot = ds.slope.hvplot.image(\n", - " x=\"x\", y=\"y\", aggregator=\"max\", rasterize=True, cmap=\"terrain\", crs=crs, projection=crs, title=\"Slope\"\n", + " x=\"x\", y=\"y\", aggregator=\"max\", rasterize=True, data_aspect=1, crs=crs, projection=crs, title=\"Slope\"\n", ")\n", "dem_plot + tpi_plot + slope_plot" ] diff --git a/notebooks/test-e2e.ipynb b/notebooks/test-e2e.ipynb index 6061985..cf816c6 100644 --- a/notebooks/test-e2e.ipynb +++ b/notebooks/test-e2e.ipynb @@ -72,7 +72,7 @@ " )\n", " for z in tile.data_vars\n", " ]\n", - " return hv.Layout(var_plots).cols(ncols)\n" + " return hv.Layout(var_plots).cols(ncols)" ] }, { @@ -82,13 +82,13 @@ "outputs": [], "source": [ "cache_file = DATA_ROOT / \"intermediate\" / f\"planet_{fpath.stem}.nc\"\n", - "force = True\n", + "force = False\n", "slc = {\"x\": slice(0, 1000), \"y\": slice(7000, 8000)}\n", "if cache_file.exists() and not force:\n", - " tile = xr.open_dataset(cache_file, engine=\"h5netcdf\", mask_and_scale=False).isel(slc)\n", + " tile = xr.open_dataset(cache_file, engine=\"h5netcdf\", mask_and_scale=False).set_coords(\"spatial_ref\")\n", "else:\n", " optical = load_planet_scene(fpath).isel(slc)\n", - " arcticdem = load_arcticdem_tile(optical, arcticdem_dir, resolution=2)\n", + " arcticdem = load_arcticdem_tile(optical.odc.geobox, arcticdem_dir, resolution=2)\n", " tcvis = load_tcvis(optical, None if force else cache_dir)\n", " data_masks = load_planet_masks(fpath).isel(slc)\n", " tile = preprocess_legacy_fast(optical, arcticdem, tcvis, data_masks)\n", diff --git a/pyproject.toml b/pyproject.toml index 5421b1b..225621f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ dependencies = [ "folium>=0.18.0", "bokeh>=3.5.2", "jupyter-bokeh>=4.0.5", + "setuptools>=75.5.0", ] readme = "README.md" requires-python = ">= 3.11" @@ -56,14 +57,18 @@ cpu = ["torch==2.2.0+cpu", "torchvision==0.17.0+cpu"] cuda11 = [ "torch==2.2.0+cu118", "torchvision==0.17.0+cu118", - "cupy-cuda11x>=13.3.0", "cucim-cu11>=24.8.0", + "cupy>=13.3.0", + "cupy-xarray>=0.1.4", + "cuda-python>=12.6.2.post1", ] cuda12 = [ "torch==2.2.0+cu121", "torchvision==0.17.0+cu121", - "cupy-cuda12x>=13.3.0", "cucim-cu12==24.8.*", + "cupy>=13.3.0", + "cupy-xarray>=0.1.4", + "cuda-python>=12.6.2.post1", ] gdal393 = ["gdal==3.9.3"] gdal39 = ["gdal==3.9.2"]