Skip to content

Add scalar diffusion solver #940

@brendancol

Description

@brendancol

Author of Proposal: @brendan

Reason or Problem

There's no way to model how a scalar field spreads across a raster over time. The hydrology tools handle directed flow (steepest descent), but diffusion is different: it spreads in all directions proportional to local gradients. This comes up in heat transfer, pollutant dispersion, moisture migration, and similar problems where you want to answer "the HVAC unit at cell (50, 30) failed, what does the temperature field look like in 2 hours?"

Proposal

Add a diffuse() function that runs an explicit finite-difference solver for the 2D diffusion equation on a raster.

Design:

The diffusion equation:

dU/dt = alpha * laplacian(U)

With a standard 5-point Laplacian stencil and forward Euler time-stepping, each step updates every cell:

u_new[i,j] = u[i,j] + dt * alpha[i,j] * (u[i-1,j] + u[i+1,j] + u[i,j-1] + u[i,j+1] - 4*u[i,j]) / dx^2

The function takes:

  • agg: initial scalar field (xr.DataArray, 2D)
  • diffusivity: float or xr.DataArray (spatially varying)
  • steps: number of time steps to run
  • dt: time step size (default: auto-computed from stability limit)
  • boundary: edge handling ('nan', 'nearest', 'reflect', 'wrap')

When dt is not given, the solver picks the largest stable step: dt = 0.25 * dx^2 / max(alpha) (CFL condition for the explicit scheme).

All four backends (numpy, cupy, dask+numpy, dask+cupy) supported via ArrayTypeFunctionMapping. The numpy path uses @ngjit for the inner loop. The cupy path uses a @cuda.jit kernel. Dask paths use map_overlap with depth=1 since the stencil is 3x3.

Usage:

from xrspatial.diffusion import diffuse

# uniform diffusivity
result = diffuse(temperature, diffusivity=0.01, steps=100)

# spatially varying diffusivity
result = diffuse(concentration, diffusivity=alpha_raster, steps=50, dt=0.1)

The existing tools don't cover this case: focal tools do single-pass neighborhood ops, hydrology tools do directed flow. Diffusion is the missing piece for thermal modeling, pollutant transport, and spatially-aware smoothing.

Stakeholders and Impacts

Anyone doing environmental modeling, building simulation, or contaminant transport. New module (xrspatial/diffusion.py), no changes to existing code.

Drawbacks

The explicit Euler scheme has a stability limit on dt, so large diffusivity values need many small steps. An implicit solver would remove this but requires a sparse linear solver, which is a lot more work. The explicit scheme covers most practical cases.

Alternatives

  • focal.mean() repeated gives a diffusion-like effect but you can't control the diffusivity or time step.
  • Scipy's ndimage.convolve with a Laplacian kernel works for numpy but doesn't extend to the other backends.

Unresolved Questions

None.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions