tidalgeopro/upl.py at master · ogourgue/tidalgeopro · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
""" Unchanneled Path Length.

Compute the unchanneled path length (UPL), i.e., the shortest distance to a
    channel node (triangular grids).

Author: Olivier Gourgue (University of Antwerp & Boston University).

"""

import numpy as np
from scipy import spatial


################################################################################
# Unchanneled path length. #####################################################
################################################################################

def upl(x, y, chn, mask = None):
    """Compute the unchanneled path length (UPL).

    Args:
        x, y (NumPy arrays): Node coordinates.
        chn (NumPy array, boolean): True for channel nodes, False otherwise, for
            one (1D array) or several time steps (2D array, second dimension for
            time).
        mask (NumPy array, boolean): True at grid nodes where UPL is not
            computed (same shape as chn for one time step; default to None, that
            is, no mask).

    Returns:
        NumPy array: Unchanneled path length (same shape as chn).
    """

    # Reshape chn as a 2D array, if needed.
    if chn.ndim == 1:
        chn = chn.reshape((-1, 1))

    # Initialize.
    upl = np.zeros(chn.shape)

    # Number of nodes.
    n = chn.shape[0]

    # Number of time steps.
    nt = chn.shape[1]

    # Default mask.
    if mask is None:
        mask = np.zeros((n, 1), dtype = bool)

    # Area of interest.
    not_mask = np.logical_not(mask)

    # Loop over time steps.
    for i in range(nt):

        # Channel node indices and coordinates.
        ind_chn = np.flatnonzero(chn[:, i] * not_mask)
        xy_chn = np.array([x[ind_chn], y[ind_chn]]).T

        # Platform node indices and coordinates.
        ind_plt = np.flatnonzero(np.logical_not(chn[:, i]) * not_mask)
        xy_plt = np.array([x[ind_plt], y[ind_plt]]).T

        # UPL.
        if len(ind_chn) > 0:
            tree = spatial.KDTree(xy_chn)
            upl_plt, ind = tree.query(xy_plt)
            upl[ind_plt, i] = upl_plt

    # Reshape as a 1D array, if needed.
    if nt == 1:
        upl = upl.reshape(-1)

    return upl