flow_accumulation.py

"""D8 flow accumulation: count of upstream cells draining through each cell.

Computes the number of upstream cells (including itself) that drain
through each cell in a D8 flow direction grid.  This is a global
graph traversal (topological sort of the D8 flow forest), not a
local window operation.

Algorithm
---------
CPU : Kahn's BFS topological sort — O(N).
GPU : iterative frontier peeling with pull-based kernels.
Dask: iterative tile sweep with boundary propagation (one tile in
      RAM at a time), following the ``cost_distance.py`` pattern.
"""

from __future__ import annotations

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

try:
    import cupy
except ImportError:
    class cupy:  # type: ignore[no-redef]
        ndarray = False

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

from xrspatial.utils import (
    _validate_raster,
    cuda_args,
    has_cuda_and_cupy,
    is_cupy_array,
    is_dask_cupy,
    ngjit,
)
from xrspatial._boundary_store import BoundaryStore
from xrspatial.dataset_support import supports_dataset


# =====================================================================
# Direction helpers
# =====================================================================

@ngjit
def _code_to_offset(code):
    """Return (dy, dx) row/col offset for a D8 direction code."""
    c = int(code)
    if c == 1:
        return 0, 1
    elif c == 2:
        return 1, 1
    elif c == 4:
        return 1, 0
    elif c == 8:
        return 1, -1
    elif c == 16:
        return 0, -1
    elif c == 32:
        return -1, -1
    elif c == 64:
        return -1, 0
    elif c == 128:
        return -1, 1
    return 0, 0


def _code_to_offset_py(code):
    """Pure-Python version for non-numba contexts."""
    c = int(code)
    _map = {1: (0, 1), 2: (1, 1), 4: (1, 0), 8: (1, -1),
            16: (0, -1), 32: (-1, -1), 64: (-1, 0), 128: (-1, 1)}
    return _map.get(c, (0, 0))


# =====================================================================
# CPU kernel
# =====================================================================

@ngjit
def _flow_accum_cpu(flow_dir, height, width):
    """Kahn's BFS topological sort for flow accumulation."""
    accum = np.empty((height, width), dtype=np.float64)
    in_degree = np.zeros((height, width), dtype=np.int32)
    valid = np.zeros((height, width), dtype=np.int8)

    # Pass 1: initialise
    for r in range(height):
        for c in range(width):
            v = flow_dir[r, c]
            if v == v:  # not NaN
                valid[r, c] = 1
                accum[r, c] = 1.0
            else:
                accum[r, c] = np.nan

    # Pass 2: compute in-degrees
    for r in range(height):
        for c in range(width):
            if valid[r, c] == 0:
                continue
            dy, dx = _code_to_offset(flow_dir[r, c])
            if dy == 0 and dx == 0:
                continue
            nr = r + dy
            nc = c + dx
            if 0 <= nr < height and 0 <= nc < width and valid[nr, nc] == 1:
                in_degree[nr, nc] += 1

    # BFS queue (flat arrays with head/tail pointers)
    queue_r = np.empty(height * width, dtype=np.int64)
    queue_c = np.empty(height * width, dtype=np.int64)
    head = np.int64(0)
    tail = np.int64(0)

    for r in range(height):
        for c in range(width):
            if valid[r, c] == 1 and in_degree[r, c] == 0:
                queue_r[tail] = r
                queue_c[tail] = c
                tail += 1

    while head < tail:
        r = queue_r[head]
        c = queue_c[head]
        head += 1

        dy, dx = _code_to_offset(flow_dir[r, c])
        if dy == 0 and dx == 0:
            continue
        nr = r + dy
        nc = c + dx
        if 0 <= nr < height and 0 <= nc < width and valid[nr, nc] == 1:
            accum[nr, nc] += accum[r, c]
            in_degree[nr, nc] -= 1
            if in_degree[nr, nc] == 0:
                queue_r[tail] = nr
                queue_c[tail] = nc
                tail += 1

    return accum


# =====================================================================
# GPU kernels
# =====================================================================

@cuda.jit
def _init_accum_indegree(flow_dir, accum, in_degree, state, H, W):
    """Initialise accum, in_degree and state arrays on GPU."""
    i, j = cuda.grid(2)
    if i >= H or j >= W:
        return

    v = flow_dir[i, j]
    if v != v:  # NaN
        state[i, j] = 0
        accum[i, j] = 0.0
        return

    state[i, j] = 1
    accum[i, j] = 1.0

    # Decode direction (inline — can't call @ngjit from @cuda.jit)
    code = int(v)
    dy = 0
    dx = 0
    if code == 1:
        dy, dx = 0, 1
    elif code == 2:
        dy, dx = 1, 1
    elif code == 4:
        dy, dx = 1, 0
    elif code == 8:
        dy, dx = 1, -1
    elif code == 16:
        dy, dx = 0, -1
    elif code == 32:
        dy, dx = -1, -1
    elif code == 64:
        dy, dx = -1, 0
    elif code == 128:
        dy, dx = -1, 1

    if dy == 0 and dx == 0:
        return  # pit

    ni = i + dy
    nj = j + dx
    if 0 <= ni < H and 0 <= nj < W:
        cuda.atomic.add(in_degree, (ni, nj), 1)


@cuda.jit
def _find_ready_and_finalize(in_degree, state, changed, H, W):
    """Finalize previous frontier (2→3), mark new frontier (1→2)."""
    i, j = cuda.grid(2)
    if i >= H or j >= W:
        return

    if state[i, j] == 2:
        state[i, j] = 3

    if state[i, j] == 1 and in_degree[i, j] == 0:
        state[i, j] = 2
        cuda.atomic.add(changed, 0, 1)


@cuda.jit
def _pull_from_frontier(flow_dir, accum, in_degree, state, H, W):
    """Active cells pull accumulation from frontier neighbours."""
    i, j = cuda.grid(2)
    if i >= H or j >= W:
        return

    if state[i, j] != 1:
        return

    # Check all 8 neighbours
    for k in range(8):
        if k == 0:
            dy, dx = 0, 1
        elif k == 1:
            dy, dx = 1, 1
        elif k == 2:
            dy, dx = 1, 0
        elif k == 3:
            dy, dx = 1, -1
        elif k == 4:
            dy, dx = 0, -1
        elif k == 5:
            dy, dx = -1, -1
        elif k == 6:
            dy, dx = -1, 0
        else:
            dy, dx = -1, 1

        ni = i + dy
        nj = j + dx
        if ni < 0 or ni >= H or nj < 0 or nj >= W:
            continue

        if state[ni, nj] != 2:
            continue

        # Check if neighbour's flow_dir points to me
        nv = flow_dir[ni, nj]
        ncode = int(nv)
        ndy = 0
        ndx = 0
        if ncode == 1:
            ndy, ndx = 0, 1
        elif ncode == 2:
            ndy, ndx = 1, 1
        elif ncode == 4:
            ndy, ndx = 1, 0
        elif ncode == 8:
            ndy, ndx = 1, -1
        elif ncode == 16:
            ndy, ndx = 0, -1
        elif ncode == 32:
            ndy, ndx = -1, -1
        elif ncode == 64:
            ndy, ndx = -1, 0
        elif ncode == 128:
            ndy, ndx = -1, 1

        if ni + ndy == i and nj + ndx == j:
            accum[i, j] += accum[ni, nj]
            in_degree[i, j] -= 1


def _flow_accum_cupy(flow_dir_data):
    """GPU driver: iterative frontier peeling."""
    import cupy as cp

    H, W = flow_dir_data.shape
    flow_dir_f64 = flow_dir_data.astype(cp.float64)

    accum = cp.zeros((H, W), dtype=cp.float64)
    in_degree = cp.zeros((H, W), dtype=cp.int32)
    state = cp.zeros((H, W), dtype=cp.int32)
    changed = cp.zeros(1, dtype=cp.int32)

    griddim, blockdim = cuda_args((H, W))

    _init_accum_indegree[griddim, blockdim](
        flow_dir_f64, accum, in_degree, state, H, W)

    max_iter = H * W
    for _ in range(max_iter):
        changed[0] = 0
        _find_ready_and_finalize[griddim, blockdim](
            in_degree, state, changed, H, W)

        if int(changed[0]) == 0:
            break

        _pull_from_frontier[griddim, blockdim](
            flow_dir_f64, accum, in_degree, state, H, W)

    # Convert invalid cells to NaN
    accum = cp.where(state == 0, cp.nan, accum)
    return accum


# =====================================================================
# Tile kernel for dask iterative path
# =====================================================================

@ngjit
def _flow_accum_tile_kernel(flow_dir, h, w,
                            seed_top, seed_bottom, seed_left, seed_right,
                            seed_tl, seed_tr, seed_bl, seed_br):
    """Seeded BFS flow accumulation for a single tile.

    Same as ``_flow_accum_cpu`` but adds external seeds to boundary
    cells before BFS.
    """
    accum = np.empty((h, w), dtype=np.float64)
    in_degree = np.zeros((h, w), dtype=np.int32)
    valid = np.zeros((h, w), dtype=np.int8)

    # Initialise
    for r in range(h):
        for c in range(w):
            v = flow_dir[r, c]
            if v == v:
                valid[r, c] = 1
                accum[r, c] = 1.0
            else:
                accum[r, c] = np.nan

    # Add external seeds to boundary cells
    for c in range(w):
        if valid[0, c] == 1:
            accum[0, c] += seed_top[c]
        if valid[h - 1, c] == 1:
            accum[h - 1, c] += seed_bottom[c]
    for r in range(h):
        if valid[r, 0] == 1:
            accum[r, 0] += seed_left[r]
        if valid[r, w - 1] == 1:
            accum[r, w - 1] += seed_right[r]

    # Corner seeds
    if valid[0, 0] == 1:
        accum[0, 0] += seed_tl
    if valid[0, w - 1] == 1:
        accum[0, w - 1] += seed_tr
    if valid[h - 1, 0] == 1:
        accum[h - 1, 0] += seed_bl
    if valid[h - 1, w - 1] == 1:
        accum[h - 1, w - 1] += seed_br

    # Compute in-degrees
    for r in range(h):
        for c in range(w):
            if valid[r, c] == 0:
                continue
            dy, dx = _code_to_offset(flow_dir[r, c])
            if dy == 0 and dx == 0:
                continue
            nr = r + dy
            nc = c + dx
            if 0 <= nr < h and 0 <= nc < w and valid[nr, nc] == 1:
                in_degree[nr, nc] += 1

    # BFS
    queue_r = np.empty(h * w, dtype=np.int64)
    queue_c = np.empty(h * w, dtype=np.int64)
    head = np.int64(0)
    tail = np.int64(0)

    for r in range(h):
        for c in range(w):
            if valid[r, c] == 1 and in_degree[r, c] == 0:
                queue_r[tail] = r
                queue_c[tail] = c
                tail += 1

    while head < tail:
        r = queue_r[head]
        c = queue_c[head]
        head += 1

        dy, dx = _code_to_offset(flow_dir[r, c])
        if dy == 0 and dx == 0:
            continue
        nr = r + dy
        nc = c + dx
        if 0 <= nr < h and 0 <= nc < w and valid[nr, nc] == 1:
            accum[nr, nc] += accum[r, c]
            in_degree[nr, nc] -= 1
            if in_degree[nr, nc] == 0:
                queue_r[tail] = nr
                queue_c[tail] = nc
                tail += 1

    return accum


# =====================================================================
# Dask iterative tile sweep
# =====================================================================

def _preprocess_tiles(flow_dir_da, chunks_y, chunks_x):
    """Extract boundary flow-direction strips into a BoundaryStore."""
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    flow_bdry = BoundaryStore(chunks_y, chunks_x, fill_value=np.nan)

    for iy in range(n_tile_y):
        for ix in range(n_tile_x):
            chunk = flow_dir_da.blocks[iy, ix].compute()
            flow_bdry.set('top', iy, ix,
                          np.asarray(chunk[0, :], dtype=np.float64))
            flow_bdry.set('bottom', iy, ix,
                          np.asarray(chunk[-1, :], dtype=np.float64))
            flow_bdry.set('left', iy, ix,
                          np.asarray(chunk[:, 0], dtype=np.float64))
            flow_bdry.set('right', iy, ix,
                          np.asarray(chunk[:, -1], dtype=np.float64))

    return flow_bdry


def _compute_seeds(iy, ix, boundaries, flow_bdry,
                   chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Compute seed arrays for tile (iy, ix) from neighbour boundaries.

    For each boundary cell of the current tile, checks whether cells
    in adjacent tiles flow INTO this tile.  If so, adds the adjacent
    cell's boundary accum value as a seed.

    Returns (seed_top, seed_bottom, seed_left, seed_right,
             seed_tl, seed_tr, seed_bl, seed_br).
    """
    tile_h = chunks_y[iy]
    tile_w = chunks_x[ix]

    seed_top = np.zeros(tile_w, dtype=np.float64)
    seed_bottom = np.zeros(tile_w, dtype=np.float64)
    seed_left = np.zeros(tile_h, dtype=np.float64)
    seed_right = np.zeros(tile_h, dtype=np.float64)
    seed_tl = 0.0
    seed_tr = 0.0
    seed_bl = 0.0
    seed_br = 0.0

    # --- Top edge: bottom of tile above ---
    if iy > 0:
        nb_fdir = flow_bdry.get('bottom', iy - 1, ix)
        nb_accum = boundaries.get('bottom', iy - 1, ix)
        w = len(nb_fdir)
        # Cardinal S (code 4): nb cell j → our (0, j)
        mask = (nb_fdir == 4)
        seed_top += np.where(mask, nb_accum, 0.0)
        if w > 1:
            # Diagonal SE (code 2): nb cell j → our (0, j+1)
            mask = (nb_fdir[:-1] == 2)
            seed_top[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal SW (code 8): nb cell j → our (0, j-1)
            mask = (nb_fdir[1:] == 8)
            seed_top[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Bottom edge: top of tile below ---
    if iy < n_tile_y - 1:
        nb_fdir = flow_bdry.get('top', iy + 1, ix)
        nb_accum = boundaries.get('top', iy + 1, ix)
        w = len(nb_fdir)
        # Cardinal N (code 64)
        mask = (nb_fdir == 64)
        seed_bottom += np.where(mask, nb_accum, 0.0)
        if w > 1:
            # Diagonal NE (code 128): nb cell j → our (h-1, j+1)
            mask = (nb_fdir[:-1] == 128)
            seed_bottom[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal NW (code 32): nb cell j → our (h-1, j-1)
            mask = (nb_fdir[1:] == 32)
            seed_bottom[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Left edge: right column of tile to the left ---
    if ix > 0:
        nb_fdir = flow_bdry.get('right', iy, ix - 1)
        nb_accum = boundaries.get('right', iy, ix - 1)
        h = len(nb_fdir)
        # Cardinal E (code 1)
        mask = (nb_fdir == 1)
        seed_left += np.where(mask, nb_accum, 0.0)
        if h > 1:
            # Diagonal SE (code 2): nb cell i → our (i+1, 0)
            mask = (nb_fdir[:-1] == 2)
            seed_left[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal NE (code 128): nb cell i → our (i-1, 0)
            mask = (nb_fdir[1:] == 128)
            seed_left[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Right edge: left column of tile to the right ---
    if ix < n_tile_x - 1:
        nb_fdir = flow_bdry.get('left', iy, ix + 1)
        nb_accum = boundaries.get('left', iy, ix + 1)
        h = len(nb_fdir)
        # Cardinal W (code 16)
        mask = (nb_fdir == 16)
        seed_right += np.where(mask, nb_accum, 0.0)
        if h > 1:
            # Diagonal SW (code 8): nb cell i → our (i+1, w-1)
            mask = (nb_fdir[:-1] == 8)
            seed_right[1:] += np.where(mask, nb_accum[:-1], 0.0)
            # Diagonal NW (code 32): nb cell i → our (i-1, w-1)
            mask = (nb_fdir[1:] == 32)
            seed_right[:-1] += np.where(mask, nb_accum[1:], 0.0)

    # --- Diagonal corner seeds ---
    # TL: bottom-right of (iy-1, ix-1) flows SE (code 2)
    if iy > 0 and ix > 0:
        fdir = flow_bdry.get('bottom', iy - 1, ix - 1)[-1]
        if fdir == 2:
            seed_tl = float(boundaries.get('bottom', iy - 1, ix - 1)[-1])
    # TR: bottom-left of (iy-1, ix+1) flows SW (code 8)
    if iy > 0 and ix < n_tile_x - 1:
        fdir = flow_bdry.get('bottom', iy - 1, ix + 1)[0]
        if fdir == 8:
            seed_tr = float(boundaries.get('bottom', iy - 1, ix + 1)[0])
    # BL: top-right of (iy+1, ix-1) flows NE (code 128)
    if iy < n_tile_y - 1 and ix > 0:
        fdir = flow_bdry.get('top', iy + 1, ix - 1)[-1]
        if fdir == 128:
            seed_bl = float(boundaries.get('top', iy + 1, ix - 1)[-1])
    # BR: top-left of (iy+1, ix+1) flows NW (code 32)
    if iy < n_tile_y - 1 and ix < n_tile_x - 1:
        fdir = flow_bdry.get('top', iy + 1, ix + 1)[0]
        if fdir == 32:
            seed_br = float(boundaries.get('top', iy + 1, ix + 1)[0])

    return (seed_top, seed_bottom, seed_left, seed_right,
            seed_tl, seed_tr, seed_bl, seed_br)


def _process_tile(iy, ix, flow_dir_da, boundaries, flow_bdry,
                  chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Run seeded BFS on one tile; update boundaries in-place.

    Returns the maximum absolute boundary change (float).
    """
    chunk = np.asarray(
        flow_dir_da.blocks[iy, ix].compute(), dtype=np.float64)
    h, w = chunk.shape

    seeds = _compute_seeds(
        iy, ix, boundaries, flow_bdry,
        chunks_y, chunks_x, n_tile_y, n_tile_x)

    accum = _flow_accum_tile_kernel(chunk, h, w, *seeds)

    # Extract new boundary strips
    new_top = accum[0, :].copy()
    new_bottom = accum[-1, :].copy()
    new_left = accum[:, 0].copy()
    new_right = accum[:, -1].copy()

    # Compute max absolute change
    change = 0.0
    for side, new in (('top', new_top), ('bottom', new_bottom),
                      ('left', new_left), ('right', new_right)):
        old = boundaries.get(side, iy, ix)
        with np.errstate(invalid='ignore'):
            diff = np.abs(new - old)
        diff = np.where(np.isnan(diff), 0.0, diff)
        m = float(np.max(diff))
        if m > change:
            change = m

    # Store updated boundaries
    boundaries.set('top', iy, ix, new_top)
    boundaries.set('bottom', iy, ix, new_bottom)
    boundaries.set('left', iy, ix, new_left)
    boundaries.set('right', iy, ix, new_right)

    return change


def _flow_accum_dask_iterative(flow_dir_da):
    """Iterative boundary-propagation for arbitrarily large dask arrays.

    Memory usage is O(tile_size + boundary_strips) per iteration.
    """
    chunks_y = flow_dir_da.chunks[0]
    chunks_x = flow_dir_da.chunks[1]
    n_tile_y = len(chunks_y)
    n_tile_x = len(chunks_x)

    # Phase 0: extract boundary flow dirs
    flow_bdry = _preprocess_tiles(flow_dir_da, chunks_y, chunks_x)
    flow_bdry = flow_bdry.snapshot()  # read-only from here; release temp files

    # Phase 1: initialise boundary accum to 0
    boundaries = BoundaryStore(chunks_y, chunks_x, fill_value=0.0)

    # Phase 2: iterative forward/backward sweeps
    max_iterations = max(n_tile_y, n_tile_x) + 10

    for _iteration in range(max_iterations):
        max_change = 0.0

        # Forward sweep (top-left -> bottom-right)
        for iy in range(n_tile_y):
            for ix in range(n_tile_x):
                c = _process_tile(iy, ix, flow_dir_da, boundaries,
                                  flow_bdry, chunks_y, chunks_x,
                                  n_tile_y, n_tile_x)
                if c > max_change:
                    max_change = c

        # Backward sweep (bottom-right -> top-left)
        for iy in reversed(range(n_tile_y)):
            for ix in reversed(range(n_tile_x)):
                c = _process_tile(iy, ix, flow_dir_da, boundaries,
                                  flow_bdry, chunks_y, chunks_x,
                                  n_tile_y, n_tile_x)
                if c > max_change:
                    max_change = c

        if max_change == 0.0:
            break

    # Snapshot converged boundaries before assembly (releases temp files)
    boundaries = boundaries.snapshot()

    # Phase 3: lazy assembly via da.map_blocks
    return _assemble_result(flow_dir_da, boundaries, flow_bdry,
                            chunks_y, chunks_x, n_tile_y, n_tile_x)


def _assemble_result(flow_dir_da, boundaries, flow_bdry,
                     chunks_y, chunks_x, n_tile_y, n_tile_x):
    """Build a lazy dask array by re-running each tile with converged seeds."""

    def _tile_fn(flow_dir_block, block_info=None):
        if block_info is None or 0 not in block_info:
            return np.full(flow_dir_block.shape, np.nan, dtype=np.float64)
        iy, ix = block_info[0]['chunk-location']
        h, w = flow_dir_block.shape
        seeds = _compute_seeds(
            iy, ix, boundaries, flow_bdry,
            chunks_y, chunks_x, n_tile_y, n_tile_x)
        return _flow_accum_tile_kernel(
            np.asarray(flow_dir_block, dtype=np.float64), h, w, *seeds)

    return da.map_blocks(
        _tile_fn,
        flow_dir_da,
        dtype=np.float64,
        meta=np.array((), dtype=np.float64),
    )


def _flow_accum_dask_cupy(flow_dir_da):
    """Dask+CuPy: convert to numpy, run CPU iterative path, convert back."""
    import cupy as cp

    flow_dir_np = flow_dir_da.map_blocks(
        lambda b: b.get(), dtype=flow_dir_da.dtype,
        meta=np.array((), dtype=flow_dir_da.dtype),
    )
    result = _flow_accum_dask_iterative(flow_dir_np)
    return result.map_blocks(
        cp.asarray, dtype=result.dtype,
        meta=cp.array((), dtype=result.dtype),
    )


# =====================================================================
# Public API
# =====================================================================

@supports_dataset
def flow_accumulation(flow_dir: xr.DataArray,
                      name: str = 'flow_accumulation') -> xr.DataArray:
    """Compute D8 flow accumulation from a flow direction grid.

    For each cell, counts how many upstream cells drain through it
    (including itself).  The input must be a D8 flow direction grid
    (codes 0/1/2/4/8/16/32/64/128; NaN for nodata).

    Parameters
    ----------
    flow_dir : xarray.DataArray or xr.Dataset
        2D NumPy, CuPy, NumPy-backed Dask, or CuPy-backed Dask
        xarray DataArray of D8 flow direction codes.
        If a Dataset is passed, the operation is applied to each
        data variable independently.
    name : str, default='flow_accumulation'
        Name of output DataArray.

    Returns
    -------
    xarray.DataArray or xr.Dataset
        2D float64 array of flow accumulation values.  Each cell
        contains the count of upstream cells (including itself) that
        drain through it.  Cells with NaN flow direction produce NaN.

    References
    ----------
    Jenson, S.K. and Domingue, J.O. (1988). Extracting Topographic
    Structure from Digital Elevation Data for Geographic Information
    System Analysis. Photogrammetric Engineering and Remote Sensing,
    54(11), 1593-1600.
    """
    _validate_raster(flow_dir, func_name='flow_accumulation', name='flow_dir')

    data = flow_dir.data

    if isinstance(data, np.ndarray):
        out = _flow_accum_cpu(data.astype(np.float64), *data.shape)

    elif has_cuda_and_cupy() and is_cupy_array(data):
        out = _flow_accum_cupy(data)

    elif has_cuda_and_cupy() and is_dask_cupy(flow_dir):
        out = _flow_accum_dask_cupy(data)

    elif da is not None and isinstance(data, da.Array):
        out = _flow_accum_dask_iterative(data)

    else:
        raise TypeError(f"Unsupported array type: {type(data)}")

    return xr.DataArray(out,
                        name=name,
                        coords=flow_dir.coords,
                        dims=flow_dir.dims,
                        attrs=flow_dir.attrs)