Repository: opendatacube/odc-geo
Pull request #163 (2024-06-18 06:38:11)

refactor: nodata handling #162
- adding new types
  - ~`MaybeAutoNodata`~ `SomeNodata` -- `None|int|float|str|"auto"`
  - `Nodata` is now `None|float|int`
  - `MaybeNodata` is now ~`None|float|int|str`~  have been removed
  - "auto" replaces what used to be `None`
  - `None` now means "no nodata value"
- `resolve_nodata()` is used to handle nodata options consistently across the library
- `.odc.nodata` is used to extract `nodata` from xarrays
- fixes in reproject/overview generation for float data. It is assumed to have `nan` values, and GDAL needs `nodata=nan` to handle it correctly.

### nodata types explained

- ~`MaybeAutoNodata`~`SomeNodata` should be used in APIs that allow configuring nodata and nodata overrides
- `Nodata` is a type suitable to pass to GDAL `{src,dst)_nodata` it can be `None` which means expect no missing values in the data.
- `FillValue` is a type suitable for `np.full` and is resolved from `Nodata` + `dtype`
   - `NaN` for not-configured floating point images
   - `0` for not-configured integer images
   - value of `nodata` if set (`NaN` is fine here too) 

```
SomeNodata + dtype + [fallback] -> Nodata + dtype -> FillValue
```


github-actions[bot]: <!-- NETLIFY DEPLOY COMMENT GENERATED BY ACTIONS_NETLIFY - APP ID SHA256: 2516d918dc82767c089255cf7129f75d9f959e53c07db8d3368972c4656b3dce -->
🚀 Deployed on https://66712bd4a1ffecbbd5cd5a8f--odc-geo-docs.netlify.app
SpacemanPaul: Thanks for the swift action on this Kirill - I will review tomorrow.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/opendatacube/odc-geo/pull/163?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) Report
Attention: Patch coverage is `98.59155%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 95.46%. Comparing base [(`6af5d0c`)](https://app.codecov.io/gh/opendatacube/odc-geo/commit/6af5d0cce76d073e83ade6dca51b0377ca7d4e57?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) to head [(`44c956a`)](https://app.codecov.io/gh/opendatacube/odc-geo/commit/44c956ab580dce4581f02bff5a86cbbf4105389b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).
> Report is 33 commits behind head on develop.

| [Files](https://app.codecov.io/gh/opendatacube/odc-geo/pull/163?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) | Patch % | Lines |
|---|---|---|
| [odc/geo/math.py](https://app.codecov.io/gh/opendatacube/odc-geo/pull/163?src=pr&el=tree&filepath=odc%2Fgeo%2Fmath.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube#diff-b2RjL2dlby9tYXRoLnB5) | 96.55% | [1 Missing :warning: ](https://app.codecov.io/gh/opendatacube/odc-geo/pull/163?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop     #163      +/-   ##
===========================================
+ Coverage    95.26%   95.46%   +0.20%     
===========================================
  Files           31       31              
  Lines         5323     5474     +151     
===========================================
+ Hits          5071     5226     +155     
+ Misses         252      248       -4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/opendatacube/odc-geo/pull/163?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).

robbibt: Thanks heaps for this @Kirill888 - the test cases appear to match the functionality I was hoping for nicely, but I'll do a "user" test today and verify that TIFFs exported using the updated code work as intended (e.g. in ESRI etc).
Kirill888: @robbibt thanks, BTW pleas use `pip install odc-geo==0.4.7rc1` for your test (that release is not on conda, only pypi)
robbibt: I think this is working perfectly. For example, we have a dataset with an Xarray `nodata: nan` attribute:
``` 
import datacube
import odc.geo.xr
from datacube.utils.cog import write_cog

dc = datacube.Datacube()

query_params = dict(
    x=(142.13223, 142.65461),
    y=(-32.17591, -32.54618),
    time=("2022", "2022"),
)

ds = dc.load(product="ga_ls8cls9c_gm_cyear_3", measurements=["edev"], **query_params)
da = ds.edev.squeeze()
```
![image](https://github.com/opendatacube/odc-geo/assets/17680388/692eab3d-3ed1-4f57-9b8e-2abc50508cfb)


This gets written out with a GeoTIFF nodata flag `nodata=nan` using both `datacube` and `odc-geo` tooling:
```
da.odc.write_cog("nodata_nan_odcgeo.tif")
write_cog(da, "nodata_nan_datacube.tif")
```
![image](https://github.com/opendatacube/odc-geo/assets/17680388/de99ba3c-5141-41c2-a8bf-dd7e72af0271)
![image](https://github.com/opendatacube/odc-geo/assets/17680388/1afaf412-d076-4151-86e4-77a18b3105a4)


However, if we write out data without a Xarray `nodata: nan` attribute, `datacube` doesn't include a GeoTIFF nodata flag, but now `odc-geo` does!
```
del da.attrs["nodata"]

da.odc.write_cog("nodata_missing_odcgeo.tif")
write_cog(da, "nodata_missing_datacube.tif")
```
![image](https://github.com/opendatacube/odc-geo/assets/17680388/1022abd2-d061-43b4-9e48-ad364781a913)
![image](https://github.com/opendatacube/odc-geo/assets/17680388/eafa6872-b51b-4932-84a1-eecd88468f5a)

I can though return a true missing nodata value like this:
```
da.odc.write_cog("nodata_truenone_odcgeo.tif", nodata=None)
```
![image](https://github.com/opendatacube/odc-geo/assets/17680388/95b61f05-c63a-4080-8f25-e5a21cdcbcee)



robbibt: I do want to do a quick test of the COG overviews stuff, so will approve this as soon as I've finished that.
robbibt: > I think some of the naming is confusing, but the approach is sound.

I also find "MaybeAutoNodata" a bit confusing - took me a while to work out what it meant.
Kirill888: > > I think some of the naming is confusing, but the approach is sound.
> 
> I also find "MaybeAutoNodata" a bit confusing - took me a while to work out what it meant.

What about `SomeNodata` instead of `MaybeAutoNodata`, @SpacemanPaul @robbibt 
SpacemanPaul: > > > I think some of the naming is confusing, but the approach is sound.
> > 
> > 
> > I also find "MaybeAutoNodata" a bit confusing - took me a while to work out what it meant.
> 
> What about `SomeNodata` instead of `MaybeAutoNodata`, @SpacemanPaul @robbibt

Yes `SomeNodata` or even `AnyNodata` works for me.
Kirill888: Updated,

- Removed no longer relevant `MaybeNodata` type, if anyone was using that type, they were probably also using `Nodata` type too, and that have changed, so...

- use `SomeNodata` instead of `MaybeAutoNodata` name



Patch:
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index bd87e5a..d3ae2ca 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -222,6 +222,7 @@ jobs:
         with:
           fail_ci_if_error: false
           verbose: false
+          token: ${{ secrets.CODECOV_TOKEN }}
 
   test-wheels:
     runs-on: ubuntu-latest
@@ -306,7 +307,7 @@ jobs:
       - name: Deploy to Netlify
         id: netlify
         if: github.event_name == 'pull_request'
-        uses: nwtgck/actions-netlify@v2
+        uses: nwtgck/actions-netlify@v3
         with:
           production-branch: "main"
           publish-dir: "docs/_build/html"
@@ -320,7 +321,7 @@ jobs:
           NETLIFY_SITE_ID: ${{ secrets.NETLIFY_SITE_ID }}
 
       - name: Print Notice
-        uses: actions/github-script@v6
+        uses: actions/github-script@v7
         if: github.event_name == 'pull_request'
         env:
           NETLIFY_URL: ${{ steps.netlify.outputs.deploy-url }}
diff --git a/odc/geo/_dask.py b/odc/geo/_dask.py
index 004583c..5b69c68 100644
--- a/odc/geo/_dask.py
+++ b/odc/geo/_dask.py
@@ -9,13 +9,9 @@ from dask.highlevelgraph import HighLevelGraph
 from ._blocks import BlockAssembler
 from .gcp import GCPGeoBox
 from .geobox import GeoBox, GeoboxTiles
-from .warp import (
-    Nodata,
-    Resampling,
-    _rio_reproject,
-    resampling_s2rio,
-    resolve_fill_value,
-)
+from .math import resolve_fill_value
+from .types import Nodata
+from .warp import Resampling, _rio_reproject, resampling_s2rio
 
 
 def _do_chunked_reproject(
diff --git a/odc/geo/_rgba.py b/odc/geo/_rgba.py
index f99efb2..5d7c486 100644
--- a/odc/geo/_rgba.py
+++ b/odc/geo/_rgba.py
@@ -8,6 +8,7 @@ import numpy as np
 import xarray as xr
 
 from ._interop import is_dask_collection
+from .types import Nodata
 
 # pylint: disable=import-outside-toplevel
 
@@ -59,7 +60,7 @@ def _np_to_rgba(
     r: np.ndarray,
     g: np.ndarray,
     b: np.ndarray,
-    nodata: Optional[float],
+    nodata: Nodata,
     vmin: float,
     vmax: float,
 ) -> np.ndarray:
@@ -67,7 +68,7 @@ def _np_to_rgba(
 
     if r.dtype.kind == "f":
         valid = ~np.isnan(r)
-        if nodata is not None:
+        if nodata is not None and not np.isnan(nodata):
             valid = valid * (r != nodata)
     elif nodata is not None:
         valid = r != nodata
@@ -130,7 +131,7 @@ def to_rgba(
     assert vmax is not None
 
     _b = ds[bands[0]]
-    nodata = getattr(_b, "nodata", None)
+    nodata = _b.odc.nodata
     dims = (*_b.dims, "band")
 
     r, g, b = (ds[name].data for name in bands)
@@ -171,12 +172,20 @@ def _np_colorize(x, cmap, clip):
     return cmap[x]
 
 
-def _matplotlib_colorize(x, cmap, vmin=None, vmax=None, nodata=None, robust=False):
+def _matplotlib_colorize(
+    x,
+    cmap,
+    vmin=None,
+    vmax=None,
+    nodata: Nodata = None,
+    robust=False,
+):
     from matplotlib import colormaps
     from matplotlib.colors import Normalize
 
     if cmap is None or isinstance(cmap, str):
-        cmap = colormaps.get_cmap(cmap)
+        # None is a valid input, maps to default cmap
+        cmap = colormaps.get_cmap(cmap)  # type: ignore
 
     if nodata is not None:
         x = np.where(x == nodata, np.float32("nan"), x)
@@ -234,8 +243,11 @@ def colorize(
     :param clip: If ``True`` clip values from ``x`` to be in the safe range for ``cmap``.
     """
     # pylint: disable=too-many-locals
+    from ._xr_interop import ODCExtensionDa
 
     assert isinstance(x, xr.DataArray)
+    assert isinstance(x.odc, ODCExtensionDa)
+
     _is_dask = is_dask_collection(x.data)
 
     if isinstance(cmap, np.ndarray):
@@ -263,7 +275,7 @@ def colorize(
             _matplotlib_colorize,
             vmin=vmin,
             vmax=vmax,
-            nodata=getattr(x, "nodata", None),
+            nodata=x.odc.nodata,
             robust=robust,
         )
         nc, cmap_dtype = 4, "uint8"
diff --git a/odc/geo/_version.py b/odc/geo/_version.py
index 5eefbda..15ba425 100644
--- a/odc/geo/_version.py
+++ b/odc/geo/_version.py
@@ -1,3 +1,3 @@
 """version information only."""
 
-__version__ = "0.4.6"
+__version__ = "0.4.7rc1"
diff --git a/odc/geo/_xr_interop.py b/odc/geo/_xr_interop.py
index 45ffeec..40dc37d 100644
--- a/odc/geo/_xr_interop.py
+++ b/odc/geo/_xr_interop.py
@@ -39,13 +39,15 @@ from .math import (
     affine_from_axis,
     approx_equal_affine,
     is_affine_st,
+    is_nodata_empty,
     maybe_int,
     resolution_from_affine,
+    resolve_fill_value,
+    resolve_nodata,
 )
 from .overlap import compute_output_geobox
 from .roi import roi_is_empty
-from .types import Resolution, SomeResolution, SomeShape, xy_
-from .warp import resolve_fill_value
+from .types import Nodata, Resolution, SomeNodata, SomeResolution, SomeShape, xy_
 
 # pylint: disable=import-outside-toplevel
 # pylint: disable=too-many-lines
@@ -656,7 +658,7 @@ def xr_reproject(
     how: Union[SomeCRS, GeoBox],
     *,
     resampling: Union[str, int] = "nearest",
-    dst_nodata: Optional[float] = None,
+    dst_nodata: SomeNodata = "auto",
     dtype=None,
     resolution: Union[SomeResolution, Literal["auto", "fit", "same"]] = "auto",
     shape: Union[SomeShape, int, None] = None,
@@ -753,7 +755,7 @@ def _xr_reproject_ds(
     how: Union[SomeCRS, GeoBox],
     *,
     resampling: Union[str, int] = "nearest",
-    dst_nodata: Optional[float] = None,
+    dst_nodata: SomeNodata = "auto",
     dtype=None,
     **kw,
 ) -> xarray.Dataset:
@@ -797,7 +799,7 @@ def _xr_reproject_da(
     how: Union[SomeCRS, GeoBox],
     *,
     resampling: Union[str, int] = "nearest",
-    dst_nodata: Optional[float] = None,
+    dst_nodata: SomeNodata = "auto",
     dtype=None,
     **kw,
 ) -> xarray.DataArray:
@@ -828,9 +830,8 @@ def _xr_reproject_da(
     assert ydim + 1 == src.odc.xdim
     dst_shape = (*src.shape[:ydim], *dst_geobox.shape, *src.shape[ydim + 2 :])
 
-    src_nodata = kw.pop("src_nodata", None)
-    if src_nodata is None:
-        src_nodata = src.odc.nodata
+    src_nodata = resolve_nodata(kw.pop("src_nodata", "auto"), src.dtype, src.odc.nodata)
+    dst_nodata = resolve_nodata(dst_nodata, dtype, src_nodata)
 
     fill_value = resolve_fill_value(dst_nodata, src_nodata, dtype)
 
@@ -865,10 +866,9 @@ def _xr_reproject_da(
         )
 
     attrs = {k: v for k, v in src.attrs.items() if k not in REPROJECT_SKIP_ATTRS}
-    if numpy.isfinite(fill_value) and (
-        dst_nodata is not None or src_nodata is not None
-    ):
-        attrs.update({k: maybe_int(float(fill_value), 1e-6) for k in NODATA_ATTRIBUTES})
+    if not is_nodata_empty(dst_nodata):
+        assert dst_nodata is not None
+        attrs.update({k: maybe_int(float(dst_nodata), 1e-6) for k in NODATA_ATTRIBUTES})
 
     # new set of coords (replace x,y dims)
     # discard all coords that reference spatial dimensions
@@ -997,16 +997,39 @@ class ODCExtensionDa(ODCExtension):
         return assign_crs(self._xx, crs=crs, crs_coord_name=crs_coord_name)
 
     @property
-    def nodata(self) -> Optional[float]:
+    def nodata(self) -> Nodata:
         """Extract ``nodata/_FillValue`` attribute if set."""
         attrs = self._xx.attrs
+        encoding = self._xx.encoding
+
         for k in ["nodata", "_FillValue"]:
-            nodata = attrs.get(k, None)
-            if nodata is not None:
-                return float(nodata)
+            nodata = attrs.get(k, ())
+            if nodata == ():
+                nodata = encoding.get(k, ())
+
+            if nodata == ():
+                continue
+
+            if nodata is None:
+                return None
+
+            return float(nodata)
 
         return None
 
+    @nodata.setter
+    def nodata(self, value: Nodata):
+        nodata = resolve_nodata(value, self._xx.dtype)
+
+        if nodata is None:
+            for k in ["nodata", "_FillValue"]:
+                self._xx.attrs.pop(k, None)
+                self._xx.encoding.pop(k, None)
+            return
+
+        self._xx.attrs["nodata"] = nodata
+        self._xx.encoding["_FillValue"] = nodata
+
     colorize = _wrap_op(colorize)
 
     if have.rasterio:
@@ -1076,7 +1099,7 @@ def wrap_xr(
     gbox: SomeGeoBox,
     *,
     time=None,
-    nodata=None,
+    nodata: SomeNodata = "auto",
     crs_coord_name: Optional[str] = _DEFAULT_CRS_COORD_NAME,
     always_yx: bool = False,
     dims: Optional[Tuple[str, ...]] = None,
@@ -1159,8 +1182,9 @@ def wrap_xr(
                 [f"b{i}" for i in range(nb)], dims=(dim,), name=dim
             )
 
-    if nodata is not None:
-        attrs = {"nodata": nodata, **attrs}
+    _nodata = resolve_nodata(nodata, im.dtype)
+    if not is_nodata_empty(_nodata) or nodata != "auto":
+        attrs = {"nodata": _nodata, **attrs}
 
     out = xarray.DataArray(im, coords=coords, dims=dims, attrs=attrs)
     if crs_coord_name is not None:
diff --git a/odc/geo/cog/_rio.py b/odc/geo/cog/_rio.py
index b8b6dd7..97d8cef 100644
--- a/odc/geo/cog/_rio.py
+++ b/odc/geo/cog/_rio.py
@@ -19,7 +19,8 @@ import xarray as xr
 from rasterio.shutil import copy as rio_copy  # pylint: disable=no-name-in-module
 
 from ..geobox import GeoBox
-from ..types import MaybeNodata, SomeShape, shape_, wh_
+from ..math import resolve_nodata
+from ..types import Nodata, SomeNodata, SomeShape, shape_, wh_
 from ..warp import resampling_s2rio
 from ._shared import adjust_blocksize
 
@@ -123,7 +124,7 @@ def _write_cog(
     pix: np.ndarray,
     geobox: GeoBox,
     fname: Union[Path, str],
-    nodata: MaybeNodata = None,
+    nodata: Nodata = None,
     overwrite: bool = False,
     blocksize: Optional[int] = None,
     overview_resampling: Optional[str] = None,
@@ -283,6 +284,7 @@ def write_cog(
     use_windowed_writes: bool = False,
     intermediate_compression: Union[bool, str, Dict[str, Any]] = False,
     tags: Optional[Dict[str, Any]] = None,
+    nodata: SomeNodata = "auto",
     **extra_rio_opts,
 ) -> Union[Path, bytes]:
     """
@@ -298,7 +300,7 @@ def write_cog(
     :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32],
                             to disable overviews supply empty list ``[]``
     :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is
-                   read from the attributes of the input array (``geo_im.attrs['nodata']``).
+                   read from the attributes of the input array (``geo_im.odc.nodata``).
     :param use_windowed_writes: Write image block by block (might need this for large images)
     :param intermediate_compression: Configure compression settings for first pass write
                     , default is no compression
@@ -322,6 +324,8 @@ def write_cog(
 
        This means that this function will use about 1.5 to 2 times memory taken by ``geo_im``.
     """
+    nodata = resolve_nodata(nodata, geo_im.dtype, geo_im.odc.nodata)
+
     if overviews is not None:
         layers = [geo_im, *overviews]
         result = write_cog_layers(
@@ -333,6 +337,7 @@ def write_cog(
             use_windowed_writes=use_windowed_writes,
             intermediate_compression=intermediate_compression,
             tags=tags,
+            nodata=nodata,
             **extra_rio_opts,
         )
         assert result is not None
@@ -340,9 +345,6 @@ def write_cog(
 
     pix = geo_im.data
     geobox = geo_im.odc.geobox
-    nodata = extra_rio_opts.pop("nodata", None)
-    if nodata is None:
-        nodata = geo_im.attrs.get("nodata", None)
 
     if geobox is None:
         raise ValueError("Need geo-registered array on input")
@@ -448,6 +450,7 @@ def write_cog_layers(
     intermediate_compression: Union[bool, str, Dict[str, Any]] = False,
     use_windowed_writes: bool = False,
     tags: Optional[Dict[str, Any]] = None,
+    nodata: Nodata = None,
     **extra_rio_opts,
 ) -> Union[Path, bytes, None]:
     """
@@ -475,14 +478,14 @@ def write_cog_layers(
         blocksize=blocksize,
         shape=gbox.shape,
         is_float=pix.dtype.kind == "f",
-        nodata=pix.attrs.get("nodata", None),
+        nodata=nodata,
     )
     rio_opts.update(extra_rio_opts)
 
     first_pass_cfg: Dict[str, Any] = {
         "num_threads": "ALL_CPUS",
         "blocksize": blocksize,
-        "nodata": rio_opts.get("nodata", None),
+        "nodata": nodata,
         "use_windowed_writes": use_windowed_writes,
         "gdal_metadata": _get_gdal_metadata(xx, tags),
         **_norm_compression_opts(intermediate_compression),
diff --git a/odc/geo/cog/_shared.py b/odc/geo/cog/_shared.py
index 7f6622d..c495913 100644
--- a/odc/geo/cog/_shared.py
+++ b/odc/geo/cog/_shared.py
@@ -12,7 +12,7 @@ import numpy as np
 
 from ..geobox import GeoBox
 from ..math import align_down_pow2, align_up
-from ..types import MaybeNodata, Shape2d, SomeShape, shape_, wh_
+from ..types import Nodata, Shape2d, SomeShape, shape_, wh_
 
 # pylint: disable=too-many-locals,too-many-branches,too-many-arguments,too-many-statements,too-many-instance-attributes
 
@@ -63,7 +63,7 @@ class CogMeta:
     compressionargs: Dict[str, Any] = field(default_factory=dict, repr=False)
     gbox: Optional[GeoBox] = None
     overviews: Tuple["CogMeta", ...] = field(default=(), repr=False)
-    nodata: MaybeNodata = None
+    nodata: Nodata = None
 
     def _pix_shape(self, shape: Shape2d) -> Tuple[int, ...]:
         if self.axis == "YX":
diff --git a/odc/geo/cog/_tifffile.py b/odc/geo/cog/_tifffile.py
index c35aec5..60bd1fa 100644
--- a/odc/geo/cog/_tifffile.py
+++ b/odc/geo/cog/_tifffile.py
@@ -18,7 +18,8 @@ import xarray as xr
 
 from .._interop import have
 from ..geobox import GeoBox
-from ..types import MaybeNodata, Shape2d, Unset, shape_
+from ..math import resolve_nodata
+from ..types import Shape2d, SomeNodata, Unset, shape_
 from ._mpu import mpu_write
 from ._mpu_fs import MPUFileSink
 from ._s3 import MultiPartUpload, s3_parse_url
@@ -120,7 +121,7 @@ def _make_empty_cog(
     dtype: Any,
     gbox: Optional[GeoBox] = None,
     *,
-    nodata: MaybeNodata = None,
+    nodata: SomeNodata = "auto",
     gdal_metadata: Optional[str] = None,
     compression: Union[str, Unset] = Unset(),
     compressionargs: Any = None,
@@ -140,6 +141,8 @@ def _make_empty_cog(
         enumarg,
     )
 
+    nodata = resolve_nodata(nodata, dtype)
+
     predictor, compression, compressionargs = _norm_compression_tifffile(
         dtype,
         predictor,
@@ -773,7 +776,7 @@ def save_cog_with_dask(
 
 def geotiff_metadata(
     geobox: GeoBox,
-    nodata: MaybeNodata = None,
+    nodata: SomeNodata = "auto",
     gdal_metadata: Optional[str] = None,
 ) -> Tuple[List[Tuple[int, int, int, Any]], Dict[str, Any]]:
     """
diff --git a/odc/geo/math.py b/odc/geo/math.py
index 60c0484..ac70992 100644
--- a/odc/geo/math.py
+++ b/odc/geo/math.py
@@ -29,7 +29,10 @@ from numpy.polynomial.polynomial import polygrid2d, polyval2d
 from .types import (
     XY,
     AnchorEnum,
+    FillValue,
+    Nodata,
     Resolution,
+    SomeNodata,
     SomeResolution,
     SomeShape,
     res_,
@@ -181,6 +184,49 @@ def is_almost_int(x: float, tol: float) -> bool:
     return x < tol
 
 
+def resolve_fill_value(dst_nodata: Nodata, src_nodata: Nodata, dtype) -> FillValue:
+    dtype = np.dtype(dtype)
+
+    if dst_nodata is not None:
+        return dtype.type(dst_nodata)
+    if np.issubdtype(dtype, np.floating):
+        return dtype.type("nan")
+    if src_nodata is not None:
+        return dtype.type(src_nodata)
+    return dtype.type(0)
+
+
+def resolve_nodata(
+    nodata: SomeNodata,
+    dtype=None,
+    xr_nodata=None,
+) -> Nodata:
+    # pylint: disable=too-many-return-statements
+    if nodata is None:
+        return None
+
+    if nodata == "auto":
+        if xr_nodata is not None:
+            return xr_nodata
+        if dtype is None:
+            return None
+        if np.issubdtype(dtype, np.floating):
+            return np.nan
+        return None
+
+    if isinstance(nodata, str):
+        return float(nodata)
+    return nodata
+
+
+def is_nodata_empty(nodata: Nodata) -> bool:
+    if nodata is None:
+        return True
+    if isinstance(nodata, float) and np.isnan(nodata):
+        return True
+    return False
+
+
 def _snap_edge_pos(x0: float, x1: float, res: float, tol: float) -> Tuple[float, int]:
     assert res > 0
     assert x1 >= x0
diff --git a/odc/geo/types.py b/odc/geo/types.py
index 773f32a..c13c161 100644
--- a/odc/geo/types.py
+++ b/odc/geo/types.py
@@ -21,8 +21,9 @@ from typing import (
 
 MaybeInt = Optional[int]
 MaybeFloat = Optional[float]
-Nodata = Union[float, int, str]
-MaybeNodata = Optional[Nodata]
+FillValue = Union[float, int]
+Nodata = Union[float, int, None]
+SomeNodata = Union[float, int, str, None, Literal["auto"]]
 T = TypeVar("T")
 T1 = TypeVar("T1")
 T2 = TypeVar("T2")
diff --git a/odc/geo/warp.py b/odc/geo/warp.py
index b441d01..7798e2f 100644
--- a/odc/geo/warp.py
+++ b/odc/geo/warp.py
@@ -10,19 +10,18 @@ from affine import Affine
 
 from .gcp import GCPGeoBox
 from .geobox import GeoBox
-from .types import wh_
+from .math import resolve_fill_value, resolve_nodata
+from .types import Nodata, SomeNodata, wh_
 
 # pylint: disable=invalid-name, too-many-arguments
 Resampling = Union[str, int, rasterio.warp.Resampling]
-Nodata = Optional[Union[int, float]]
 _WRP_CRS = "epsg:3857"
 
 __all__ = [
     "resampling_s2rio",
     "is_resampling_nn",
     "resolve_fill_value",
-    "warp_affine",
-    "warp_affine_rio",
+    "rio_warp_affine",
     "rio_reproject",
 ]
 
@@ -47,25 +46,13 @@ def is_resampling_nn(resampling: Resampling) -> bool:
     return resampling == rasterio.warp.Resampling.nearest
 
 
-def resolve_fill_value(dst_nodata, src_nodata, dtype):
-    dtype = np.dtype(dtype)
-
-    if dst_nodata is not None:
-        return dtype.type(dst_nodata)
-    if np.issubdtype(dtype, np.floating):
-        return dtype.type("nan")
-    if src_nodata is not None:
-        return dtype.type(src_nodata)
-    return dtype.type(0)
-
-
-def warp_affine_rio(
+def rio_warp_affine(
     src: np.ndarray,
     dst: np.ndarray,
     A: Affine,
     resampling: Resampling,
-    src_nodata: Nodata = None,
-    dst_nodata: Nodata = None,
+    src_nodata: SomeNodata = "auto",
+    dst_nodata: SomeNodata = "auto",
     **kwargs,
 ) -> np.ndarray:
     """
@@ -89,36 +76,13 @@ def warp_affine_rio(
 
     s_gbox = GeoBox(wh_(sw, sh), Affine.identity(), _WRP_CRS)
     d_gbox = GeoBox(wh_(dw, dh), A, _WRP_CRS)
-    return _rio_reproject(
-        src, dst, s_gbox, d_gbox, resampling, src_nodata, dst_nodata, **kwargs
-    )
-
-
-def warp_affine(
-    src: np.ndarray,
-    dst: np.ndarray,
-    A: Affine,
-    resampling: Resampling,
-    src_nodata: Nodata = None,
-    dst_nodata: Nodata = None,
-    **kwargs,
-) -> np.ndarray:
-    """
-    Perform Affine warp using best available backend (GDAL via rasterio is the only one so far).
 
-    :param        src: image as ndarray
-    :param        dst: image as ndarray
-    :param          A: Affine transformm, maps from dst_coords to src_coords
-    :param resampling: str resampling strategy
-    :param src_nodata: Value representing "no data" in the source image
-    :param dst_nodata: Value to represent "no data" in the destination image
-
-    :param     kwargs: any other args to pass to implementation
+    src_nodata = resolve_nodata(src_nodata, src.dtype)
+    dst_nodata = resolve_nodata(dst_nodata, dst.dtype)
+    fill_value = resolve_fill_value(dst_nodata, src_nodata, dst.dtype)
 
-    :returns: dst
-    """
-    return warp_affine_rio(
-        src, dst, A, resampling, src_nodata=src_nodata, dst_nodata=dst_nodata, **kwargs
+    return _rio_reproject(
+        src, dst, s_gbox, d_gbox, resampling, src_nodata, fill_value, **kwargs
     )
 
 
@@ -128,8 +92,8 @@ def rio_reproject(
     s_gbox: Union[GeoBox, GCPGeoBox],
     d_gbox: GeoBox,
     resampling: Resampling,
-    src_nodata: Nodata = None,
-    dst_nodata: Nodata = None,
+    src_nodata: SomeNodata = "auto",
+    dst_nodata: SomeNodata = "auto",
     ydim: Optional[int] = None,
     **kwargs,
 ) -> np.ndarray:
@@ -151,13 +115,22 @@ def rio_reproject(
     """
     assert src.ndim == dst.ndim
 
+    src_nodata = resolve_nodata(src_nodata, src.dtype)
+    dst_nodata = resolve_nodata(dst_nodata, dst.dtype)
+    fill_value = resolve_fill_value(dst_nodata, src_nodata, dst.dtype)
+
     if src.ndim == 2:
         return _rio_reproject(
-            src, dst, s_gbox, d_gbox, resampling, src_nodata, dst_nodata, **kwargs
+            src,
+            dst,
+            s_gbox,
+            d_gbox,
+            resampling=resampling,
+            src_nodata=src_nodata,
+            dst_nodata=fill_value,
+            **kwargs,
         )
 
-    fill_value = resolve_fill_value(dst_nodata, src_nodata, dst.dtype)
-
     if ydim is None:
         # Assume last two dimensions are Y/X
         ydim = src.ndim - 2


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index aecb750..0e62389 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import numpy as np
 import pytest
 import xarray as xr
 
@@ -7,6 +8,8 @@ from odc.geo.data import country_geom, ocean_geom
 from odc.geo.geobox import GeoBox
 from odc.geo.xr import rasterize
 
+# pylint: disable=protected-access,import-outside-toplevel,redefined-outer-name
+
 
 @pytest.fixture(scope="session")
 def data_dir():
@@ -45,4 +48,18 @@ def crs():
 
 @pytest.fixture()
 def country(iso3, crs):
-    return country_geom(iso3, crs=crs)
+    yield country_geom(iso3, crs=crs)
+
+
+@pytest.fixture()
+def country_raster(country, resolution):
+    geobox = GeoBox.from_geopolygon(country, resolution=resolution, tight=True)
+    yield rasterize(country, geobox)
+
+
+@pytest.fixture()
+def country_raster_f32(country, resolution):
+    geobox = GeoBox.from_geopolygon(country, resolution=resolution, tight=True)
+    xx = rasterize(country, geobox)
+    xx = xr.where(xx, np.random.uniform(0, 100, xx.shape).astype("float32"), 0)
+    yield xx
diff --git a/tests/test_cog.py b/tests/test_cog.py
index 4f2a23e..1d0991c 100644
--- a/tests/test_cog.py
+++ b/tests/test_cog.py
@@ -345,7 +345,7 @@ def test_norm_compress():
         _gbox.center_pixel.pad(3),
     ],
 )
-@pytest.mark.parametrize("nodata", [None, float("nan"), 0, -999])
+@pytest.mark.parametrize("nodata", ["auto", None, float("nan"), 0, -999])
 @pytest.mark.parametrize("gdal_metadata", [None, "<GDALMetadata></GDALMetadata>"])
 def test_geotiff_metadata(gbox: GeoBox, nodata, gdal_metadata: Optional[str]):
     assert gbox.crs is not None
diff --git a/tests/test_map.py b/tests/test_map.py
index 57e8e9d..8b2ad9e 100644
--- a/tests/test_map.py
+++ b/tests/test_map.py
@@ -5,6 +5,8 @@ import xarray as xr
 from odc.geo._interop import have
 from odc.geo.xr import ODCExtensionDa
 
+# pylint: disable=protected-access,import-outside-toplevel
+
 cmap = np.asarray(
     [
         [153, 153, 102, 255],
diff --git a/tests/test_math.py b/tests/test_math.py
index 7e5829c..55ba516 100644
--- a/tests/test_math.py
+++ b/tests/test_math.py
@@ -2,6 +2,7 @@ import math
 from typing import Tuple
 
 import numpy as np
+import numpy.testing as npt
 import pytest
 from affine import Affine
 
@@ -17,10 +18,13 @@ from odc.geo.math import (
     apply_affine,
     data_resolution_and_offset,
     is_almost_int,
+    is_nodata_empty,
     maybe_int,
     maybe_zero,
     quasi_random_r2,
     resolution_from_affine,
+    resolve_fill_value,
+    resolve_nodata,
     snap_affine,
     snap_grid,
     snap_scale,
@@ -29,6 +33,8 @@ from odc.geo.math import (
 )
 from odc.geo.testutils import mkA
 
+NaN = float("nan")
+
 
 def test_math_ops():
     assert align_up(32, 16) == 32
@@ -391,3 +397,54 @@ def test_align_down_pow2(x: int):
     assert isinstance(y, int)
     assert y <= x
     assert 2 ** int(math.log2(y)) == y
+
+
+@pytest.mark.parametrize(
+    "nodata, dtype, expect",
+    [
+        (None, None, None),
+        (None, "float32", None),
+        ("auto", None, None),
+        ("auto", "float32", NaN),
+        ("auto", np.dtype("float64"), NaN),
+        ("auto", "uint32", None),
+        ("auto", "int16", None),
+        ("auto", "uint8", None),
+        ("auto", "bool", None),
+        ("nan", "float32", NaN),
+        (NaN, "float32", NaN),
+        (np.nan, "float32", NaN),
+        (33, "int8", 33),
+        (-99.99, "float64", -99.99),
+    ],
+)
+def test_resolve_nodata(nodata, dtype, expect):
+    npt.assert_equal(resolve_nodata(nodata, dtype), expect)
+    assert resolve_nodata("auto", dtype, 13) == 13
+
+
+@pytest.mark.parametrize(
+    "dst_nodata, src_nodata, dtype, expect",
+    [
+        (None, None, "uint16", 0),
+        (None, 3, "uint16", 3),
+        (4, 3, "uint16", 4),
+        (10, 0, "uint16", 10),
+        (NaN, 0, "float32", NaN),
+        (None, 0, "float32", NaN),
+        (None, None, "bool", False),
+        (None, None, "uint8", 0),
+    ],
+)
+def test_resolve_fill_value(dst_nodata, src_nodata, dtype, expect):
+    npt.assert_equal(resolve_fill_value(dst_nodata, src_nodata, dtype), expect)
+
+
+def test_empty():
+    assert is_nodata_empty(None) is True
+    assert is_nodata_empty(NaN) is True
+    assert is_nodata_empty(np.nan) is True
+    assert is_nodata_empty(0) is False
+    assert is_nodata_empty(-1) is False
+    assert is_nodata_empty(-1.0) is False
+    assert is_nodata_empty(np.uint8(0)) is False
diff --git a/tests/test_warp.py b/tests/test_warp.py
new file mode 100644
index 0000000..4d5ee5e
--- /dev/null
+++ b/tests/test_warp.py
@@ -0,0 +1,68 @@
+import numpy as np
+import numpy.testing as npt
+import pytest
+import xarray as xr
+from affine import Affine
+
+from odc.geo import MaybeCRS
+from odc.geo.warp import resampling_s2rio, rio_reproject, rio_warp_affine
+
+NaN = float("nan")
+
+
+@pytest.mark.parametrize(
+    "iso3, crs, resolution",
+    [
+        ("AUS", "epsg:4326", 0.1),
+        ("AUS", "epsg:3577", 10_000),
+        ("AUS", "epsg:3857", 10_000),
+        ("NZL", "epsg:3857", 5_000),
+    ],
+)
+@pytest.mark.parametrize("resampling", ["nearest", "bilinear", "average", "sum"])
+def test_warp_nan(country_raster_f32: xr.DataArray, crs: MaybeCRS, resampling: str):
+    xx = country_raster_f32
+    assert isinstance(xx, xr.DataArray)
+    assert xx.odc.crs == crs
+    assert xx.odc.nodata is None
+    assert xx.dtype == "float32"
+
+    mid = xx.shape[0] // 2
+    xx.data[mid, :] = NaN
+    xx.data[:, -10] = NaN
+
+    assert resampling_s2rio(resampling) is not None
+    assert np.isnan(xx.data).sum() > 0
+
+    src_gbox = xx.odc.geobox
+    dst_gbox = src_gbox.zoom_to(shape=100).pad(10)
+
+    yy1 = np.full(dst_gbox.shape, -333, dtype=xx.dtype)
+    yy2 = np.full(dst_gbox.shape, -333, dtype=xx.dtype)
+
+    assert rio_reproject(xx.data, yy1, src_gbox, dst_gbox, resampling=resampling) is yy1
+    assert (
+        rio_reproject(
+            xx.data,
+            yy2,
+            src_gbox,
+            dst_gbox,
+            resampling=resampling,
+            src_nodata=NaN,
+            dst_nodata=NaN,
+        )
+        is yy2
+    )
+
+    npt.assert_array_equal(yy1, yy2)
+
+    # make sure all pixels were replaced
+    assert (yy1 == -333).sum() == 0
+
+    # expect to see NaNs in the output
+    assert np.isnan(yy2).sum() > 0
+
+    A = Affine.identity()
+    xx_ = xx.data.copy() * 0
+    assert rio_warp_affine(xx.data, xx_, A, resampling) is xx_
+    npt.assert_array_equal(xx.data, xx_)
diff --git a/tests/test_xr_interop.py b/tests/test_xr_interop.py
index 1b8272b..bec8693 100644
--- a/tests/test_xr_interop.py
+++ b/tests/test_xr_interop.py
@@ -32,6 +32,8 @@ TEST_GEOBOXES_SMALL_AXIS_ALIGNED = [
     GeoBox.from_bbox((-10, -2, 5, 4), "epsg:3857", tight=True, resolution=resxy_(1, 2)),
 ]
 
+NaN = float("nan")
+
 
 @pytest.fixture
 def geobox_epsg4326():
@@ -163,6 +165,50 @@ def test_purge_crs_info(xx_epsg4326: xr.DataArray):
     assert xx.encoding == {}
 
 
+def test_set_nodata(xx_epsg4326: xr.DataArray):
+    xx = xx_epsg4326
+
+    assert xx.odc.nodata is None
+    assert xx.attrs == {}
+
+    for new_nodata in [0, 999, 3]:
+        xx.odc.nodata = new_nodata
+        assert xx.odc.nodata == new_nodata
+        assert "_FillValue" in xx.encoding
+        assert xx.encoding["_FillValue"] == new_nodata
+        assert xx.attrs["nodata"] == new_nodata
+
+        xx.odc.nodata = None
+        assert xx.odc.nodata is None
+        assert "_FillValue" not in xx.encoding
+        assert xx.attrs == {}
+
+
+def test_nodata_nan(xx_epsg4326: xr.DataArray):
+    xx = xx_epsg4326.astype("float32")
+    assert isinstance(xx.odc, ODCExtensionDa)
+
+    assert xx.odc.nodata is None
+    assert xx.dtype.kind == "f"
+
+    # set to NaN from nothing
+    xx.odc.nodata = NaN
+    assert isinstance(xx.odc.nodata, float)
+    assert np.isnan(xx.odc.nodata)
+    assert np.isnan(xx.encoding["_FillValue"])
+    assert np.isnan(xx.attrs["nodata"])
+
+    xx.odc.nodata = -9999
+    assert xx.odc.nodata == -9999
+
+    # set to NaN from non-NaN, and use a string
+    xx.odc.nodata = "nan"  # type: ignore
+    assert isinstance(xx.odc.nodata, float)
+    assert np.isnan(xx.odc.nodata)
+    assert np.isnan(xx.encoding["_FillValue"])
+    assert np.isnan(xx.attrs["nodata"])
+
+
 def test_odc_extension(xx_epsg4326: xr.DataArray, geobox_epsg4326: GeoBox):
     gbox = geobox_epsg4326
     xx = xx_epsg4326
@@ -361,7 +407,7 @@ def test_wrap_xr():
     assert wrap_xr(data, gbox, always_yx=True).dims == ("y", "x")
     assert wrap_xr(data, gbox, dims=("Y", "X")).dims == ("Y", "X")
 
-    xx = wrap_xr(data, gbox, nodata=None)
+    xx = wrap_xr(data, gbox)
     assert xx.attrs == {}
 
     xx = wrap_xr(data, gbox, nodata=10, some_flag=3)
@@ -385,6 +431,9 @@ def test_wrap_xr():
     assert xx.shape == (*gbox.shape, 1)
     assert xx.band.data.tolist() == ["b0"]
 
+    assert wrap_xr(data.astype("float32"), gbox).odc.nodata is None
+    assert np.isnan(wrap_xr(data.astype("float32"), gbox, nodata=NaN).odc.nodata)
+
 
 @pytest.mark.parametrize("gbox", TEST_GEOBOXES_SMALL_AXIS_ALIGNED)
 @pytest.mark.parametrize("nprefix", [0, 1, 2])
@@ -422,6 +471,9 @@ def test_xr_reproject(xx_epsg4326: xr.DataArray):
     assert xx_epsg4326.odc.nodata is None
     dst_gbox = xx_epsg4326.odc.geobox.zoom_out(1.3)
     xx = xx_epsg4326.odc.reproject(dst_gbox)
+    assert isinstance(xx, xr.DataArray)
+    assert isinstance(xx.odc, ODCExtensionDa)
+
     assert xx.odc.geobox == dst_gbox
     assert xx.encoding["grid_mapping"] == "spatial_ref"
     assert "crs" not in xx.attrs