erosion.py

from __future__ import annotations

import numpy as np
import xarray as xr
from numba import jit

try:
    import cupy
    from numba import cuda
except ImportError:
    class cupy(object):
        ndarray = False
    cuda = None

try:
    import dask.array as da
except ImportError:
    da = None

from xrspatial.utils import (
    ArrayTypeFunctionMapping,
    _validate_scalar,
    has_cuda_and_cupy,
    is_cupy_array,
    is_dask_cupy,
)

# Default erosion parameters
_DEFAULT_PARAMS = dict(
    inertia=0.05,
    capacity=4.0,
    deposition=0.3,
    erosion=0.3,
    evaporation=0.01,
    gravity=4.0,
    min_slope=0.01,
    radius=3,
    max_lifetime=30,
)

# Upper bounds on user-controlled iteration / kernel-size parameters.
# These cap the worst-case CPU work and the worst-case host/device
# allocation regardless of raster shape.
#
# - _MAX_ITERATIONS sizes the (iterations, 2) random_pos buffer
#   (16 bytes/particle on the host, copied to GPU).  10**8 is ~1.6 GB,
#   well above any realistic erosion run (default is 50_000).
# - _MAX_RADIUS bounds the brush size: (2r+1)**2 cells, three arrays.
#   At r=1024 the brush arrays total ~50 MB.
# - _MAX_LIFETIME bounds the inner per-particle loop in the JIT kernel.
#   At 100_000 with iterations=10**8 the total JIT step count is 10**13,
#   which is still excessive; the iteration cap is the binding limit.
_MAX_ITERATIONS = 10**8
_MAX_RADIUS = 1024
_MAX_LIFETIME = 100_000


def _build_brush(radius):
    """Precompute brush offsets and weights for the erosion kernel."""
    offsets_y = []
    offsets_x = []
    weights = []
    for dy in range(-radius, radius + 1):
        for dx in range(-radius, radius + 1):
            dist2 = dx * dx + dy * dy
            if dist2 <= radius * radius:
                w = max(0.0, radius - dist2 ** 0.5)
                offsets_y.append(dy)
                offsets_x.append(dx)
                weights.append(w)
    weight_sum = sum(weights)
    boy = np.array(offsets_y, dtype=np.int32)
    box = np.array(offsets_x, dtype=np.int32)
    bw = np.array([w / weight_sum for w in weights], dtype=np.float64)
    return boy, box, bw


@jit(nopython=True, nogil=True)
def _erode_cpu(heightmap, random_pos, boy, box, bw,
               inertia, capacity, deposition, erosion_rate,
               evaporation, gravity, min_slope, max_lifetime):
    """Particle-based hydraulic erosion on a 2D heightmap (CPU).

    random_pos : float64 array of shape (n_iterations, 2) with pre-generated
                 random starting positions in [0, 1).
    """
    height, width = heightmap.shape
    n_iterations = random_pos.shape[0]
    n_brush = bw.shape[0]

    for iteration in range(n_iterations):
        pos_x = random_pos[iteration, 0] * (width - 3) + 1
        pos_y = random_pos[iteration, 1] * (height - 3) + 1
        dir_x = 0.0
        dir_y = 0.0
        speed = 1.0
        water = 1.0
        sediment = 0.0

        for step in range(max_lifetime):
            node_x = int(pos_x)
            node_y = int(pos_y)

            if node_x < 1 or node_x >= width - 2 or node_y < 1 or node_y >= height - 2:
                break

            fx = pos_x - node_x
            fy = pos_y - node_y

            h00 = heightmap[node_y, node_x]
            h10 = heightmap[node_y, node_x + 1]
            h01 = heightmap[node_y + 1, node_x]
            h11 = heightmap[node_y + 1, node_x + 1]

            grad_x = (h10 - h00) * (1 - fy) + (h11 - h01) * fy
            grad_y = (h01 - h00) * (1 - fx) + (h11 - h10) * fx

            dir_x = dir_x * inertia - grad_x * (1 - inertia)
            dir_y = dir_y * inertia - grad_y * (1 - inertia)

            dir_len = (dir_x * dir_x + dir_y * dir_y) ** 0.5
            if dir_len < 1e-10:
                break
            dir_x /= dir_len
            dir_y /= dir_len

            new_x = pos_x + dir_x
            new_y = pos_y + dir_y

            if new_x < 1 or new_x >= width - 2 or new_y < 1 or new_y >= height - 2:
                break

            h_old = h00 * (1 - fx) * (1 - fy) + h10 * fx * (1 - fy) + \
                    h01 * (1 - fx) * fy + h11 * fx * fy

            new_node_x = int(new_x)
            new_node_y = int(new_y)
            new_fx = new_x - new_node_x
            new_fy = new_y - new_node_y
            h_new = (heightmap[new_node_y, new_node_x] * (1 - new_fx) * (1 - new_fy) +
                     heightmap[new_node_y, new_node_x + 1] * new_fx * (1 - new_fy) +
                     heightmap[new_node_y + 1, new_node_x] * (1 - new_fx) * new_fy +
                     heightmap[new_node_y + 1, new_node_x + 1] * new_fx * new_fy)

            h_diff = h_new - h_old

            sed_capacity = max(-h_diff, min_slope) * speed * water * capacity

            if sediment > sed_capacity or h_diff > 0:
                if h_diff > 0:
                    amount = min(h_diff, sediment)
                else:
                    amount = (sediment - sed_capacity) * deposition

                sediment -= amount

                heightmap[node_y, node_x] += amount * (1 - fx) * (1 - fy)
                heightmap[node_y, node_x + 1] += amount * fx * (1 - fy)
                heightmap[node_y + 1, node_x] += amount * (1 - fx) * fy
                heightmap[node_y + 1, node_x + 1] += amount * fx * fy
            else:
                amount = min((sed_capacity - sediment) * erosion_rate, -h_diff)

                for k in range(n_brush):
                    ey = node_y + boy[k]
                    ex = node_x + box[k]
                    if 0 <= ey < height and 0 <= ex < width:
                        heightmap[ey, ex] -= amount * bw[k]

                sediment += amount

            speed_sq = speed * speed + h_diff * gravity
            speed = speed_sq ** 0.5 if speed_sq > 0 else 0.0
            water *= (1 - evaporation)

            pos_x = new_x
            pos_y = new_y

    return heightmap


# ---------------------------------------------------------------------------
# GPU (CuPy / CUDA) implementation
# ---------------------------------------------------------------------------

if cuda is not None:

    @cuda.jit
    def _erode_gpu_kernel(
        heightmap, random_pos, boy, box, bw,
        inertia, capacity, deposition, erosion_rate,
        evaporation, gravity, min_slope, max_lifetime,
    ):
        """One CUDA thread per particle.

        Particles within a single launch are independent.  Conflicts at
        shared heightmap cells are resolved via ``cuda.atomic.add``.
        """
        tid = cuda.grid(1)
        n_iterations = random_pos.shape[0]
        if tid >= n_iterations:
            return

        height = heightmap.shape[0]
        width = heightmap.shape[1]
        n_brush = bw.shape[0]

        pos_x = random_pos[tid, 0] * (width - 3) + 1
        pos_y = random_pos[tid, 1] * (height - 3) + 1
        dir_x = 0.0
        dir_y = 0.0
        speed = 1.0
        water = 1.0
        sediment = 0.0

        for step in range(max_lifetime):
            node_x = int(pos_x)
            node_y = int(pos_y)

            if node_x < 1 or node_x >= width - 2:
                return
            if node_y < 1 or node_y >= height - 2:
                return

            fx = pos_x - node_x
            fy = pos_y - node_y

            h00 = heightmap[node_y, node_x]
            h10 = heightmap[node_y, node_x + 1]
            h01 = heightmap[node_y + 1, node_x]
            h11 = heightmap[node_y + 1, node_x + 1]

            grad_x = (h10 - h00) * (1 - fy) + (h11 - h01) * fy
            grad_y = (h01 - h00) * (1 - fx) + (h11 - h10) * fx

            dir_x = dir_x * inertia - grad_x * (1 - inertia)
            dir_y = dir_y * inertia - grad_y * (1 - inertia)

            dir_len = (dir_x * dir_x + dir_y * dir_y) ** 0.5
            if dir_len < 1e-10:
                return
            dir_x /= dir_len
            dir_y /= dir_len

            new_x = pos_x + dir_x
            new_y = pos_y + dir_y

            if new_x < 1 or new_x >= width - 2:
                return
            if new_y < 1 or new_y >= height - 2:
                return

            h_old = (h00 * (1 - fx) * (1 - fy) + h10 * fx * (1 - fy) +
                     h01 * (1 - fx) * fy + h11 * fx * fy)

            new_node_x = int(new_x)
            new_node_y = int(new_y)
            new_fx = new_x - new_node_x
            new_fy = new_y - new_node_y
            h_new = (heightmap[new_node_y, new_node_x] * (1 - new_fx) * (1 - new_fy) +
                     heightmap[new_node_y, new_node_x + 1] * new_fx * (1 - new_fy) +
                     heightmap[new_node_y + 1, new_node_x] * (1 - new_fx) * new_fy +
                     heightmap[new_node_y + 1, new_node_x + 1] * new_fx * new_fy)

            h_diff = h_new - h_old

            neg_h_diff = -h_diff
            if neg_h_diff < min_slope:
                neg_h_diff = min_slope
            sed_capacity = neg_h_diff * speed * water * capacity

            if sediment > sed_capacity or h_diff > 0:
                if h_diff > 0:
                    amount = h_diff if h_diff < sediment else sediment
                else:
                    amount = (sediment - sed_capacity) * deposition

                sediment -= amount

                cuda.atomic.add(heightmap, (node_y, node_x),
                                amount * (1 - fx) * (1 - fy))
                cuda.atomic.add(heightmap, (node_y, node_x + 1),
                                amount * fx * (1 - fy))
                cuda.atomic.add(heightmap, (node_y + 1, node_x),
                                amount * (1 - fx) * fy)
                cuda.atomic.add(heightmap, (node_y + 1, node_x + 1),
                                amount * fx * fy)
            else:
                neg_h = -h_diff
                sc_minus_sed = (sed_capacity - sediment) * erosion_rate
                amount = sc_minus_sed if sc_minus_sed < neg_h else neg_h

                for k in range(n_brush):
                    ey = node_y + boy[k]
                    ex = node_x + box[k]
                    if 0 <= ey < height and 0 <= ex < width:
                        cuda.atomic.add(heightmap, (ey, ex),
                                        -(amount * bw[k]))

                sediment += amount

            speed_sq = speed * speed + h_diff * gravity
            if speed_sq > 0:
                speed = speed_sq ** 0.5
            else:
                speed = 0.0
            water *= (1 - evaporation)

            pos_x = new_x
            pos_y = new_y


def _available_memory_bytes():
    """Best-effort estimate of available host memory in bytes."""
    try:
        with open('/proc/meminfo', 'r') as f:
            for line in f:
                if line.startswith('MemAvailable:'):
                    return int(line.split()[1]) * 1024  # kB -> bytes
    except (OSError, ValueError, IndexError):
        pass
    try:
        import psutil
        return psutil.virtual_memory().available
    except (ImportError, AttributeError):
        pass
    return 2 * 1024 ** 3


def _check_erosion_memory(shape, dtype, iterations, n_brush):
    """Raise MemoryError if the projected working set would exceed RAM.

    Covers four allocations:
      - input copy as float64 (numpy path) or float32 output
      - random_pos buffer of shape (iterations, 2), float64
      - brush arrays of length n_brush (two int32, one float64)
      - output cast back to float32

    The check fires before any of these are allocated, so callers see
    a clean MemoryError instead of an opaque OOM from numpy/cupy.
    """
    raster_bytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
    # float64 working copy + float32 output ~ 3x raster_bytes for float32
    # input (8/4 = 2 for the copy and another 1 for the output cast).
    raster_working = raster_bytes * 3
    rng_bytes = int(iterations) * 2 * 8  # float64 (iterations, 2)
    brush_bytes = int(n_brush) * (4 + 4 + 8)  # two int32 + one float64
    required = raster_working + rng_bytes + brush_bytes
    avail = _available_memory_bytes()
    if required > 0.8 * avail:
        raise MemoryError(
            f"erode() needs ~{required / 1e9:.1f} GB of working memory "
            f"(raster ~{raster_working / 1e9:.2f} GB, "
            f"random_pos ~{rng_bytes / 1e9:.2f} GB, "
            f"brush ~{brush_bytes / 1e9:.2f} GB) but only "
            f"~{avail / 1e9:.1f} GB is available. "
            f"Reduce `iterations`, `params['radius']`, or downsample "
            f"the input."
        )


def _erode_cupy(data, random_pos_np, boy, box, bw, params):
    """Run erosion on a CuPy array using the CUDA kernel."""
    _check_erosion_memory(data.shape, data.dtype,
                          random_pos_np.shape[0], bw.shape[0])
    hm = data.astype(cupy.float64)

    # Transfer brush and random positions to device
    random_pos_d = cupy.asarray(random_pos_np)
    boy_d = cupy.asarray(boy)
    box_d = cupy.asarray(box)
    bw_d = cupy.asarray(bw)

    n_particles = random_pos_np.shape[0]
    threads_per_block = 256
    blocks = (n_particles + threads_per_block - 1) // threads_per_block

    _erode_gpu_kernel[blocks, threads_per_block](
        hm, random_pos_d, boy_d, box_d, bw_d,
        params['inertia'], params['capacity'], params['deposition'],
        params['erosion'], params['evaporation'], params['gravity'],
        params['min_slope'], int(params['max_lifetime']),
    )
    cuda.synchronize()

    return hm.astype(cupy.float32)


def _erode_numpy(data, random_pos, boy, box, bw, params):
    """Run erosion on a NumPy array using the CPU kernel."""
    _check_erosion_memory(data.shape, data.dtype,
                          random_pos.shape[0], bw.shape[0])
    hm = data.astype(np.float64).copy()

    hm = _erode_cpu(
        hm, random_pos, boy, box, bw,
        params['inertia'], params['capacity'], params['deposition'],
        params['erosion'], params['evaporation'], params['gravity'],
        params['min_slope'], int(params['max_lifetime']),
    )

    return hm.astype(np.float32)


def _erode_dask_numpy(data, random_pos, boy, box, bw, params):
    """Run erosion on a dask+numpy array.

    Erosion is a global operation (particles traverse the full grid),
    so we materialize to numpy, run the CPU kernel, then re-wrap.
    """
    np_data = data.compute()
    result = _erode_numpy(np_data, random_pos, boy, box, bw, params)
    return da.from_array(result, chunks=data.chunksize)


def _erode_dask_cupy(data, random_pos, boy, box, bw, params):
    """Run erosion on a dask+cupy array.

    Materializes to a single CuPy array, runs the GPU kernel, then
    re-wraps as dask.
    """
    cp_data = data.compute()  # CuPy ndarray
    result = _erode_cupy(cp_data, random_pos, boy, box, bw, params)
    return da.from_array(result, chunks=data.chunksize,
                         meta=cupy.array((), dtype=cupy.float32))


def erode(agg, iterations=50000, seed=42, params=None):
    """Apply particle-based hydraulic erosion to a terrain DataArray.

    Erosion is a global operation that cannot be chunked, so dask arrays
    are materialized before processing and re-wrapped afterwards.

    Parameters
    ----------
    agg : xr.DataArray
        2D terrain heightmap.
    iterations : int
        Number of water droplets to simulate. Capped at
        ``_MAX_ITERATIONS`` (1e8) to keep host and device allocations
        bounded.
    seed : int
        Random seed for droplet placement.
    params : dict, optional
        Override default erosion constants. Keys:
        inertia, capacity, deposition, erosion, evaporation,
        gravity, min_slope, radius, max_lifetime.
        ``radius`` is capped at ``_MAX_RADIUS`` (1024) and
        ``max_lifetime`` at ``_MAX_LIFETIME`` (1e5).

    Returns
    -------
    xr.DataArray
        Eroded terrain.

    Raises
    ------
    ValueError
        If `iterations`, `params['radius']`, or `params['max_lifetime']`
        is outside the allowed range.
    MemoryError
        If the projected working set exceeds available memory.
    """
    _validate_scalar(
        iterations, func_name='erode', name='iterations',
        dtype=int, min_val=1, max_val=_MAX_ITERATIONS,
    )

    p = dict(_DEFAULT_PARAMS)
    if params is not None:
        p.update(params)

    _validate_scalar(
        p['radius'], func_name='erode', name="params['radius']",
        dtype=int, min_val=1, max_val=_MAX_RADIUS,
    )
    _validate_scalar(
        p['max_lifetime'], func_name='erode', name="params['max_lifetime']",
        dtype=int, min_val=1, max_val=_MAX_LIFETIME,
    )

    # Precompute brush and random positions on the host
    boy, box, bw = _build_brush(int(p['radius']))
    rng = np.random.RandomState(seed)
    random_pos = rng.random((iterations, 2))

    mapper = ArrayTypeFunctionMapping(
        numpy_func=lambda d: _erode_numpy(d, random_pos, boy, box, bw, p),
        cupy_func=lambda d: _erode_cupy(d, random_pos, boy, box, bw, p),
        dask_func=lambda d: _erode_dask_numpy(d, random_pos, boy, box, bw, p),
        dask_cupy_func=lambda d: _erode_dask_cupy(d, random_pos, boy, box, bw, p),
    )

    result_data = mapper(agg)(agg.data)

    return xr.DataArray(result_data, dims=agg.dims, coords=agg.coords,
                        attrs=agg.attrs, name=agg.name)