rastr 0.4.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

rastr/raster.py

@@ -4,17 +4,20 @@ from __future__ import annotations
 
 import importlib.util
 import warnings
+from collections.abc import Collection
 from contextlib import contextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, overload
 
 import numpy as np
 import numpy.ma
+import rasterio.features
 import rasterio.plot
 import rasterio.sample
 import rasterio.transform
 import skimage.measure
 from pydantic import BaseModel, InstanceOf, field_validator
+from pyproj import Transformer
 from pyproj.crs.crs import CRS
 from rasterio.enums import Resampling
 from rasterio.io import MemoryFile
@@ -29,10 +32,14 @@ if TYPE_CHECKING:
     from collections.abc import Callable, Generator
 
     import geopandas as gpd
+    from affine import Affine
+    from branca.colormap import LinearColormap as BrancaLinearColormap
     from folium import Map
     from matplotlib.axes import Axes
+    from matplotlib.image import AxesImage
     from numpy.typing import ArrayLike, NDArray
     from rasterio.io import BufferedDatasetWriter, DatasetReader, DatasetWriter
+    from shapely import MultiPolygon
     from typing_extensions import Self
 
 try:
@@ -45,17 +52,28 @@ FOLIUM_INSTALLED = importlib.util.find_spec("folium") is not None
 BRANCA_INSTALLED = importlib.util.find_spec("branca") is not None
 MATPLOTLIB_INSTALLED = importlib.util.find_spec("matplotlib") is not None
 
+CONTOUR_PERTURB_EPS = 1e-10
+
 
 class RasterCellArrayShapeError(ValueError):
     """Custom error for invalid raster cell array shapes."""
 
 
-class RasterModel(BaseModel):
+class Raster(BaseModel):
     """2-dimensional raster and metadata."""
 
     arr: InstanceOf[np.ndarray]
     raster_meta: RasterMeta
 
+    @field_validator("arr")
+    @classmethod
+    def check_2d_array(cls, v: NDArray) -> NDArray:
+        """Validator to ensure the cell array is 2D."""
+        if v.ndim != 2:
+            msg = "Cell array must be 2D"
+            raise RasterCellArrayShapeError(msg)
+        return v
+
     @property
     def meta(self) -> RasterMeta:
         """Alias for raster_meta."""
@@ -80,6 +98,26 @@ class RasterModel(BaseModel):
         """Set the CRS via meta."""
         self.meta.crs = value
 
+    @property
+    def transform(self) -> Affine:
+        """Convenience property to access the transform via meta."""
+        return self.meta.transform
+
+    @transform.setter
+    def transform(self, value: Affine) -> None:
+        """Set the transform via meta."""
+        self.meta.transform = value
+
+    @property
+    def cell_size(self) -> float:
+        """Convenience property to access the cell size via meta."""
+        return self.meta.cell_size
+
+    @cell_size.setter
+    def cell_size(self, value: float) -> None:
+        """Set the cell size via meta."""
+        self.meta.cell_size = value
+
     def __init__(
         self,
         *,
@@ -108,14 +146,25 @@ class RasterModel(BaseModel):
         super().__init__(arr=arr, raster_meta=raster_meta)
 
     def __eq__(self, other: object) -> bool:
-        """Check equality of two RasterModel objects."""
-        if not isinstance(other, RasterModel):
+        """Check equality of two Raster objects."""
+        if not isinstance(other, Raster):
             return NotImplemented
         return (
             np.array_equal(self.arr, other.arr)
             and self.raster_meta == other.raster_meta
         )
 
+    def is_like(self, other: Raster) -> bool:
+        """Check if two Raster objects have the same metadata and shape.
+
+        Args:
+            other: Another Raster to compare with.
+
+        Returns:
+            True if both rasters have the same meta and shape attributes.
+        """
+        return self.meta == other.meta and self.shape == other.shape
+
     __hash__ = BaseModel.__hash__
 
     def __add__(self, other: float | Self) -> Self:
@@ -123,7 +172,7 @@ class RasterModel(BaseModel):
         if isinstance(other, float | int):
             new_arr = self.arr + other
             return cls(arr=new_arr, raster_meta=self.raster_meta)
-        elif isinstance(other, RasterModel):
+        elif isinstance(other, Raster):
             if self.raster_meta != other.raster_meta:
                 msg = (
                     "Rasters must have the same metadata (e.g. CRS, cell size, etc.) "
@@ -149,7 +198,7 @@ class RasterModel(BaseModel):
         if isinstance(other, float | int):
             new_arr = self.arr * other
             return cls(arr=new_arr, raster_meta=self.raster_meta)
-        elif isinstance(other, RasterModel):
+        elif isinstance(other, Raster):
             if self.raster_meta != other.raster_meta:
                 msg = (
                     "Rasters must have the same metadata (e.g. CRS, cell size, etc.) "
@@ -172,7 +221,7 @@ class RasterModel(BaseModel):
         if isinstance(other, float | int):
             new_arr = self.arr / other
             return cls(arr=new_arr, raster_meta=self.raster_meta)
-        elif isinstance(other, RasterModel):
+        elif isinstance(other, Raster):
             if self.raster_meta != other.raster_meta:
                 msg = (
                     "Rasters must have the same metadata (e.g. CRS, cell size, etc.) "
@@ -222,7 +271,7 @@ class RasterModel(BaseModel):
         """Create a rasterio in-memory dataset from the Raster object.
 
         Example:
-            >>> raster = RasterModel.example()
+            >>> raster = Raster.example()
             >>> with raster.to_rasterio_dataset() as dataset:
             >>>     ...
         """
@@ -248,12 +297,30 @@ class RasterModel(BaseModel):
         finally:
             memfile.close()
 
+    @overload
+    def sample(
+        self,
+        xy: Collection[tuple[float, float]] | Collection[Point] | ArrayLike,
+        *,
+        na_action: Literal["raise", "ignore"] = "raise",
+    ) -> NDArray: ...
+    @overload
     def sample(
         self,
-        xy: list[tuple[float, float]] | list[Point] | ArrayLike,
+        xy: tuple[float, float] | Point,
         *,
         na_action: Literal["raise", "ignore"] = "raise",
-    ) -> NDArray[np.float64]:
+    ) -> float: ...
+    def sample(
+        self,
+        xy: Collection[tuple[float, float]]
+        | Collection[Point]
+        | ArrayLike
+        | tuple[float, float]
+        | Point,
+        *,
+        na_action: Literal["raise", "ignore"] = "raise",
+    ) -> NDArray | float:
         """Sample raster values at GeoSeries locations and return sampled values.
 
         Args:
@@ -270,13 +337,30 @@ class RasterModel(BaseModel):
         # https://rdrn.me/optimising-sampling/
 
         # Convert shapely Points to coordinate tuples if needed
-        if isinstance(xy, (list, tuple)):
-            xy = [_get_xy_tuple(point) for point in xy]
+        if isinstance(xy, Point):
+            xy = [(xy.x, xy.y)]
+            singleton = True
+        elif (
+            isinstance(xy, Collection)
+            and len(xy) > 0
+            and isinstance(next(iter(xy)), Point)
+        ):
+            xy = [(point.x, point.y) for point in xy]  # pyright: ignore[reportAttributeAccessIssue]
+            singleton = False
+        elif (
+            isinstance(xy, tuple)
+            and len(xy) == 2
+            and isinstance(next(iter(xy)), (float, int))
+        ):
+            xy = [xy]  # pyright: ignore[reportAssignmentType]
+            singleton = True
+        else:
+            singleton = False
 
         xy = np.asarray(xy, dtype=float)
 
-        # Short-circuit
         if len(xy) == 0:
+            # Short-circuit
             return np.array([], dtype=float)
 
         # Create in-memory rasterio dataset from the incumbent Raster object
@@ -326,6 +410,10 @@ class RasterModel(BaseModel):
                     axis=0,
                 )
 
+        if singleton:
+            (raster_value,) = raster_values
+            return raster_value
+
         return raster_values
 
     @property
@@ -354,13 +442,16 @@ class RasterModel(BaseModel):
             ]
         )
 
-    def explore(
+    def explore(  # noqa: PLR0913 c.f. geopandas.explore which also has many input args
         self,
         *,
         m: Map | None = None,
         opacity: float = 1.0,
-        colormap: str = "viridis",
+        colormap: str
+        | Callable[[float], tuple[float, float, float, float]] = "viridis",
         cbar_label: str | None = None,
+        vmin: float | None = None,
+        vmax: float | None = None,
     ) -> Map:
         """Display the raster on a folium map."""
         if not FOLIUM_INSTALLED or not MATPLOTLIB_INSTALLED:
@@ -368,39 +459,44 @@ class RasterModel(BaseModel):
             raise ImportError(msg)
 
         import folium.raster_layers
-        import geopandas as gpd
         import matplotlib as mpl
 
         if m is None:
             m = folium.Map()
 
-        rgba_map: Callable[[float], tuple[float, float, float, float]] = mpl.colormaps[
-            colormap
-        ]
+        if vmin is not None and vmax is not None and vmax <= vmin:
+            msg = "'vmin' must be less than 'vmax'."
+            raise ValueError(msg)
+
+        if isinstance(colormap, str):
+            colormap = mpl.colormaps[colormap]
 
-        # Cast to GDF to facilitate converting bounds to WGS84
+        # Transform bounds to WGS84 using pyproj directly
         wgs84_crs = CRS.from_epsg(4326)
-        gdf = gpd.GeoDataFrame(geometry=[self.bbox], crs=self.raster_meta.crs).to_crs(
-            wgs84_crs
+        transformer = Transformer.from_crs(
+            self.raster_meta.crs, wgs84_crs, always_xy=True
         )
-        xmin, ymin, xmax, ymax = gdf.total_bounds
 
-        arr = np.array(self.arr)
+        # Get the corner points of the bounding box
+        raster_xmin, raster_ymin, raster_xmax, raster_ymax = self.bounds
+        corner_points = [
+            (raster_xmin, raster_ymin),
+            (raster_xmin, raster_ymax),
+            (raster_xmax, raster_ymax),
+            (raster_xmax, raster_ymin),
+        ]
 
-        # Normalize the data to the range [0, 1] as this is the cmap range
-        with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore",
-                message="All-NaN slice encountered",
-                category=RuntimeWarning,
-            )
-            min_val = np.nanmin(arr)
-            max_val = np.nanmax(arr)
+        # Transform all corner points to WGS84
+        transformed_points = [transformer.transform(x, y) for x, y in corner_points]
 
-        if max_val > min_val:  # Prevent division by zero
-            arr = (arr - min_val) / (max_val - min_val)
-        else:
-            arr = np.zeros_like(arr)  # In case all values are the same
+        # Find the bounding box of the transformed points
+        transformed_xs, transformed_ys = zip(*transformed_points, strict=True)
+        xmin, xmax = min(transformed_xs), max(transformed_xs)
+        ymin, ymax = min(transformed_ys), max(transformed_ys)
+
+        # Normalize the array to [0, 1] for colormap mapping
+        _vmin, _vmax = _get_vmin_vmax(self, vmin=vmin, vmax=vmax)
+        arr = self.normalize(vmin=_vmin, vmax=_vmax).arr
 
         # Finally, need to determine whether to flip the image based on negative Affine
         # coefficients
@@ -411,11 +507,12 @@ class RasterModel(BaseModel):
         if flip_y:
             arr = np.flip(arr, axis=0)
 
+        bounds = [[ymin, xmin], [ymax, xmax]]
         img = folium.raster_layers.ImageOverlay(
             image=arr,
-            bounds=[[ymin, xmin], [ymax, xmax]],
+            bounds=bounds,
             opacity=opacity,
-            colormap=rgba_map,
+            colormap=colormap,
             mercator_project=True,
         )
 
@@ -423,26 +520,39 @@ class RasterModel(BaseModel):
 
         # Add a colorbar legend
         if BRANCA_INSTALLED:
-            from branca.colormap import LinearColormap as BrancaLinearColormap
-            from matplotlib.colors import to_hex
-
-            # Determine legend data range in original units
-            vmin = float(min_val) if np.isfinite(min_val) else 0.0
-            vmax = float(max_val) if np.isfinite(max_val) else 1.0
-            if vmax <= vmin:
-                vmax = vmin + 1.0
-
-            sample_points = np.linspace(0, 1, rgba_map.N)
-            colors = [to_hex(rgba_map(x)) for x in sample_points]
-            legend = BrancaLinearColormap(colors=colors, vmin=vmin, vmax=vmax)
+            cbar = _map_colorbar(colormap=colormap, vmin=_vmin, vmax=_vmax)
             if cbar_label:
-                legend.caption = cbar_label
-            legend.add_to(m)
+                cbar.caption = cbar_label
+            cbar.add_to(m)
 
-        m.fit_bounds([[ymin, xmin], [ymax, xmax]])
+        m.fit_bounds(bounds)
 
         return m
 
+    def normalize(
+        self, *, vmin: float | None = None, vmax: float | None = None
+    ) -> Self:
+        """Normalize the raster values to the range [0, 1].
+
+        If custom vmin and vmax values are provided, values below vmin will be set to 0,
+        and values above vmax will be set to 1.
+
+        Args:
+            vmin: Minimum value for normalization. Values below this will be set to 0.
+                  If None, the minimum value in the array is used.
+            vmax: Maximum value for normalization. Values above this will be set to 1.
+                  If None, the maximum value in the array is used.
+        """
+        _vmin, _vmax = _get_vmin_vmax(self, vmin=vmin, vmax=vmax)
+
+        arr = self.arr.copy()
+        if _vmax > _vmin:
+            arr = (arr - _vmin) / (_vmax - _vmin)
+            arr = np.clip(arr, 0, 1)
+        else:
+            arr = np.zeros_like(arr)
+        return self.__class__(arr=arr, raster_meta=self.raster_meta)
+
     def to_clipboard(self) -> None:
         """Copy the raster cell array to the clipboard."""
         import pandas as pd
@@ -456,8 +566,22 @@ class RasterModel(BaseModel):
         cbar_label: str | None = None,
         basemap: bool = False,
         cmap: str = "viridis",
+        suppressed: Collection[float] | float = tuple(),
+        **kwargs: Any,
     ) -> Axes:
-        """Plot the raster on a matplotlib axis."""
+        """Plot the raster on a matplotlib axis.
+
+        Args:
+            ax: A matplotlib axes object to plot on. If None, a new figure will be
+                created.
+            cbar_label: Label for the colorbar. If None, no label is added.
+            basemap: Whether to add a basemap. Currently not implemented.
+            cmap: Colormap to use for the plot.
+            suppressed: Values to suppress from the plot (i.e. not display). This can be
+                        useful for zeroes especially.
+            **kwargs: Additional keyword arguments to pass to `rasterio.plot.show()`.
+                      This includes parameters like `alpha` for transparency.
+        """
         if not MATPLOTLIB_INSTALLED:
             msg = "The 'matplotlib' package is required for 'plot()'."
             raise ImportError(msg)
@@ -465,6 +589,8 @@ class RasterModel(BaseModel):
         from matplotlib import pyplot as plt
         from mpl_toolkits.axes_grid1 import make_axes_locatable
 
+        suppressed = np.array(suppressed)
+
         if ax is None:
             _, _ax = plt.subplots()
             _ax: Axes
@@ -474,33 +600,34 @@ class RasterModel(BaseModel):
             msg = "Basemap plotting is not yet implemented."
             raise NotImplementedError(msg)
 
-        arr = self.arr.copy()
+        model = self.model_copy()
+        model.arr = model.arr.copy()
 
-        # Get extent of the non-zero values in array index coordinates
-        (x_nonzero,) = np.nonzero(arr.any(axis=0))
-        (y_nonzero,) = np.nonzero(arr.any(axis=1))
+        # Get extent of the unsuppressed values in array index coordinates
+        suppressed_mask = np.isin(model.arr, suppressed)
+        (x_unsuppressed,) = np.nonzero((~suppressed_mask).any(axis=0))
+        (y_unsuppressed,) = np.nonzero((~suppressed_mask).any(axis=1))
 
-        if len(x_nonzero) == 0 or len(y_nonzero) == 0:
-            msg = "Raster contains no non-zero values; cannot plot."
+        if len(x_unsuppressed) == 0 or len(y_unsuppressed) == 0:
+            msg = "Raster contains no unsuppressed values; cannot plot."
             raise ValueError(msg)
 
-        min_x_nonzero = np.min(x_nonzero)
-        max_x_nonzero = np.max(x_nonzero)
-        min_y_nonzero = np.min(y_nonzero)
-        max_y_nonzero = np.max(y_nonzero)
+        # N.B. these are array index coordinates, so np.min and np.max are safe since
+        # they cannot encounter NaN values.
+        min_x_unsuppressed = np.min(x_unsuppressed)
+        max_x_unsuppressed = np.max(x_unsuppressed)
+        min_y_unsuppressed = np.min(y_unsuppressed)
+        max_y_unsuppressed = np.max(y_unsuppressed)
 
         # Transform to raster CRS
-        x1, y1 = self.raster_meta.transform * (min_x_nonzero, min_y_nonzero)  # type: ignore[reportAssignmentType] overloaded tuple size in affine
-        x2, y2 = self.raster_meta.transform * (max_x_nonzero, max_y_nonzero)  # type: ignore[reportAssignmentType]
+        x1, y1 = self.raster_meta.transform * (min_x_unsuppressed, min_y_unsuppressed)  # type: ignore[reportAssignmentType] overloaded tuple size in affine
+        x2, y2 = self.raster_meta.transform * (max_x_unsuppressed, max_y_unsuppressed)  # type: ignore[reportAssignmentType]
         xmin, xmax = sorted([x1, x2])
         ymin, ymax = sorted([y1, y2])
 
-        arr[arr == 0] = np.nan
+        model.arr[suppressed_mask] = np.nan
 
-        with self.to_rasterio_dataset() as dataset:
-            img, *_ = rasterio.plot.show(
-                dataset, with_bounds=True, ax=ax, cmap=cmap
-            ).get_images()
+        img, *_ = model.rio_show(ax=ax, cmap=cmap, with_bounds=True, **kwargs)
 
         ax.set_xlim(xmin, xmax)
         ax.set_ylim(ymin, ymax)
@@ -516,6 +643,20 @@ class RasterModel(BaseModel):
             fig.colorbar(img, label=cbar_label, cax=cax)
         return ax
 
+    def rio_show(self, **kwargs: Any) -> list[AxesImage]:
+        """Plot the raster using rasterio's built-in plotting function.
+
+        This is useful for lower-level access to rasterio's plotting capabilities.
+        Generally, the `plot()` method is preferred for most use cases.
+
+        Args:
+            **kwargs: Keyword arguments to pass to `rasterio.plot.show()`. This includes
+            parameters like `alpha` for transparency, and `with_bounds` to control
+            whether to plot in spatial coordinates or array index coordinates.
+        """
+        with self.to_rasterio_dataset() as dataset:
+            return rasterio.plot.show(dataset, **kwargs).get_images()
+
     def as_geodataframe(self, name: str = "value") -> gpd.GeoDataFrame:
         """Create a GeoDataFrame representation of the raster."""
         import geopandas as gpd
@@ -532,8 +673,15 @@ class RasterModel(BaseModel):
 
         return raster_gdf
 
-    def to_file(self, path: Path | str) -> None:
-        """Write the raster to a GeoTIFF file."""
+    def to_file(self, path: Path | str, **kwargs: Any) -> None:
+        """Write the raster to a GeoTIFF file.
+
+        Args:
+            path: Path to output file.
+            **kwargs: Additional keyword arguments to pass to `rasterio.open()`. If
+                      `nodata` is provided, NaN values in the raster will be replaced
+                      with the nodata value.
+        """
 
         path = Path(path)
 
@@ -548,6 +696,15 @@ class RasterModel(BaseModel):
             msg = f"Unsupported file extension: {suffix}"
             raise ValueError(msg)
 
+        # Handle nodata: use provided value or default to np.nan
+        if "nodata" in kwargs:
+            # Replace NaN values with the nodata value
+            nodata_value = kwargs.pop("nodata")
+            arr_to_write = np.where(np.isnan(self.arr), nodata_value, self.arr)
+        else:
+            nodata_value = np.nan
+            arr_to_write = self.arr
+
         with rasterio.open(
             path,
             "w",
@@ -558,10 +715,11 @@ class RasterModel(BaseModel):
             dtype=self.arr.dtype,
             crs=self.raster_meta.crs,
             transform=self.raster_meta.transform,
-            nodata=np.nan,
+            nodata=nodata_value,
+            **kwargs,
         ) as dst:
             try:
-                dst.write(self.arr, 1)
+                dst.write(arr_to_write, 1)
             except CPLE_BaseError as err:
                 msg = f"Failed to write raster to file: {err}"
                 raise OSError(msg) from err
@@ -576,7 +734,7 @@ class RasterModel(BaseModel):
 
     @classmethod
     def example(cls) -> Self:
-        """Create an example RasterModel."""
+        """Create an example Raster."""
         # Peaks dataset style example
         n = 256
         x = np.linspace(-3, 3, n)
@@ -588,6 +746,34 @@ class RasterModel(BaseModel):
         raster_meta = RasterMeta.example()
         return cls(arr=arr, raster_meta=raster_meta)
 
+    @classmethod
+    def full_like(cls, other: Raster, *, fill_value: float) -> Self:
+        """Create a raster with the same metadata as another but filled with a constant.
+
+        Args:
+            other: The raster to copy metadata from.
+            fill_value: The constant value to fill all cells with.
+
+        Returns:
+            A new raster with the same shape and metadata as `other`, but with all cells
+            set to `fill_value`.
+        """
+        arr = np.full(other.shape, fill_value, dtype=np.float32)
+        return cls(arr=arr, raster_meta=other.raster_meta)
+
+    @classmethod
+    def read_file(cls, filename: Path | str, crs: CRS | str | None = None) -> Self:
+        """Read raster data from a file and return an in-memory Raster object.
+
+        Args:
+            filename: Path to the raster file.
+            crs: Optional coordinate reference system to override the file's CRS.
+        """
+        # Import here to avoid circular import (rastr.io imports Raster)
+        from rastr.io import read_raster_inmem  # noqa: PLC0415
+
+        return read_raster_inmem(filename, crs=crs, cls=cls)
+
     @overload
     def apply(
         self,
@@ -625,6 +811,59 @@ class RasterModel(BaseModel):
         new_raster.arr = np.asarray(new_arr)
         return new_raster
 
+    def max(self) -> float:
+        """Get the maximum value in the raster, ignoring NaN values.
+
+        Returns:
+            The maximum value in the raster. Returns NaN if all values are NaN.
+        """
+        return float(np.nanmax(self.arr))
+
+    def min(self) -> float:
+        """Get the minimum value in the raster, ignoring NaN values.
+
+        Returns:
+            The minimum value in the raster. Returns NaN if all values are NaN.
+        """
+        return float(np.nanmin(self.arr))
+
+    def mean(self) -> float:
+        """Get the mean value in the raster, ignoring NaN values.
+
+        Returns:
+            The mean value in the raster. Returns NaN if all values are NaN.
+        """
+        return float(np.nanmean(self.arr))
+
+    def std(self) -> float:
+        """Get the standard deviation of values in the raster, ignoring NaN values.
+
+        Returns:
+            The standard deviation of the raster. Returns NaN if all values are NaN.
+        """
+        return float(np.nanstd(self.arr))
+
+    def quantile(self, q: float) -> float:
+        """Get the specified quantile value in the raster, ignoring NaN values.
+
+        Args:
+            q: Quantile to compute, must be between 0 and 1 inclusive.
+
+        Returns:
+            The quantile value. Returns NaN if all values are NaN.
+        """
+        return float(np.nanquantile(self.arr, q))
+
+    def median(self) -> float:
+        """Get the median value in the raster, ignoring NaN values.
+
+        This is equivalent to quantile(0.5).
+
+        Returns:
+            The median value in the raster. Returns NaN if all values are NaN.
+        """
+        return float(np.nanmedian(self.arr))
+
     def fillna(self, value: float) -> Self:
         """Fill NaN values in the raster with a specified value.
 
@@ -635,6 +874,16 @@ class RasterModel(BaseModel):
         new_raster.arr = filled_arr
         return new_raster
 
+    def copy(self) -> Self:  # type: ignore[override]
+        """Create a copy of the raster.
+
+        This method wraps `model_copy()` for convenience.
+
+        Returns:
+            A new Raster instance.
+        """
+        return self.model_copy(deep=True)
+
     def get_xy(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
         """Get the x and y coordinates of the raster cell centres in meshgrid format.
 
@@ -651,7 +900,7 @@ class RasterModel(BaseModel):
         return coords[:, :, 0], coords[:, :, 1]
 
     def contour(
-        self, levels: list[float] | NDArray, *, smoothing: bool = True
+        self, levels: Collection[float] | NDArray, *, smoothing: bool = True
     ) -> gpd.GeoDataFrame:
         """Create contour lines from the raster data, optionally with smoothing.
 
@@ -664,8 +913,8 @@ class RasterModel(BaseModel):
         contouring, to denoise the contours.
 
         Args:
-            levels: A list or array of contour levels to generate. The contour lines
-                    will be generated for each level in this sequence.
+            levels: A collection or array of contour levels to generate. The contour
+                    lines will be generated for each level in this sequence.
             smoothing: Defaults to true, which corresponds to applying a smoothing
                        algorithm to the contour lines. At the moment, this is the
                        Catmull-Rom spline algorithm. If set to False, the raw
@@ -676,9 +925,17 @@ class RasterModel(BaseModel):
         all_levels = []
         all_geoms = []
         for level in levels:
+            # If this is the maximum or minimum level, perturb it ever-so-slightly to
+            # ensure we get contours at the edges of the raster
+            perturbed_level = level
+            if level == self.max():
+                perturbed_level -= CONTOUR_PERTURB_EPS
+            elif level == self.min():
+                perturbed_level += CONTOUR_PERTURB_EPS
+
             contours = skimage.measure.find_contours(
                 self.arr,
-                level=level,
+                level=perturbed_level,
             )
 
             # Construct shapely LineString objects
@@ -713,19 +970,40 @@ class RasterModel(BaseModel):
         # Dissolve contours by level to merge all contour lines of the same level
         return contour_gdf.dissolve(by="level", as_index=False)
 
-    def blur(self, sigma: float) -> Self:
+    def blur(self, sigma: float, *, preserve_nan: bool = True) -> Self:
         """Apply a Gaussian blur to the raster data.
 
         Args:
             sigma: Standard deviation for Gaussian kernel, in units of geographic
                    coordinate distance (e.g. meters). A larger sigma results in a more
                    blurred image.
+            preserve_nan: If True, applies NaN-safe blurring by extrapolating NaN values
+                          before blurring and restoring them afterwards. This prevents
+                          NaNs from spreading into valid data during the blur operation.
         """
         from scipy.ndimage import gaussian_filter
 
         cell_sigma = sigma / self.raster_meta.cell_size
 
-        blurred_array = gaussian_filter(self.arr, sigma=cell_sigma)
+        if preserve_nan:
+            # Save the original NaN mask
+            nan_mask = np.isnan(self.arr)
+
+            # If there are no NaNs, just apply regular blur
+            if not np.any(nan_mask):
+                blurred_array = gaussian_filter(self.arr, sigma=cell_sigma)
+            else:
+                # Extrapolate to fill NaN values temporarily
+                extrapolated_arr = fillna_nearest_neighbours(arr=self.arr)
+
+                # Apply blur to the extrapolated array
+                blurred_array = gaussian_filter(extrapolated_arr, sigma=cell_sigma)
+
+                # Restore original NaN values
+                blurred_array = np.where(nan_mask, np.nan, blurred_array)
+        else:
+            blurred_array = gaussian_filter(self.arr, sigma=cell_sigma)
+
         new_raster = self.model_copy()
         new_raster.arr = blurred_array
         return new_raster
@@ -751,9 +1029,75 @@ class RasterModel(BaseModel):
 
         return raster
 
+    def pad(self, width: float, *, value: float = np.nan) -> Self:
+        """Extend the raster by adding a constant fill value around the edges.
+
+        By default, the padding value is NaN, but this can be changed via the
+        `value` parameter.
+
+        This grows the raster by adding padding around all edges. New cells are
+        filled with the constant `value`.
+
+        If the width is not an exact multiple of the cell size, the padding may be
+        slightly larger than the specified width, i.e. the value is rounded up to
+        the nearest whole number of cells.
+
+        Args:
+            width: The width of the padding, in the same units as the raster CRS
+                   (e.g. meters). This defines how far from the edge the padding
+                   extends.
+            value: The constant value to use for padding. Default is NaN.
+        """
+        cell_size = self.raster_meta.cell_size
+
+        # Calculate number of cells to pad in each direction
+        pad_cells = int(np.ceil(width / cell_size))
+
+        # Get current bounds
+        xmin, ymin, xmax, ymax = self.bounds
+
+        # Calculate new bounds with padding
+        new_xmin = xmin - (pad_cells * cell_size)
+        new_ymin = ymin - (pad_cells * cell_size)
+        new_xmax = xmax + (pad_cells * cell_size)
+        new_ymax = ymax + (pad_cells * cell_size)
+
+        # Create padded array
+        new_height = self.arr.shape[0] + 2 * pad_cells
+        new_width = self.arr.shape[1] + 2 * pad_cells
+
+        # Create new array filled with the padding value
+        padded_arr = np.full((new_height, new_width), value, dtype=self.arr.dtype)
+
+        # Copy original array into the center of the padded array
+        padded_arr[
+            pad_cells : pad_cells + self.arr.shape[0],
+            pad_cells : pad_cells + self.arr.shape[1],
+        ] = self.arr
+
+        # Create new transform for the padded raster
+        new_transform = rasterio.transform.from_bounds(
+            west=new_xmin,
+            south=new_ymin,
+            east=new_xmax,
+            north=new_ymax,
+            width=new_width,
+            height=new_height,
+        )
+
+        # Create new raster metadata
+        new_meta = RasterMeta(
+            cell_size=cell_size,
+            crs=self.raster_meta.crs,
+            transform=new_transform,
+        )
+
+        return self.__class__(arr=padded_arr, raster_meta=new_meta)
+
     def crop(
         self,
         bounds: tuple[float, float, float, float],
+        *,
         strategy: Literal["underflow", "overflow"] = "underflow",
     ) -> Self:
         """Crop the raster to the specified bounds as (minx, miny, maxx, maxy).
@@ -767,7 +1111,7 @@ class RasterModel(BaseModel):
                         remains covered with cells.
 
         Returns:
-            A new RasterModel instance cropped to the specified bounds.
+            A new Raster instance cropped to the specified bounds.
         """
 
         minx, miny, maxx, maxy = bounds
@@ -827,6 +1171,159 @@ class RasterModel(BaseModel):
         )
         return cls(arr=cropped_arr, raster_meta=new_meta)
 
+    def taper_border(self, width: float, *, limit: float = 0.0) -> Self:
+        """Taper values to a limiting value around the border of the raster.
+
+        By default, the borders are tapered to zero, but this can be changed via the
+        `limit` parameter.
+
+        This keeps the raster size the same, overwriting values in the border area.
+        To instead grow the raster, consider using `pad()` followed by `taper_border()`.
+
+        The tapering is linear from the cell centres around the border of the raster,
+        so the value at the edge of the raster will be equal to `limit`.
+
+        Args:
+            width: The width of the taper, in the same units as the raster CRS
+                   (e.g. meters). This defines how far from the edge the tapering
+                   starts.
+            limit: The limiting value to taper to at the edges. Default is zero.
+        """
+
+        # Determine the width in cell units (possibly fractional)
+        cell_size = self.raster_meta.cell_size
+        width_in_cells = width / cell_size
+
+        # Calculate the distance from the edge in cell units
+        arr_height, arr_width = self.arr.shape
+        y_indices, x_indices = np.indices((int(arr_height), int(arr_width)))
+        dist_from_left = x_indices
+        dist_from_right = arr_width - 1 - x_indices
+        dist_from_top = y_indices
+        dist_from_bottom = arr_height - 1 - y_indices
+        dist_from_edge = np.minimum.reduce(
+            [dist_from_left, dist_from_right, dist_from_top, dist_from_bottom]
+        )
+
+        # Mask the arrays to only the area within the width from the edge, rounding up
+        mask = dist_from_edge < np.ceil(width_in_cells)
+        masked_dist_arr = np.where(mask, dist_from_edge, np.nan)
+        masked_arr = np.where(mask, self.arr, np.nan)
+
+        # Calculate the tapering factor based on the distance from the edge
+        taper_factor = np.clip(masked_dist_arr / width_in_cells, 0.0, 1.0)
+        tapered_values = limit + (masked_arr - limit) * taper_factor
+
+        # Create the new raster array
+        new_arr = self.arr.copy()
+        new_arr[mask] = tapered_values[mask]
+        new_raster = self.model_copy()
+        new_raster.arr = new_arr
+
+        return new_raster
+
+    def clip(
+        self,
+        polygon: Polygon | MultiPolygon,
+        *,
+        strategy: Literal["centres"] = "centres",
+    ) -> Self:
+        """Clip the raster to the specified polygon, replacing cells outside with NaN.
+
+        The clipping strategy determines how to handle cells that are partially
+        within the polygon. Currently, only the 'centres' strategy is supported, which
+        retains cells whose centres fall within the polygon.
+
+        Args:
+            polygon: A shapely Polygon or MultiPolygon defining the area to clip to.
+            strategy: The clipping strategy to use. Currently only 'centres' is
+                      supported, which retains cells whose centres fall within the
+                      polygon.
+
+        Returns:
+            A new Raster with cells outside the polygon set to NaN.
+        """
+        if strategy != "centres":
+            msg = f"Unsupported clipping strategy: {strategy}"
+            raise NotImplementedError(msg)
+
+        raster = self.model_copy()
+
+        mask = rasterio.features.rasterize(
+            [(polygon, 1)],
+            fill=0,
+            out_shape=self.shape,
+            transform=self.meta.transform,
+            dtype=np.uint8,
+        )
+
+        raster.arr = np.where(mask, raster.arr, np.nan)
+
+        return raster
+
+    def _trim_value(self, *, value_mask: NDArray[np.bool_], value_name: str) -> Self:
+        """Crop the raster by trimming away slices matching the mask at the edges.
+
+        Args:
+            value_mask: Boolean mask where True indicates values to trim
+            value_name: Name of the value type for error messages (e.g., 'NaN', 'zero')
+        """
+        arr = self.arr
+
+        # Check if the entire array matches the mask
+        if np.all(value_mask):
+            msg = f"Cannot crop raster: all values are {value_name}"
+            raise ValueError(msg)
+
+        # Find rows and columns that are not all matching the mask
+        row_mask = np.all(value_mask, axis=1)
+        col_mask = np.all(value_mask, axis=0)
+
+        # Find the bounding indices
+        (row_indices,) = np.where(~row_mask)
+        (col_indices,) = np.where(~col_mask)
+
+        min_row, max_row = row_indices[0], row_indices[-1]
+        min_col, max_col = col_indices[0], col_indices[-1]
+
+        # Crop the array
+        cropped_arr = arr[min_row : max_row + 1, min_col : max_col + 1]
+
+        # Shift the transform by the number of pixels cropped (min_col, min_row)
+        new_transform = (
+            self.raster_meta.transform
+            * rasterio.transform.Affine.translation(min_col, min_row)
+        )
+
+        # Create new metadata
+        new_meta = RasterMeta(
+            cell_size=self.raster_meta.cell_size,
+            crs=self.raster_meta.crs,
+            transform=new_transform,
+        )
+
+        return self.__class__(arr=cropped_arr, raster_meta=new_meta)
+
+    def trim_nan(self) -> Self:
+        """Crop the raster by trimming away all-NaN slices at the edges.
+
+        This effectively trims the raster to the smallest bounding box that contains all
+        of the non-NaN values. Note that this does not guarantee no NaN values at all
+        around the edges, only that there won't be entire edges which are all-NaN.
+
+        Consider using `.extrapolate()` for further cleanup of NaN values.
+        """
+        return self._trim_value(value_mask=np.isnan(self.arr), value_name="NaN")
+
+    def trim_zeros(self) -> Self:
+        """Crop the raster by trimming away all-zero slices at the edges.
+
+        This effectively trims the raster to the smallest bounding box that contains all
+        of the non-zero values. Note that this does not guarantee no zero values at all
+        around the edges, only that there won't be entire edges which are all-zero.
+        """
+        return self._trim_value(value_mask=(self.arr == 0), value_name="zero")
+
     def resample(
         self, new_cell_size: float, *, method: Literal["bilinear"] = "bilinear"
     ) -> Self:
@@ -874,26 +1371,55 @@ class RasterModel(BaseModel):
 
             return cls(arr=new_arr, raster_meta=new_raster_meta)
 
-    @field_validator("arr")
-    @classmethod
-    def check_2d_array(cls, v: NDArray) -> NDArray:
-        """Validator to ensure the cell array is 2D."""
-        if v.ndim != 2:
-            msg = "Cell array must be 2D"
-            raise RasterCellArrayShapeError(msg)
-        return v
 
+def _map_colorbar(
+    *,
+    colormap: Callable[[float], tuple[float, float, float, float]],
+    vmin: float,
+    vmax: float,
+) -> BrancaLinearColormap:
+    from branca.colormap import LinearColormap as BrancaLinearColormap
+    from matplotlib.colors import ListedColormap, to_hex
+
+    # Determine legend data range in original units
+    vmin = float(vmin) if np.isfinite(vmin) else 0.0
+    vmax = float(vmax) if np.isfinite(vmax) else 1.0
+    if vmax <= vmin:
+        vmax = vmin + 1.0
+
+    if isinstance(colormap, ListedColormap):
+        n = colormap.N
+    else:
+        n = 256
+
+    sample_points = np.linspace(0, 1, n)
+    colors = [to_hex(colormap(x)) for x in sample_points]
+    return BrancaLinearColormap(colors=colors, vmin=vmin, vmax=vmax)
 
-def _get_xy_tuple(xy: Any) -> tuple[float, float]:
-    """Convert Point or coordinate tuple to coordinate tuple.
 
-    Args:
-        xy: Either a coordinate tuple or a shapely Point object.
+def _get_vmin_vmax(
+    raster: Raster, *, vmin: float | None = None, vmax: float | None = None
+) -> tuple[float, float]:
+    """Get maximum and minimum values from a raster array, ignoring NaNs.
 
-    Returns:
-        A coordinate tuple (x, y).
+    Allows for custom over-ride vmin and vmax values to be provided.
     """
-    if isinstance(xy, Point):
-        return (xy.x, xy.y)
-    x, y = xy
-    return (float(x), float(y))
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message="All-NaN slice encountered",
+            category=RuntimeWarning,
+        )
+        if vmin is None:
+            _vmin = float(raster.min())
+        else:
+            _vmin = vmin
+        if vmax is None:
+            _vmax = float(raster.max())
+        else:
+            _vmax = vmax
+
+    return _vmin, _vmax
+
+
+RasterModel = Raster