Merge branch 'dev'

raster_tools/clipping.py

@@ -7,8 +7,12 @@
 from raster_tools.exceptions import RasterNoDataError
 from raster_tools.general import band_concat
 from raster_tools.masking import get_default_null_value
-from raster_tools.raster import Raster, get_raster
-from raster_tools.utils import make_raster_ds
+from raster_tools.raster import (
+    Raster,
+    dataarray_to_xr_raster_ds,
+    get_raster,
+    xr_where_with_meta,
+)
 from raster_tools.vector import get_vector
 
 
@@ -21,54 +25,56 @@ def _clip(
     envelope=False,
 ):
     feat = get_vector(feature)
-    rs = get_raster(data_raster)
-    if rs.crs is None:
+    data_raster = get_raster(data_raster)
+    if data_raster.crs is None:
         raise ValueError("Data raster has no CRS")
     if bounds is not None and len(bounds) != 4:
         raise ValueError("Invalid bounds. Must be a size 4 array or tuple.")
     if envelope and invert:
         raise ValueError("Envelope and invert cannot both be true")
 
-    feat = feat.to_crs(rs.crs)
-
+    crs = data_raster.crs
+    feat = feat.to_crs(crs)
     if trim:
         if bounds is None:
             (bounds,) = dask.compute(feat.bounds)
         else:
-            bounds = np.atleast_1d(bounds)
+            bounds = np.atleast_1d(bounds).ravel()
         try:
-            rs = clip_box(rs, bounds)
+            data_raster = clip_box(data_raster, bounds)
         except RasterNoDataError as err:
             raise RuntimeError(
                 "No data in given bounds. Make sure that the bounds are in the"
                 " same CRS as the data raster."
             ) from err
 
-    feat_rs = feat.to_raster(rs, mask=True)
+    feature_raster = feat.to_raster(data_raster, mask=True)
     if not envelope:
         if invert:
-            clip_mask = feat_rs.to_null_mask()
+            clip_mask = feature_raster.to_null_mask()
         else:
-            clip_mask = ~feat_rs.to_null_mask()
+            clip_mask = ~feature_raster.to_null_mask()
     else:
         if invert:
-            clip_mask = zeros_like(feat_rs, dtype=bool)
+            clip_mask = zeros_like(feature_raster, dtype=bool)
         else:
-            clip_mask = ones_like(feat_rs, dtype=bool)
-
-    nv = rs.null_value if rs._masked else get_default_null_value(rs.dtype)
+            clip_mask = ones_like(feature_raster, dtype=bool)
 
-    if rs.nbands > 1:
-        clip_mask = band_concat([clip_mask] * rs.nbands)
-    xdata_out = xr.where(clip_mask.xdata, rs.xdata, nv)
+    nv = (
+        data_raster.null_value
+        if data_raster._masked
+        else get_default_null_value(data_raster.dtype)
+    )
+    if data_raster.nbands > 1:
+        clip_mask = band_concat([clip_mask] * data_raster.nbands)
+    xdata_out = xr_where_with_meta(
+        clip_mask.xdata, data_raster.xdata, nv, crs=crs, nv=nv
+    )
     xmask_out = ~clip_mask.xdata
-
-    if rs._masked:
-        xmask_out |= rs.xmask
+    if data_raster._masked:
+        xmask_out |= data_raster.xmask
         xdata_out = xdata_out.rio.write_nodata(nv)
-    ds_out = make_raster_ds(xdata_out, xmask_out)
-    if rs.crs is not None:
-        ds_out = ds_out.rio.write_crs(rs.crs)
+    ds_out = dataarray_to_xr_raster_ds(xdata_out, xmask=xmask_out, crs=crs)
     return Raster(ds_out, _fast_path=True)
 
 
@@ -212,20 +218,18 @@ def clip_box(raster, bounds):
         The raster clipped to the given bounds.
 
     """
-    rs = get_raster(raster)
+    raster = get_raster(raster)
     if len(bounds) != 4:
         raise ValueError("Invalid bounds. Must be a size 4 array or tuple.")
     try:
-        xrs = rs.xdata.rio.clip_box(*bounds, auto_expand=True)
+        xdata = raster.xdata.rio.clip_box(*bounds, auto_expand=True)
     except rxr.exceptions.NoDataInBounds as err:
         raise RasterNoDataError(
             "No data found within provided bounds"
         ) from err
-    if rs._masked:
-        xmask = rs.xmask.rio.clip_box(*bounds, auto_expand=True)
+    if raster._masked:
+        xmask = raster.xmask.rio.clip_box(*bounds, auto_expand=True)
     else:
-        xmask = xr.zeros_like(xrs, dtype=bool)
-    ds = make_raster_ds(xrs, xmask)
-    if rs.crs is not None:
-        ds = ds.rio.write_crs(rs.crs)
+        xmask = xr.zeros_like(xdata, dtype=bool)
+    ds = dataarray_to_xr_raster_ds(xdata, xmask=xmask, crs=raster.crs)
     return Raster(ds, _fast_path=True)

raster_tools/creation.py

@@ -2,12 +2,10 @@
 
 import dask.array as da
 import numpy as np
-import xarray as xr
 
 from raster_tools.dtypes import is_int, is_scalar
 from raster_tools.masking import get_default_null_value
-from raster_tools.raster import Raster, get_raster
-from raster_tools.utils import make_raster_ds
+from raster_tools.raster import data_to_raster_like, get_raster
 
 __all__ = [
     "constant_raster",
@@ -60,28 +58,28 @@ def _get_dtype(dtype):
 
 def _copy_mask(template, out_bands):
     if out_bands == template.nbands:
-        return template.xmask.copy()
-    mask = template.get_bands(1).xmask
-    mask = xr.concat([mask] * out_bands, "band")
-    mask["band"] = np.arange(out_bands) + 1
+        return template.mask.copy()
+    mask = template.mask[0]
+    mask = da.stack([mask] * out_bands, axis=0)
     return mask
 
 
-def _build_result(template, xdata, nbands, copy_mask, copy_nv):
+def _build_result(template, data, nbands, copy_mask, copy_nv):
     if template._masked and copy_mask:
-        xmask = _copy_mask(template, nbands)
+        mask = _copy_mask(template, nbands)
         nv = (
             template.null_value
             if copy_nv
-            else get_default_null_value(xdata.dtype)
+            else get_default_null_value(data.dtype)
         )
-        xdata = xdata.rio.write_nodata(nv)
+        burn_mask = True
     else:
-        xmask = xr.zeros_like(xdata, dtype=bool)
-    ds = make_raster_ds(xdata, xmask)
-    if template.crs is not None:
-        ds = ds.rio.write_crs(template.crs)
-    return Raster(ds, _fast_path=True).burn_mask()
+        mask = None
+        nv = None
+        burn_mask = False
+    return data_to_raster_like(
+        data, template, mask=mask, nv=nv, burn=burn_mask
+    )
 
 
 def random_raster(
@@ -136,7 +134,7 @@ def random_raster(
         The resulting raster of random values pulled from the distribution.
 
     """
-    rst = get_raster(raster_template)
+    raster_template = get_raster(raster_template)
     bands = _get_bands(bands)
     if not isinstance(params, Sequence):
         try:
@@ -148,8 +146,8 @@ def random_raster(
     else:
         params = list(params)
 
-    shape = (bands,) + rst.shape[1:]
-    chunks = ((1,) * bands,) + rst.data.chunks[1:]
+    shape = (bands,) + raster_template.shape[1:]
+    chunks = ((1,) * bands,) + raster_template.data.chunks[1:]
 
     dist = distribution.lower()
     if dist not in _VALID_RANDOM_DISTRIBUTIONS:
@@ -183,11 +181,7 @@ def random_raster(
         ndata = da.random.random(size=shape, chunks=chunks)
     # TODO: add more distributions
 
-    cband = np.arange(bands) + 1
-    xdata = xr.DataArray(
-        ndata, coords=(cband, rst.y, rst.x), dims=("band", "y", "x")
-    )
-    return _build_result(rst, xdata, bands, copy_mask, False)
+    return _build_result(raster_template, ndata, bands, copy_mask, False)
 
 
 def empty_like(raster_template, bands=1, dtype=None, copy_mask=False):
@@ -219,13 +213,8 @@ def empty_like(raster_template, bands=1, dtype=None, copy_mask=False):
     shape = (bands,) + rst.shape[1:]
     chunks = ((1,) * bands,) + rst.data.chunks[1:]
     ndata = da.empty(shape, chunks=chunks, dtype=dtype)
-    xdata = xr.DataArray(
-        ndata,
-        coords=(np.arange(bands) + 1, rst.y, rst.x),
-        dims=("band", "y", "x"),
-    )
     copy_null = dtype is None or np.dtype(dtype) == rst.dtype
-    return _build_result(rst, xdata, bands, copy_mask, copy_null)
+    return _build_result(rst, ndata, bands, copy_mask, copy_null)
 
 
 def full_like(raster_template, value, bands=1, dtype=None, copy_mask=False):
@@ -266,13 +255,8 @@ def full_like(raster_template, value, bands=1, dtype=None, copy_mask=False):
     shape = (bands,) + rst.shape[1:]
     chunks = ((1,) * bands,) + rst.data.chunks[1:]
     ndata = da.full(shape, value, chunks=chunks, dtype=dtype)
-    xdata = xr.DataArray(
-        ndata,
-        coords=(np.arange(bands) + 1, rst.y, rst.x),
-        dims=("band", "y", "x"),
-    )
     copy_null = dtype is None or np.dtype(dtype) == rst.dtype
-    return _build_result(rst, xdata, bands, copy_mask, copy_null)
+    return _build_result(rst, ndata, bands, copy_mask, copy_null)
 
 
 def constant_raster(raster_template, value=1, bands=1, copy_mask=False):

raster_tools/distance/cost_distance.py

@@ -20,8 +20,7 @@
     is_scalar,
     is_str,
 )
-from raster_tools.raster import Raster
-from raster_tools.utils import make_raster_ds
+from raster_tools.raster import Raster, dataarray_to_xr_raster_ds
 
 JIT_KWARGS = {"nopython": True, "nogil": True}
 ngjit = nb.jit(**JIT_KWARGS)
@@ -575,13 +574,13 @@ def cost_distance_analysis(costs, sources, elevation=None):
     # Add 1 to match ESRI 0-8 scale
     xtr += 1
 
-    cd_ds = make_raster_ds(
+    cd_ds = dataarray_to_xr_raster_ds(
         xcd.rio.write_nodata(costs.null_value), costs.xmask.copy()
     )
-    tr_ds = make_raster_ds(
+    tr_ds = dataarray_to_xr_raster_ds(
         xtr.rio.write_nodata(_TRACEBACK_NOT_REACHED + 1), costs.xmask.copy()
     )
-    al_ds = make_raster_ds(
+    al_ds = dataarray_to_xr_raster_ds(
         xal.rio.write_nodata(sources_null_value), costs.xmask.copy()
     )
     if costs.crs is not None:

raster_tools/distance/proximity.py

@@ -1,11 +1,10 @@
 import dask.array as da
 import numba as nb
 import numpy as np
-import xarray as xr
 
 from raster_tools.dtypes import F16, F32, F64
 from raster_tools.masking import get_default_null_value
-from raster_tools.raster import Raster, get_raster
+from raster_tools.raster import Raster, data_to_raster_like, get_raster
 from raster_tools.utils import single_band_mappable
 
 __all__ = [
@@ -620,14 +619,7 @@ def _proximity_analysis(
         nodata=nodata,
         mode=mode,
     )
-    xout = xr.DataArray(
-        out_data,
-        coords=raster.xdata.coords,
-        dims=raster.xdata.dims,
-        attrs=raster.xdata.attrs,
-    ).rio.write_nodata(nodata)
-    rs_out = Raster(xout)
-    return rs_out
+    return data_to_raster_like(out_data, raster, nv=nodata)
 
 
 def pa_proximity(

raster_tools/focal.py

@@ -3,7 +3,6 @@
 import dask.array as da
 import numba as nb
 import numpy as np
-import xarray as xr
 from dask_image import ndfilters
 
 from raster_tools.dtypes import (
@@ -18,7 +17,7 @@
     promote_data_dtype,
     promote_dtype_to_float,
 )
-from raster_tools.raster import Raster, get_raster
+from raster_tools.raster import Raster, data_to_xr_raster_ds_like, get_raster
 from raster_tools.stat_common import (
     nan_unique_count_jit,
     nanasm_jit,
@@ -32,7 +31,7 @@
     nansum_jit,
     nanvar_jit,
 )
-from raster_tools.utils import make_raster_ds, single_band_mappable
+from raster_tools.utils import single_band_mappable
 
 __all__ = [
     "check_kernel",
@@ -417,24 +416,16 @@ def focal(raster, focal_type, width_or_radius, height=None, ignore_null=False):
                     break
             data = data.astype(unq_dtype)
 
-    xdata = xr.DataArray(
-        data, coords=raster.xdata.coords, dims=raster.xdata.dims
-    )
-    xmask = raster.xmask.copy()
-    # Nan values will only appear in the result data if there were null values
-    # present in the input. Thus we only need to worry about updating the mask
-    # if the input was masked.
     if raster._masked:
-        nan_mask = np.isnan(xdata)
-        if ignore_null:
-            xmask = nan_mask
-        else:
-            xmask |= nan_mask
-        xdata = xdata.rio.write_nodata(raster.null_value)
-    ds = make_raster_ds(xdata, xmask)
-    if raster.crs is not None:
-        ds = ds.rio.write_crs(raster.crs)
-    return Raster(ds, _fast_path=True).burn_mask()
+        nv = np.nan
+        mask = raster.mask.copy()
+        mask = mask if ignore_null else mask | np.isnan(data)
+        ds = data_to_xr_raster_ds_like(
+            data, raster.xdata, mask=mask, nv=nv, burn=True
+        )
+    else:
+        ds = data_to_xr_raster_ds_like(data, raster.xdata)
+    return Raster(ds, _fast_path=True)
 
 
 def get_focal_window(width_or_radius, height=None):
@@ -587,14 +578,15 @@ def correlate(raster, kernel, mode="constant", cval=0.0):
     if upcast:
         data = data.astype(final_dtype)
 
-    xdata = xr.DataArray(
-        data, coords=raster.xdata.coords, dims=raster.xdata.dims
-    ).rio.write_nodata(raster.null_value)
-    xmask = raster.xmask.copy()
-    ds = make_raster_ds(xdata, xmask)
-    if raster.crs is not None:
-        ds = ds.rio.write_crs(raster.crs)
-    return Raster(ds, _fast_path=True).burn_mask()
+    if raster._masked:
+        nv = raster.null_value
+        mask = raster.mask.copy()
+        ds = data_to_xr_raster_ds_like(
+            data, raster.xdata, mask=mask, nv=nv, burn=True
+        )
+    else:
+        ds = data_to_xr_raster_ds_like(data, raster.xdata)
+    return Raster(ds, _fast_path=True)
 
 
 def convolve(raster, kernel, mode="constant", cval=0.0):