datagenerator.meanderpy

   1import numpy as np
   2import matplotlib.pyplot as plt
   3import scipy.interpolate
   4from scipy.spatial import distance
   5from scipy import ndimage
   6from PIL import Image, ImageDraw
   7from skimage import measure
   8from skimage import morphology
   9
  10# from matplotlib.colors import LinearSegmentedColormap
  11# import time, sys
  12import numba
  13import matplotlib.colors as mcolors
  14from matplotlib import cm
  15from tqdm import trange
  16
  17
  18class Channel:
  19    """class for Channel objects"""
  20
  21    def __init__(self, x, y, z, W, D):
  22        """Initialize Channel object
  23
  24        :param x: x-coordinate of centerline
  25        :param y: y-coordinate of centerline
  26        :param z: z-coordinate of centerline
  27        :param W: channel width
  28        :param D: channel depth"""
  29
  30        self.x = x
  31        self.y = y
  32        self.z = z
  33        self.W = W
  34        self.D = D
  35
  36
  37class Cutoff:
  38    """class for Cutoff objects"""
  39
  40    def __init__(self, x, y, z, W, D):
  41        """Initialize Cutoff object
  42
  43        :param x: x-coordinate of centerline
  44        :param y: y-coordinate of centerline
  45        :param z: z-coordinate of centerline
  46        :param W: channel width
  47        :param D: channel depth"""
  48
  49        self.x = x
  50        self.y = y
  51        self.z = z
  52        self.W = W
  53        self.D = D
  54
  55
  56class ChannelBelt3D:
  57    """class for 3D models of channel belts"""
  58
  59    def __init__(self, model_type, topo, strat, facies, facies_code, dx, channels):
  60        """initialize ChannelBelt3D object
  61
  62        :param model_type: type of model to be built; can be either 'fluvial' or 'submarine'
  63        :param topo: set of topographic surfaces (3D numpy array)
  64        :param strat: set of stratigraphic surfaces (3D numpy array)
  65        :param facies: facies volume (3D numpy array)
  66        :param facies_code: dictionary of facies codes, e.g. {0:'oxbow', 1:'point bar', 2:'levee'}
  67        :param dx: gridcell size (m)
  68        :param channels: list of channel objects that form 3D model"""
  69
  70        self.model_type = model_type
  71        self.topo = topo
  72        self.strat = strat
  73        self.facies = facies
  74        self.facies_code = facies_code
  75        self.dx = dx
  76        self.channels = channels
  77
  78    def plot_xsection(self, xsec, colors, ve):
  79        """method for plotting a cross section through a 3D model; also plots map of
  80        basal erosional surface and map of final geomorphic surface
  81
  82        :param xsec: location of cross section along the x-axis (in pixel/ voxel coordinates)
  83        :param colors: list of RGB values that define the colors for different facies
  84        :param ve: vertical exaggeration
  85        :return: handles to the three figures"""
  86
  87        strat = self.strat
  88        dx = self.dx
  89        fig1 = plt.figure(figsize=(20, 5))
  90        ax1 = fig1.add_subplot(111)
  91        r, c, ts = np.shape(strat)
  92        Xv = dx * np.arange(0, r)
  93        for i in range(0, ts - 1, 3):
  94            X1 = np.concatenate((Xv, Xv[::-1]))
  95            Y1 = np.concatenate((strat[:, xsec, i], strat[::-1, xsec, i + 1]))
  96            Y2 = np.concatenate((strat[:, xsec, i + 1], strat[::-1, xsec, i + 2]))
  97            Y3 = np.concatenate((strat[:, xsec, i + 2], strat[::-1, xsec, i + 3]))
  98            if self.model_type == "submarine":
  99                ax1.fill(
 100                    X1, Y1, facecolor=colors[2], linewidth=0.5, edgecolor=[0, 0, 0]
 101                )  # oxbow mud
 102                ax1.fill(
 103                    X1, Y2, facecolor=colors[0], linewidth=0.5, edgecolor=[0, 0, 0]
 104                )  # point bar sand
 105                ax1.fill(X1, Y3, facecolor=colors[1], linewidth=0.5)  # levee mud
 106            if self.model_type == "fluvial":
 107                ax1.fill(
 108                    X1, Y1, facecolor=colors[0], linewidth=0.5, edgecolor=[0, 0, 0]
 109                )  # levee mud
 110                ax1.fill(
 111                    X1, Y2, facecolor=colors[1], linewidth=0.5, edgecolor=[0, 0, 0]
 112                )  # oxbow mud
 113                ax1.fill(X1, Y3, facecolor=colors[2], linewidth=0.5)  # channel sand
 114        ax1.set_xlim(0, dx * (r - 1))
 115        ax1.set_aspect(ve, adjustable="datalim")
 116        fig2 = plt.figure()
 117        ax2 = fig2.add_subplot(111)
 118        ax2.contourf(strat[:, :, ts - 1], 100, cmap="viridis")
 119        ax2.contour(
 120            strat[:, :, ts - 1],
 121            100,
 122            colors="k",
 123            linestyles="solid",
 124            linewidths=0.1,
 125            alpha=0.4,
 126        )
 127        ax2.plot([xsec, xsec], [0, r], "k", linewidth=2)
 128        ax2.axis([0, c, 0, r])
 129        ax2.set_aspect("equal", adjustable="box")
 130        ax2.set_title("final geomorphic surface")
 131        ax2.tick_params(
 132            bottom=False,
 133            top=False,
 134            left=False,
 135            right=False,
 136            labelbottom=False,
 137            labelleft=False,
 138        )
 139        fig3 = plt.figure()
 140        ax3 = fig3.add_subplot(111)
 141        ax3.contourf(strat[:, :, 0], 100, cmap="viridis")
 142        ax3.contour(
 143            strat[:, :, 0],
 144            100,
 145            colors="k",
 146            linestyles="solid",
 147            linewidths=0.1,
 148            alpha=0.4,
 149        )
 150        ax3.plot([xsec, xsec], [0, r], "k", linewidth=2)
 151        ax3.axis([0, c, 0, r])
 152        ax3.set_aspect("equal", adjustable="box")
 153        ax3.set_title("basal erosional surface")
 154        ax3.tick_params(
 155            bottom=False,
 156            top=False,
 157            left=False,
 158            right=False,
 159            labelbottom=False,
 160            labelleft=False,
 161        )
 162        return fig1, fig2, fig3
 163
 164
 165class ChannelBelt:
 166    """class for ChannelBelt objects"""
 167
 168    def __init__(self, channels, cutoffs, cl_times, cutoff_times):
 169        """initialize ChannelBelt object
 170
 171        :param channels: list of Channel objects
 172        :param cutoffs: list of Cutoff objects
 173        :param cl_times: list of ages of Channel objects (in years)
 174        :param cutoff_times: list of ages of Cutoff objects"""
 175
 176        self.channels = channels
 177        self.cutoffs = cutoffs
 178        self.cl_times = cl_times
 179        self.cutoff_times = cutoff_times
 180
 181    def migrate(
 182        self,
 183        nit,
 184        saved_ts,
 185        deltas,
 186        pad,
 187        crdist,
 188        Cf,
 189        kl,
 190        kv,
 191        dt,
 192        dens,
 193        t1,
 194        t2,
 195        t3,
 196        aggr_factor,
 197        *D
 198    ):
 199        """method for computing migration rates along channel centerlines and moving the centerlines accordingly
 200
 201        :param nit: number of iterations
 202        :param saved_ts: which time steps will be saved; e.g., if saved_ts = 10, every tenth time step will be saved
 203        :param deltas: distance between nodes on centerline
 204        :param pad: padding (number of nodepoints along centerline)
 205        :param crdist: threshold distance at which cutoffs occur
 206        :param Cf: dimensionless Chezy friction factor
 207        :param kl: migration rate constant (m/s)
 208        :param kv: vertical slope-dependent erosion rate constant (m/s)
 209        :param dt: time step (s)
 210        :param dens: density of fluid (kg/m3)
 211        :param t1: time step when incision starts
 212        :param t2: time step when lateral migration starts
 213        :param t3: time step when aggradation starts
 214        :param aggr_factor: aggradation factor
 215        :param D: channel depth (m)"""
 216
 217        channel = self.channels[
 218            -1
 219        ]  # first channel is the same as last channel of input
 220        x = channel.x
 221        y = channel.y
 222        z = channel.z
 223        W = channel.W
 224        if len(D) == 0:
 225            D = channel.D
 226        else:
 227            D = D[0]
 228        k = 1.0  # constant in HK equation
 229        xc = []  # initialize cutoff coordinates
 230        # determine age of last channel:
 231        if len(self.cl_times) > 0:
 232            last_cl_time = self.cl_times[-1]
 233        else:
 234            last_cl_time = 0
 235        dx, dy, dz, ds, s = compute_derivatives(x, y, z)
 236        slope = np.gradient(z) / ds
 237        # padding at the beginning can be shorter than padding at the downstream end:
 238        pad1 = int(pad / 10.0)
 239        if pad1 < 5:
 240            pad1 = 5
 241        omega = (
 242            -1.0
 243        )  # constant in migration rate calculation (Howard and Knutson, 1984)
 244        gamma = 2.5  # from Ikeda et al., 1981 and Howard and Knutson, 1984
 245        for itn in trange(nit):  # main loop
 246            # update_progress(itn/nit)
 247            x, y = migrate_one_step(
 248                x, y, z, W, kl, dt, k, Cf, D, pad, pad1, omega, gamma
 249            )
 250            # x, y = migrate_one_step_w_bias(x,y,z,W,kl,dt,k,Cf,D,pad,pad1,omega,gamma)
 251            x, y, z, xc, yc, zc = cut_off_cutoffs(
 252                x, y, z, s, crdist, deltas
 253            )  # find and execute cutoffs
 254            x, y, z, dx, dy, dz, ds, s = resample_centerline(
 255                x, y, z, deltas
 256            )  # resample centerline
 257            slope = np.gradient(z) / ds
 258            # for itn<=t1, z is unchanged; for itn>t1:
 259            if (itn > t1) & (itn <= t2):  # incision
 260                if np.min(np.abs(slope)) != 0:  # if slope is not zero
 261                    z = z + kv * dens * 9.81 * D * slope * dt
 262                else:
 263                    z = z - kv * dens * 9.81 * D * dt * 0.05  # if slope is zero
 264            if (itn > t2) & (itn <= t3):  # lateral migration
 265                if np.min(np.abs(slope)) != 0:  # if slope is not zero
 266                    # use the median slope to counterbalance incision:
 267                    z = (
 268                        z
 269                        + kv * dens * 9.81 * D * slope * dt
 270                        - kv * dens * 9.81 * D * np.median(slope) * dt
 271                    )
 272                else:
 273                    z = z  # no change in z
 274            if itn > t3:  # aggradation
 275                if np.min(np.abs(slope)) != 0:  # if slope is not zero
 276                    # 'aggr_factor' should be larger than 1 so that this leads to overall aggradation:
 277                    z = (
 278                        z
 279                        + kv * dens * 9.81 * D * slope * dt
 280                        - aggr_factor * kv * dens * 9.81 * D * np.mean(slope) * dt
 281                    )
 282                else:
 283                    z = z + aggr_factor * dt
 284            if len(xc) > 0:  # save cutoff data
 285                self.cutoff_times.append(
 286                    last_cl_time + (itn + 1) * dt / (365 * 24 * 60 * 60.0)
 287                )
 288                cutoff = Cutoff(xc, yc, zc, W, D)  # create cutoff object
 289                self.cutoffs.append(cutoff)
 290            # saving centerlines:
 291            if np.mod(itn, saved_ts) == 0:
 292                self.cl_times.append(
 293                    last_cl_time + (itn + 1) * dt / (365 * 24 * 60 * 60.0)
 294                )
 295                channel = Channel(x, y, z, W, D)  # create channel object
 296                self.channels.append(channel)
 297
 298    def plot(self, plot_type, pb_age, ob_age, end_time, n_channels):
 299        """method  for plotting ChannelBelt object
 300
 301        :param plot_type: can be either 'strat' (for stratigraphic plot), 'morph' (for morphologic plot), or 'age' (for age plot)
 302        :param pb_age: age of point bars (in years) at which they get covered by vegetation
 303        :param ob_age: age of oxbow lakes (in years) at which they get covered by vegetation
 304        :param end_time: age of last channel to be plotted (in years)
 305        :param n_channels: total number of channels (used in 'age' plots; can be larger than number of channels being plotted)
 306        :return: handle to figure"""
 307
 308        cot = np.array(self.cutoff_times)
 309        sclt = np.array(self.cl_times)
 310        if end_time > 0:
 311            cot = cot[cot <= end_time]
 312            sclt = sclt[sclt <= end_time]
 313        times = np.sort(np.hstack((cot, sclt)))
 314        times = np.unique(times)
 315        order = 0  # variable for ordering objects in plot
 316        # set up min and max x and y coordinates of the plot:
 317        xmin = np.min(self.channels[0].x)
 318        xmax = np.max(self.channels[0].x)
 319        ymax = 0
 320        for i in range(len(self.channels)):
 321            ymax = max(ymax, np.max(np.abs(self.channels[i].y)))
 322        ymax = ymax + 2 * self.channels[0].W  # add a bit of space on top and bottom
 323        ymin = -1 * ymax
 324        # size figure so that its size matches the size of the model:
 325        fig = plt.figure(figsize=(20, (ymax - ymin) * 20 / (xmax - xmin)))
 326        if plot_type == "morph":
 327            pb_crit = len(times[times < times[-1] - pb_age]) / float(len(times))
 328            ob_crit = len(times[times < times[-1] - ob_age]) / float(len(times))
 329            green = (106 / 255.0, 159 / 255.0, 67 / 255.0)  # vegetation color
 330            pb_color = (189 / 255.0, 153 / 255.0, 148 / 255.0)  # point bar color
 331            ob_color = (15 / 255.0, 58 / 255.0, 65 / 255.0)  # oxbow color
 332            pb_cmap = make_colormap(
 333                [green, green, pb_crit, green, pb_color, 1.0, pb_color]
 334            )  # colormap for point bars
 335            ob_cmap = make_colormap(
 336                [green, green, ob_crit, green, ob_color, 1.0, ob_color]
 337            )  # colormap for oxbows
 338            plt.fill(
 339                [xmin, xmax, xmax, xmin],
 340                [ymin, ymin, ymax, ymax],
 341                color=(106 / 255.0, 159 / 255.0, 67 / 255.0),
 342            )
 343        if plot_type == "age":
 344            age_cmap = cm.get_cmap("magma", n_channels)
 345        for i in range(0, len(times)):
 346            if times[i] in sclt:
 347                ind = np.where(sclt == times[i])[0][0]
 348                x1 = self.channels[ind].x
 349                y1 = self.channels[ind].y
 350                W = self.channels[ind].W
 351                xm, ym = get_channel_banks(x1, y1, W)
 352                if plot_type == "morph":
 353                    if times[i] > times[-1] - pb_age:
 354                        plt.fill(
 355                            xm,
 356                            ym,
 357                            facecolor=pb_cmap(i / float(len(times) - 1)),
 358                            edgecolor="k",
 359                            linewidth=0.2,
 360                        )
 361                    else:
 362                        plt.fill(xm, ym, facecolor=pb_cmap(i / float(len(times) - 1)))
 363                if plot_type == "strat":
 364                    order += 1
 365                    plt.fill(
 366                        xm,
 367                        ym,
 368                        "xkcd:light tan",
 369                        edgecolor="k",
 370                        linewidth=0.25,
 371                        zorder=order,
 372                    )
 373                if plot_type == "age":
 374                    order += 1
 375                    plt.fill(
 376                        xm,
 377                        ym,
 378                        facecolor=age_cmap(i / float(n_channels - 1)),
 379                        edgecolor="k",
 380                        linewidth=0.1,
 381                        zorder=order,
 382                    )
 383            if times[i] in cot:
 384                ind = np.where(cot == times[i])[0][0]
 385                for j in range(0, len(self.cutoffs[ind].x)):
 386                    x1 = self.cutoffs[ind].x[j]
 387                    y1 = self.cutoffs[ind].y[j]
 388                    xm, ym = get_channel_banks(x1, y1, self.cutoffs[ind].W)
 389                    if plot_type == "morph":
 390                        plt.fill(xm, ym, color=ob_cmap(i / float(len(times) - 1)))
 391                    if plot_type == "strat":
 392                        order = order + 1
 393                        plt.fill(
 394                            xm,
 395                            ym,
 396                            "xkcd:ocean blue",
 397                            edgecolor="k",
 398                            linewidth=0.25,
 399                            zorder=order,
 400                        )
 401                    if plot_type == "age":
 402                        order += 1
 403                        plt.fill(
 404                            xm,
 405                            ym,
 406                            "xkcd:sea blue",
 407                            edgecolor="k",
 408                            linewidth=0.1,
 409                            zorder=order,
 410                        )
 411
 412        x1 = self.channels[len(sclt) - 1].x
 413        y1 = self.channels[len(sclt) - 1].y
 414        xm, ym = get_channel_banks(x1, y1, self.channels[len(sclt) - 1].W)
 415        order = order + 1
 416        if plot_type == "age":
 417            plt.fill(
 418                xm,
 419                ym,
 420                color="xkcd:sea blue",
 421                zorder=order,
 422                edgecolor="k",
 423                linewidth=0.1,
 424            )
 425        else:
 426            plt.fill(
 427                xm, ym, color=(16 / 255.0, 73 / 255.0, 90 / 255.0), zorder=order
 428            )  # ,edgecolor='k')
 429        plt.axis("equal")
 430        plt.xlim(xmin, xmax)
 431        plt.ylim(ymin, ymax)
 432        return fig
 433
 434    def create_movie(
 435        self,
 436        xmin,
 437        xmax,
 438        plot_type,
 439        filename,
 440        dirname,
 441        pb_age,
 442        ob_age,
 443        end_time,
 444        n_channels,
 445    ):
 446        """method for creating movie frames (PNG files) that capture the plan-view evolution of a channel belt through time
 447        movie has to be assembled from the PNG file after this method is applied
 448
 449        :param xmin: value of x coodinate on the left side of frame
 450        :param xmax: value of x coordinate on right side of frame
 451        :param plot_type: plot type; can be either 'strat' (for stratigraphic plot) or 'morph' (for morphologic plot)
 452        :param filename: first few characters of the output filenames
 453        :param dirname: name of directory where output files should be written
 454        :param pb_age: age of point bars (in years) at which they get covered by vegetation (if the 'morph' option is used for 'plot_type')
 455        :param ob_age: age of oxbow lakes (in years) at which they get covered by vegetation (if the 'morph' option is used for 'plot_type')
 456        :param scale: scaling factor (e.g., 2) that determines how many times larger you want the frame to be, compared to the default scaling of the figure
 457        :param end_time: time at which simulation should be stopped
 458        :param n_channels: total number of channels + cutoffs for which simulation is run (usually it is len(chb.cutoffs) + len(chb.channels)). Used when plot_type = 'age'"""
 459
 460        sclt = np.array(self.cl_times)
 461        if type(end_time) != list:
 462            sclt = sclt[sclt <= end_time]
 463        channels = self.channels[: len(sclt)]
 464        ymax = 0
 465        for i in range(len(channels)):
 466            ymax = max(ymax, np.max(np.abs(channels[i].y)))
 467        ymax = ymax + 2 * channels[0].W  # add a bit of space on top and bottom
 468        ymin = -1 * ymax
 469        for i in range(0, len(sclt)):
 470            fig = self.plot(plot_type, pb_age, ob_age, sclt[i], n_channels)
 471            scale = 1
 472            fig_height = scale * fig.get_figheight()
 473            fig_width = (xmax - xmin) * fig_height / (ymax - ymin)
 474            fig.set_figwidth(fig_width)
 475            fig.set_figheight(fig_height)
 476            fig.gca().set_xlim(xmin, xmax)
 477            fig.gca().set_xticks([])
 478            fig.gca().set_yticks([])
 479            plt.plot(
 480                [xmin + 200, xmin + 200 + 5000],
 481                [ymin + 200, ymin + 200],
 482                "k",
 483                linewidth=2,
 484            )
 485            plt.text(xmin + 200 + 2000, ymin + 200 + 100, "5 km", fontsize=14)
 486            fname = dirname + filename + "%03d.png" % (i)
 487            fig.savefig(fname, bbox_inches="tight")
 488            plt.close()
 489
 490    def build_3d_model(
 491        self,
 492        model_type,
 493        h_mud,
 494        levee_width,
 495        h,
 496        w,
 497        bth,
 498        dcr,
 499        dx,
 500        delta_s,
 501        starttime,
 502        endtime,
 503        xmin,
 504        xmax,
 505        ymin,
 506        ymax,
 507    ):
 508        """method for building 3D model from set of centerlines (that are part of a ChannelBelt object)
 509
 510        :param model_type: model type ('fluvial' or 'submarine')
 511        :param h_mud: maximum thickness of overbank mud
 512        :param levee_width: width of overbank mud
 513        :param h: channel depth
 514        :param w: channel width
 515        :param bth: thickness of channel sand (only used in submarine models)
 516        :param dcr: critical channel depth where sand thickness goes to zero (only used in submarine models)
 517        :param dx: cell size in x and y directions
 518        :param delta_s: sampling distance alogn centerlines
 519        :param starttime: age of centerline that will be used as the first centerline in the model
 520        :param endtime: age of centerline that will be used as the last centerline in the model
 521        :param xmin: minimum x coordinate that defines the model domain; if xmin is set to zero,
 522        a plot of the centerlines is generated and the model domain has to be defined by clicking its upper left and lower right corners
 523        :param xmax: maximum x coordinate that defines the model domain
 524        :param ymin: minimum y coordinate that defines the model domain
 525        :param ymax: maximum y coordinate that defines the model domain
 526        :return chb_3d: a ChannelBelt3D object
 527        :return xmin, xmax, ymin, ymax: x and y coordinates that define the model domain (so that they can be reused later)"""
 528
 529        sclt = np.array(self.cl_times)
 530        ind1 = np.where(sclt >= starttime)[0][0]
 531        ind2 = np.where(sclt <= endtime)[0][-1]
 532        sclt = sclt[ind1 : ind2 + 1]
 533        channels = self.channels[ind1 : ind2 + 1]
 534        cot = np.array(self.cutoff_times)
 535        if (
 536            (len(cot) > 0)
 537            & (len(np.where(cot >= starttime)[0]) > 0)
 538            & (len(np.where(cot <= endtime)[0]) > 0)
 539        ):
 540            cfind1 = np.where(cot >= starttime)[0][0]
 541            cfind2 = np.where(cot <= endtime)[0][-1]
 542            cot = cot[cfind1 : cfind2 + 1]
 543            cutoffs = self.cutoffs[cfind1 : cfind2 + 1]
 544        else:
 545            cot = []
 546            cutoffs = []
 547        n_steps = len(sclt)  # number of events
 548        if xmin == 0:  # plot centerlines and define model domain
 549            plt.figure(figsize=(15, 4))
 550            maxX, minY, maxY = 0, 0, 0
 551            for i in range(n_steps):  # plot centerlines
 552                plt.plot(channels[i].x, channels[i].y, "k")
 553                maxX = max(maxX, np.max(channels[i].x))
 554                maxY = max(maxY, np.max(channels[i].y))
 555                minY = min(minY, np.min(channels[i].y))
 556            plt.axis([0, maxX, minY - 10 * w, maxY + 10 * w])
 557            plt.gca().set_aspect("equal", adjustable="box")
 558            plt.tight_layout()
 559            pts = np.zeros((2, 2))
 560            for i in range(0, 2):
 561                pt = np.asarray(plt.ginput(1))
 562                pts[i, :] = pt
 563                plt.scatter(pt[0][0], pt[0][1])
 564            plt.plot(
 565                [pts[0, 0], pts[1, 0], pts[1, 0], pts[0, 0], pts[0, 0]],
 566                [pts[0, 1], pts[0, 1], pts[1, 1], pts[1, 1], pts[0, 1]],
 567                "r",
 568            )
 569            xmin = min(pts[0, 0], pts[1, 0])
 570            xmax = max(pts[0, 0], pts[1, 0])
 571            ymin = min(pts[0, 1], pts[1, 1])
 572            ymax = max(pts[0, 1], pts[1, 1])
 573        iwidth = int((xmax - xmin) / dx)
 574        iheight = int((ymax - ymin) / dx)
 575        topo = np.zeros(
 576            (iheight, iwidth, 4 * n_steps)
 577        )  # array for storing topographic surfaces
 578        dists = np.zeros((iheight, iwidth, n_steps))
 579        zmaps = np.zeros((iheight, iwidth, n_steps))
 580        facies = np.zeros((4 * n_steps, 1))
 581        # create initial topography:
 582        x1 = np.linspace(0, iwidth - 1, iwidth)
 583        y1 = np.linspace(0, iheight - 1, iheight)
 584        xv, yv = np.meshgrid(x1, y1)
 585        z1 = channels[0].z
 586        z1 = z1[(channels[0].x > xmin) & (channels[0].x < xmax)]
 587        topoinit = (
 588            z1[0] - ((z1[0] - z1[-1]) / (xmax - xmin)) * xv * dx
 589        )  # initial (sloped) topography
 590        topo[:, :, 0] = topoinit.copy()
 591        surf = topoinit.copy()
 592        facies[0] = np.NaN
 593        # generate surfaces:
 594        channels3D = []
 595        for i in trange(n_steps):
 596            x = channels[i].x
 597            y = channels[i].y
 598            z = channels[i].z
 599            cutoff_ind = []
 600            # check if there were cutoffs during the last time step and collect indices in an array:
 601            for j in range(len(cot)):
 602                if (cot[j] >= sclt[i - 1]) & (cot[j] < sclt[i]):
 603                    cutoff_ind = np.append(cutoff_ind, j)
 604            # create distance map:
 605            cl_dist, x_pix, y_pix, z_pix, s_pix, z_map, x1, y1, z1 = dist_map(
 606                x, y, z, xmin, xmax, ymin, ymax, dx, delta_s
 607            )
 608            if i == 0:
 609                cl_dist_prev = cl_dist
 610            # erosion:
 611            surf = np.minimum(surf, erosion_surface(h, w / dx, cl_dist, z_map))
 612            topo[:, :, 4 * i] = surf  # erosional surface
 613            dists[:, :, i] = cl_dist
 614            zmaps[:, :, i] = z_map
 615            facies[4 * i] = np.NaN
 616
 617            if model_type == "fluvial":
 618                pb = point_bar_surface(cl_dist, z_map, h, w / dx)
 619                th = np.maximum(surf, pb) - surf
 620                th_oxbows = th.copy()
 621                # setting sand thickness to zero at cutoff locations:
 622                if len(cutoff_ind) > 0:
 623                    cutoff_dists = 1e10 * np.ones(
 624                        np.shape(th)
 625                    )  # initialize cutoff_dists with a large number
 626                    for j in range(len(cutoff_ind)):
 627                        cutoff_dist, cfx_pix, cfy_pix = cl_dist_map(
 628                            cutoffs[int(cutoff_ind[j])].x[0],
 629                            cutoffs[int(cutoff_ind[j])].y[0],
 630                            cutoffs[int(cutoff_ind[j])].z[0],
 631                            xmin,
 632                            xmax,
 633                            ymin,
 634                            ymax,
 635                            dx,
 636                        )
 637                        cutoff_dists = np.minimum(cutoff_dists, cutoff_dist)
 638                    th_oxbows[
 639                        cutoff_dists >= 0.9 * w / dx
 640                    ] = 0  # set oxbow fill thickness to zero outside of oxbows
 641                    th[
 642                        cutoff_dists < 0.9 * w / dx
 643                    ] = 0  # set point bar thickness to zero inside of oxbows
 644                else:  # no cutoffs
 645                    th_oxbows = np.zeros(np.shape(th))
 646                th[th < 0] = 0  # eliminate negative th values
 647                surf = (
 648                    surf + th_oxbows
 649                )  # update topographic surface with oxbow deposit thickness
 650                topo[:, :, 4 * i + 1] = surf  # top of oxbow mud
 651                facies[4 * i + 1] = 0
 652                surf = surf + th  # update topographic surface with sand thickness
 653                topo[:, :, 4 * i + 2] = surf  # top of sand
 654                facies[4 * i + 2] = 1
 655                surf = surf + mud_surface(
 656                    h_mud, levee_width / dx, cl_dist, w / dx, z_map, surf
 657                )  # mud/levee deposition
 658                topo[:, :, 4 * i + 3] = surf  # top of levee
 659                facies[4 * i + 3] = 2
 660                channels3D.append(Channel(x1 - xmin, y1 - ymin, z1, w, h))
 661
 662            if model_type == "submarine":
 663                surf = surf + mud_surface(
 664                    h_mud[i], levee_width / dx, cl_dist, w / dx, z_map, surf
 665                )  # mud/levee deposition
 666                topo[:, :, 4 * i + 1] = surf  # top of levee
 667                facies[4 * i + 1] = 2
 668                # sand thickness:
 669                th, relief = sand_surface(surf, bth, dcr, z_map, h)
 670                th[th < 0] = 0  # eliminate negative th values
 671                th[cl_dist > 1.0 * w / dx] = 0  # eliminate sand outside of channel
 672                th_oxbows = th.copy()
 673                # setting sand thickness to zero at cutoff locations:
 674                if len(cutoff_ind) > 0:
 675                    cutoff_dists = 1e10 * np.ones(
 676                        np.shape(th)
 677                    )  # initialize cutoff_dists with a large number
 678                    for j in range(len(cutoff_ind)):
 679                        cutoff_dist, cfx_pix, cfy_pix = cl_dist_map(
 680                            cutoffs[int(cutoff_ind[j])].x[0],
 681                            cutoffs[int(cutoff_ind[j])].y[0],
 682                            cutoffs[int(cutoff_ind[j])].z[0],
 683                            xmin,
 684                            xmax,
 685                            ymin,
 686                            ymax,
 687                            dx,
 688                        )
 689                        cutoff_dists = np.minimum(cutoff_dists, cutoff_dist)
 690                    th_oxbows[
 691                        cutoff_dists >= 0.9 * w / dx
 692                    ] = 0  # set oxbow fill thickness to zero outside of oxbows
 693                    th[
 694                        cutoff_dists < 0.9 * w / dx
 695                    ] = 0  # set point bar thickness to zero inside of oxbows
 696                    # adding back sand near the channel axis (submarine only):
 697                    # th[cl_dist<0.5*w/dx] = bth*(1 - relief[cl_dist<0.5*w/dx]/dcr)
 698                else:  # no cutoffs
 699                    th_oxbows = np.zeros(np.shape(th))
 700                surf = (
 701                    surf + th_oxbows
 702                )  # update topographic surface with oxbow deposit thickness
 703                topo[:, :, 4 * i + 2] = surf  # top of oxbow mud
 704                facies[4 * i + 2] = 0
 705                surf = surf + th  # update topographic surface with sand thickness
 706                topo[:, :, 4 * i + 3] = surf  # top of sand
 707                facies[4 * i + 3] = 1
 708                channels3D.append(Channel(x1 - xmin, y1 - ymin, z1, w, h))
 709
 710            cl_dist_prev = cl_dist.copy()
 711        topo = np.concatenate(
 712            (np.reshape(topoinit, (iheight, iwidth, 1)), topo), axis=2
 713        )  # add initial topography to array
 714        strat = topostrat(topo)  # create stratigraphic surfaces
 715        strat = np.delete(
 716            strat, np.arange(4 * n_steps + 1)[1::4], 2
 717        )  # get rid of unnecessary stratigraphic surfaces (duplicates)
 718        facies = np.delete(
 719            facies, np.arange(4 * n_steps)[::4]
 720        )  # get rid of unnecessary facies layers (NaNs)
 721        if model_type == "fluvial":
 722            facies_code = {0: "oxbow", 1: "point bar", 2: "levee"}
 723        if model_type == "submarine":
 724            facies_code = {0: "oxbow", 1: "channel sand", 2: "levee"}
 725        chb_3d = ChannelBelt3D(
 726            model_type, topo, strat, facies, facies_code, dx, channels3D
 727        )
 728        return chb_3d, xmin, xmax, ymin, ymax  # , dists, zmaps
 729
 730
 731def resample_centerline(x, y, z, deltas):
 732    """resample centerline so that 'deltas' is roughly constant, using parametric
 733    spline representation of curve; note that there is *no* smoothing
 734
 735    :param x: x-coordinates of centerline
 736    :param y: y-coordinates of centerline
 737    :param z: z-coordinates of centerline
 738    :param deltas: distance between points on centerline
 739    :return x: x-coordinates of resampled centerline
 740    :return y: y-coordinates of resampled centerline
 741    :return z: z-coordinates of resampled centerline
 742    :return dx: dx of resampled centerline
 743    :return dy: dy of resampled centerline
 744    :return dz: dz of resampled centerline
 745    :return s: s-coordinates of resampled centerline"""
 746
 747    dx, dy, dz, ds, s = compute_derivatives(x, y, z)  # compute derivatives
 748    tck, u = scipy.interpolate.splprep([x, y, z], s=0)
 749    unew = np.linspace(0, 1, 1 + int(round(s[-1] / deltas)))  # vector for resampling
 750    out = scipy.interpolate.splev(unew, tck)  # resampling
 751    x, y, z = out[0], out[1], out[2]  # assign new coordinate values
 752    dx, dy, dz, ds, s = compute_derivatives(x, y, z)  # recompute derivatives
 753    return x, y, z, dx, dy, dz, ds, s
 754
 755
 756def migrate_one_step(x, y, z, W, kl, dt, k, Cf, D, pad, pad1, omega, gamma):
 757    """migrate centerline during one time step, using the migration computed as in Howard & Knutson (1984)
 758
 759    :param x: x-coordinates of centerline
 760    :param y: y-coordinates of centerline
 761    :param z: z-coordinates of centerline
 762    :param W: channel width
 763    :param kl: migration rate (or erodibility) constant (m/s)
 764    :param dt: duration of time step (s)
 765    :param k: constant for calculating the exponent alpha (= 1.0)
 766    :param Cf: dimensionless Chezy friction factor
 767    :param D: channel depth
 768    :param omega: constant in Howard & Knutson equation (= -1.0)
 769    :param gamma: constant in Howard & Knutson equation (= 2.5)
 770    :return x: new x-coordinates of centerline after migration
 771    :return y: new y-coordinates of centerline after migration
 772    """
 773
 774    ns = len(x)
 775    curv = compute_curvature(x, y)
 776    dx, dy, dz, ds, s = compute_derivatives(x, y, z)
 777    sinuosity = s[-1] / (x[-1] - x[0])
 778    curv = W * curv  # dimensionless curvature
 779    R0 = (
 780        kl * curv
 781    )  # simple linear relationship between curvature and nominal migration rate
 782    alpha = k * 2 * Cf / D  # exponent for convolution function G
 783    R1 = compute_migration_rate(pad, ns, ds, alpha, omega, gamma, R0)
 784    R1 = sinuosity ** (-2 / 3.0) * R1
 785    # calculate new centerline coordinates:
 786    dy_ds = dy[pad1 : ns - pad + 1] / ds[pad1 : ns - pad + 1]
 787    dx_ds = dx[pad1 : ns - pad + 1] / ds[pad1 : ns - pad + 1]
 788    # adjust x and y coordinates (this *is* the migration):
 789    x[pad1 : ns - pad + 1] = (
 790        x[pad1 : ns - pad + 1] + R1[pad1 : ns - pad + 1] * dy_ds * dt
 791    )
 792    y[pad1 : ns - pad + 1] = (
 793        y[pad1 : ns - pad + 1] - R1[pad1 : ns - pad + 1] * dx_ds * dt
 794    )
 795    return x, y
 796
 797
 798def migrate_one_step_w_bias(x, y, z, W, kl, dt, k, Cf, D, pad, pad1, omega, gamma):
 799    ns = len(x)
 800    curv = compute_curvature(x, y)
 801    dx, dy, dz, ds, s = compute_derivatives(x, y, z)
 802    sinuosity = s[-1] / (x[-1] - x[0])
 803    curv = W * curv  # dimensionless curvature
 804    R0 = (
 805        kl * curv
 806    )  # simple linear relationship between curvature and nominal migration rate
 807    alpha = k * 2 * Cf / D  # exponent for convolution function G
 808    R1 = compute_migration_rate(pad, ns, ds, alpha, omega, gamma, R0)
 809    R1 = sinuosity ** (-2 / 3.0) * R1
 810    pad = -1
 811    # calculate new centerline coordinates:
 812    dy_ds = dy[pad1 : ns - pad + 1] / ds[pad1 : ns - pad + 1]
 813    dx_ds = dx[pad1 : ns - pad + 1] / ds[pad1 : ns - pad + 1]
 814    tilt_factor = 0.2
 815    T = kl * tilt_factor * np.ones(np.shape(x))
 816    angle = 90.0
 817    # adjust x and y coordinates (this *is* the migration):
 818    x[pad1 : ns - pad + 1] = (
 819        x[pad1 : ns - pad + 1]
 820        + R1[pad1 : ns - pad + 1] * dy_ds * dt
 821        + T[pad1 : ns - pad + 1]
 822        * dy_ds
 823        * dt
 824        * (np.sin(np.deg2rad(angle)) * dx_ds + np.cos(np.deg2rad(angle)) * dy_ds)
 825    )
 826    y[pad1 : ns - pad + 1] = (
 827        y[pad1 : ns - pad + 1]
 828        - R1[pad1 : ns - pad + 1] * dx_ds * dt
 829        - T[pad1 : ns - pad + 1]
 830        * dx_ds
 831        * dt
 832        * (np.sin(np.deg2rad(angle)) * dx_ds + np.cos(np.deg2rad(angle)) * dy_ds)
 833    )
 834    return x, y
 835
 836
 837def generate_initial_channel(W, D, Sl, deltas, pad, n_bends):
 838    """generate straight Channel object with some noise added that can serve
 839    as input for initializing a ChannelBelt object
 840    W - channel width
 841    D - channel depth
 842    Sl - channel gradient
 843    deltas - distance between nodes on centerline
 844    pad - padding (number of nodepoints along centerline)
 845    n_bends - approximate number of bends to be simulated"""
 846    noisy_len = n_bends * 10 * W / 2.0  # length of noisy part of initial centerline
 847    pad1 = int(
 848        pad / 10.0
 849    )  # padding at upstream end can be shorter than padding on downstream end
 850    if pad1 < 5:
 851        pad1 = 5
 852    x = np.linspace(
 853        0, noisy_len + (pad + pad1) * deltas, int(noisy_len / deltas + pad + pad1) + 1
 854    )  # x coordinate
 855    y = 10.0 * (
 856        2
 857        * np.random.random_sample(
 858            int(noisy_len / deltas) + 1,
 859        )
 860        - 1
 861    )
 862    y = np.hstack(
 863        (
 864            np.zeros(
 865                (pad1),
 866            ),
 867            y,
 868            np.zeros(
 869                (pad),
 870            ),
 871        )
 872    )  # y coordinate
 873    deltaz = Sl * deltas * (len(x) - 1)
 874    z = np.linspace(0, deltaz, len(x))[::-1]  # z coordinate
 875    return Channel(x, y, z, W, D)
 876
 877
 878@numba.jit(nopython=True)  # use Numba to speed up the heaviest computation
 879def compute_migration_rate(pad, ns, ds, alpha, omega, gamma, R0):
 880    """compute migration rate as weighted sum of upstream curvatures
 881    pad - padding (number of nodepoints along centerline)
 882    ns - number of points in centerline
 883    ds - distances between points in centerline
 884    omega - constant in HK model
 885    gamma - constant in HK model
 886    R0 - nominal migration rate (dimensionless curvature * migration rate constant)"""
 887    R1 = np.zeros(ns)  # preallocate adjusted channel migration rate
 888    pad1 = int(
 889        pad / 10.0
 890    )  # padding at upstream end can be shorter than padding on downstream end
 891    if pad1 < 5:
 892        pad1 = 5
 893    for i in range(pad1, ns - pad):
 894        si2 = np.hstack(
 895            (np.array([0]), np.cumsum(ds[i - 1 :: -1]))
 896        )  # distance along centerline, backwards from current point
 897        G = np.exp(-alpha * si2)  # convolution vector
 898        R1[i] = omega * R0[i] + gamma * np.sum(R0[i::-1] * G) / np.sum(
 899            G
 900        )  # main equation
 901    return R1
 902
 903
 904def compute_derivatives(x, y, z):
 905    """function for computing first derivatives of a curve (centerline)
 906    x,y are cartesian coodinates of the curve
 907    outputs:
 908    dx - first derivative of x coordinate
 909    dy - first derivative of y coordinate
 910    ds - distances between consecutive points along the curve
 911    s - cumulative distance along the curve"""
 912    dx = np.gradient(x)  # first derivatives
 913    dy = np.gradient(y)
 914    dz = np.gradient(z)
 915    ds = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
 916    s = np.hstack((0, np.cumsum(ds[1:])))
 917    return dx, dy, dz, ds, s
 918
 919
 920def compute_curvature(x, y):
 921    """function for computing first derivatives and curvature of a curve (centerline)
 922    x,y are cartesian coodinates of the curve
 923    outputs:
 924    dx - first derivative of x coordinate
 925    dy - first derivative of y coordinate
 926    ds - distances between consecutive points along the curve
 927    s - cumulative distance along the curve
 928    curvature - curvature of the curve (in 1/units of x and y)"""
 929    dx = np.gradient(x)  # first derivatives
 930    dy = np.gradient(y)
 931    ddx = np.gradient(dx)  # second derivatives
 932    ddy = np.gradient(dy)
 933    curvature = (dx * ddy - dy * ddx) / ((dx ** 2 + dy ** 2) ** 1.5)
 934    return curvature
 935
 936
 937def make_colormap(seq):
 938    """Return a LinearSegmentedColormap
 939    seq: a sequence of floats and RGB-tuples. The floats should be increasing
 940    and in the interval (0,1).
 941    [from: https://stackoverflow.com/questions/16834861/create-own-colormap-using-matplotlib-and-plot-color-scale]
 942    """
 943    seq = [(None,) * 3, 0.0] + list(seq) + [1.0, (None,) * 3]
 944    cdict = {"red": [], "green": [], "blue": []}
 945    for i, item in enumerate(seq):
 946        if isinstance(item, float):
 947            r1, g1, b1 = seq[i - 1]
 948            r2, g2, b2 = seq[i + 1]
 949            cdict["red"].append([item, r1, r2])
 950            cdict["green"].append([item, g1, g2])
 951            cdict["blue"].append([item, b1, b2])
 952    return mcolors.LinearSegmentedColormap("CustomMap", cdict)
 953
 954
 955def kth_diag_indices(a, k):
 956    """function for finding diagonal indices with k offset
 957    [from https://stackoverflow.com/questions/10925671/numpy-k-th-diagonal-indices]"""
 958    rows, cols = np.diag_indices_from(a)
 959    if k < 0:
 960        return rows[:k], cols[-k:]
 961    elif k > 0:
 962        return rows[k:], cols[:-k]
 963    else:
 964        return rows, cols
 965
 966
 967def find_cutoffs(x, y, crdist, deltas):
 968    """function for identifying locations of cutoffs along a centerline
 969    and the indices of the segments that will become part of the oxbows
 970    x,y - coordinates of centerline
 971    crdist - critical cutoff distance
 972    deltas - distance between neighboring points along the centerline"""
 973    diag_blank_width = int((crdist + 20 * deltas) / deltas)
 974    # distance matrix for centerline points:
 975    dist = distance.cdist(np.array([x, y]).T, np.array([x, y]).T)
 976    dist[
 977        dist > crdist
 978    ] = np.NaN  # set all values that are larger than the cutoff threshold to NaN
 979    # set matrix to NaN along the diagonal zone:
 980    for k in range(-diag_blank_width, diag_blank_width + 1):
 981        rows, cols = kth_diag_indices(dist, k)
 982        dist[rows, cols] = np.NaN
 983    i1, i2 = np.where(~np.isnan(dist))
 984    ind1 = i1[np.where(i1 < i2)[0]]  # get rid of unnecessary indices
 985    ind2 = i2[np.where(i1 < i2)[0]]  # get rid of unnecessary indices
 986    return ind1, ind2  # return indices of cutoff points and cutoff coordinates
 987
 988
 989def cut_off_cutoffs(x, y, z, s, crdist, deltas):
 990    """function for executing cutoffs - removing oxbows from centerline and storing cutoff coordinates
 991    x,y - coordinates of centerline
 992    crdist - critical cutoff distance
 993    deltas - distance between neighboring points along the centerline
 994    outputs:
 995    x,y,z - updated coordinates of centerline
 996    xc, yc, zc - lists with coordinates of cutoff segments"""
 997    xc = []
 998    yc = []
 999    zc = []
1000    ind1, ind2 = find_cutoffs(x, y, crdist, deltas)  # initial check for cutoffs
1001    while len(ind1) > 0:
1002        xc.append(x[ind1[0] : ind2[0] + 1])  # x coordinates of cutoff
1003        yc.append(y[ind1[0] : ind2[0] + 1])  # y coordinates of cutoff
1004        zc.append(z[ind1[0] : ind2[0] + 1])  # z coordinates of cutoff
1005        x = np.hstack((x[: ind1[0] + 1], x[ind2[0] :]))  # x coordinates after cutoff
1006        y = np.hstack((y[: ind1[0] + 1], y[ind2[0] :]))  # y coordinates after cutoff
1007        z = np.hstack((z[: ind1[0] + 1], z[ind2[0] :]))  # z coordinates after cutoff
1008        ind1, ind2 = find_cutoffs(x, y, crdist, deltas)
1009    return x, y, z, xc, yc, zc
1010
1011
1012def get_channel_banks(x, y, W):
1013    """function for finding coordinates of channel banks, given a centerline and a channel width
1014    x,y - coordinates of centerline
1015    W - channel width
1016    outputs:
1017    xm, ym - coordinates of channel banks (both left and right banks)"""
1018    x1 = x.copy()
1019    y1 = y.copy()
1020    x2 = x.copy()
1021    y2 = y.copy()
1022    ns = len(x)
1023    dx = np.diff(x)
1024    dy = np.diff(y)
1025    ds = np.sqrt(dx ** 2 + dy ** 2)
1026    x1[:-1] = x[:-1] + 0.5 * W * np.diff(y) / ds
1027    y1[:-1] = y[:-1] - 0.5 * W * np.diff(x) / ds
1028    x2[:-1] = x[:-1] - 0.5 * W * np.diff(y) / ds
1029    y2[:-1] = y[:-1] + 0.5 * W * np.diff(x) / ds
1030    x1[ns - 1] = x[ns - 1] + 0.5 * W * (y[ns - 1] - y[ns - 2]) / ds[ns - 2]
1031    y1[ns - 1] = y[ns - 1] - 0.5 * W * (x[ns - 1] - x[ns - 2]) / ds[ns - 2]
1032    x2[ns - 1] = x[ns - 1] - 0.5 * W * (y[ns - 1] - y[ns - 2]) / ds[ns - 2]
1033    y2[ns - 1] = y[ns - 1] + 0.5 * W * (x[ns - 1] - x[ns - 2]) / ds[ns - 2]
1034    xm = np.hstack((x1, x2[::-1]))
1035    ym = np.hstack((y1, y2[::-1]))
1036    return xm, ym
1037
1038
1039def dist_map(x, y, z, xmin, xmax, ymin, ymax, dx, delta_s):
1040    """function for centerline rasterization and distance map calculation
1041    inputs:
1042    x,y,z - coordinates of centerline
1043    xmin, xmax, ymin, ymax - x and y coordinates that define the area of interest
1044    dx - gridcell size (m)
1045    delta_s - distance between points along centerline (m)
1046    returns:
1047    cl_dist - distance map (distance from centerline)
1048    x_pix, y_pix, z_pix - x,y, and z pixel coordinates of the centerline
1049    s_pix - along-channel distance in pixels
1050    z_map - map of reference channel thalweg elevation (elevation of closest point along centerline)
1051    x, y, z - x,y,z centerline coordinates clipped to the 3D model domain"""
1052    y = y[(x > xmin) & (x < xmax)]
1053    z = z[(x > xmin) & (x < xmax)]
1054    x = x[(x > xmin) & (x < xmax)]
1055    dummy, dy, dz, ds, s = compute_derivatives(x, y, z)
1056    if len(np.where(ds > 2 * delta_s)[0]) > 0:
1057        inds = np.where(ds > 2 * delta_s)[0]
1058        inds = np.hstack((0, inds, len(x)))
1059        lengths = np.diff(inds)
1060        long_segment = np.where(lengths == max(lengths))[0][0]
1061        start_ind = inds[long_segment] + 1
1062        end_ind = inds[long_segment + 1]
1063        if end_ind < len(x):
1064            x = x[start_ind:end_ind]
1065            y = y[start_ind:end_ind]
1066            z = z[start_ind:end_ind]
1067        else:
1068            x = x[start_ind:]
1069            y = y[start_ind:]
1070            z = z[start_ind:]
1071    xdist = xmax - xmin
1072    ydist = ymax - ymin
1073    iwidth = int((xmax - xmin) / dx)
1074    iheight = int((ymax - ymin) / dx)
1075    xratio = iwidth / xdist
1076    # create list with pixel coordinates:
1077    pixels = []
1078    for i in range(0, len(x)):
1079        px = int(iwidth - (xmax - x[i]) * xratio)
1080        py = int(iheight - (ymax - y[i]) * xratio)
1081        pixels.append((px, py))
1082    # create image and numpy array:
1083    img = Image.new("RGB", (iwidth, iheight), "white")
1084    draw = ImageDraw.Draw(img)
1085    draw.line(pixels, fill="rgb(0, 0, 0)")  # draw centerline as black line
1086    pix = np.array(img)
1087    cl = pix[:, :, 0]
1088    cl[cl == 255] = 1  # set background to 1 (centerline is 0)
1089    y_pix, x_pix = np.where(cl == 0)
1090    x_pix, y_pix = order_cl_pixels(x_pix, y_pix)
1091    # This next block of code is kind of a hack. Looking for, and eliminating, 'bad' pixels.
1092    img = np.array(img)
1093    img = img[:, :, 0]
1094    img[img == 255] = 1
1095    img1 = morphology.binary_dilation(img, morphology.square(2)).astype(np.uint8)
1096    if len(np.where(img1 == 0)[0]) > 0:
1097        x_pix, y_pix = eliminate_bad_pixels(img, img1)
1098        x_pix, y_pix = order_cl_pixels(x_pix, y_pix)
1099    img1 = morphology.binary_dilation(
1100        img, np.array([[1, 0, 1], [1, 1, 1]], dtype=np.uint8)
1101    ).astype(np.uint8)
1102    if len(np.where(img1 == 0)[0]) > 0:
1103        x_pix, y_pix = eliminate_bad_pixels(img, img1)
1104        x_pix, y_pix = order_cl_pixels(x_pix, y_pix)
1105    img1 = morphology.binary_dilation(
1106        img, np.array([[1, 0, 1], [0, 1, 0], [1, 0, 1]], dtype=np.uint8)
1107    ).astype(np.uint8)
1108    if len(np.where(img1 == 0)[0]) > 0:
1109        x_pix, y_pix = eliminate_bad_pixels(img, img1)
1110        x_pix, y_pix = order_cl_pixels(x_pix, y_pix)
1111    # redo the distance calculation (because x_pix and y_pix do not always contain all the points in cl):
1112    cl[cl == 0] = 1
1113    cl[y_pix, x_pix] = 0
1114    cl_dist, inds = ndimage.distance_transform_edt(cl, return_indices=True)
1115    dx, dy, dz, ds, s = compute_derivatives(x, y, z)
1116    dx_pix = np.diff(x_pix)
1117    dy_pix = np.diff(y_pix)
1118    ds_pix = np.sqrt(dx_pix ** 2 + dy_pix ** 2)
1119    s_pix = np.hstack((0, np.cumsum(ds_pix)))
1120    f = scipy.interpolate.interp1d(s, z)
1121    snew = s_pix * s[-1] / s_pix[-1]
1122    if snew[-1] > s[-1]:
1123        snew[-1] = s[-1]
1124    snew[snew < s[0]] = s[0]
1125    z_pix = f(snew)
1126    # create z_map:
1127    z_map = np.zeros(np.shape(cl_dist))
1128    z_map[y_pix, x_pix] = z_pix
1129    xinds = inds[1, :, :]
1130    yinds = inds[0, :, :]
1131    for i in range(0, len(x_pix)):
1132        z_map[(xinds == x_pix[i]) & (yinds == y_pix[i])] = z_pix[i]
1133    return cl_dist, x_pix, y_pix, z_pix, s_pix, z_map, x, y, z
1134
1135
1136def erosion_surface(h, w, cl_dist, z):
1137    """function for creating a parabolic erosional surface
1138    inputs:
1139    h - geomorphic channel depth (m)
1140    w - geomorphic channel width (in pixels, as cl_dist is also given in pixels)
1141    cl_dist - distance map (distance from centerline)
1142    z - reference elevation (m)
1143    returns:
1144    surf - map of the quadratic erosional surface (m)
1145    """
1146    surf = z + (4 * h / w ** 2) * (cl_dist + w * 0.5) * (cl_dist - w * 0.5)
1147    return surf
1148
1149
1150def point_bar_surface(cl_dist, z, h, w):
1151    """function for creating a Gaussian-based point bar surface
1152    used in 3D fluvial model
1153    inputs:
1154    cl_dist - distance map (distance from centerline)
1155    z - reference elevation (m)
1156    h - channel depth (m)
1157    w - channel width, in pixels, as cl_dist is also given in pixels
1158    returns:
1159    pb - map of the Gaussian surface that can be used to from a point bar deposit (m)"""
1160    pb = z - h * np.exp(-(cl_dist ** 2) / (2 * (w * 0.33) ** 2))
1161    return pb
1162
1163
1164def sand_surface(surf, bth, dcr, z_map, h):
1165    """function for creating the top horizontal surface sand-rich deposit in the bottom of the channel
1166    used in 3D submarine channel models
1167    inputs:
1168    surf - current geomorphic surface
1169    bth - thickness of sand deposit in axis of channel (m)
1170    dcr - critical channel depth, above which there is no sand deposition (m)
1171    z_map - map of reference channel thalweg elevation (elevation of closest point along centerline)
1172    h - channel depth (m)
1173    returns:
1174    th - thickness map of sand deposit (m)
1175    relief - map of channel relief (m)"""
1176    relief = np.abs(surf - z_map + h)
1177    relief = np.abs(relief - np.amin(relief))
1178    th = bth * (1 - relief / dcr)  # bed thickness inversely related to relief
1179    th[th < 0] = 0.0  # set negative th values to zero
1180    return th, relief
1181
1182
1183def mud_surface(h_mud, levee_width, cl_dist, w, z_map, topo):
1184    """function for creating a map of overbank deposit thickness
1185    inputs:
1186    h_mud - maximum thickness of overbank deposit (m)
1187    levee_width - half-width of overbank deposit (m)
1188    cl_dist - distance map (distance from centerline)
1189    w - channel width (in pixels, as cl_dist is also given in pixels)
1190    z_map - map of reference channel thalweg elevation (elevation of closest point along centerline)
1191    topo - current geomorphic surface
1192    returns:
1193    surf - map of overbank deposit thickness (m)"""
1194    # create a surface that thins linearly away from the channel centerline:
1195    surf1 = (-2 * h_mud / levee_width) * cl_dist + h_mud
1196    surf2 = (2 * h_mud / levee_width) * cl_dist + h_mud
1197    surf = np.minimum(surf1, surf2)
1198    # surface for 'eroding' the central part of the mud layer:
1199    surf3 = h_mud + (4 * 1.5 * h_mud / w ** 2) * (cl_dist + w * 0.5) * (
1200        cl_dist - w * 0.5
1201    )
1202    surf = np.minimum(surf, surf3)
1203    surf[surf < 0] = 0
1204    # eliminate negative thicknesses
1205    return surf
1206
1207
1208def topostrat(topo):
1209    """function for converting a stack of geomorphic surfaces into stratigraphic surfaces
1210    inputs:
1211    topo - 3D numpy array of geomorphic surfaces
1212    returns:
1213    strat - 3D numpy array of stratigraphic surfaces
1214    """
1215    r, c, ts = np.shape(topo)
1216    strat = np.copy(topo)
1217    for i in range(0, ts):
1218        strat[:, :, i] = np.amin(topo[:, :, i:], axis=2)
1219    return strat
1220
1221
1222def cl_dist_map(x, y, z, xmin, xmax, ymin, ymax, dx):
1223    """function for centerline rasterization and distance map calculation (does not return zmap)
1224    used for cutoffs only
1225    inputs:
1226    x,y,z - coordinates of centerline
1227    xmin, xmax, ymin, ymax - x and y coordinates that define the area of interest
1228    dx - gridcell size (m)
1229    returns:
1230    cl_dist - distance map (distance from centerline)
1231    x_pix, y_pix, - x and y pixel coordinates of the centerline
1232    """
1233    y = y[(x > xmin) & (x < xmax)]
1234    z = z[(x > xmin) & (x < xmax)]
1235    x = x[(x > xmin) & (x < xmax)]
1236    xdist = xmax - xmin
1237    ydist = ymax - ymin
1238    iwidth = int((xmax - xmin) / dx)
1239    iheight = int((ymax - ymin) / dx)
1240    xratio = iwidth / xdist
1241    # create list with pixel coordinates:
1242    pixels = []
1243    for i in range(0, len(x)):
1244        px = int(iwidth - (xmax - x[i]) * xratio)
1245        py = int(iheight - (ymax - y[i]) * xratio)
1246        pixels.append((px, py))
1247    # create image and numpy array:
1248    img = Image.new("RGB", (iwidth, iheight), "white")
1249    draw = ImageDraw.Draw(img)
1250    draw.line(pixels, fill="rgb(0, 0, 0)")  # draw centerline as black line
1251    pix = np.array(img)
1252    cl = pix[:, :, 0]
1253    cl[cl == 255] = 1  # set background to 1 (centerline is 0)
1254    # calculate Euclidean distance map:
1255    cl_dist, inds = ndimage.distance_transform_edt(cl, return_indices=True)
1256    y_pix, x_pix = np.where(cl == 0)
1257    return cl_dist, x_pix, y_pix
1258
1259
1260def eliminate_bad_pixels(img, img1):
1261    x_ind = np.where(img1 == 0)[1][0]
1262    y_ind = np.where(img1 == 0)[0][0]
1263    img[y_ind : y_ind + 2, x_ind : x_ind + 2] = np.ones(
1264        1,
1265    ).astype(np.uint8)
1266    all_labels = measure.label(img, background=1, connectivity=2)
1267    cl = all_labels.copy()
1268    cl[cl == 2] = 0
1269    cl[cl > 0] = 1
1270    y_pix, x_pix = np.where(cl == 1)
1271    return x_pix, y_pix
1272
1273
1274def order_cl_pixels(x_pix, y_pix):
1275    dist = distance.cdist(np.array([x_pix, y_pix]).T, np.array([x_pix, y_pix]).T)
1276    dist[np.diag_indices_from(dist)] = 100.0
1277    ind = np.argmin(x_pix)  # select starting point on left side of image
1278    clinds = [ind]
1279    count = 0
1280    while count < len(x_pix):
1281        t = dist[ind, :].copy()
1282        if len(clinds) > 2:
1283            t[clinds[-2]] = t[clinds[-2]] + 100.0
1284            t[clinds[-3]] = t[clinds[-3]] + 100.0
1285        ind = np.argmin(t)
1286        clinds.append(ind)
1287        count = count + 1
1288    x_pix = x_pix[clinds]
1289    y_pix = y_pix[clinds]
1290    return x_pix, y_pix