We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Feature request: write_cog to explicitly define no_data_values and dtypes
### Expected behaviour
`write_cog(ds, "output_file.tif",  dtype='float32', nodata=nan)`

### Actual behaviour
`write_cog(ds, "output_file.tif")`
Super lovely, but a lot of our products have problems as the nodata output defaults to zero. Eg. with NDVI, 0 means water not nodata. Explicitly stating the nodata default in the doco would help; additionally making it obvious that you can set the dtype would help. I know that rasterio arguments can be passed, but it took me about a month of using the function until I had the brains to look for this in the rasterio.open doco.....


Happy to help with this one. 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `datacube/utils/cog.py`
Content:
```
1 import warnings
2 import rasterio
3 from rasterio.shutil import copy as rio_copy
4 import numpy as np
5 import xarray as xr
6 import dask
7 from dask.delayed import Delayed
8 from pathlib import Path
9 from typing import Union, Optional, List, Any
10 
11 from .io import check_write_path
12 from .geometry import GeoBox
13 from .geometry.tools import align_up
14 
15 __all__ = ("write_cog", "to_cog")
16 
17 
18 def _adjust_blocksize(block, dim):
19     if block > dim:
20         return align_up(dim, 16)
21     return align_up(block, 16)
22 
23 
24 def _write_cog(pix: np.ndarray,
25                geobox: GeoBox,
26                fname: Union[Path, str],
27                nodata: Optional[float] = None,
28                overwrite: bool = False,
29                blocksize: Optional[int] = None,
30                overview_resampling: Optional[str] = None,
31                overview_levels: Optional[List[int]] = None,
32                ovr_blocksize: Optional[int] = None,
33                **extra_rio_opts) -> Union[Path, bytes]:
34     """Write geo-registered ndarray to GeoTiff file or RAM.
35 
36     :param pix: ``xarray.DataArray`` with crs or (ndarray, geobox, nodata) triple
37     :param fname:  Output file or ":mem:"
38     :param nodata: Set ``nodata`` flag to this value if supplied
39     :param overwrite: True -- replace existing file, False -- abort with IOError exception
40     :param blocksize: Size of internal tiff tiles (512x512 pixels)
41     :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)
42     :param overview_resampling: Use this resampling when computing overviews
43     :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]
44                             to disable overviews supply empty list ``[]``
45     :param extra_rio_opts: Any other option is passed to ``rasterio.open``
46 
47     When fname=":mem:" write COG to memory rather than to a file and return it
48     as memoryview object.
49 
50     NOTE: about memory requirements
51 
52     This function generates temporary in memory tiff file without compression
53     to speed things up. It then adds overviews to this file and only then
54     copies it to the final destination with requested compression settings.
55     This is necessary to produce compliant COG, since COG standard demands
56     overviews to be placed before native resolution data and double pass is the
57     only way to achieve this currently.
58 
59     This means that this function will use about 1.5 to 2 times memory taken by `pix`.
60     """
61     # pylint: disable=too-many-locals
62     if blocksize is None:
63         blocksize = 512
64     if ovr_blocksize is None:
65         ovr_blocksize = blocksize
66     if overview_levels is None:
67         overview_levels = [2 ** i for i in range(1, 6)]
68     if overview_resampling is None:
69         overview_resampling = "nearest"
70 
71     if pix.ndim == 2:
72         h, w = pix.shape
73         nbands = 1
74         band = 1  # type: Any
75     elif pix.ndim == 3:
76         if pix.shape[:2] == geobox.shape:
77             pix = pix.transpose([2, 0, 1])
78         elif pix.shape[-2:] != geobox.shape:
79             raise ValueError('GeoBox shape does not match image shape')
80 
81         nbands, h, w = pix.shape
82         band = tuple(i for i in range(1, nbands + 1))
83     else:
84         raise ValueError("Need 2d or 3d ndarray on input")
85 
86     assert geobox.shape == (h, w)
87 
88     if fname != ":mem:":
89         path = check_write_path(
90             fname, overwrite
91         )  # aborts if overwrite=False and file exists already
92 
93     resampling = rasterio.enums.Resampling[overview_resampling]
94 
95     if (blocksize % 16) != 0:
96         warnings.warn(f"Block size must be a multiple of 16, will be adjusted")
97 
98     rio_opts = dict(
99         width=w,
100         height=h,
101         count=nbands,
102         dtype=pix.dtype.name,
103         crs=str(geobox.crs),
104         transform=geobox.transform,
105         tiled=True,
106         blockxsize=_adjust_blocksize(blocksize, w),
107         blockysize=_adjust_blocksize(blocksize, h),
108         zlevel=6,
109         predictor=3 if pix.dtype.kind == "f" else 2,
110         compress="DEFLATE",
111     )
112 
113     if nodata is not None:
114         rio_opts.update(nodata=nodata)
115 
116     rio_opts.update(extra_rio_opts)
117 
118     # Deal efficiently with "no overviews needed case"
119     if len(overview_levels) == 0:
120         if fname == ":mem:":
121             with rasterio.MemoryFile() as mem:
122                 with mem.open(driver="GTiff", **rio_opts) as dst:
123                     dst.write(pix, band)
124                 return bytes(mem.getbuffer())
125         else:
126             with rasterio.open(path, mode='w', driver='GTiff', **rio_opts) as dst:
127                 dst.write(pix, band)
128             return path
129 
130     # copy re-compresses anyway so skip compression for temp image
131     tmp_opts = rio_opts.copy()
132     tmp_opts.pop("compress")
133     tmp_opts.pop("predictor")
134     tmp_opts.pop("zlevel")
135 
136     with rasterio.Env(GDAL_TIFF_OVR_BLOCKSIZE=ovr_blocksize):
137         with rasterio.MemoryFile() as mem:
138             with mem.open(driver="GTiff", **tmp_opts) as tmp:
139                 tmp.write(pix, band)
140                 tmp.build_overviews(overview_levels, resampling)
141 
142                 if fname == ":mem:":
143                     with rasterio.MemoryFile() as mem2:
144                         rio_copy(
145                             tmp,
146                             mem2.name,
147                             driver="GTiff",
148                             copy_src_overviews=True,
149                             **rio_opts
150                         )
151                         return bytes(mem2.getbuffer())
152 
153                 rio_copy(
154                     tmp, path, driver="GTiff", copy_src_overviews=True, **rio_opts
155                 )
156 
157     return path
158 
159 
160 _delayed_write_cog_to_mem = dask.delayed(  # pylint: disable=invalid-name
161     _write_cog,
162     name="compress-cog", pure=True, nout=1
163 )
164 
165 _delayed_write_cog_to_file = dask.delayed(   # pylint: disable=invalid-name
166     _write_cog,
167     name="save-cog", pure=False, nout=1
168 )
169 
170 
171 def write_cog(geo_im: xr.DataArray,
172               fname: Union[str, Path],
173               blocksize: Optional[int] = None,
174               ovr_blocksize: Optional[int] = None,
175               overview_resampling: Optional[str] = None,
176               overview_levels: Optional[List[int]] = None,
177               **extra_rio_opts) -> Union[Path, bytes, Delayed]:
178     """
179     Save ``xarray.DataArray`` to a file in Cloud Optimized GeoTiff format.
180 
181     This function is "Dask aware". If ``geo_im`` is a Dask array, then the
182     output of this function is also a Dask Delayed object. This allows us to
183     save multiple images concurrently across a Dask cluster. If you are not
184     familiar with Dask this can be confusing, as no operation is performed until
185     ``.compute()`` method is called, so if you call this function with Dask
186     array it will return immediately without writing anything to disk.
187 
188     If you are using Dask to speed up data loading, follow example below:
189 
190     .. code-block:: python
191 
192        # Example: save red band from first time slice to file "red.tif"
193        xx = dc.load(.., dask_chunks=dict(x=1024, y=1024))
194        write_cog(xx.isel(time=0).red, "red.tif").compute()
195        # or compute input first instead
196        write_cog(xx.isel(time=0).red.compute(), "red.tif")
197 
198     :param geo_im: ``xarray.DataArray`` with crs
199     :param fname: Output path or ``":mem:"`` in which case compress to RAM and return bytes
200     :param blocksize: Size of internal tiff tiles (512x512 pixels)
201     :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)
202     :param overview_resampling: Use this resampling when computing overviews
203     :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32],
204                             to disable overviews supply empty list ``[]``
205     :param extra_rio_opts: Any other option is passed to ``rasterio.open``
206 
207     :returns: Path to which output was written
208     :returns: Bytes if ``fname=":mem:"``
209     :returns: ``dask.Delayed`` object if input is a Dask array
210 
211     .. note ::
212 
213        **memory requirements**
214 
215        This function generates temporary in memory tiff file without
216        compression to speed things up. It then adds overviews to this file and
217        only then copies it to the final destination with requested compression
218        settings. This is necessary to produce compliant COG, since COG standard
219        demands overviews to be placed before native resolution data and double
220        pass is the only way to achieve this currently.
221 
222        This means that this function will use about 1.5 to 2 times memory taken by ``geo_im``.
223     """
224     pix = geo_im.data
225     geobox = getattr(geo_im, 'geobox', None)
226     nodata = extra_rio_opts.pop('nodata', None)
227     if nodata is None:
228         nodata = geo_im.attrs.get("nodata", None)
229 
230     if geobox is None:
231         raise ValueError("Need geo-registered array on input")
232 
233     if dask.is_dask_collection(pix):
234         real_op = _delayed_write_cog_to_mem if fname == ":mem:" else _delayed_write_cog_to_file
235     else:
236         real_op = _write_cog
237 
238     return real_op(
239         pix,
240         geobox,
241         fname,
242         nodata=nodata,
243         blocksize=blocksize,
244         ovr_blocksize=ovr_blocksize,
245         overview_resampling=overview_resampling,
246         overview_levels=overview_levels,
247         **extra_rio_opts)
248 
249 
250 def to_cog(geo_im: xr.DataArray,
251            blocksize: Optional[int] = None,
252            ovr_blocksize: Optional[int] = None,
253            overview_resampling: Optional[str] = None,
254            overview_levels: Optional[List[int]] = None,
255            **extra_rio_opts) -> Union[bytes, Delayed]:
256     """
257     Compress ``xarray.DataArray`` into Cloud Optimized GeoTiff bytes in memory.
258 
259     This function doesn't write to disk, it compresses in RAM, this is useful
260     for saving data to S3 or other cloud object stores.
261 
262     This function is "Dask aware". If ``geo_im`` is a Dask array, then the
263     output of this function is also a Dask Delayed object. This allows us to
264     compress multiple images concurrently across a Dask cluster. If you are not
265     familiar with Dask this can be confusing, as no operation is performed until
266     ``.compute()`` method is called, so if you call this function with Dask
267     array it will return immediately without compressing any data.
268 
269     :param geo_im: ``xarray.DataArray`` with crs
270     :param blocksize: Size of internal tiff tiles (512x512 pixels)
271     :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)
272     :param overview_resampling: Use this resampling when computing overviews
273     :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]
274     :param extra_rio_opts: Any other option is passed to ``rasterio.open``
275 
276     :returns: In-memory GeoTiff file as bytes
277     :returns: ``dask.Delayed`` object if input is a Dask array
278 
279     Also see :py:meth:`~datacube.utils.cog.write_cog`
280 
281     """
282     bb = write_cog(geo_im,
283                    ":mem:",
284                    blocksize=blocksize,
285                    ovr_blocksize=ovr_blocksize,
286                    overview_resampling=overview_resampling,
287                    overview_levels=overview_levels,
288                    **extra_rio_opts)
289 
290     assert isinstance(bb, (bytes, Delayed))  # for mypy sake for :mem: output it bytes or delayed bytes
291     return bb
292
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/datacube/utils/cog.py b/datacube/utils/cog.py
--- a/datacube/utils/cog.py
+++ b/datacube/utils/cog.py
@@ -202,6 +202,8 @@
     :param overview_resampling: Use this resampling when computing overviews
     :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']``).
     :param extra_rio_opts: Any other option is passed to ``rasterio.open``
 
     :returns: Path to which output was written
@@ -271,6 +273,8 @@
     :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)
     :param overview_resampling: Use this resampling when computing overviews
     :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]
+    :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']``).
     :param extra_rio_opts: Any other option is passed to ``rasterio.open``
 
     :returns: In-memory GeoTiff file as bytes

{"golden_diff": "diff --git a/datacube/utils/cog.py b/datacube/utils/cog.py\n--- a/datacube/utils/cog.py\n+++ b/datacube/utils/cog.py\n@@ -202,6 +202,8 @@\n     :param overview_resampling: Use this resampling when computing overviews\n     :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32],\n                             to disable overviews supply empty list ``[]``\n+    :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is\n+                   read from the attributes of the input array (``geo_im.attrs['nodata']``).\n     :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n \n     :returns: Path to which output was written\n@@ -271,6 +273,8 @@\n     :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)\n     :param overview_resampling: Use this resampling when computing overviews\n     :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]\n+    :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is\n+                   read from the attributes of the input array (``geo_im.attrs['nodata']``).\n     :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n \n     :returns: In-memory GeoTiff file as bytes\n", "issue": "Feature request: write_cog to explicitly define no_data_values and dtypes\n### Expected behaviour\r\n`write_cog(ds, \"output_file.tif\",  dtype='float32', nodata=nan)`\r\n\r\n### Actual behaviour\r\n`write_cog(ds, \"output_file.tif\")`\r\nSuper lovely, but a lot of our products have problems as the nodata output defaults to zero. Eg. with NDVI, 0 means water not nodata. Explicitly stating the nodata default in the doco would help; additionally making it obvious that you can set the dtype would help. I know that rasterio arguments can be passed, but it took me about a month of using the function until I had the brains to look for this in the rasterio.open doco.....\r\n\r\n\r\nHappy to help with this one. \n", "before_files": [{"content": "import warnings\nimport rasterio\nfrom rasterio.shutil import copy as rio_copy\nimport numpy as np\nimport xarray as xr\nimport dask\nfrom dask.delayed import Delayed\nfrom pathlib import Path\nfrom typing import Union, Optional, List, Any\n\nfrom .io import check_write_path\nfrom .geometry import GeoBox\nfrom .geometry.tools import align_up\n\n__all__ = (\"write_cog\", \"to_cog\")\n\n\ndef _adjust_blocksize(block, dim):\n    if block > dim:\n        return align_up(dim, 16)\n    return align_up(block, 16)\n\n\ndef _write_cog(pix: np.ndarray,\n               geobox: GeoBox,\n               fname: Union[Path, str],\n               nodata: Optional[float] = None,\n               overwrite: bool = False,\n               blocksize: Optional[int] = None,\n               overview_resampling: Optional[str] = None,\n               overview_levels: Optional[List[int]] = None,\n               ovr_blocksize: Optional[int] = None,\n               **extra_rio_opts) -> Union[Path, bytes]:\n    \"\"\"Write geo-registered ndarray to GeoTiff file or RAM.\n\n    :param pix: ``xarray.DataArray`` with crs or (ndarray, geobox, nodata) triple\n    :param fname:  Output file or \":mem:\"\n    :param nodata: Set ``nodata`` flag to this value if supplied\n    :param overwrite: True -- replace existing file, False -- abort with IOError exception\n    :param blocksize: Size of internal tiff tiles (512x512 pixels)\n    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)\n    :param overview_resampling: Use this resampling when computing overviews\n    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]\n                            to disable overviews supply empty list ``[]``\n    :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n\n    When fname=\":mem:\" write COG to memory rather than to a file and return it\n    as memoryview object.\n\n    NOTE: about memory requirements\n\n    This function generates temporary in memory tiff file without compression\n    to speed things up. It then adds overviews to this file and only then\n    copies it to the final destination with requested compression settings.\n    This is necessary to produce compliant COG, since COG standard demands\n    overviews to be placed before native resolution data and double pass is the\n    only way to achieve this currently.\n\n    This means that this function will use about 1.5 to 2 times memory taken by `pix`.\n    \"\"\"\n    # pylint: disable=too-many-locals\n    if blocksize is None:\n        blocksize = 512\n    if ovr_blocksize is None:\n        ovr_blocksize = blocksize\n    if overview_levels is None:\n        overview_levels = [2 ** i for i in range(1, 6)]\n    if overview_resampling is None:\n        overview_resampling = \"nearest\"\n\n    if pix.ndim == 2:\n        h, w = pix.shape\n        nbands = 1\n        band = 1  # type: Any\n    elif pix.ndim == 3:\n        if pix.shape[:2] == geobox.shape:\n            pix = pix.transpose([2, 0, 1])\n        elif pix.shape[-2:] != geobox.shape:\n            raise ValueError('GeoBox shape does not match image shape')\n\n        nbands, h, w = pix.shape\n        band = tuple(i for i in range(1, nbands + 1))\n    else:\n        raise ValueError(\"Need 2d or 3d ndarray on input\")\n\n    assert geobox.shape == (h, w)\n\n    if fname != \":mem:\":\n        path = check_write_path(\n            fname, overwrite\n        )  # aborts if overwrite=False and file exists already\n\n    resampling = rasterio.enums.Resampling[overview_resampling]\n\n    if (blocksize % 16) != 0:\n        warnings.warn(f\"Block size must be a multiple of 16, will be adjusted\")\n\n    rio_opts = dict(\n        width=w,\n        height=h,\n        count=nbands,\n        dtype=pix.dtype.name,\n        crs=str(geobox.crs),\n        transform=geobox.transform,\n        tiled=True,\n        blockxsize=_adjust_blocksize(blocksize, w),\n        blockysize=_adjust_blocksize(blocksize, h),\n        zlevel=6,\n        predictor=3 if pix.dtype.kind == \"f\" else 2,\n        compress=\"DEFLATE\",\n    )\n\n    if nodata is not None:\n        rio_opts.update(nodata=nodata)\n\n    rio_opts.update(extra_rio_opts)\n\n    # Deal efficiently with \"no overviews needed case\"\n    if len(overview_levels) == 0:\n        if fname == \":mem:\":\n            with rasterio.MemoryFile() as mem:\n                with mem.open(driver=\"GTiff\", **rio_opts) as dst:\n                    dst.write(pix, band)\n                return bytes(mem.getbuffer())\n        else:\n            with rasterio.open(path, mode='w', driver='GTiff', **rio_opts) as dst:\n                dst.write(pix, band)\n            return path\n\n    # copy re-compresses anyway so skip compression for temp image\n    tmp_opts = rio_opts.copy()\n    tmp_opts.pop(\"compress\")\n    tmp_opts.pop(\"predictor\")\n    tmp_opts.pop(\"zlevel\")\n\n    with rasterio.Env(GDAL_TIFF_OVR_BLOCKSIZE=ovr_blocksize):\n        with rasterio.MemoryFile() as mem:\n            with mem.open(driver=\"GTiff\", **tmp_opts) as tmp:\n                tmp.write(pix, band)\n                tmp.build_overviews(overview_levels, resampling)\n\n                if fname == \":mem:\":\n                    with rasterio.MemoryFile() as mem2:\n                        rio_copy(\n                            tmp,\n                            mem2.name,\n                            driver=\"GTiff\",\n                            copy_src_overviews=True,\n                            **rio_opts\n                        )\n                        return bytes(mem2.getbuffer())\n\n                rio_copy(\n                    tmp, path, driver=\"GTiff\", copy_src_overviews=True, **rio_opts\n                )\n\n    return path\n\n\n_delayed_write_cog_to_mem = dask.delayed(  # pylint: disable=invalid-name\n    _write_cog,\n    name=\"compress-cog\", pure=True, nout=1\n)\n\n_delayed_write_cog_to_file = dask.delayed(   # pylint: disable=invalid-name\n    _write_cog,\n    name=\"save-cog\", pure=False, nout=1\n)\n\n\ndef write_cog(geo_im: xr.DataArray,\n              fname: Union[str, Path],\n              blocksize: Optional[int] = None,\n              ovr_blocksize: Optional[int] = None,\n              overview_resampling: Optional[str] = None,\n              overview_levels: Optional[List[int]] = None,\n              **extra_rio_opts) -> Union[Path, bytes, Delayed]:\n    \"\"\"\n    Save ``xarray.DataArray`` to a file in Cloud Optimized GeoTiff format.\n\n    This function is \"Dask aware\". If ``geo_im`` is a Dask array, then the\n    output of this function is also a Dask Delayed object. This allows us to\n    save multiple images concurrently across a Dask cluster. If you are not\n    familiar with Dask this can be confusing, as no operation is performed until\n    ``.compute()`` method is called, so if you call this function with Dask\n    array it will return immediately without writing anything to disk.\n\n    If you are using Dask to speed up data loading, follow example below:\n\n    .. code-block:: python\n\n       # Example: save red band from first time slice to file \"red.tif\"\n       xx = dc.load(.., dask_chunks=dict(x=1024, y=1024))\n       write_cog(xx.isel(time=0).red, \"red.tif\").compute()\n       # or compute input first instead\n       write_cog(xx.isel(time=0).red.compute(), \"red.tif\")\n\n    :param geo_im: ``xarray.DataArray`` with crs\n    :param fname: Output path or ``\":mem:\"`` in which case compress to RAM and return bytes\n    :param blocksize: Size of internal tiff tiles (512x512 pixels)\n    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)\n    :param overview_resampling: Use this resampling when computing overviews\n    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32],\n                            to disable overviews supply empty list ``[]``\n    :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n\n    :returns: Path to which output was written\n    :returns: Bytes if ``fname=\":mem:\"``\n    :returns: ``dask.Delayed`` object if input is a Dask array\n\n    .. note ::\n\n       **memory requirements**\n\n       This function generates temporary in memory tiff file without\n       compression to speed things up. It then adds overviews to this file and\n       only then copies it to the final destination with requested compression\n       settings. This is necessary to produce compliant COG, since COG standard\n       demands overviews to be placed before native resolution data and double\n       pass is the only way to achieve this currently.\n\n       This means that this function will use about 1.5 to 2 times memory taken by ``geo_im``.\n    \"\"\"\n    pix = geo_im.data\n    geobox = getattr(geo_im, 'geobox', None)\n    nodata = extra_rio_opts.pop('nodata', None)\n    if nodata is None:\n        nodata = geo_im.attrs.get(\"nodata\", None)\n\n    if geobox is None:\n        raise ValueError(\"Need geo-registered array on input\")\n\n    if dask.is_dask_collection(pix):\n        real_op = _delayed_write_cog_to_mem if fname == \":mem:\" else _delayed_write_cog_to_file\n    else:\n        real_op = _write_cog\n\n    return real_op(\n        pix,\n        geobox,\n        fname,\n        nodata=nodata,\n        blocksize=blocksize,\n        ovr_blocksize=ovr_blocksize,\n        overview_resampling=overview_resampling,\n        overview_levels=overview_levels,\n        **extra_rio_opts)\n\n\ndef to_cog(geo_im: xr.DataArray,\n           blocksize: Optional[int] = None,\n           ovr_blocksize: Optional[int] = None,\n           overview_resampling: Optional[str] = None,\n           overview_levels: Optional[List[int]] = None,\n           **extra_rio_opts) -> Union[bytes, Delayed]:\n    \"\"\"\n    Compress ``xarray.DataArray`` into Cloud Optimized GeoTiff bytes in memory.\n\n    This function doesn't write to disk, it compresses in RAM, this is useful\n    for saving data to S3 or other cloud object stores.\n\n    This function is \"Dask aware\". If ``geo_im`` is a Dask array, then the\n    output of this function is also a Dask Delayed object. This allows us to\n    compress multiple images concurrently across a Dask cluster. If you are not\n    familiar with Dask this can be confusing, as no operation is performed until\n    ``.compute()`` method is called, so if you call this function with Dask\n    array it will return immediately without compressing any data.\n\n    :param geo_im: ``xarray.DataArray`` with crs\n    :param blocksize: Size of internal tiff tiles (512x512 pixels)\n    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)\n    :param overview_resampling: Use this resampling when computing overviews\n    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]\n    :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n\n    :returns: In-memory GeoTiff file as bytes\n    :returns: ``dask.Delayed`` object if input is a Dask array\n\n    Also see :py:meth:`~datacube.utils.cog.write_cog`\n\n    \"\"\"\n    bb = write_cog(geo_im,\n                   \":mem:\",\n                   blocksize=blocksize,\n                   ovr_blocksize=ovr_blocksize,\n                   overview_resampling=overview_resampling,\n                   overview_levels=overview_levels,\n                   **extra_rio_opts)\n\n    assert isinstance(bb, (bytes, Delayed))  # for mypy sake for :mem: output it bytes or delayed bytes\n    return bb\n", "path": "datacube/utils/cog.py"}], "after_files": [{"content": "import warnings\nimport rasterio\nfrom rasterio.shutil import copy as rio_copy\nimport numpy as np\nimport xarray as xr\nimport dask\nfrom dask.delayed import Delayed\nfrom pathlib import Path\nfrom typing import Union, Optional, List, Any\n\nfrom .io import check_write_path\nfrom .geometry import GeoBox\nfrom .geometry.tools import align_up\n\n__all__ = (\"write_cog\", \"to_cog\")\n\n\ndef _adjust_blocksize(block, dim):\n    if block > dim:\n        return align_up(dim, 16)\n    return align_up(block, 16)\n\n\ndef _write_cog(pix: np.ndarray,\n               geobox: GeoBox,\n               fname: Union[Path, str],\n               nodata: Optional[float] = None,\n               overwrite: bool = False,\n               blocksize: Optional[int] = None,\n               overview_resampling: Optional[str] = None,\n               overview_levels: Optional[List[int]] = None,\n               ovr_blocksize: Optional[int] = None,\n               **extra_rio_opts) -> Union[Path, bytes]:\n    \"\"\"Write geo-registered ndarray to GeoTiff file or RAM.\n\n    :param pix: ``xarray.DataArray`` with crs or (ndarray, geobox, nodata) triple\n    :param fname:  Output file or \":mem:\"\n    :param nodata: Set ``nodata`` flag to this value if supplied\n    :param overwrite: True -- replace existing file, False -- abort with IOError exception\n    :param blocksize: Size of internal tiff tiles (512x512 pixels)\n    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)\n    :param overview_resampling: Use this resampling when computing overviews\n    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]\n                            to disable overviews supply empty list ``[]``\n    :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n\n    When fname=\":mem:\" write COG to memory rather than to a file and return it\n    as memoryview object.\n\n    NOTE: about memory requirements\n\n    This function generates temporary in memory tiff file without compression\n    to speed things up. It then adds overviews to this file and only then\n    copies it to the final destination with requested compression settings.\n    This is necessary to produce compliant COG, since COG standard demands\n    overviews to be placed before native resolution data and double pass is the\n    only way to achieve this currently.\n\n    This means that this function will use about 1.5 to 2 times memory taken by `pix`.\n    \"\"\"\n    # pylint: disable=too-many-locals\n    if blocksize is None:\n        blocksize = 512\n    if ovr_blocksize is None:\n        ovr_blocksize = blocksize\n    if overview_levels is None:\n        overview_levels = [2 ** i for i in range(1, 6)]\n    if overview_resampling is None:\n        overview_resampling = \"nearest\"\n\n    if pix.ndim == 2:\n        h, w = pix.shape\n        nbands = 1\n        band = 1  # type: Any\n    elif pix.ndim == 3:\n        if pix.shape[:2] == geobox.shape:\n            pix = pix.transpose([2, 0, 1])\n        elif pix.shape[-2:] != geobox.shape:\n            raise ValueError('GeoBox shape does not match image shape')\n\n        nbands, h, w = pix.shape\n        band = tuple(i for i in range(1, nbands + 1))\n    else:\n        raise ValueError(\"Need 2d or 3d ndarray on input\")\n\n    assert geobox.shape == (h, w)\n\n    if fname != \":mem:\":\n        path = check_write_path(\n            fname, overwrite\n        )  # aborts if overwrite=False and file exists already\n\n    resampling = rasterio.enums.Resampling[overview_resampling]\n\n    if (blocksize % 16) != 0:\n        warnings.warn(f\"Block size must be a multiple of 16, will be adjusted\")\n\n    rio_opts = dict(\n        width=w,\n        height=h,\n        count=nbands,\n        dtype=pix.dtype.name,\n        crs=str(geobox.crs),\n        transform=geobox.transform,\n        tiled=True,\n        blockxsize=_adjust_blocksize(blocksize, w),\n        blockysize=_adjust_blocksize(blocksize, h),\n        zlevel=6,\n        predictor=3 if pix.dtype.kind == \"f\" else 2,\n        compress=\"DEFLATE\",\n    )\n\n    if nodata is not None:\n        rio_opts.update(nodata=nodata)\n\n    rio_opts.update(extra_rio_opts)\n\n    # Deal efficiently with \"no overviews needed case\"\n    if len(overview_levels) == 0:\n        if fname == \":mem:\":\n            with rasterio.MemoryFile() as mem:\n                with mem.open(driver=\"GTiff\", **rio_opts) as dst:\n                    dst.write(pix, band)\n                return bytes(mem.getbuffer())\n        else:\n            with rasterio.open(path, mode='w', driver='GTiff', **rio_opts) as dst:\n                dst.write(pix, band)\n            return path\n\n    # copy re-compresses anyway so skip compression for temp image\n    tmp_opts = rio_opts.copy()\n    tmp_opts.pop(\"compress\")\n    tmp_opts.pop(\"predictor\")\n    tmp_opts.pop(\"zlevel\")\n\n    with rasterio.Env(GDAL_TIFF_OVR_BLOCKSIZE=ovr_blocksize):\n        with rasterio.MemoryFile() as mem:\n            with mem.open(driver=\"GTiff\", **tmp_opts) as tmp:\n                tmp.write(pix, band)\n                tmp.build_overviews(overview_levels, resampling)\n\n                if fname == \":mem:\":\n                    with rasterio.MemoryFile() as mem2:\n                        rio_copy(\n                            tmp,\n                            mem2.name,\n                            driver=\"GTiff\",\n                            copy_src_overviews=True,\n                            **rio_opts\n                        )\n                        return bytes(mem2.getbuffer())\n\n                rio_copy(\n                    tmp, path, driver=\"GTiff\", copy_src_overviews=True, **rio_opts\n                )\n\n    return path\n\n\n_delayed_write_cog_to_mem = dask.delayed(  # pylint: disable=invalid-name\n    _write_cog,\n    name=\"compress-cog\", pure=True, nout=1\n)\n\n_delayed_write_cog_to_file = dask.delayed(   # pylint: disable=invalid-name\n    _write_cog,\n    name=\"save-cog\", pure=False, nout=1\n)\n\n\ndef write_cog(geo_im: xr.DataArray,\n              fname: Union[str, Path],\n              blocksize: Optional[int] = None,\n              ovr_blocksize: Optional[int] = None,\n              overview_resampling: Optional[str] = None,\n              overview_levels: Optional[List[int]] = None,\n              **extra_rio_opts) -> Union[Path, bytes, Delayed]:\n    \"\"\"\n    Save ``xarray.DataArray`` to a file in Cloud Optimized GeoTiff format.\n\n    This function is \"Dask aware\". If ``geo_im`` is a Dask array, then the\n    output of this function is also a Dask Delayed object. This allows us to\n    save multiple images concurrently across a Dask cluster. If you are not\n    familiar with Dask this can be confusing, as no operation is performed until\n    ``.compute()`` method is called, so if you call this function with Dask\n    array it will return immediately without writing anything to disk.\n\n    If you are using Dask to speed up data loading, follow example below:\n\n    .. code-block:: python\n\n       # Example: save red band from first time slice to file \"red.tif\"\n       xx = dc.load(.., dask_chunks=dict(x=1024, y=1024))\n       write_cog(xx.isel(time=0).red, \"red.tif\").compute()\n       # or compute input first instead\n       write_cog(xx.isel(time=0).red.compute(), \"red.tif\")\n\n    :param geo_im: ``xarray.DataArray`` with crs\n    :param fname: Output path or ``\":mem:\"`` in which case compress to RAM and return bytes\n    :param blocksize: Size of internal tiff tiles (512x512 pixels)\n    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)\n    :param overview_resampling: Use this resampling when computing overviews\n    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32],\n                            to disable overviews supply empty list ``[]``\n    :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is\n                   read from the attributes of the input array (``geo_im.attrs['nodata']``).\n    :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n\n    :returns: Path to which output was written\n    :returns: Bytes if ``fname=\":mem:\"``\n    :returns: ``dask.Delayed`` object if input is a Dask array\n\n    .. note ::\n\n       **memory requirements**\n\n       This function generates temporary in memory tiff file without\n       compression to speed things up. It then adds overviews to this file and\n       only then copies it to the final destination with requested compression\n       settings. This is necessary to produce compliant COG, since COG standard\n       demands overviews to be placed before native resolution data and double\n       pass is the only way to achieve this currently.\n\n       This means that this function will use about 1.5 to 2 times memory taken by ``geo_im``.\n    \"\"\"\n    pix = geo_im.data\n    geobox = getattr(geo_im, 'geobox', None)\n    nodata = extra_rio_opts.pop('nodata', None)\n    if nodata is None:\n        nodata = geo_im.attrs.get(\"nodata\", None)\n\n    if geobox is None:\n        raise ValueError(\"Need geo-registered array on input\")\n\n    if dask.is_dask_collection(pix):\n        real_op = _delayed_write_cog_to_mem if fname == \":mem:\" else _delayed_write_cog_to_file\n    else:\n        real_op = _write_cog\n\n    return real_op(\n        pix,\n        geobox,\n        fname,\n        nodata=nodata,\n        blocksize=blocksize,\n        ovr_blocksize=ovr_blocksize,\n        overview_resampling=overview_resampling,\n        overview_levels=overview_levels,\n        **extra_rio_opts)\n\n\ndef to_cog(geo_im: xr.DataArray,\n           blocksize: Optional[int] = None,\n           ovr_blocksize: Optional[int] = None,\n           overview_resampling: Optional[str] = None,\n           overview_levels: Optional[List[int]] = None,\n           **extra_rio_opts) -> Union[bytes, Delayed]:\n    \"\"\"\n    Compress ``xarray.DataArray`` into Cloud Optimized GeoTiff bytes in memory.\n\n    This function doesn't write to disk, it compresses in RAM, this is useful\n    for saving data to S3 or other cloud object stores.\n\n    This function is \"Dask aware\". If ``geo_im`` is a Dask array, then the\n    output of this function is also a Dask Delayed object. This allows us to\n    compress multiple images concurrently across a Dask cluster. If you are not\n    familiar with Dask this can be confusing, as no operation is performed until\n    ``.compute()`` method is called, so if you call this function with Dask\n    array it will return immediately without compressing any data.\n\n    :param geo_im: ``xarray.DataArray`` with crs\n    :param blocksize: Size of internal tiff tiles (512x512 pixels)\n    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)\n    :param overview_resampling: Use this resampling when computing overviews\n    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]\n    :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is\n                   read from the attributes of the input array (``geo_im.attrs['nodata']``).\n    :param extra_rio_opts: Any other option is passed to ``rasterio.open``\n\n    :returns: In-memory GeoTiff file as bytes\n    :returns: ``dask.Delayed`` object if input is a Dask array\n\n    Also see :py:meth:`~datacube.utils.cog.write_cog`\n\n    \"\"\"\n    bb = write_cog(geo_im,\n                   \":mem:\",\n                   blocksize=blocksize,\n                   ovr_blocksize=ovr_blocksize,\n                   overview_resampling=overview_resampling,\n                   overview_levels=overview_levels,\n                   **extra_rio_opts)\n\n    assert isinstance(bb, (bytes, Delayed))  # for mypy sake for :mem: output it bytes or delayed bytes\n    return bb\n", "path": "datacube/utils/cog.py"}]}