Merge branch 'dev'

raster_tools/_compat.py

@@ -1,5 +1,9 @@
 import os
 
+import numpy as np
+
+from raster_tools.utils import version_to_tuple
+
 # Force the use of shapely 2 instead of pygeos in geopandas
 os.environ["USE_PYGEOS"] = "0"
 
@@ -10,3 +14,7 @@
 # backend for line_stats.
 if gpd.options.use_pygeos:
     gpd.options.use_pygeos = False
+
+
+# Numpy 2.0 made several changes to type promotion rules.
+NUMPY_GE_2 = version_to_tuple(np.__version__) >= (2, 0, 0)

raster_tools/focal.py

@@ -11,6 +11,7 @@
     U16,
     U32,
     U64,
+    get_dtype_min_max,
     is_bool,
     is_float,
     is_int,
@@ -411,7 +412,7 @@ def focal(raster, focal_type, width_or_radius, height=None, ignore_null=False):
             n = window.size
             unq_dtype = None
             for dt in (U8, U16, U32, U64):
-                if np.can_cast(n, dt):
+                if n <= get_dtype_min_max(dt)[1]:
                     unq_dtype = dt
                     break
             data = data.astype(unq_dtype)

raster_tools/io.py

@@ -17,6 +17,7 @@
     RasterDataError,
     RasterIOError,
 )
+from raster_tools.masking import get_default_null_value
 from raster_tools.utils import to_chunk_dict, validate_path
 
 
@@ -201,29 +202,21 @@ def _get_valid_variables(meta, ignore_too_many_dims):
     return valid
 
 
-def _build_xr_raster(path, variable, affine, crs, xarray_kwargs):
+def _build_raster(path, variable, affine, crs, xarray_kwargs):
+    from raster_tools.raster import data_to_raster
+
     if affine is None:
         affine = Affine(1, 0, 0, 0, -1, 0, 0)
     kwargs = xarray_kwargs.copy()
     kwargs["chunks"] = "auto"
     var = xr.open_dataset(path, **kwargs)[variable].squeeze()
-    var_data = var.data
-    if var.ndim == 2:
-        var_data = np.expand_dims(var_data, axis=0)
-    var_data = var_data.rechunk((1, "auto", "auto"))
-    band = np.array(list(range(var_data.shape[0])))
     x = var[var.rio.x_dim].to_numpy()
     y = var[var.rio.y_dim].to_numpy()
-    new_var = xr.DataArray(
-        var_data, dims=["band", "y", "x"], coords=(band, y, x)
-    )
-    new_var = new_var.rio.write_transform(affine)
-    if crs is not None:
-        new_var = new_var.rio.write_crs(crs)
     nv = var._FillValue if "_FillValue" in var.attrs else var.rio.nodata
-    if nv is not None:
-        new_var = new_var.rio.write_nodata(nv)
-    return new_var
+    raster = data_to_raster(var.data, x=x, y=y, affine=affine, crs=crs, nv=nv)
+    if nv is None or np.isnan(nv):
+        raster = raster.set_null_value(get_default_null_value(raster.dtype))
+    return raster
 
 
 def _get_affine(ds):
@@ -287,8 +280,6 @@ def open_dataset(
         raster.
 
     """
-    from raster_tools.raster import Raster
-
     if xarray_kwargs is None:
         xarray_kwargs = {}
     xarray_kwargs["decode_coords"] = "all"
@@ -299,6 +290,5 @@ def open_dataset(
     tmp_ds = None
     ds = {}
     for v in data_vars:
-        var = _build_xr_raster(path, v, affine, crs, xarray_kwargs)
-        ds[v] = Raster(var)
+        ds[v] = _build_raster(path, v, affine, crs, xarray_kwargs)
     return ds

raster_tools/masking.py

@@ -18,6 +18,7 @@
     U16,
     U32,
     U64,
+    get_dtype_min_max,
     is_bool,
     is_float,
 )
@@ -69,12 +70,25 @@ def reconcile_nullvalue_with_dtype(null_value, dtype, warn=False):
         return get_default_null_value(dtype)
 
     dtype = np.dtype(dtype)
+    null_value = np.min_scalar_type(null_value).type(null_value)
     if np.can_cast(null_value, dtype):
         return null_value
 
     original_nv = null_value
+    # TODO: stop checking if the value is an integer. Too much complexity. If
+    # it is a float, treat it as a float.
     if is_float(original_nv) and float(original_nv).is_integer():
-        nv = int(original_nv)
+        if any(
+            np.allclose(original_nv, v)
+            for v in (*get_dtype_min_max(F32), *get_dtype_min_max(F64))
+        ):
+            nv = original_nv
+        else:
+            nv = int(original_nv)
+            min_type = np.min_scalar_type(nv)
+            nv = min_type.type(nv)
+            if min_type.kind == "O":
+                nv = original_nv
     else:
         nv = original_nv
     if np.can_cast(nv, dtype) or (is_bool(dtype) and nv in (0, 1)):

raster_tools/raster.py

@@ -17,6 +17,7 @@
 from odc.geo.geobox import GeoBox
 from shapely.geometry import box
 
+from raster_tools._compat import NUMPY_GE_2
 from raster_tools.dask_utils import (
     chunks_to_array_locations,
     dask_nanmax,
@@ -165,6 +166,18 @@ def _normalize_ufunc_other(other, this):
 
 
 def _apply_ufunc(ufunc, this, left, right=None, kwargs=None, out=None):
+    ufname = ufunc.__name__
+    if NUMPY_GE_2:
+        # If the ufunc is ldexp and the left arg is a python scalar, casting
+        # causes the resulting dtype to be different from the analongous
+        # (scalar X numpy-array) combination. So we disable casting for this
+        # special case.
+        if type(left) in (int, float) and ufname != "ldexp":
+            left = np.int64(left) if type(left) is int else np.float64(left)
+        if type(right) in (int, float):
+            right = (
+                np.int64(right) if type(right) is int else np.float64(right)
+            )
     args = [left]
     if right is not None:
         args.append(right)
@@ -179,7 +192,6 @@ def _apply_ufunc(ufunc, this, left, right=None, kwargs=None, out=None):
     elif isinstance(other, Raster) and other.crs is not None:
         out_crs = other.crs
 
-    ufname = ufunc.__name__
     if ufname.startswith("bitwise") and any(is_float(t) for t in types):
         raise TypeError(
             "Bitwise operations are not compatible with float dtypes. You may"
@@ -309,7 +321,27 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             right = other
         else:
             left = other
-        return _apply_ufunc(ufunc, self, left, right=right, out=out)
+
+        try:
+            return _apply_ufunc(ufunc, self, left, right=right, out=out)
+        except TypeError as err:
+            if right is not None and str(err).startswith(
+                "operand type(s) all returned NotImplemented from "
+                "__array_ufunc__"
+            ):
+                msgs = []
+                for obj in (left, right):
+                    if isinstance(obj, Raster):
+                        msg = f"Raster<{obj.dtype}>"
+                    else:
+                        msg = type(obj)
+                    msgs.append(msg)
+                left_msg, right_msg = msgs
+                raise TypeError(
+                    f"Could not apply {ufunc} to types "
+                    f"{left_msg} and {right_msg}"
+                ) from err
+            raise err
 
     def __array__(self, dtype=None):
         return self._ds.raster.__array__(dtype)

raster_tools/rasterize.py

@@ -428,9 +428,20 @@ def _rasterize_spatial_aware(
         like = like.get_bands(1)
 
     sparts = dgdf.spatial_partitions.to_frame("geometry")
+    # The null value can be different from the fill value when mask=True so set
+    # a separate variable.
+    nv = fill
     if mask:
-        fill = 1 if mask_invert else 0
         target_dtype = U8
+        # The cells that will eventually be set to null values will always have
+        # 0 in them after the rasterize operation when mask=True. final_nv,
+        # below, will then be used to set the actual null values.
+        nv = target_dtype.type(0)
+        final_nv = None
+        if fill is not None:
+            # Store the original value to set later
+            final_nv = fill
+        fill = target_dtype.type(1 if mask_invert else 0)
         # Only need geometry for masking
         dgdf = dgdf.geometry.to_frame("geometry")
     else:
@@ -475,7 +486,10 @@ def _rasterize_spatial_aware(
         mask=mask,
         mask_invert=mask_invert,
     )
-    return data_to_raster_like(out_data, like, nv=fill)
+    raster = data_to_raster_like(out_data, like, nv=nv)
+    if mask and final_nv is not None:
+        raster = raster.set_null_value(final_nv)
+    return raster
 
 
 def _rasterize_spatial_naive(

raster_tools/vector.py

@@ -666,7 +666,7 @@ def to_crs(self, crs):
         The transformed vector(s).
 
         """
-        return Vector(self._geo.to_crs(crs))
+        return Vector(self._geo.to_crs(crs), self._size)
 
     def cast_field(self, name, dtype):
         """Cast a field to the specified `dtype`.
@@ -714,7 +714,11 @@ def __getitem__(self, idx):
         # [idx], however. This breaks future iteration we reset it to [0] to
         # allow iteration over the resulting Vector.
         subgeo = self._geo.loc[[idx]]
-        subgeo.index = dask.array.from_array([0]).to_dask_dataframe()
+        subgeo = subgeo.assign(__new_idx__=0)
+        old_idx_name = subgeo.index.name
+        subgeo = subgeo.set_index("__new_idx__", divisions=[0, 1])
+        new_idx_name = "index" if old_idx_name is None else old_idx_name
+        subgeo.index = subgeo.index.rename(new_idx_name)
         return Vector(subgeo, 1)
 
     def buffer(self, *args, **kwargs):

requirements/310-ci.yml

@@ -5,14 +5,14 @@ dependencies:
   - python=3.10
   - affine
   - cfgrib
-  - dask<=2024.2.1
+  - dask
   - dask-geopandas
   - dask-image
   - fiona
   - geopandas
   - netcdf4
   - numba
-  - numpy<2
+  - numpy
   - odc-geo
   - pandas
   - pyogrio

requirements/311-ci.yml

@@ -5,14 +5,14 @@ dependencies:
   - python=3.11
   - affine
   - cfgrib
-  - dask<=2024.2.1
+  - dask
   - dask-geopandas
   - dask-image
   - fiona
   - geopandas
   - netcdf4
   - numba
-  - numpy<2
+  - numpy
   - odc-geo
   - pandas
   - pyogrio

requirements/312-ci.yml

@@ -5,14 +5,14 @@ dependencies:
   - python=3.12
   - affine
   - cfgrib
-  - dask<=2024.2.1
+  - dask
   - dask-geopandas
   - dask-image
   - fiona
   - geopandas
   - netcdf4
   - numba
-  - numpy<2
+  - numpy
   - odc-geo
   - pandas
   - pyogrio

requirements/39-ci.yml

@@ -5,15 +5,15 @@ dependencies:
   - python=3.9
   - affine
   - cfgrib
-  - dask<=2024.2.1
+  - dask
   - dask-geopandas
   - dask-image
   - fiona
   - geopandas
   - numba
-  - numpy<2
+  - numpy
   - netcdf4
-  - odc-geo<=0.4.3
+  - odc-geo
   - pandas
   - pyogrio
   - pyproj

requirements/default.txt

@@ -1,13 +1,13 @@
 affine
 cfgrib
-dask<=2024.2.1
+dask
 dask-geopandas
 dask-image
 fiona
 geopandas
 netcdf4
 numba
-numpy>=1.22,<2
+numpy>=1.22
 odc-geo<=0.4.3
 pandas
 pyogrio

requirements/dev.txt

@@ -1,14 +1,14 @@
 affine
 build
 cfgrib
-dask<=2024.2.1
+dask
 dask-geopandas
 dask-image
 fiona
 geopandas
 netcdf4
 numba
-numpy<2
+numpy
 odc-geo
 pandas
 pre-commit

requirements/dev.yml

@@ -6,7 +6,7 @@ dependencies:
   - python=3.12
   - affine
   - cfgrib
-  - dask<=2024.2.1
+  - dask
   - dask-geopandas
   - dask-image
   - fiona
@@ -15,7 +15,7 @@ dependencies:
   - matplotlib
   - netcdf4
   - numba
-  - numpy<2
+  - numpy
   - odc-geo
   - pandas
   - pre-commit

requirements/release.yml

@@ -6,14 +6,14 @@ dependencies:
   - python=3.10
   - affine
   - cfgrib
-  - dask<=2024.2.1
+  - dask
   - dask-geopandas
   - dask-image
   - fiona
   - geopandas
   - netcdf4
   - numba
-  - numpy<2
+  - numpy
   - pandas
   - pre-commit
   - pygeos