text:
### Soil erosion ###
@get_units @with_kw struct LandSediment{T}
    # number of cells
    n::Int | "-"
    ### Soil erosion part ###
    # length of cells in y direction [m]
    yl::Vector{T} | "m"
    # length of cells in x direction [m]
    xl::Vector{T} | "m"
    # Fraction of river [-]
    riverfrac::Vector{T} | "-"
    # Waterbodies areas [-]
    wbcover::Vector{T} | "-"
    # Depth of overland flow [m]
    h_land::Vector{T} | "m"
    # Canopy interception [mm]
    interception::Vector{T} | "mm"
    # Precipitation [mm]
    precipitation::Vector{T} | "mm"
    # Overland flow [m3/s]
    q_land::Vector{T} | "m3 s-1"
    # Canopy height [m]
    canopyheight::Vector{T} | "m"
    # Canopy gap fraction [mm]
    canopygapfraction::Vector{T} | "mm"
    # Coefficient for EUROSEM rainfall erosion [-]
    erosk::Vector{T} | "-"
    # Exponent for EUROSEM rainfall erosion [-]
    erosspl::Vector{T} | "-"
    # Coefficient for ANSWERS overland flow erosion [-]
    erosov::Vector{T} | "-"
    # Fraction of impervious area per grid cell [-]
    pathfrac::Vector{T} | "-"
    # Land slope [-]
    slope::Vector{T} | "-"
    # USLE crop management factor [-]
    usleC::Vector{T} | "-"
    # USLE soil erodibility factor [-]
    usleK::Vector{T} | "-"
    # Sediment eroded by rainfall [ton]
    sedspl::Vector{T} | "t"
    # Sediment eroded by overland flow [ton]
    sedov::Vector{T} | "t"
    # Total eroded soil [ton]
    soilloss::Vector{T} | "t"
    # Eroded soil per particle class [ton]
    erosclay::Vector{T} | "t"  # clay
    erossilt::Vector{T} | "t"  # silt
    erossand::Vector{T} | "t"  # sand
    erossagg::Vector{T} | "t"  # small aggregates
    eroslagg::Vector{T} | "t"  # large aggregates
    ## Interception related to leaf_area_index climatology ###
    # Specific leaf storage [mm]
    sl::Vector{T} | "mm"
    # Storage woody part of vegetation [mm]
    swood::Vector{T} | "mm"
    # Extinction coefficient [-] (to calculate canopy gap fraction)
    kext::Vector{T} | "-"
    # Leaf area index [m² m⁻²]
    leaf_area_index::Vector{T} | "m2 m-2"
    ### Transport capacity part ###
    # Drain length [m]
    dl::Vector{T} | "m"
    # Flow width [m]
    dw::Vector{T} | "m"
    # Govers transport capacity coefficients [-]
    cGovers::Vector{T} | "-"
    nGovers::Vector{T} | "-"
    # Median particle diameter of the topsoil [mm]
    D50::Vector{T} | "mm"
    # Median diameter per particle class [um]
    dmclay::Vector{T} | "µm"
    dmsilt::Vector{T} | "µm"
    dmsand::Vector{T} | "µm"
    dmsagg::Vector{T} | "µm"
    dmlagg::Vector{T} | "µm"
    # Fraction of the different particle class [-]
    fclay::Vector{T} | "-"
    fsilt::Vector{T} | "-"
    fsand::Vector{T} | "-"
    fsagg::Vector{T} | "-"
    flagg::Vector{T} | "-"
    # Density of sediment [kg/m3]
    rhos::Vector{T} | "kg m-3"
    # Filter with river cells
    rivcell::Vector{T} | "-"
    # Total transport capacity of overland flow [ton]
    TCsed::Vector{T} | "t"
    # Transport capacity of overland flow per particle class [ton]
    TCclay::Vector{T} | "t"
    TCsilt::Vector{T} | "t"
    TCsand::Vector{T} | "t"
    TCsagg::Vector{T} | "t"
    TClagg::Vector{T} | "t"

    function LandSediment{T}(args...) where {T}
        equal_size_vectors(args)
        return new(args...)
    end
end

statevars(::LandSediment) = ()

function initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)
    # Initialize parameters for the soil loss part
    n = length(inds)
    do_river = get(config.model, "runrivermodel", false)::Bool
    # Reservoir / lake
    do_reservoirs = get(config.model, "doreservoir", false)::Bool
    do_lakes = get(config.model, "dolake", false)::Bool
    # Rainfall erosion equation: ["answers", "eurosem"]
    rainerosmethod = get(config.model, "rainerosmethod", "answers")::String
    # Overland flow transport capacity method: ["yalinpart", "govers", "yalin"]
    landtransportmethod = get(config.model, "landtransportmethod", "yalinpart")::String

    altitude =
        ncread(nc, param(config, "input.vertical.altitude"); sel = inds, type = Float)
    canopyheight = ncread(
        nc,
        param(config, "input.vertical.canopyheight", nothing);
        sel = inds,
        defaults = 3.0,
        type = Float,
    )
    erosk = ncread(
        nc,
        param(config, "input.vertical.erosk", nothing);
        sel = inds,
        defaults = 0.6,
        type = Float,
    )
    erosspl = ncread(
        nc,
        param(config, "input.vertical.erosspl", nothing);
        sel = inds,
        defaults = 2.0,
        type = Float,
    )
    erosov = ncread(
        nc,
        param(config, "input.vertical.erosov", nothing);
        sel = inds,
        defaults = 0.9,
        type = Float,
    )
    pathfrac = ncread(
        nc,
        param(config, "input.vertical.pathfrac", nothing);
        sel = inds,
        defaults = 0.01,
        type = Float,
    )
    slope = ncread(
        nc,
        param(config, "input.vertical.slope", nothing);
        sel = inds,
        defaults = 0.01,
        type = Float,
    )
    usleC = ncread(
        nc,
        param(config, "input.vertical.usleC", nothing);
        sel = inds,
        defaults = 0.01,
        type = Float,
    )
    usleK = ncread(
        nc,
        param(config, "input.vertical.usleK", nothing);
        sel = inds,
        defaults = 0.1,
        type = Float,
    )

    cmax, _, canopygapfraction, sl, swood, kext = initialize_canopy(nc, config, inds)

    # Initialise parameters for the transport capacity part
    βₗ = slope
    clamp!(βₗ, 0.00001, Inf)
    ldd_2d = ncread(nc, param(config, "input.ldd"); allow_missing = true)
    ldd = ldd_2d[inds]
    dl = fill(mv, n)
    dw = fill(mv, n)

    for i = 1:n
        dl[i] = detdrainlength(ldd[i], xl[i], yl[i])
        dw[i] = detdrainwidth(ldd[i], xl[i], yl[i])
    end
    dmclay = ncread(
        nc,
        param(config, "input.vertical.dmclay", nothing);
        sel = inds,
        defaults = 2.0,
        type = Float,
    )
    dmsilt = ncread(
        nc,
        param(config, "input.vertical.dmsilt", nothing);
        sel = inds,
        defaults = 10.0,
        type = Float,
    )
    dmsand = ncread(
        nc,
        param(config, "input.vertical.dmsand", nothing);
        sel = inds,
        defaults = 200.0,
        type = Float,
    )
    dmsagg = ncread(
        nc,
        param(config, "input.vertical.dmsagg", nothing);
        sel = inds,
        defaults = 30.0,
        type = Float,
    )
    dmlagg = ncread(
        nc,
        param(config, "input.vertical.dmlagg", nothing);
        sel = inds,
        defaults = 500.0,
        type = Float,
    )
    pclay = ncread(
        nc,
        param(config, "input.vertical.pclay", nothing);
        sel = inds,
        defaults = 0.1,
        type = Float,
    )
    psilt = ncread(
        nc,
        param(config, "input.vertical.psilt", nothing);
        sel = inds,
        defaults = 0.1,
        type = Float,
    )
    rhos = ncread(
        nc,
        param(config, "input.vertical.rhosed", nothing);
        sel = inds,
        defaults = 2650.0,
        type = Float,
    )

    ### Initialize transport capacity variables ###
    rivcell = float(river)
    # Percent Sand
    psand = 100 .- pclay .- psilt
    # Govers coefficient for transport capacity
    if landtransportmethod != "yalinpart"
        # Calculation of D50 and fraction of fine and very fine sand (fvfs) from Fooladmand et al, 2006
        psand999 = psand .* ((999 - 25) / (1000 - 25))
        vd50 = log.((1 ./ (0.01 .* (pclay .+ psilt)) .- 1) ./ (1 ./ (0.01 .* pclay) .- 1))
        wd50 =
            log.(
                (1 ./ (0.01 .* (pclay .+ psilt .+ psand999)) .- 1) ./
                (1 ./ (0.01 .* pclay) .- 1),
            )
        ad50 = 1 / log((25 - 1) / (999 - 1))
        bd50 = ad50 ./ log((25 - 1) / 1)
        cd50 = ad50 .* log.(vd50 ./ wd50)
        ud50 = (.-vd50) .^ (1 .- bd50) ./ ((.-wd50) .^ (.-bd50))
        D50 = 1 .+ (-1 ./ ud50 .* log.(1 ./ (1 ./ (0.01 .* pclay) .- 1))) .^ (1 ./ cd50) #[um]
        D50 = D50 ./ 1000 # [mm]
    else
        D50 = fill(mv, n)
    end
    if landtransportmethod == "govers"
        cGovers = ((D50 .* 1000 .+ 5) ./ 0.32) .^ (-0.6)
        nGovers = ((D50 .* 1000 .+ 5) ./ 300) .^ (0.25)
    else
        cGovers = fill(mv, n)
        nGovers = fill(mv, n)
    end
    if do_river || landtransportmethod == "yalinpart"
        # Determine sediment size distribution, estimated from primary particle size distribution (Foster et al., 1980)
        fclay = 0.20 .* pclay ./ 100
        fsilt = 0.13 .* psilt ./ 100
        fsand = 0.01 .* psand .* (1 .- 0.01 .* pclay) .^ (2.4)
        fsagg = 0.28 .* (0.01 .* pclay .- 0.25) .+ 0.5
        for i = 1:n
            if pclay[i] > 50.0
                fsagg[i] = 0.57
            elseif pclay[i] < 25
                fsagg[i] = 2.0 * 0.01 * pclay[i]
            end
        end
        flagg = 1.0 .- fclay .- fsilt .- fsand .- fsagg
    else
        fclay = fill(mv, n)
        fsilt = fill(mv, n)
        fsand = fill(mv, n)
        fsagg = fill(mv, n)
        flagg = fill(mv, n)
    end

    # Reservoir and lakes
    wbcover = zeros(Float, n)
    if do_reservoirs
        rescoverage_2d = ncread(
            nc,
            param(config, "input.vertical.resareas");
            sel = inds,
            type = Float,
            fill = 0.0,
        )
        wbcover = wbcover .+ rescoverage_2d
    end
    if do_lakes
        lakecoverage_2d = ncread(
            nc,
            param(config, "input.vertical.lakeareas");
            sel = inds,
            type = Float,
            fill = 0.0,
        )
        wbcover = wbcover .+ lakecoverage_2d
    end

    eros = LandSediment{Float}(
        n = n,
        yl = yl,
        xl = xl,
        riverfrac = riverfrac,
        wbcover = wbcover,
        ### Soil erosion part ###
        # Forcing
        interception = fill(mv, n),
        h_land = fill(mv, n),
        precipitation = fill(mv, n),
        q_land = fill(mv, n),
        # Parameters
        canopyheight = canopyheight,
        canopygapfraction = canopygapfraction,
        erosk = erosk,
        erosspl = erosspl,
        erosov = erosov,
        pathfrac = pathfrac,
        slope = slope,
        usleC = usleC,
        usleK = usleK,
        # Interception related to climatology (leaf_area_index)
        sl = sl,
        swood = swood,
        kext = kext,
        leaf_area_index = fill(mv, n),
        # Outputs
        sedspl = fill(mv, n),
        sedov = fill(mv, n),
        soilloss = fill(mv, n),
        erosclay = fill(mv, n),
        erossilt = fill(mv, n),
        erossand = fill(mv, n),
        erossagg = fill(mv, n),
        eroslagg = fill(mv, n),
        ### Transport capacity part ###
        # Parameters
        dl = dl,
        dw = dw,
        cGovers = cGovers,
        D50 = D50,
        dmclay = dmclay,
        dmsilt = dmsilt,
        dmsand = dmsand,
        dmsagg = dmsagg,
        dmlagg = dmlagg,
        fclay = fclay,
        fsilt = fsilt,
        fsand = fsand,
        fsagg = fsagg,
        flagg = flagg,
        nGovers = nGovers,
        rhos = rhos,
        rivcell = rivcell,
        # Outputs
        TCsed = fill(mv, n),
        TCclay = fill(mv, n),
        TCsilt = fill(mv, n),
        TCsand = fill(mv, n),
        TCsagg = fill(mv, n),
        TClagg = fill(mv, n),
    )

    return eros
end

# Soil erosion
function update_until_ols(eros::LandSediment, config)
    # Options from config
    do_lai = haskey(config.input.vertical, "leaf_area_index")
    rainerosmethod = get(config.model, "rainerosmethod", "answers")::String
    Δt = Second(config.timestepsecs)
    ts = Float(Δt.value)

    for i = 1:eros.n

        ### Splash / Rainfall erosion ###
        # ANSWERS method
        if rainerosmethod == "answers"
            # calculate rainfall intensity [mm/min]
            rintnsty = eros.precipitation[i] / (ts / 60)
            # splash erosion [kg/min]
            sedspl =
                0.108 * eros.usleC[i] * eros.usleK[i] * eros.xl[i] * eros.yl[i] * rintnsty^2
            # [ton/timestep]
            sedspl = sedspl * (ts / 60) * 1e-3
        end
        # TODO check eurosem method output values (too high!)
        if rainerosmethod == "eurosem"
            if do_lai
                cmax = eros.sl[i] * eros.leaf_area_index[i] + eros.swood[i]
                canopygapfraction = exp(-eros.kext[i] * eros.leaf_area_index[i])
            end
            # calculate rainfall intensity [mm/h]
            rintnsty = eros.precipitation[i] / (ts / 3600)
            # Kinetic energy of direct throughfall [J/m2/mm]
            # kedir = max(11.87 + 8.73 * log10(max(0.0001, rintnsty)),0.0) #basis used in USLE
            kedir = max(8.95 + 8.44 * log10(max(0.0001, rintnsty)), 0.0) #variant used in most distributed mdoels
            # Kinetic energy of leaf drainage [J/m2/mm]
            pheff = 0.5 * eros.canopyheight[i]
            keleaf = max((15.8 * pheff^0.5) - 5.87, 0.0)

            #Depths of rainfall (total, leaf drianage, direct) [mm]
            rdtot = eros.precipitation[i]
            rdleaf = rdtot * 0.1 * canopygapfraction #stemflow
            rddir = max(rdtot - rdleaf - eros.interception[i], 0.0) #throughfall

            #Total kinetic energy by rainfall [J/m2]
            ketot = (rddir * kedir + rdleaf * keleaf) * 0.001
            # Rainfall / splash erosion [g/m2]
            sedspl = eros.erosk[i] * ketot * exp(-eros.erosspl[i] * eros.h_land[i])
            sedspl = sedspl * eros.xl[i] * eros.yl[i] * 1e-6 # ton/cell
        end

        # Remove the impervious area
        sedspl = sedspl * (1.0 - eros.pathfrac[i])

        ### Overland flow erosion ###
        # ANWERS method
        # Overland flow rate [m2/min]
        qr_land = eros.q_land[i] * 60 / ((eros.xl[i] + eros.yl[i]) / 2)
        # Sine of the slope
        sinslope = sin(atan(eros.slope[i]))

        # Overland flow erosion [kg/min]
        # For a wide range of slope, it is better to use the sine of slope rather than tangeant
        sedov =
            eros.erosov[i] *
            eros.usleC[i] *
            eros.usleK[i] *
            eros.xl[i] *
            eros.yl[i] *
            sinslope *
            qr_land
        # [ton/timestep]
        sedov = sedov * (ts / 60) * 1e-3
        # Remove the impervious area
        sedov = sedov * (1.0 - eros.pathfrac[i])

        # Assume no erosion in reservoir/lake areas
        if eros.wbcover[i] > 0
            sedspl = 0.0
            sedov = 0.0
        end

        # Total soil loss [ton/cell/timestep]
        soilloss = sedspl + sedov

        # update the outputs and states
        eros.sedspl[i] = sedspl
        eros.sedov[i] = sedov
        eros.soilloss[i] = soilloss
        # Eroded amount per particle class
        eros.erosclay[i] = soilloss * eros.fclay[i]
        eros.erossilt[i] = soilloss * eros.fsilt[i]
        eros.erossand[i] = soilloss * eros.fsand[i]
        eros.erossagg[i] = soilloss * eros.fsagg[i]
        eros.eroslagg[i] = soilloss * eros.flagg[i]
    end

end

### Sediment transport capacity in overland flow ###
function update_until_oltransport(ols::LandSediment, config::Config)

    do_river = get(config.model, "runrivermodel", false)::Bool
    tcmethod = get(config.model, "landtransportmethod", "yalinpart")::String
    Δt = Second(config.timestepsecs)
    ts = Float(Δt.value)

    for i = 1:ols.n
        sinslope = sin(atan(ols.slope[i]))

        if !do_river
            # Total transport capacity without particle differentiation
            if tcmethod == "govers"
                TC = tc_govers(ols, i, sinslope, ts)
            end

            if tcmethod == "yalin"
                TC = tc_yalin(ols, i, sinslope, ts)
            end

            # Filter TC land for river cells (0 in order for sediment from land to stop when entering the river)
            if ols.rivcell[i] == 1.0
                TC = 0.0
            end

            # Set particle TC to 0
            TCclay = 0.0
            TCsilt = 0.0
            TCsand = 0.0
            TCsagg = 0.0
            TClagg = 0.0
        end

        if do_river || tcmethod == "yalinpart"
            TC, TCclay, TCsilt, TCsand, TCsagg, TClagg = tc_yalinpart(ols, i, sinslope, ts)
        end

        # update the outputs and states
        ols.TCsed[i] = TC
        ols.TCclay[i] = TCclay
        ols.TCsilt[i] = TCsilt
        ols.TCsand[i] = TCsand
        ols.TCsagg[i] = TCsagg
        ols.TClagg[i] = TClagg
    end

end

function tc_govers(ols::LandSediment, i::Int, sinslope::Float, ts::Float)::Float
    # Transport capacity from govers 1990
    # Unit stream power
    if ols.h_land[i] > 0.0
        velocity = ols.q_land[i] / (ols.dw[i] * ols.h_land[i])
    else
        velocity = 0.0
    end
    omega = 10 * sinslope * 100 * velocity #cm/s
    if omega > 0.4
        TCf = ols.cGovers[i] * (omega - 0.4)^(ols.nGovers[i]) * 2650 #kg/m3
    else
        TCf = 0.0
    end
    TC = TCf * ols.q_land[i] * ts * 1e-3 #[ton/cell]

    return TC
end

function tc_yalin(ols::LandSediment, i::Int, sinslope::Float, ts::Float)::Float
    # Transport capacity from Yalin without particle differentiation
    delta = max(
        (
            ols.h_land[i] * sinslope / (ols.D50[i] * 0.001 * (ols.rhos[i] / 1000 - 1)) /
            0.06 - 1
        ),
        0.0,
    )
    alphay = delta * 2.45 / (0.001 * ols.rhos[i])^0.4 * 0.06^(0.5)
    if ols.q_land[i] > 0.0 && alphay != 0.0
        TC = (
            ols.dw[i] / ols.q_land[i] *
            (ols.rhos[i] - 1000) *
            ols.D50[i] *
            0.001 *
            (9.81 * ols.h_land[i] * sinslope) *
            0.635 *
            delta *
            (1 - log(1 + alphay) / (alphay))
        ) # [kg/m3]
        TC = TC * ols.q_land[i] * ts * 1e-3 #[ton]
    else
        TC = 0.0
    end
    return TC
end

function tc_yalinpart(ols::LandSediment, i::Int, sinslope::Float, ts::Float)
    # Transport capacity from Yalin with particle differentiation
    # Delta parameter of Yalin for each particle class
    delta = ols.h_land[i] * sinslope / (1e-6 * (ols.rhos[i] / 1000 - 1)) / 0.06
    dclay = max(1 / ols.dmclay[i] * delta - 1, 0.0)
    dsilt = max(1 / ols.dmsilt[i] * delta - 1, 0.0)
    dsand = max(1 / ols.dmsand[i] * delta - 1, 0.0)
    dsagg = max(1 / ols.dmsagg[i] * delta - 1, 0.0)
    dlagg = max(1 / ols.dmlagg[i] * delta - 1, 0.0)
    # Total transportability
    dtot = dclay + dsilt + dsand + dsagg + dlagg

    # Yalin transport capacity of overland flow for each particle class
    if ols.q_land[i] > 0.0
        TCa =
            ols.dw[i] / ols.q_land[i] *
            (ols.rhos[i] - 1000) *
            1e-6 *
            (9.81 * ols.h_land[i] * sinslope)
    else
        TCa = 0.0
    end
    TCb = 2.45 / (ols.rhos[i] / 1000)^0.4 * 0.06^0.5
    if dtot != 0.0 && dclay != 0.0
        TCclay =
            TCa * ols.dmclay[i] * dclay / dtot *
            0.635 *
            dclay *
            (1 - log(1 + dclay * TCb) / dclay * TCb) # [kg/m3]
        TCclay = TCclay * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCclay = 0.0
    end
    if dtot != 0.0 && dsilt != 0.0
        TCsilt =
            TCa * ols.dmsilt[i] * dsilt / dtot *
            0.635 *
            dsilt *
            (1 - log(1 + dsilt * TCb) / dsilt * TCb) # [kg/m3]
        TCsilt = TCsilt * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCsilt = 0.0
    end
    if dtot != 0.0 && dsand != 0.0
        TCsand =
            TCa * ols.dmsand[i] * dsand / dtot *
            0.635 *
            dsand *
            (1 - log(1 + dsand * TCb) / dsand * TCb) # [kg/m3]
        TCsand = TCsand * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCsand = 0.0
    end
    if dtot != 0.0 && dsagg != 0.0
        TCsagg =
            TCa * ols.dmsagg[i] * dsagg / dtot *
            0.635 *
            dsagg *
            (1 - log(1 + dsagg * TCb) / dsagg * TCb) # [kg/m3]
        TCsagg = TCsagg * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCsagg = 0.0
    end
    if dtot != 0.0 && dlagg != 0.0
        TClagg =
            TCa * ols.dmlagg[i] * dlagg / dtot *
            0.635 *
            dlagg *
            (1 - log(1 + dlagg * TCb) / dlagg * TCb) # [kg/m3]
        TClagg = TClagg * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TClagg = 0.0
    end

    # Assume that ols all reach the river in reservoir/lake areas (very high TC)
    if ols.wbcover[i] > 0
        TC = 1e9
        TCclay = 1e9
        TCsilt = 1e9
        TCsand = 1e9
        TCsagg = 1e9
        TClagg = 1e9
    end

    # Filter TC land for river cells (0 in order for sediment from land to stop when entering the river)
    if ols.rivcell[i] == 1.0
        TCclay = 0.0
        TCsilt = 0.0
        TCsand = 0.0
        TCsagg = 0.0
        TClagg = 0.0
    end

    # Set total TC to 0
    TC = 0.0

    return TC, TCclay, TCsilt, TCsand, TCsagg, TClagg
end

### Sediment transport in overland flow ###
@get_units @with_kw struct OverlandFlowSediment{T}
    # number of cells
    n::Int | "-"
    # Filter with river cells
    rivcell::Vector{T} | "-"
    # Total eroded soil [ton]
    soilloss::Vector{T} | "t"
    # Eroded soil per particle class [ton]
    erosclay::Vector{T} | "t"
    erossilt::Vector{T} | "t"
    erossand::Vector{T} | "t"
    erossagg::Vector{T} | "t"
    eroslagg::Vector{T} | "t"
    # Total transport capacity of overland flow [ton]
    TCsed::Vector{T} | "t"
    # Transport capacity of overland flow per particle class [ton]
    TCclay::Vector{T} | "t"
    TCsilt::Vector{T} | "t"
    TCsand::Vector{T} | "t"
    TCsagg::Vector{T} | "t"
    TClagg::Vector{T} | "t"
    # Sediment flux in overland flow [ton]
    olsed::Vector{T} | "t"
    olclay::Vector{T} | "t"
    olsilt::Vector{T} | "t"
    olsand::Vector{T} | "t"
    olsagg::Vector{T} | "t"
    ollagg::Vector{T} | "t"
    # Sediment reaching the river with overland flow [ton]
    inlandsed::Vector{T} | "t"
    inlandclay::Vector{T} | "t"
    inlandsilt::Vector{T} | "t"
    inlandsand::Vector{T} | "t"
    inlandsagg::Vector{T} | "t"
    inlandlagg::Vector{T} | "t"

    function OverlandFlowSediment{T}(args...) where {T}
        equal_size_vectors(args)
        return new(args...)
    end
end

statevars(::OverlandFlowSediment) = ()

function partial_update!(inland, rivcell, eroded)
    no_erosion = zero(eltype(eroded))
    for i in eachindex(inland)
        inland[i] = rivcell[i] == 1 ? eroded[i] : no_erosion
    end
    return inland
end

function update(ols::OverlandFlowSediment, network, config)
    do_river = get(config.model, "runrivermodel", false)::Bool
    tcmethod = get(config.model, "landtransportmethod", "yalinpart")::String
    zeroarr = fill(0.0, ols.n)

    # transport sediment down the network
    if do_river || tcmethod == "yalinpart"
        # clay
        accucapacityflux!(ols.olclay, ols.erosclay, network, ols.TCclay)
        partial_update!(ols.inlandclay, ols.rivcell, ols.erosclay)
        # silt
        accucapacityflux!(ols.olsilt, ols.erossilt, network, ols.TCsilt)
        partial_update!(ols.inlandsilt, ols.rivcell, ols.erossilt)
        # sand
        accucapacityflux!(ols.olsand, ols.erossand, network, ols.TCsand)
        partial_update!(ols.inlandsand, ols.rivcell, ols.erossand)
        # small aggregates
        accucapacityflux!(ols.olsagg, ols.erossagg, network, ols.TCsagg)
        partial_update!(ols.inlandsagg, ols.rivcell, ols.erossagg)
        # large aggregates
        accucapacityflux!(ols.ollagg, ols.eroslagg, network, ols.TClagg)
        partial_update!(ols.inlandlagg, ols.rivcell, ols.eroslagg)

        # total sediment, all particle classes
        ols.olsed .= ols.olclay .+ ols.olsilt .+ ols.olsand .+ ols.olsagg .+ ols.ollagg
        ols.inlandsed .=
            ols.inlandclay .+ ols.inlandsilt .+ ols.inlandsand .+ ols.inlandsagg .+
            ols.inlandlagg
    else
        accucapacityflux!(ols.olsed, ols.soilloss, network, ols.TCsed)
        ols.inlandsed .= ifelse.(ols.rivcell .== 1, ols.soilloss, zeroarr)
    end
end

### River transport and processes ###
@get_units @with_kw struct RiverSediment{T}
    # number of cells
    n::Int | "-"
    # Timestep [s]
    Δt::T | "s"
    # River geometry (slope [-], length [m], width [m])
    sl::Vector{T} | "m"
    dl::Vector{T} | "m"
    width::Vector{T} | "m"
    # Sediment mean diameter in the river bed [mm]
    d50::Vector{T}
    # Particle mean diameter [mm]
    dmclay::Vector{T}
    dmsilt::Vector{T}
    dmsand::Vector{T}
    dmsagg::Vector{T}
    dmlagg::Vector{T}
    dmgrav::Vector{T}
    # River bed and bank particle fraction composition [-]
    fclayriv::Vector{T} | "-"
    fsiltriv::Vector{T} | "-"
    fsandriv::Vector{T} | "-"
    fgravriv::Vector{T} | "-"
    # Sediment mean diameter for Engelund and Hansen transport equation [mm]
    d50engelund::Vector{T}
    # Parameters for Bagnold transport equation
    cbagnold::Vector{T} | "-"
    ebagnold::Vector{T} | "-"
    # Parameters for Kodatie transport equation
    ak::Vector{T} | "-"
    bk::Vector{T} | "-"
    ck::Vector{T} | "-"
    dk::Vector{T} | "-"
    # Critical bed and bank shear stress [N/m2]
    TCrbank::Vector{T} | "N m-2"
    TCrbed::Vector{T} | "N m-2"
    # Bed and bank erodibilities [m3/N.s]
    kdbank::Vector{T} | "m3 N-1 s-1"
    kdbed::Vector{T} | "m3 N-1 s-1"
    # Sediment density [kg/m3]
    rhos::Vector{T} | "kg m-3"
    # River water level [m]
    h_riv::Vector{T} | "m"
    # River discharge [m3/s]
    q_riv::Vector{T} | "m3 s-1"
    # Sediment input from land erosion [ton]
    inlandclay::Vector{T} | "t"
    inlandsilt::Vector{T} | "t"
    inlandsand::Vector{T} | "t"
    inlandsagg::Vector{T} | "t"
    inlandlagg::Vector{T} | "t"
    inlandsed::Vector{T} | "t"
    # Sediment / particle left in the cell [ton]
    sedload::Vector{T} | "t"
    clayload::Vector{T} | "t"
    siltload::Vector{T} | "t"
    sandload::Vector{T} | "t"
    saggload::Vector{T} | "t"
    laggload::Vector{T} | "t"
    gravload::Vector{T} | "t"
    # Sediment / particle stored on the river bed after deposition [ton]
    sedstore::Vector{T} | "t"
    claystore::Vector{T} | "t"
    siltstore::Vector{T} | "t"
    sandstore::Vector{T} | "t"
    saggstore::Vector{T} | "t"
    laggstore::Vector{T} | "t"
    gravstore::Vector{T} | "t"
    # Sediment / particle flux [ton]
    outsed::Vector{T} | "t"
    outclay::Vector{T} | "t"
    outsilt::Vector{T} | "t"
    outsand::Vector{T} | "t"
    outsagg::Vector{T} | "t"
    outlagg::Vector{T} | "t"
    outgrav::Vector{T} | "t"
    # Sediment concentrations [kg/m3]
    Sedconc::Vector{T} | "kg m-3"
    SSconc::Vector{T} | "kg m-3"
    Bedconc::Vector{T} | "kg m-3"
    # River transport capacity
    maxsed::Vector{T} | "t"
    # Eroded sediment (total, bank and bed)
    erodsed::Vector{T} | "t"
    erodsedbank::Vector{T} | "t"
    erodsedbed::Vector{T} | "t"
    # Deposited sediment
    depsed::Vector{T} | "t"
    # Sediment in
    insed::Vector{T} | "t"
    # Reservoir and lakes
    wbcover::Vector{T} | "-"
    wblocs::Vector{T} | "-"
    wbarea::Vector{T} | "m2"
    wbtrap::Vector{T} | "-"

    # function RiverSediment{T}(args...) where {T}
    #     equal_size_vectors(args)
    #     return new(args...)
    # end
end

statevars(::RiverSediment) = (
    :clayload,
    :siltload,
    :sandload,
    :saggload,
    :laggload,
    :gravload,
    :claystore,
    :siltstore,
    :sandstore,
    :saggstore,
    :laggstore,
    :gravstore,
    :outclay,
    :outsilt,
    :outsand,
    :outsagg,
    :outlagg,
    :outgrav,
)

function initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)
    # Initialize river parameters
    nriv = length(inds_riv)
    # River flow transport capacity method: ["bagnold", "engelund", "yang", "kodatie", "molinas"]
    tcmethodriv = get(config.model, "rivtransportmethod", "bagnold")::String
    Δt = Second(config.timestepsecs)
    # Reservoir / lakes
    do_reservoirs = get(config.model, "doreservoir", false)::Bool
    do_lakes = get(config.model, "dolake", false)::Bool
    wbcover = zeros(Float, nriv)
    wblocs = zeros(Float, nriv)
    wbarea = zeros(Float, nriv)
    wbtrap = zeros(Float, nriv)

    if do_reservoirs
        reslocs = ncread(
            nc,
            param(config, "input.lateral.river.reslocs");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        rescoverage_2d = ncread(
            nc,
            param(config, "input.lateral.river.resareas");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        resarea = ncread(
            nc,
            param(config, "input.lateral.river.resarea");
            sel = inds_riv,
            type = Float,
            fill = 0.0,
        )
        restrapefficiency = ncread(
            nc,
            param(config, "input.lateral.river.restrapeff");
            sel = inds_riv,
            type = Float,
            defaults = 1.0,
            fill = 0.0,
        )

        wbcover = wbcover .+ rescoverage_2d
        wblocs = wblocs .+ reslocs
        wbarea = wbarea .+ resarea
        wbtrap = wbtrap .+ restrapefficiency
    end

    if do_lakes
        lakelocs = ncread(
            nc,
            param(config, "input.lateral.river.lakelocs");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        lakecoverage_2d = ncread(
            nc,
            param(config, "input.lateral.river.lakeareas");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        lakearea = ncread(
            nc,
            param(config, "input.lateral.river.lakearea");
            sel = inds_riv,
            type = Float,
            fill = 0.0,
        )

        wbcover = wbcover .+ lakecoverage_2d
        wblocs = wblocs .+ lakelocs
        wbarea = wbarea .+ lakearea
    end

    riverslope =
        ncread(nc, param(config, "input.lateral.river.slope"); sel = inds_riv, type = Float)
    clamp!(riverslope, 0.00001, Inf)
    rhos = ncread(
        nc,
        param(config, "input.lateral.river.rhosed", nothing);
        sel = inds_riv,
        defaults = 2650.0,
        type = Float,
    )
    dmclay = ncread(
        nc,
        param(config, "input.lateral.river.dmclay", nothing);
        sel = inds_riv,
        defaults = 2.0,
        type = Float,
    )
    dmsilt = ncread(
        nc,
        param(config, "input.lateral.river.dmsilt", nothing);
        sel = inds_riv,
        defaults = 10.0,
        type = Float,
    )
    dmsand = ncread(
        nc,
        param(config, "input.lateral.river.dmsand", nothing);
        sel = inds_riv,
        defaults = 200.0,
        type = Float,
    )
    dmsagg = ncread(
        nc,
        param(config, "input.lateral.river.dmsagg", nothing);
        sel = inds_riv,
        defaults = 30.0,
        type = Float,
    )
    dmlagg = ncread(
        nc,
        param(config, "input.lateral.river.dmlagg", nothing);
        sel = inds_riv,
        defaults = 500.0,
        type = Float,
    )
    dmgrav = ncread(
        nc,
        param(config, "input.lateral.river.dmgrav", nothing);
        sel = inds_riv,
        defaults = 2000.0,
        type = Float,
    )
    fclayriv = ncread(
        nc,
        param(config, "input.lateral.river.fclayriv");
        sel = inds_riv,
        type = Float,
    )
    fsiltriv = ncread(
        nc,
        param(config, "input.lateral.river.fsiltriv");
        sel = inds_riv,
        type = Float,
    )
    fsandriv = ncread(
        nc,
        param(config, "input.lateral.river.fsandriv");
        sel = inds_riv,
        type = Float,
    )
    fgravriv = ncread(
        nc,
        param(config, "input.lateral.river.fgravriv");
        sel = inds_riv,
        type = Float,
    )
    d50riv =
        ncread(nc, param(config, "input.lateral.river.d50"); sel = inds_riv, type = Float)
    d50engelund = ncread(
        nc,
        param(config, "input.lateral.river.d50engelund");
        sel = inds_riv,
        type = Float,
    )
    cbagnold = ncread(
        nc,
        param(config, "input.lateral.river.cbagnold");
        sel = inds_riv,
        type = Float,
    )
    ebagnold = ncread(
        nc,
        param(config, "input.lateral.river.ebagnold");
        sel = inds_riv,
        type = Float,
    )

    # Initialisation of parameters for Kodatie transport capacity
    ak = zeros(Float, nriv)
    bk = zeros(Float, nriv)
    ck = zeros(Float, nriv)
    dk = zeros(Float, nriv)
    if tcmethodriv == "kodatie"
        for i = 1:nriv
            if d50riv[i] <= 0.05
                ak[i] = 281.4
                bk[i] = 2.622
                ck[i] = 0.182
                dk[i] = 0.0
            elseif d50riv[i] <= 0.25
                ak[i] = 2829.6
                bk[i] = 3.646
                ck[i] = 0.406
                dk[i] = 0.412
            elseif d50riv[i] <= 2.0
                ak[i] = 2123.4
                bk[i] = 3.3
                ck[i] = 0.468
                dk[i] = 0.613
            else
                ak[i] = 431884.8
                bk[i] = 1.0
                ck[i] = 1.0
                dk[i] = 2.0
            end
        end
    end
    # Initialisation of parameters for river erosion
    # Bed and Bank from Shields diagram, Da Silva & Yalin (2017)
    E_ = (2.65 - 1) * 9.81
    E = (E_ .* (d50riv .* 1e-3) .^ 3 ./ 1e-12) .^ 0.33
    TCrbed = @. Float(
        E_ *
        d50riv *
        (0.13 * E^(-0.392) * exp(-0.015 * E^2) + 0.045 * (1 - exp(-0.068 * E))),
    )
    TCrbank = TCrbed
    # kd from Hanson & Simon 2001
    kdbank = @. Float(0.2 * TCrbank^(-0.5) * 1e-6)
    kdbed = @. Float(0.2 * TCrbed^(-0.5) * 1e-6)

    rs = RiverSediment(
        n = nriv,
        Δt = Float(Δt.value),
        # Parameters
        sl = riverslope,
        dl = riverlength,
        width = riverwidth,
        dmclay = dmclay,
        dmsilt = dmsilt,
        dmsand = dmsand,
        dmsagg = dmsagg,
        dmlagg = dmlagg,
        dmgrav = dmgrav,
        fclayriv = fclayriv,
        fsiltriv = fsiltriv,
        fsandriv = fsandriv,
        fgravriv = fgravriv,
        d50 = d50riv,
        d50engelund = d50engelund,
        cbagnold = cbagnold,
        ebagnold = ebagnold,
        ak = ak,
        bk = bk,
        ck = ck,
        dk = dk,
        kdbank = kdbank,
        kdbed = kdbed,
        TCrbank = TCrbank,
        TCrbed = TCrbed,
        rhos = rhos,
        # Forcing
        h_riv = fill(mv, nriv),
        q_riv = fill(mv, nriv),
        # Input from land
        inlandclay = zeros(Float, nriv),
        inlandsilt = zeros(Float, nriv),
        inlandsand = zeros(Float, nriv),
        inlandsagg = zeros(Float, nriv),
        inlandlagg = zeros(Float, nriv),
        inlandsed = zeros(Float, nriv),
        # States
        sedload = zeros(Float, nriv),
        clayload = zeros(Float, nriv),
        siltload = zeros(Float, nriv),
        sandload = zeros(Float, nriv),
        saggload = zeros(Float, nriv),
        laggload = zeros(Float, nriv),
        gravload = zeros(Float, nriv),
        sedstore = zeros(Float, nriv),
        claystore = zeros(Float, nriv),
        siltstore = zeros(Float, nriv),
        sandstore = zeros(Float, nriv),
        saggstore = zeros(Float, nriv),
        laggstore = zeros(Float, nriv),
        gravstore = zeros(Float, nriv),
        outsed = zeros(Float, nriv),
        outclay = zeros(Float, nriv),
        outsilt = zeros(Float, nriv),
        outsand = zeros(Float, nriv),
        outsagg = zeros(Float, nriv),
        outlagg = zeros(Float, nriv),
        outgrav = zeros(Float, nriv),
        # Outputs
        Sedconc = zeros(Float, nriv),
        SSconc = zeros(Float, nriv),
        Bedconc = zeros(Float, nriv),
        maxsed = zeros(Float, nriv),
        erodsed = zeros(Float, nriv),
        erodsedbank = zeros(Float, nriv),
        erodsedbed = zeros(Float, nriv),
        depsed = zeros(Float, nriv),
        insed = zeros(Float, nriv),
        # Reservoir / lake
        wbcover = wbcover,
        wblocs = wblocs,
        wbarea = wbarea,
        wbtrap = wbtrap,
    )

    return rs
end

function update(rs::RiverSediment, network, config)
    @unpack graph, order = network
    tcmethod = get(config.model, "rivtransportmethod", "bagnold")::String

    # River sediment loads are separated into different particle class.
    # Clay, silt and sand can both come from land, resuspension or river channel erosion.
    # Small and large aggregates only come from land erosion or resuspension.
    # Gravel only comes from resuspension or river channel erosion.

    for v in order
        ### Sediment input in the cell (left from previous timestep + from land + from upstream outflux) ###
        upstream_nodes = inneighbors(graph, v)

        inrivclay = 0.0
        inrivsilt = 0.0
        inrivsand = 0.0
        inrivsagg = 0.0
        inrivlagg = 0.0
        inrivgrav = 0.0
        if !isempty(upstream_nodes)
            for i in upstream_nodes
                if rs.outclay[i] >= 0.0 # avoid NaN from upstream non-river cells
                    inrivclay += rs.outclay[i]
                    inrivsilt += rs.outsilt[i]
                    inrivsand += rs.outsand[i]
                    inrivsagg += rs.outsagg[i]
                    inrivlagg += rs.outlagg[i]
                    inrivgrav += rs.outgrav[i]
                end
            end
        end

        inclay = rs.clayload[v] + rs.inlandclay[v] + inrivclay
        insilt = rs.siltload[v] + rs.inlandsilt[v] + inrivsilt
        insand = rs.sandload[v] + rs.inlandsand[v] + inrivsand
        insagg = rs.saggload[v] + rs.inlandsagg[v] + inrivsagg
        inlagg = rs.laggload[v] + rs.inlandlagg[v] + inrivlagg
        ingrav = rs.gravload[v] + inrivgrav

        insed = inclay + insilt + insand + insagg + inlagg + ingrav
        rs.insed[v] = insed
        rs.inlandsed[v] =
            rs.inlandclay[v] +
            rs.inlandsilt[v] +
            rs.inlandsand[v] +
            rs.inlandsagg[v] +
            rs.inlandlagg[v]

        ### Transport capacity of the flow ###
        # Hydraulic radius of the river [m] (rectangular channel)
        hydrad = rs.h_riv[v] * rs.width[v] / (rs.width[v] + 2 * rs.h_riv[v])

        # Engelund and Hansen transport formula
        if tcmethod == "engelund"
            vmean =
                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
            vshear = (9.81 * hydrad * rs.sl[v])^0.5
            # Concentration by weight
            cw = ifelse(
                hydrad > 0.0,
                (
                    rs.rhos[v] / 1000 * 0.05 * vmean * vshear^3 /
                    ((rs.rhos[v] / 1000 - 1)^2 * 9.81^2 * rs.d50engelund[v] * hydrad)
                ),
                0.0,
            )
            cw = min(1.0, cw)
            # Transport capacity [tons/m3]
            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
        elseif tcmethod == "bagnold"
            maxsed =
                rs.cbagnold[v] * (rs.q_riv[v] / (rs.h_riv[v] * rs.width[v]))^rs.ebagnold[v]
        elseif tcmethod == "kodatie"
            vmean =
                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
            maxsed = rs.ak[v] * vmean^rs.bk[v] * rs.h_riv[v]^rs.ck[v] * rs.sl[v]^rs.dk[v]
            # Transport capacity [tons/m3]
            maxsed =
                ifelse(rs.q_riv[v] > 0.0, maxsed * rs.width[v] / (rs.q_riv[v] * rs.Δt), 0.0)
        elseif tcmethod == "yang"
            ws = 411 * rs.d50[v]^2 / 3600
            vshear = (9.81 * hydrad * rs.sl[v])^0.5
            var1 = vshear * rs.d50[v] / 1000 / (1.16 * 1e-6)
            var2 = ws * rs.d50[v] / 1000 / (1.16 * 1e-6)
            vcr = min(
                0.0,
                ifelse(var1 >= 70.0, 2.05 * ws, ws * (2.5 / (log10(var1) - 0.06) + 0.66)),
            )
            # Sand equation
            if (rs.width[v] * rs.h_riv[v]) >= vcr && rs.d50[v] < 2.0
                logcppm = (
                    5.435 - 0.286 * log10(var2) - 0.457 * log10(vshear / ws) + 1.799 -
                    0.409 * log10(var2) -
                    0.314 *
                    log10(vshear / ws) *
                    log10(
                        (rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]) - vcr) * rs.sl[v] / ws,
                    )
                )
                # Gravel equation
            elseif (rs.width[v] * rs.h_riv[v]) >= vcr && rs.d50[v] < 2.0
                logcppm = (
                    6.681 - 0.633 * log10(var2) - 4.816 * log10(vshear / ws) + 2.784 -
                    0.305 * log10(var2) -
                    0.282 *
                    log10(vshear / ws) *
                    log10(
                        (rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]) - vcr) * rs.sl[v] / ws,
                    )
                )
            else
                logcppm = 0.0
            end
            # Sediment concentration by weight
            cw = 10^logcppm * 1e-6
            # Transport capacity [ton/m3]
            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
        elseif tcmethod == "molinas"
            ws = 411 * rs.d50[v]^2 / 3600
            vmean =
                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
            if rs.h_riv[v] > 0.0
                psi = (
                    vmean^3 / (
                        (rs.rhos[v] / 1000 - 1) *
                        9.81 *
                        rs.h_riv[v] *
                        ws *
                        log10(1000 * rs.h_riv[v] / rs.d50[v])^2
                    )
                )
            else
                psi = 0.0
            end
            # Concentration by weight
            cw = 1430 * (0.86 + psi^0.5) * psi^1.5 / (0.016 + psi) * 1e-6
            # Transport capacity [ton/m3]
            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
        end

        # 1285 g/L: boundary between streamflow and debris flow (Costa, 1988)
        maxsed = min(maxsed, 1.285)
        # Transport capacity [ton]
        maxsed = maxsed * (rs.h_riv[v] * rs.width[v] * rs.dl[v] + rs.q_riv[v] * rs.Δt)
        rs.maxsed[v] = maxsed

        ### River erosion ###
        # Erosion only if the load is below the transport capacity of the flow.
        sedex = max(maxsed - insed, 0.0)
        # No erosion in lake and reservoir cells
        if rs.wbcover[v] > 0.0
            sedex = 0.0
        end
        # Bed and bank are eroded only if the previously deposited material is not enough
        rs.sedstore[v] =
            rs.claystore[v] +
            rs.siltstore[v] +
            rs.sandstore[v] +
            rs.saggstore[v] +
            rs.laggstore[v] +
            rs.gravstore[v]
        if sedex > 0.0 && sedex > rs.sedstore[v]
            # Effective sediment needed fom river bed and bank erosion [ton]
            effsedex = sedex - rs.sedstore[v]

            # Repartition of the effective shear stress between the bank and the bed from Knight et al. 1984 [%]
            SFbank = ifelse(
                rs.h_riv[v] > 0.0,
                exp(-3.23 * log10(rs.width[v] / rs.h_riv[v] + 3) + 6.146),
                0.0,
            )
            # Effective shear stress on river bed and banks [N/m2]
            TEffbank = ifelse(
                rs.h_riv[v] > 0.0,
                1000 * 9.81 * hydrad * rs.sl[v] * SFbank / 100 *
                (1 + rs.width[v] / (2 * rs.h_riv[v])),
                0.0,
            )
            TEffbed =
                1000 *
                9.81 *
                hydrad *
                rs.sl[v] *
                (1 - SFbank / 100) *
                (1 + 2 * rs.h_riv[v] / rs.width[v])
            # Potential erosion rates of the bed and bank [t/cell/timestep]
            #(assuming only one bank is eroding)
            Tex = max(TEffbank - rs.TCrbank[v], 0.0)
            # 1.4 is bank default bulk density
            ERbank = max(0.0, rs.kdbank[v] * Tex * rs.dl[v] * rs.h_riv[v] * 1.4 * rs.Δt)
            # 1.5 is bed default bulk density
            ERbed = max(
                0.0,
                rs.kdbed[v] *
                (TEffbed - rs.TCrbed[v]) *
                rs.dl[v] *
                rs.width[v] *
                1.5 *
                rs.Δt,
            )
            # Relative potential erosion rates of the bed and the bank [-]
            RTEbank = ifelse(ERbank + ERbed > 0.0, ERbank / (ERbank + ERbed), 0.0)
            RTEbed = 1.0 - RTEbank

            # Bank erosion (difference between effective and potential erosion) [ton]
            sedbank = ifelse(effsedex * RTEbank <= ERbank, effsedex * RTEbank, ERbank)
            claybank = rs.fclayriv[v] * sedbank
            siltbank = rs.fsiltriv[v] * sedbank
            sandbank = rs.fsandriv[v] * sedbank
            gravbank = rs.fgravriv[v] * sedbank

            # Bed erosion [ton]
            sedbed = ifelse(effsedex * RTEbed <= ERbed, effsedex * RTEbed, ERbed)
            claybed = rs.fclayriv[v] * sedbed
            siltbed = rs.fsiltriv[v] * sedbed
            sandbed = rs.fsandriv[v] * sedbed
            gravbed = rs.fgravriv[v] * sedbed
        else
            sedbank = 0.0
            claybank = 0.0
            siltbank = 0.0
            sandbank = 0.0
            gravbank = 0.0

            sedbed = 0.0
            claybed = 0.0
            siltbed = 0.0
            sandbed = 0.0
            gravbed = 0.0
        end

        # Erosion/degradation of the previously deposited sediment (from clay to gravel) [ton]
        if sedex > 0.0
            degstoreclay = ifelse(rs.claystore[v] >= sedex, sedex, rs.claystore[v])
            rs.claystore[v] = rs.claystore[v] - degstoreclay
            #Update amount of sediment that need to be degraded
            sedex = sedex - degstoreclay
            degstoresilt = ifelse(rs.siltstore[v] >= sedex, sedex, rs.siltstore[v])
            rs.siltstore[v] = rs.siltstore[v] - degstoresilt
            sedex = max(0.0, sedex - degstoresilt)
            degstoresagg = ifelse(rs.saggstore[v] >= sedex, sedex, rs.saggstore[v])
            rs.saggstore[v] = rs.saggstore[v] - degstoresagg
            sedex = max(0.0, sedex - degstoresagg)
            degstoresand = ifelse(rs.sandstore[v] >= sedex, sedex, rs.sandstore[v])
            rs.sandstore[v] = rs.sandstore[v] - degstoresand
            sedex = max(0.0, sedex - degstoresand)
            degstorelagg = ifelse(rs.laggstore[v] >= sedex, sedex, rs.laggstore[v])
            rs.laggstore[v] = rs.laggstore[v] - degstorelagg
            sedex = max(0.0, sedex - degstorelagg)
            degstoregrav = ifelse(rs.gravstore[v] >= sedex, sedex, rs.gravstore[v])
            rs.gravstore[v] = rs.gravstore[v] - degstoregrav
            sedex = max(0.0, sedex - degstoregrav)
            degstoresed =
                degstoreclay +
                degstoresilt +
                degstoresagg +
                degstoresand +
                degstorelagg +
                degstoregrav
        else
            degstoreclay = 0.0
            degstoresilt = 0.0
            degstoresagg = 0.0
            degstoresand = 0.0
            degstorelagg = 0.0
            degstoregrav = 0.0
            degstoresed = 0.0
        end

        # Sum all erosion sources per particle class
        erodsed = sedbank + sedbed + degstoresed
        erodclay = claybank + claybed + degstoreclay
        erodsilt = siltbank + siltbed + degstoresilt
        erodsand = sandbank + sandbed + degstoresand
        erodsagg = degstoresagg
        erodlagg = degstorelagg
        erodgrav = gravbank + gravbed + degstoregrav

        rs.erodsed[v] = erodsed
        rs.erodsedbank[v] = sedbank
        rs.erodsedbed[v] = sedbed

        ### Deposition / settling ###
        # Fractions of deposited particles in river cells from the Einstein formula [-]
        # Particle fall velocity [m/s] from Stokes
        xs = ifelse(rs.q_riv[v] > 0.0, 1.055 * rs.dl[v] / (rs.q_riv[v] / rs.width[v]), 0.0)
        xclay = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmclay[v] / 1000)^2 / 3600)))
        xsilt = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsilt[v] / 1000)^2 / 3600)))
        xsand = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsand[v] / 1000)^2 / 3600)))
        xsagg = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsagg[v] / 1000)^2 / 3600)))
        xlagg = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmlagg[v] / 1000)^2 / 3600)))
        xgrav = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmgrav[v] / 1000)^2 / 3600)))

        # Sediment deposited in the channel [ton]
        # From natural settling with Einstein formula (density controlled).
        settclay = xclay * (inclay + erodclay)
        settsilt = xsilt * (insilt + erodsilt)
        settsand = xsand * (insand + erodsand)
        settsagg = xsagg * (insagg + erodsagg)
        settlagg = xlagg * (inlagg + erodlagg)
        settgrav = xgrav * (ingrav + erodgrav)

        # Sediment deposited in the channel (from gravel to clay) [ton]
        # From transport capacity exceedance (insed > maxsed)
        insedex = max(insed - maxsed, 0.0)
        if insedex > 0.0
            depgrav = ifelse(ingrav >= insedex, insedex, ingrav)
            insedex = max(insedex - depgrav, 0.0)
            deplagg = ifelse(inlagg >= insedex, insedex, inlagg)
            insedex = max(insedex - deplagg, 0.0)
            depsand = ifelse(insand >= insedex, insedex, insand)
            insedex = max(insedex - depsand, 0.0)
            depsagg = ifelse(insagg >= insedex, insedex, insagg)
            insedex = max(insedex - depsagg, 0.0)
            depsilt = ifelse(insilt >= insedex, insedex, insilt)
            insedex = max(insedex - depsilt, 0.0)
            depclay = ifelse(inclay >= insedex, insedex, inclay)
            insedex = max(insedex - depclay, 0.0)
        else
            depclay = settclay
            depsilt = settsilt
            depsand = settsand
            depsagg = settsagg
            deplagg = settlagg
            depgrav = settgrav
        end

        # No deposition in regular lake and reservoir cells, only at the outlet
        # Deposition in lake/reservoir from Camp 1945
        # Extra trapping of large particles for dams
        if rs.wbcover[v] > 0.0 && rs.wblocs[v] > 0.0
            # Compute deposition
            vcres = rs.q_riv[v] / rs.wbarea[v]
            DCres = 411 / 3600 / vcres
            depclay = (inclay + erodclay) * min(1.0, (DCres * (rs.dmclay[v] / 1000)^2))
            depsilt = (insilt + erodsilt) * min(1.0, (DCres * (rs.dmsilt[v] / 1000)^2))
            depsand = (insand + erodsand) * min(1.0, (DCres * (rs.dmsand[v] / 1000)^2))
            depsagg = (insagg + erodsagg) * min(1.0, (DCres * (rs.dmsagg[v] / 1000)^2))
            deplagg = (inlagg + erodlagg) * min(1.0, (DCres * (rs.dmlagg[v] / 1000)^2))
            depgrav = (ingrav + erodgrav) * min(1.0, (DCres * (rs.dmgrav[v] / 1000)^2))
            # Trapping of large particles
            # Use the rouse number for sagg (suspension or bedload)
            depsand = max(depsand, rs.wbtrap[v] * (insand + erodsand))
            deplagg = max(deplagg, rs.wbtrap[v] * (inlagg + erodlagg))
            depgrav = max(depgrav, rs.wbtrap[v] * (ingrav + erodgrav))
            # threshold diameter between suspended load and mixed load using Rouse number
            dsuspf =
                1e3 * (1.2 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
            depsagg = ifelse(
                rs.dmsagg[v] > dsuspf,
                depsagg,
                max(depsagg, rs.wbtrap[v] * (insagg + erodsagg)),
            )
        elseif rs.wbcover[v] > 0.0
            depsed = 0.0
            depclay = 0.0
            depsilt = 0.0
            depsand = 0.0
            depsagg = 0.0
            deplagg = 0.0
            depgrav = 0.0
        end

        depsed = depclay + depsilt + depsand + depsagg + deplagg + depgrav
        rs.depsed[v] = depsed

        # Update the river deposited sediment storage
        rs.sedstore[v] = rs.sedstore[v] + depsed
        rs.claystore[v] = rs.claystore[v] + depclay
        rs.siltstore[v] = rs.siltstore[v] + depsilt
        rs.sandstore[v] = rs.sandstore[v] + depsand
        rs.saggstore[v] = rs.saggstore[v] + depsagg
        rs.laggstore[v] = rs.laggstore[v] + deplagg
        rs.gravstore[v] = rs.gravstore[v] + depgrav

        ### Ouput loads ###
        # Sediment transported out of the cell during the timestep [ton]
        # 0 in case all sediment are deposited in the cell
        # Reduce the fraction so that there is still some sediment staying in the river cell
        fwaterout = min(rs.q_riv[v] * rs.Δt / (rs.h_riv[v] * rs.width[v] * rs.dl[v]), 1.0)
        rs.outsed[v] = fwaterout * (insed + erodsed - depsed)
        rs.outclay[v] = fwaterout * (inclay + erodclay - depclay)
        rs.outsilt[v] = fwaterout * (insilt + erodsilt - depsilt)
        rs.outsand[v] = fwaterout * (insand + erodsand - depsand)
        rs.outsagg[v] = fwaterout * (insagg + erodsagg - depsagg)
        rs.outlagg[v] = fwaterout * (inlagg + erodlagg - deplagg)
        rs.outgrav[v] = fwaterout * (ingrav + erodgrav - depgrav)

        ### Mass balance ###
        # Sediment left in the cell [ton]
        rs.sedload[v] = insed + erodsed - depsed - rs.outsed[v]
        rs.clayload[v] = inclay + erodclay - depclay - rs.outclay[v]
        rs.siltload[v] = insilt + erodsilt - depsilt - rs.outsilt[v]
        rs.sandload[v] = insand + erodsand - depsand - rs.outsand[v]
        rs.saggload[v] = insagg + erodsagg - depsagg - rs.outsagg[v]
        rs.gravload[v] = ingrav + erodgrav - depgrav - rs.outgrav[v]

        ### Concentrations and suspended sediments ###
        # Conversion from load [ton] to concentration for rivers [mg/L]
        toconc = ifelse(rs.q_riv[v] > 0.0, 1e6 / (rs.q_riv[v] * rs.Δt), 0.0)
        rs.Sedconc[v] = rs.outsed[v] * toconc

        # Differentiation of bed and suspended load using Rouse number for suspension
        # threshold diameter between bed load and mixed load using Rouse number
        dbedf = 1e3 * (2.5 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
        # threshold diameter between suspended load and mixed load using Rouse number
        dsuspf = 1e3 * (1.2 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
        # Rouse with diameter
        SSclay = ifelse(
            rs.dmclay[v] <= dsuspf,
            rs.outclay[v],
            ifelse(rs.dmclay[v] <= dbedf, rs.outclay[v] / 2, 0.0),
        )
        SSsilt = ifelse(
            rs.dmsilt[v] <= dsuspf,
            rs.outsilt[v],
            ifelse(rs.dmsilt[v] <= dbedf, rs.outsilt[v] / 2, 0.0),
        )
        SSsagg = ifelse(
            rs.dmsagg[v] <= dsuspf,
            rs.outsagg[v],
            ifelse(rs.dmsagg[v] <= dbedf, rs.outsagg[v] / 2, 0.0),
        )
        SSsand = ifelse(
            rs.dmsand[v] <= dsuspf,
            rs.outsand[v],
            ifelse(rs.dmsand[v] <= dbedf, rs.outsand[v] / 2, 0.0),
        )
        SSlagg = ifelse(
            rs.dmlagg[v] <= dsuspf,
            rs.outlagg[v],
            ifelse(rs.dmlagg[v] <= dbedf, rs.outlagg[v] / 2, 0.0),
        )
        SSgrav = ifelse(
            rs.dmgrav[v] <= dsuspf,
            rs.outgrav[v],
            ifelse(rs.dmgrav[v] <= dbedf, rs.outgrav[v] / 2, 0.0),
        )

        SS = SSclay + SSsilt + SSsagg + SSsand + SSlagg + SSgrav
        Bed = rs.outsed[v] - SS

        rs.SSconc[v] = SS * toconc
        rs.Bedconc[v] = Bed * toconc

    end

end


text:
# add! a whole structure to the log.
function add!(log::HDF5Logger.Log, slug::String, t::Real, data, n::Int)
    function add_structured!(log, slug, data, n)
        if isbits(data) || (isa(data, Array) && isbitstype(eltype(data)))
            HDF5Logger.add!(log, slug, data, n) # HDF5Loggeres only know scalars, vectors, and matrices.
        else
            for field in fieldnames(typeof(data))
                add_structured!(log, slug * "/" * string(field), getfield(data, field), n)
            end
        end
    end
    HDF5Logger.add!(log, slug * "time", t, n)
    add_structured!(log, slug * "data", data, n)
end

# log! a whole structure.
function log!(log::HDF5Logger.Log, slug::String, t::Real, data)
    function log_structured!(log, slug, data,)
        if isbits(data) || (isa(data, Array) && isbitstype(eltype(data)))
            HDF5Logger.log!(log, slug, data)
        else
            for field in fieldnames(typeof(data))
                log_structured!(log, slug * "/" * string(field), getfield(data, field))
            end
        end
    end
    HDF5Logger.log!(log, slug * "time", t)
    log_structured!(log, slug * "data", data)
end

ext: jl


text:
"""
    AbstractRDArray{T,N} <: DenseArray{T,N}

`N`-dimensional resizable dense array with elements of type `T` with some
pre-defined array methods.
"""
abstract type AbstractRDArray{T,N} <: DenseArray{T,N} end

"""
    AbstractRNArray{T,N} <: AbstractArray{T,N}

`N`-dimensional resizable (no-dense) array with elements of type `T` with some
pre-defined array methods.
"""
abstract type AbstractRNArray{T,N} <: AbstractArray{T,N} end

const ResizableArray{T,N} = Union{AbstractRDArray{T,N},AbstractRNArray{T,N}}

Base.size(A::ResizableArray{T,N}) where {T,N} = convert(Dims{N}, getsize(A))
Base.size(A::ResizableArray, d::Integer) =
    d < 1 ? throw(BoundsError()) : d > ndims(A) ? 1 : @inbounds getsize(A, d)
# DenseArray only
Base.unsafe_convert(::Type{Ptr{T}}, A::AbstractRDArray{T}) where {T} =
    Base.unsafe_convert(Ptr{T}, parent(A))
Base.elsize(::Type{T}) where {T<:AbstractRDArray} = Base.elsize(parent_type(T))

ArrayInterface.strides(A::AbstractRDArray) = (StaticInt(1), Base.tail(strides(A))...)
ArrayInterface.strides(A::AbstractRDArray, dim::Integer) = stride(A, Int(dim))

ArrayInterface.stride_rank(::Type{T}) where {T<:AbstractRDArray} =
    Static.nstatic(Val(ndims(T)))
ArrayInterface.dense_dims(::Type{T}) where {T<:AbstractRDArray} =
    ntuple(_ -> True(), Val(ndims(T)))
ArrayInterface.axes_types(::Type{T}) where {T<:AbstractRDArray} =
    NTuple{ndims(T), Base.OneTo{Int}}


text:
const WKB_2D = Dict{Symbol,UInt32}(
    :Point=>1,
    :LineString=>2,
    :Polygon=>3,
    :MultiPoint=>4,
    :MultiLineString=>5,
    :MultiPolygon=>6,
    :GeometryCollection=>7,
)

const WKB_Z = Dict{Symbol,UInt32}(
    :Point=>1001,
    :LineString=>1002,
    :Polygon=>1003,
    :MultiPoint=>1004,
    :MultiLineString=>1005,
    :MultiPolygon=>1006,
    :GeometryCollection=>1007,
)

const WKB_M = Dict{Symbol,UInt32}(
    :Point=>2001,
    :LineString=>2002,
    :Polygon=>2003,
    :MultiPoint=>2004,
    :MultiLineString=>2005,
    :MultiPolygon=>2006,
    :GeometryCollection=>2007,
)

const WKB_ZM = Dict{Symbol,UInt32}(
    :Point=>3001,
    :LineString=>3002,
    :Polygon=>3003,
    :MultiPoint=>3004,
    :MultiLineString=>3005,
    :MultiPolygon=>3006,
    :GeometryCollection=>3007,
)

const _WKB = Dict(
    "2D"=> WKB_2D,
    "Z"=> WKB_Z,
    "M"=> WKB_M,
    "ZM"=> WKB_ZM
)

const BINARY_TO_GEOM_TYPE = begin
    d = Dict{UInt32, Symbol}()
    for wkb_map in values(_WKB)
        for k in keys(wkb_map)
            d[wkb_map[k]] = k
        end
    end
    d
end

INT_TO_DIM_LABEL = Dict(2=> "2D", 3=> "Z", 4=> "ZM")


text:
include("./read_data.jl")
include("./init.jl")
include("./functions.jl")

possible_tasks=[]
tasks_by_position=Dict{Int, Array}()
prec=Dict{Int, Array}()
tasks_time = zeros(Int8, P, C)

for p = 1:P
    l=[]
    for t=1:CP
        #create array with all tasks
        if p==P
            push!(possible_tasks, Task(cpij[t,1]+1, cpij[t,2]+1, cpij[t,3]))
            tasks_time[cpij[t,1]+1, cpij[t,2]+1] = cpij[t,3]
        end
        #create dict with the tasks sorted by position
        if cpij[t,1]+1==p
            push!(l,Task(cpij[t,1]+1, cpij[t,2]+1, cpij[t,3]))
        end
    end
    tasks_by_position[p]=l

    #create dict with the position precedences
    s=[]
    for j = 1:P
        if prejj[p,j]==1
            push!(s,j)
        end
    end
    prec[p]=s

end

max_stars_repo_licenses: Zlib MIT

max_issues_repo_licenses: Zlib MIT

max_forks_repo_licenses: Zlib MIT


text:
using FirstOrderSolvers: HSDEMatrixQ, HSDEMatrix
using SparseArrays
using ProximalOperators: IndAffine, prox!
import LinearAlgebra: mul!

function getQ1Q2(A)
    m, n = size(A)
    b = randn(m)
    c = randn(n)

    Q1 = [zeros(n,n) A'         c;
        -A         zeros(m,m) b;
        -c'        -b'        0]

    Q2 = HSDEMatrixQ(A, b, c)
    return Q1, Q2
end

function getHSDEMatrix(Q1,Q2)
    M1 = [I           sparse(Q1') ;
          sparse(Q1)  -I          ]
    M2 = HSDEMatrix(Q2)
    return M1, M2
end

function testHSDEQ_A_mul_B(Q1, Q2, m, n)
    rhs1 = randn(m+n+1)
    rhs2 = copy(rhs1)

    y1, y2 = randn(m+n+1), randn(m+n+1)

    mul!(y1, Q1, rhs1)
    mul!(y2, Q2, rhs2)

    @test rhs1 == rhs2
    @test y1 ≈ y2

    mul!(y1, transpose(Q1), rhs1)
    mul!(y2, transpose(Q2), rhs2)

    @test rhs1 == rhs2
    @test y1 ≈ y2
end

function testHSDEMatrix_A_mul_B(M1, M2, m, n)
    rhs1 = randn(2m+2n+2)
    rhs2 = copy(rhs1)

    y1, y2 = randn(2m+2n+2), randn(2m+2n+2)

    mul!(y1, M1, rhs1)
    mul!(y2, M2, rhs2)

    @test rhs1 == rhs2
    @test y1 ≈ y2

    mul!(y1, transpose(M1), rhs1)
    mul!(y2, transpose(M2), rhs2)

    @test rhs1 == rhs2
    @test y1 ≈ y2
end

function testHSDE(A, m, n)
    Q1, Q2 = getQ1Q2(A)
    testHSDEQ_A_mul_B(Q1, Q2, m, n)

    M1, M2 = getHSDEMatrix(Q1, Q2)
    testHSDEMatrix_A_mul_B(M1, M2, m, n)

    b = randn(size(M2)[1])
    S1 = IndAffine([sparse(Q1) -I], zeros(size(Q1,1)))
    y1 = similar(b)
    y2 = similar(b)
    FirstOrderSolvers.prox!(y2, M2, b)
    prox!(y1, S1, b)

    y3 = M1\b
    y3[(size(Q1,1)+1):end] .= Q1*y3[1:size(Q1,1)]
    @test y1 ≈ y2
    @test y2 ≈ y3
end

Random.seed!(1)
ma,na = 10,20
A = randn(10*ma,10*na)
testHSDE(A, size(A)...)

A = sprandn(100*ma,100*na,0.001)
testHSDE(A, size(A)...)


text:
# Particularly useful or beautiful grammar of graphics invocations.

# A convenience plot function for quickly plotting functions or expressions.
#
# Args:
#   fs: An array in which each object is either a single argument function or an
#       expression computing some value on x.
#   a: Lower bound on x.
#   b: Upper bound on x.
#   elements: One ore more grammar elements.
#
# Returns:
#   A plot objects.
#
function plot(fs::Vector{T}, a::Number, b::Number, elements::ElementOrFunction...; mapping...) where T <: Base.Callable
    if isempty(elements)
        elements = ElementOrFunction[]
    elseif isa(elements, Tuple)
        elements = ElementOrFunction[elements...]
    end

    element_types = Set(map(typeof, elements))

    if !in(Guide.xlabel, element_types)
        push!(elements, Guide.xlabel("x"))
    end

    if !in(Guide.ylabel, element_types)
        push!(elements, Guide.ylabel("f(x)"))
    end

    if b < a
        push!(elements, Coord.cartesian(xflip=true))
    end

    mappingdict = Dict{Symbol, Any}(:y => fs, :xmin => [a], :xmax => [b])
    for (k, v) in mapping
        mappingdict[k] = v
    end

    plot(Stat.func, Geom.line, elements...; mappingdict...)
end

# Plot a single function.
plot(f::Function, a::Number, b::Number, elements::ElementOrFunction...; mapping...) =
        plot(Function[f], a, b, elements...; mapping...)

# Plot a single function using a contour plot
function plot(f::Function, xmin::Number, xmax::Number, ymin::Number, ymax::Number,
              elements::ElementOrFunction...; mapping...)
    default_elements = ElementOrFunction[]
    element_types = Set(map(typeof, elements))

    if !in(Guide.xlabel, element_types)
        push!(default_elements, Guide.xlabel("x"))
    end

    if !in(Guide.ylabel, element_types)
        push!(default_elements, Guide.ylabel("y"))
    end

    if !in(Guide.ColorKey, element_types) && !in(Guide.ManualColorKey, element_types)
        push!(default_elements, Guide.colorkey(title="f(x,y)"))
    end

    push!(default_elements, Coord.cartesian(xflip=xmin > xmax, yflip=ymin > ymax))

    plot(layer(f, xmin, xmax, ymin, ymax,
               elements...; mapping...), default_elements...)
end

function layer(f::Function, xmin::Number, xmax::Number, ymin::Number, ymax::Number,
               elements::ElementOrFunction...; mapping...)
    if isempty(elements)
        elements = ElementOrFunction[]
    elseif isa(elements, Tuple)
        elements = ElementOrFunction[elements...]
    end

    mappingdict = Dict{Symbol, Any}(:z    => f, :xmin => [xmin], :xmax => [xmax],
                                            :ymin => [ymin], :ymax => [ymax])
    for (k, v) in mapping
        mappingdict[k] = v
    end

    layer(Geom.contour, elements...; mappingdict...)
end

# Create a layer from a list of functions or expressions.
layer(fs::Array, a::Number, b::Number, elements::ElementOrFunction...) =
        layer(y=fs, xmin=[a], xmax=[b], Stat.func, Geom.line, elements...)

# Create a layer from a single function.
layer(f::Function, a::Number, b::Number, elements::ElementOrFunction...) =
        layer(Function[f], a, b, elements...)

# Simple heatmap plots of matrices.
#
# It is a wrapper around the `plot()` function using the `rectbin` geometry.
# It also applies a sane set of defaults to make sure that the plots look nice
# by default. Specifically
#   - the aspect ratio of the coordinate system is fixed Coord.cartesian(fixed=true),
#     so that the rectangles become squares
#   - the axes run from 0.5 to N+0.5, because the first row/column is drawn to
#     (0.5, 1.5) and the last one to (N-0.5, N+0.5).
#   - the y-direction is flipped, so that the [1,1] of a matrix is in the top
#     left corner, as is customary
#   - NaNs are not drawn. `spy` leaves "holes" instead into the heatmap.
#
# Args:
#   M: A matrix.
#
# Returns:
#   A plot object.
#
# Known bugs:
#   - If the matrix is only NaNs, then it throws an `ArgumentError`, because
#     an empty collection gets passed to the `plot` function / `rectbin` geometry.
#
"""
```
spy(M::AbstractMatrix, elements::ElementOrFunction...; mapping...)
```
Simple heatmap plots of matrices.

It is a wrapper around the `plot()` function using the `rectbin` geometry.
It also applies a sane set of defaults to make sure that the plots look nice
by default. Specifically
- the aspect ratio of the coordinate system is fixed Coord.cartesian(fixed=true),
so that the rectangles become squares
- the axes run from 0.5 to N+0.5, because the first row/column is drawn to
(0.5, 1.5) and the last one to (N-0.5, N+0.5).
- the y-direction is flipped, so that the [1,1] of a matrix is in the top
left corner, as is customary
- NaNs are not drawn. `spy` leaves "holes" instead into the heatmap.

### Args:
* M: A matrix.

### Returns:
A plot object.

#### Known bugs:
   - If the matrix is only NaNs, then it throws an `ArgumentError`, because
     an empty collection gets passed to the `plot` function / `rectbin` geometry.
"""

function spy(M::AbstractMatrix, elements::ElementOrFunction...; mapping...)
    is, js, values = _findnz(x->!isnan(x), M)
    n,m = size(M)
    df = DataFrames.DataFrame(i=is, j=js, value=values)
    plot(df, x=:j, y=:i, color=:value,
        Coord.cartesian(yflip=true, fixed=true, xmin=0.5, xmax=m+.5, ymin=0.5, ymax=n+.5),
        Scale.color_continuous,
        Geom.rectbin,
        Scale.x_continuous,
        Scale.y_continuous,
        elements...; mapping...)
end

# Finds the subscripts and values where the predicate returns true.
#
# It takes a predicate (`testf`), a matrix (`A`) and returns a tuple `(is, js, zs)`,
# where `is`,'js' and 'zs' are arrays of subscripts and values where the predicate
# returned true.
#
# This function is used by spy()
# Hopefully at some point something like this will be in the standard library
# and then this can be removed (https://github.com/JuliaLang/julia/pull/9340).
#
function _findnz(testf::Function, A::AbstractMatrix{T}) where T
    N = Base.count(testf, A)
    is = zeros(Int, N)
    js = zeros(Int, N)
    zs = Array{T}(N)
    if N == 0
        return (is, js, zs)
    end
    count = 1
    for j=1:size(A,2), i=1:size(A,1)
        Aij = A[i,j]
        if testf(Aij)
            is[count] = i
            js[count] = j
            zs[count] = Aij
            count += 1
        end
    end
    return (is, js, zs)
end


text:
using Random, Distributions, Printf
using LinearAlgebra: dot, ⋅

numactions = 2

mutable struct Node
    regretsum::Vector{AbstractFloat}
    strategy::Vector{AbstractFloat}
    strategysum::Vector{AbstractFloat}
end

function Base.show(io::IO, n::Node)
    sprintarr(a) = map(x -> @sprintf("%0.2f, ", x), a) |> join
    @printf("r: %s s: %s ss: %s\n", sprintarr(n.regretsum),
        sprintarr(n.strategy), sprintarr(n.strategysum))
end

struct Infoset
    hand::Vector{UInt8}
    history::Vector{Char}
end

Base.show(io::IO, i::Infoset) = print(i.hand[1], String(i.history))
Base.hash(a::Infoset, h::UInt) = hash(a.history, hash(a.hand, hash(:Infoset, h)))
Base.isequal(i1::Infoset, i2::Infoset) = i1.hand == i2.hand && i1.history == i2.history
Base.isless(i1::Infoset, i2::Infoset) = i1.hand < i2.hand

function Base.show(io::IO, d::Dict{Infoset, Node})
    for (k, v) in d
        print(k, " => ", v)
    end
end

pos(x) = x > 0.0 ? x : 0.0

issomething(s) = !isnothing(s)

function normalize(xs)
    total = sum(xs)
    n = length(xs)
    total > 0 ? xs / total : fill(1.0 / n, n)
end

function getstrategy(node, realizationweight)
    node.strategy = map(pos, node.regretsum) |> normalize
    node.strategysum += realizationweight * node.strategy
    node.strategy
end

function getnode(policys, infoset)
    if !haskey(policys, infoset)
        node = Node(zeros(numactions), zeros(numactions), zeros(numactions))
        policys[infoset] = node
    end
    policys[infoset]
end

function terminal(history, cards)
    plays = length(history)
    if plays > 1
        player = plays % 2 + 1
        opponent = player % 2 + 1
        playercardhigher = cards[player] > cards[opponent]
        if history[end] == 'p' # terminal pass
            if history == ['p', 'p']
                return playercardhigher ? 1 : -1
            else
                return 1
            end
        elseif history[end-1:end] == ['b', 'b'] # double bet
            return playercardhigher ? 2 : -2
        end
    end
    return nothing
end

function cfr(policys::Dict{Infoset, Node}, cards, history, p1, p2)
    terminalutility = terminal(history, cards)
    if issomething(terminalutility)
        return terminalutility
    end

    player = length(history) % 2 + 1
    prob = [p1, p2][player]
    infoset = Infoset([cards[player]], history)
    node = getnode(policys, infoset)
    strategy = getstrategy(node, prob)
    util = zeros(numactions)
    for a in 1:numactions
        nexthistory = copy(history)
        push!(nexthistory, ['p', 'b'][a])
        util[a] = player == 1 ?
            -cfr(policys, cards, nexthistory, p1 * strategy[a], p2) :
            -cfr(policys, cards, nexthistory, p1, p2 * strategy[a])
    end
    nodeutil = strategy ⋅ util
    regret = util .- nodeutil
    node.regretsum += [p2, p1][player] * regret
    nodeutil
end

function train(policys::Dict{Infoset, Node}, n)
    cards = [1, 2, 3]
    util = 0.0
    for _ in 1:n
        shuffle!(cards)
        util += cfr(policys, cards, [], 1.0, 1.0)
    end
    println(util / n)
    ps = Dict()
    for (k, v) in policys
        s = v.strategysum
        ps[k] = normalize(s)
    end
    ps
end

function displayresult(d)
    for (k, v) in d
        print(k, " => ")
        map(x -> @printf("%0.2f, ", x), v)
        println()
    end
end

max_stars_repo_name: jeffreyrosenbluth/LP

max_issues_repo_name: jeffreyrosenbluth/LP

max_forks_repo_name: jeffreyrosenbluth/LP


text:
module BLS12381FieldPerf

using BenchmarkTools

BenchmarkTools.DEFAULT_PARAMETERS.samples = 1000
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 10
BenchmarkTools.DEFAULT_PARAMETERS.evals = 1
BenchmarkTools.DEFAULT_PARAMETERS.gctrial = true
BenchmarkTools.DEFAULT_PARAMETERS.gcsample = false
@show BenchmarkTools.DEFAULT_PARAMETERS

using TowerFields.BLS12381Field: add, sub, mul, inv, pow, frob, powx, square
using ..TestUtils: randfq, randfq2, randfq6, randfq12

function format_trial(suite, group, res)
    a = allocs(res)
    gct = BenchmarkTools.prettytime(gctime(res))
    t = BenchmarkTools.prettytime(time(res))
    m = BenchmarkTools.prettymemory(memory(res))
    return "[$suite] $group: $t (alloc: $a, mem: $m, gc: $gct)"
end

# Add some child groups to our benchmark suite.
suite = BenchmarkGroup()

suite["FQ"] = BenchmarkGroup()
suite["FQ"]["add"] = @benchmarkable add($(randfq()), $(randfq()))
suite["FQ"]["sub"] = @benchmarkable sub($(randfq()), $(randfq()))
suite["FQ"]["mul"] = @benchmarkable mul($(randfq()), $(randfq()))
suite["FQ"]["square"] = @benchmarkable square($(randfq()))
suite["FQ"]["inv"] = @benchmarkable inv(($(randfq())))
suite["FQ"]["pow"] = @benchmarkable pow($(randfq()), $(randfq()))
suite["FQ"]["frob"] = @benchmarkable frob($(randfq()))

suite["FQ2"] = BenchmarkGroup()
suite["FQ2"]["add"] = @benchmarkable add($(randfq2()), $(randfq2()))
suite["FQ2"]["sub"] = @benchmarkable sub($(randfq2()), $(randfq2()))
suite["FQ2"]["mul"] = @benchmarkable mul($(randfq2()), $(randfq2()))
suite["FQ2"]["square"] = @benchmarkable square($(randfq2()))
suite["FQ2"]["inv"] = @benchmarkable inv(($(randfq2())))
suite["FQ2"]["pow"] = @benchmarkable pow($(randfq2()), $(randfq()))
suite["FQ2"]["powx"] = @benchmarkable powx($(randfq2()))
suite["FQ2"]["frob"] = @benchmarkable frob($(randfq2()))

suite["FQ6"] = BenchmarkGroup()
suite["FQ6"]["add"] = @benchmarkable add($(randfq6()), $(randfq6()))
suite["FQ6"]["sub"] = @benchmarkable sub($(randfq6()), $(randfq6()))
suite["FQ6"]["mul"] = @benchmarkable mul($(randfq6()), $(randfq6()))
suite["FQ6"]["square"] = @benchmarkable square($(randfq6()))
suite["FQ6"]["inv"] = @benchmarkable inv(($(randfq6())))
suite["FQ6"]["pow"] = @benchmarkable pow($(randfq6()), $(randfq()))
suite["FQ6"]["frob"] = @benchmarkable frob($(randfq6()))

suite["FQ12"] = BenchmarkGroup()
suite["FQ12"]["add"] = @benchmarkable add($(randfq12()), $(randfq12()))
suite["FQ12"]["sub"] = @benchmarkable sub($(randfq12()), $(randfq12()))
suite["FQ12"]["mul"] = @benchmarkable mul($(randfq12()), $(randfq12()))
suite["FQ12"]["square"] = @benchmarkable square($(randfq12()))
suite["FQ12"]["inv"] = @benchmarkable inv(($(randfq12())))
suite["FQ12"]["pow"] = @benchmarkable pow($(randfq12()), $(randfq()))
suite["FQ12"]["powx"] = @benchmarkable powx($(randfq12()))
suite["FQ12"]["frob"] = @benchmarkable frob($(randfq12()))

# If a cache of tuned parameters already exists, use it, otherwise, tune and cache
# the benchmark parameters. Reusing cached parameters is faster and more reliable
# than re-tuning `suite` every time the file is included.
paramspath = joinpath(@__DIR__, "params.json")
if isfile(paramspath)
    loadparams!(suite, BenchmarkTools.load(paramspath)[1], :evals);
else
    println("First run - tuning params (please be patient) ...")
    tune!(suite)
    BenchmarkTools.save(paramspath, params(suite));
end

# print the results
results = run(suite, verbose = true)
for suiteres in results
    for groupres in suiteres.second
        msg = format_trial(suiteres.first, groupres.first, groupres.second)
        println(msg)
    end
end

end


text:
using Test
using Manifolds
using Manifolds:
    _derivative,
    _derivative!,
    differential,
    differential!,
    gradient,
    gradient!,
    _gradient,
    _gradient!

using FiniteDifferences
using LinearAlgebra: Diagonal, dot

@testset "Differentiation backend" begin
    fd51 = Manifolds.FiniteDifferencesBackend()
    @testset "diff_backend" begin
        @test diff_backend() isa Manifolds.FiniteDifferencesBackend
        @test length(diff_backends()) == 1
        @test diff_backends()[1] isa Manifolds.FiniteDifferencesBackend

        @test length(fd51.method.grid) == 5
        @test fd51.method.q == 1
        fd71 = Manifolds.FiniteDifferencesBackend(central_fdm(7, 1))
        @test diff_backend!(fd71) == fd71
        @test diff_backend() == fd71
    end

    using ForwardDiff

    fwd_diff = Manifolds.ForwardDiffBackend()
    @testset "ForwardDiff" begin
        @test diff_backend() isa Manifolds.FiniteDifferencesBackend
        @test length(diff_backends()) == 2
        @test diff_backends()[1] isa Manifolds.FiniteDifferencesBackend
        @test diff_backends()[2] == fwd_diff

        @test diff_backend!(fwd_diff) == fwd_diff
        @test diff_backend() == fwd_diff
        @test diff_backend!(fd51) isa Manifolds.FiniteDifferencesBackend
        @test diff_backend() isa Manifolds.FiniteDifferencesBackend

        diff_backend!(fwd_diff)
        @test diff_backend() == fwd_diff
        diff_backend!(fd51)
    end

    using FiniteDiff

    finite_diff = Manifolds.FiniteDiffBackend()
    @testset "FiniteDiff" begin
        @test diff_backend() isa Manifolds.FiniteDifferencesBackend
        @test length(diff_backends()) == 3
        @test diff_backends()[3] == finite_diff

        @test diff_backend!(finite_diff) == finite_diff
        @test diff_backend() == finite_diff
        @test diff_backend!(fd51) isa Manifolds.FiniteDifferencesBackend
        @test diff_backend() isa Manifolds.FiniteDifferencesBackend

        diff_backend!(finite_diff)
        @test diff_backend() == finite_diff
        diff_backend!(fd51)
    end

    @testset "gradient" begin
        diff_backend!(fd51)
        r2 = Euclidean(2)

        c1(t) = [sin(t), cos(t)]
        f1(x) = x[1] + x[2]^2
        f2(x) = 3 * x[1] * x[2] + x[2]^3

        @testset "Inference" begin
            v = [-1.0, -1.0]
            @test (@inferred _derivative(c1, 0.0, Manifolds.ForwardDiffBackend())) ≈
                  [1.0, 0.0]
            @test (@inferred _derivative!(c1, v, 0.0, Manifolds.ForwardDiffBackend())) === v
            @test v ≈ [1.0, 0.0]

            @test (@inferred _derivative(c1, 0.0, finite_diff)) ≈ [1.0, 0.0]
            @test (@inferred _gradient(f1, [1.0, -1.0], finite_diff)) ≈ [1.0, -2.0]
        end

        @testset for backend in [fd51, fwd_diff, finite_diff]
            diff_backend!(backend)
            @test _derivative(c1, 0.0) ≈ [1.0, 0.0]
            v = [-1.0, -1.0]
            @test _derivative!(c1, v, 0.0) === v
            @test isapprox(v, [1.0, 0.0])
            @test _gradient(f1, [1.0, -1.0]) ≈ [1.0, -2.0]
            @test _gradient!(f1, v, [1.0, -1.0]) === v
            @test v ≈ [1.0, -2.0]
        end
        diff_backend!(Manifolds.NoneDiffBackend())
        @testset for backend in [fd51, Manifolds.ForwardDiffBackend()]
            @test _derivative(c1, 0.0, backend) ≈ [1.0, 0.0]
            @test _gradient(f1, [1.0, -1.0], backend) ≈ [1.0, -2.0]
        end

        diff_backend!(fd51)
    end
end

rb_onb_default = RiemannianONBDiffBackend(
    diff_backend(),
    Manifolds.ExponentialRetraction(),
    Manifolds.LogarithmicInverseRetraction(),
    DefaultOrthonormalBasis(),
)

rb_onb_fd51 = RiemannianONBDiffBackend(
    Manifolds.FiniteDifferencesBackend(),
    Manifolds.ExponentialRetraction(),
    Manifolds.LogarithmicInverseRetraction(),
    DefaultOrthonormalBasis(),
)

rb_onb_fwd_diff = RiemannianONBDiffBackend(
    Manifolds.ForwardDiffBackend(),
    Manifolds.ExponentialRetraction(),
    Manifolds.LogarithmicInverseRetraction(),
    DefaultOrthonormalBasis(),
)

rb_onb_finite_diff = RiemannianONBDiffBackend(
    Manifolds.FiniteDiffBackend(),
    Manifolds.ExponentialRetraction(),
    Manifolds.LogarithmicInverseRetraction(),
    DefaultOrthonormalBasis(),
)

rb_onb_default2 = RiemannianONBDiffBackend(
    diff_backend(),
    Manifolds.ExponentialRetraction(),
    Manifolds.LogarithmicInverseRetraction(),
    CachedBasis(
        DefaultOrthonormalBasis(),
        [[0.0, -1.0, 0.0], [sqrt(2) / 2, 0.0, -sqrt(2) / 2]],
    ),
)

rb_proj = Manifolds.RiemannianProjectionGradientBackend(diff_backend())

@testset "rdiff_ functions" begin
    @test Manifolds.rdifferential_backend() ===
          Manifolds._current_rdifferential_backend.backend
    @test Manifolds.rgradient_backend() === Manifolds._current_rgradient_backend.backend

    tmp_diff = Manifolds.rdifferential_backend()
    Manifolds.rdifferential_backend!(rb_onb_finite_diff)
    @test Manifolds.rdifferential_backend() === rb_onb_finite_diff
    Manifolds.rdifferential_backend!(tmp_diff)

    tmp_grad = Manifolds.rgradient_backend()
    Manifolds.rgradient_backend!(rb_onb_finite_diff)
    @test Manifolds.rgradient_backend() === rb_onb_finite_diff
    Manifolds.rgradient_backend!(tmp_grad)
end

@testset "Riemannian differentials" begin
    s2 = Sphere(2)
    p = [0.0, 0.0, 1.0]
    q = [1.0, 0.0, 0.0]
    c1(t) = geodesic(s2, q, p, t)

    Xval = [-sqrt(2) / 2, 0.0, sqrt(2) / 2]
    @test isapprox(s2, c1(π / 4), differential(s2, c1, π / 4), Xval)
    X = similar(p)
    differential!(s2, c1, X, π / 4)
    @test isapprox(s2, c1(π / 4), X, Xval)

    @testset for backend in [rb_onb_fd51, rb_onb_fwd_diff, rb_onb_finite_diff]
        @test isapprox(s2, c1(π / 4), differential(s2, c1, π / 4, backend), Xval)
        X = similar(p)
        differential!(s2, c1, X, π / 4, backend)
        @test isapprox(s2, c1(π / 4), X, Xval)
    end
end

@testset "Riemannian gradients" begin
    s2 = Sphere(2)
    f1(p) = p[1]

    q = [sqrt(2) / 2, 0, sqrt(2) / 2]
    @test isapprox(s2, q, gradient(s2, f1, q), [0.5, 0.0, -0.5])
    for backend in [rb_onb_default, rb_proj]
        @test isapprox(s2, q, gradient(s2, f1, q, backend), [0.5, 0.0, -0.5])
    end
    X = similar(q)
    gradient!(s2, f1, X, q)
    @test isapprox(s2, q, X, [0.5, 0.0, -0.5])
    for backend in [rb_onb_default, rb_proj]
        gradient!(s2, f1, X, q, backend)
        @test isapprox(s2, q, X, [0.5, 0.0, -0.5])
    end
end


text:
using SIIPExamples

using PowerSystems
using PowerSimulations

using Cbc #solver
solver = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 1, "ratioGap" => 0.5)

pkgpath = dirname(dirname(pathof(SIIPExamples)))
include(joinpath(pkgpath, "test", "2_PowerSystems_examples", "04_parse_tabulardata.jl"))

for line in get_components(Line, sys)
    if (get_base_voltage(get_from(get_arc(line))) >= 230.0) &&
       (get_base_voltage(get_to(get_arc(line))) >= 230.0)
        #if get_area(get_from(get_arc(line))) != get_area(get_to(get_arc(line)))
        @info "Changing $(get_name(line)) to MonitoredLine"
        convert_component!(MonitoredLine, line, sys)
    end
end

uc_prob =
    UnitCommitmentProblem(sys, optimizer = solver, horizon = 24, network = DCPPowerModel)

set_branch_model!(uc_prob, :L, DeviceModel(Line, StaticLineUnbounded))

construct_device!(uc_prob, :ML, DeviceModel(MonitoredLine, StaticLine))

solve!(uc_prob)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


text:
optional_BatchNorm(D, σ, isnorm; kwargs...) = isnorm ? BatchNorm(D, σ; kwargs...) : x -> σ.(x)

### DenseNet

const IntIte = Union{AbstractVector{Int},Tuple{Vararg{Int}}}

struct DenseNet <: AbstractNeuralModel
    f
end

Flux.@functor DenseNet

DenseNet(args...; kwargs...) = DenseNet(build_densenet(args...; kwargs...))

(m::DenseNet)(x::AbstractMatrix) = m.f(x)

(m::DenseNet)(x::AbstractArray) = m(reshape(x, prod(Base.front(size(x))), last(size(x))))

function build_densenet(Dhs::IntIte, σs; isnorm::Bool=false)
    @assert length(σs) == length(Dhs) - 1 "Length of `σs` should be greater than the length of `Dhs` by 1."
    layers = []
    for i in 1:length(Dhs)-2
        push!(layers, Dense(Dhs[i], Dhs[i+1]))
        push!(layers, optional_BatchNorm(Dhs[i+1], σs[i], isnorm))
    end
    push!(layers, Dense(Dhs[end-1], Dhs[end], σs[end]))
    return Chain(layers...)
end

build_densenet(Dhs::IntIte, σ::Function, σlast::Function=identity; kwargs...) =
    build_densenet(Dhs, (fill(σ, length(Dhs) - 2)..., σlast); kwargs...)

build_densenet(Din::Int, Dhs::IntIte, arg...; kwargs...) =
    build_densenet([Din, Dhs...], arg...; kwargs...)
build_densenet(Dhs::IntIte, Dout::Int, arg...; kwargs...) =
    build_densenet([Dhs..., Dout], arg...; kwargs...)
build_densenet(Din::Int, Dhs::IntIte, Dout::Int, arg...; kwargs...) =
    build_densenet([Din, Dhs..., Dout], arg...; kwargs...)
build_densenet(Din::Int, Dout::Int, arg...; kwargs...) =
    build_densenet([Din, Dout], arg...; kwargs...)

### ConvNet

struct ConvNet{Tin<:Union{Int, NTuple{3, Int}}, Tout<:Union{Int, NTuple{3, Int}}} <: AbstractNeuralModel
    Sin::Tin
    f
    Sout::Tout
end

Flux.@functor ConvNet

## Conv in

function ConvNet(WHCin::NTuple{3, Int}, Dout::Int, args...; kwargs...)
    if WHCin == (28, 28, 1)
        f = build_convin_mnist(Dout, args...; kwargs...)
    elseif WHCin == (32, 32, 3)
        f = build_convin_incifar(Dout, args...; kwargs...)
    else
        throw(ErrorException("Unsupported input and output size for `ConvNet`: `WHCin`=$WHCin, `Dout`=$Dout."))
    end
    return ConvNet(WHCin, f, Dout)
end

function (m::ConvNet{NTuple{3, Int}, Int})(x::AbstractArray{<:Real,4})
    let (W, H, C) = m.Sin
        W != size(x, 1) && throw(DimensionMismatch("`Sin[1]` ($W) != `size(x, 1)` ($(size(x, 1)))"))
        H != size(x, 2) && throw(DimensionMismatch("`Sin[2]` ($H) != `size(x, 2)` ($(size(x, 2)))"))
        C != size(x, 3) && throw(DimensionMismatch("`Sin[3]` ($C) != `size(x, 3)` ($(size(x, 3)))"))
    end
    return m.f(x)
end

(m::ConvNet{NTuple{3, Int}, Int})(x::AbstractMatrix) = m(reshape(x, m.Sin..., size(x, 2)))

function build_convin_mnist(Dout::Int, σs; isnorm::Bool=false)
    @assert length(σs) == 4 "Length of `σs` must be `4` for `build_convin_mnist`"
    return Chain(
        #    28 x 28 x  1 x B
        Conv((3, 3),  1 => 16; pad=(1, 1)), optional_BatchNorm(16, σs[1], isnorm), MaxPool((2, 2)),
        # -> 14 x 14 x 16 x B
        Conv((3, 3), 16 => 32; pad=(1, 1)), optional_BatchNorm(32, σs[2], isnorm), MaxPool((2, 2)),
        # ->  7 x  7 x 32 x B
        Conv((3, 3), 32 => 32; pad=(1, 1)), optional_BatchNorm(32, σs[3], isnorm), MaxPool((2, 2)),
        # ->  3 x  3 x 32 x B
        x -> reshape(x, :, size(x, 4)),
        # ->  288 x B
        Dense(288, Dout, σs[4])
    )
end

build_convin_mnist(Dout::Int, σ::Function, σlast::Function=identity; kwargs...) =
    build_convin_mnist(Dout, (σ, σ, σ, σlast); kwargs...)

#  function build_convin_incifar(Dout::Int, σs; isnorm::Bool=false)
#      @assert length(σs) == 5 "Length of `σs` must be `5` for `build_convin_incifar`"
#      return Chain(
#          #    32 x 32 x   3 x B
#          Conv((4, 4),  3 =>  32; stride=(2, 2), pad=(1, 1)), optional_BatchNorm( 32, σs[1], isnorm),
#          # -> 16 x 16 x  32 x B
#          Conv((4, 4), 32 =>  64; stride=(2, 2), pad=(1, 1)), optional_BatchNorm( 64, σs[2], isnorm),
#          # ->  8 x  8 x  64 x B
#          Conv((4, 4), 64 => 128; stride=(2, 2), pad=(1, 1)), optional_BatchNorm(128, σs[3], isnorm),
#          # ->  4 x  4 x 128 x B
#          x -> reshape(x, :, size(x, 4)),
#          # -> 2048 x B
#          Dense(2048, 512), optional_BatchNorm(512, σs[4], isnorm),
#          # ->  512 x B
#          Dense(512, Dout, σs[5])
#          # -> Dout x B
#      )
#  end

#  function build_convin_incifar(Dout::Int, σs; isnorm::Bool=false)
#      @assert length(σs) == 4 "Length of `σs` must be `4` for `build_convin_incifar`"
#      return Chain(
#          #    32 x 32 x   3 x B
#          Conv((4, 4),  3 =>  32; pad=(1, 1)), optional_BatchNorm( 32, σs[1], isnorm), MaxPool((2, 2)),
#          # -> 15 x 15 x  32 x B
#          Conv((4, 4), 32 =>  64; pad=(1, 1)), optional_BatchNorm( 64, σs[2], isnorm), MaxPool((2, 2)),
#          # ->  7 x  7 x  64 x B
#          Conv((4, 4), 64 => 128; pad=(1, 1)), optional_BatchNorm(128, σs[3], isnorm), MaxPool((2, 2)),
#          # ->  3 x  3 x 128 x B
#          x -> reshape(x, :, size(x, 4)),
#          # -> 1152 x B
#          Dense(1152, Dout, σs[4])
#          # -> Dout x B
#      )
#  end

# Akash version

function build_convin_incifar(Dout::Int, σs; isnorm::Bool=false)
    @assert length(σs) == 4 "Length of `σs` must be `4` for `build_convin_incifar`"
    return Chain(
        #    32 x 32 x   3 x B
        Conv((4, 4),  3 =>  32, σs[1]; stride=(2, 2), pad=(1, 1)),
        # -> 16 x 16 x  32 x B
        Conv((4, 4), 32 =>  64; stride=(2, 2), pad=(1, 1)), optional_BatchNorm( 64, σs[2], isnorm),
        # ->  8 x  8 x  64 x B
        Conv((4, 4), 64 => 128; stride=(2, 2), pad=(1, 1)), optional_BatchNorm(128, σs[3], isnorm),
        # ->  4 x  4 x 128 x B
        x -> reshape(x, :, size(x, 4)),
        # -> 2048 x B
        Dense(2048, Dout, σs[4])
        # -> Dout x B
    )
end

build_convin_incifar(Dout::Int, σ::Function, σlast::Function=identity=identity; kwargs...) =
    build_convin_incifar(Dout, (σ, σ, σ, σlast); kwargs...)

## Conv out

function ConvNet(Din::Int, WHCout::Tuple{Int,Int,Int}, args...; kwargs...)
    if WHCout == (28, 28, 1)
        f = build_convout_mnist(Din, args...; kwargs...)
    elseif WHCout == (32, 32, 3)
        f = build_convout_cifar(Din, args...; kwargs...)
    else
        throw(ErrorException("Unsupported input and output size for `ConvNet`: `Din`=$Din, `WHCout`=$WHCout."))
    end
    return ConvNet(Din, f, WHCout)
end

(m::ConvNet{Int, NTuple{3, Int}})(x) = m.f(x)

function build_convout_mnist(Din::Int, σs; isnorm::Bool=false)
    @assert length(σs) == 5 "Length of `σs` must be `5` for `build_convout_mnist`"
    return Chain(
        #     Din x B
        Dense(Din,  288), optional_BatchNorm(288, σs[1], isnorm),
        #     288 x B
        Dense(288, 1024),
        # -> 1024 x B
        x -> reshape(x, 4, 4, 64, last(size(x))), optional_BatchNorm(64, σs[2], isnorm),
        # ->    4 x  4 x 64 x B
        ConvTranspose((4, 4), 64 => 32; stride=(1, 1), pad=(0, 0)), optional_BatchNorm(32, σs[3], isnorm),
        # ->    7 x  7 x 32 x B
        ConvTranspose((4, 4), 32 => 16; stride=(2, 2), pad=(1, 1)), optional_BatchNorm(16, σs[4], isnorm),
        # ->   15 x 15 x 16 x B
        ConvTranspose((4, 4), 16 =>  1; stride=(2, 2), pad=(1, 1)), x -> σs[5].(x)
        # ->   28 x 28 x  1 x B
    )
end

build_convout_mnist(Din::Int, σ::Function, σlast::Function=identity; kwargs...) =
    build_convout_mnist(Din, (σ, σ, σ, σ, σlast); kwargs...)

# FIXME: make it 5 layers
function build_convout_cifar(Din::Int, σs; isnorm::Bool=false)
    @assert length(σs) == 4 "Length of `σs` must be `4` for `build_convnet_outcifar10`"
    return Chain(
        #     Din x B
        Dense(Din,  2048),
        # -> 2048 x B
        x -> reshape(x, 4, 4, 128, size(x, 2)),
        optional_BatchNorm(128, σs[1], isnorm; momentum=9f-1),
        # ->    4 x  4 x 128 x B
        ConvTranspose((4, 4), 128 => 64; stride=(2, 2), pad=(1, 1)),
        optional_BatchNorm(64, σs[2], isnorm; momentum=9f-1),
        # ->    8 x  8 x  64 x B
        ConvTranspose((4, 4),  64 => 32; stride=(2, 2), pad=(1, 1)),
        optional_BatchNorm(32, σs[3], isnorm; momentum=9f-1),
        # ->   16 x 16 x  64 x B
        ConvTranspose((4, 4),  32 =>  3; stride=(2, 2), pad=(1, 1)), x -> σs[4].(x)
        # ->   32 x 32 x   3 x B
    )
end

build_convout_cifar(Din::Int, σ::Function, σlast::Function; kwargs...) =
    build_convout_cifar(Din, (σ, σ, σ, σlast); kwargs...)


text:
printstyled("  parsing\n", color=:light_green)
import SeisIO: string_time, stream_float, stream_time
date_buf = BUF.date_buf

s = [ "2010",
      "2010-01-03",
      "2012-03-01",
      "2010-11-17",
      "2010-11-17T20:20:10",
      "2010-11-17T20:20:01",
      "2010-11-17T20:20:00.100000Z",
      "2010-11-17T20:20:00.010000Z",
      "2010-11-17T20:20:00.001000Z",
      "2010-11-17T20:20:00.000100Z",
      "2010-11-17T20:20:00.000010Z",
      "2010-11-17T20:20:00.000001Z" ]

t = [ 1262304000000000,
      1262476800000000,
      1330560000000000,
      1289952000000000,
      1290025210000000,
      1290025201000000,
      1290025200100000,
      1290025200010000,
      1290025200001000,
      1290025200000100,
      1290025200000010,
      1290025200000001 ]

for (i,j) in enumerate(s)
  @test string_time(j, date_buf) == t[i]
  @test stream_time(IOBuffer(j*"~"), date_buf) == t[i]
end

float_strings = ["49.981",
                 ".9183948913749817",
                 "1232.0",
                 "1232.34a",
                 "0",
                 "313,",
                 "-12345",
                 "-12345.0",
                 "3.1E19",
                 "1.23E-08",
                 "+1.23E+8",
                 "1.234E-8",
                 "1.234E+8",
                 "3.4028235f38",  # largest finite Float32
                 "-3.4028235f38", # smallest finite Float32
                 "1.0f-45"]

float_vals = Float32[49.981,
  0.91839486f0,
  1232.0,
  1232.34,
  0.0,
  313.0,
  -12345.0,
  -12345.0,
  3.1f19,
  1.23f-8,
  1.23f8,
  1.234f-8,
  1.234f8,
  3.4028235f38,
  -3.4028235f38,
  1.0f-45]

for (i,j) in enumerate(float_strings)
  # println("test ", i, ": ", j)
  b = IOBuffer(j*"~")
  @test stream_float(b, 0x00) ≈ float_vals[i]
  close(b)
end

bad_floats = ["49.9.81",
  ".9f",
 ".9fff183948913749817",
 "+-+-+-1232.0",
 "12a32.34a",
 "threeve",
 "Texa\$",
 "13.13.13.",
 "-.123efe"]

printstyled("    how bad float strings parse:\n", color=:light_green)
for (i,j) in enumerate(bad_floats)
  # println("test ", i, ": ", j, " => ", string(bad_vals[i]))
  b = IOBuffer(j*"~")
  try
    f = stream_float(b, 0x00)
    println("      ", j, " => ", string(f))
  catch err
    println("      ", j, " => error (", err, ")")
  end
  close(b)
end


text:
import HiGHS

@testset "FEASIBILITY_SENSE zeros objective" begin
    model = DiffOpt.diff_optimizer(HiGHS.Optimizer)
    MOI.set(model, MOI.Silent(), true)
    x = MOI.add_variable(model)
    MOI.add_constraint(model, x, MOI.GreaterThan(1.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.set(model, MOI.ObjectiveFunction{MOI.VariableIndex}(), x)

    MOI.optimize!(model)
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1.0 atol=ATOL rtol=RTOL

    MOI.set(model, MOI.ObjectiveSense(), MOI.FEASIBILITY_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.ObjectiveValue()) == 0.0
end


text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Composition (Modia3D/Composition/_module.jl)
#

get_eAxis(axis::Int) = axis==  1 ? Frames.Vector3D(  1.0,  0.0,  0.0) :
                       axis==  2 ? Frames.Vector3D(  0.0,  1.0,  0.0) :
                       axis==  3 ? Frames.Vector3D(  0.0,  0.0,  1.0) :
                       axis== -1 ? Frames.Vector3D( -1.0,  0.0,  0.0) :
                       axis== -2 ? Frames.Vector3D(  0.0, -1.0,  0.0) :
                       axis== -3 ? Frames.Vector3D(  0.0,  0.0, -1.0) :
                       error("Modia3D.Prismatic: axis = ", axis, " but must be 1, 2, 3, -1, -2, or -3.")

"""
    joint = Prismatic(; obj1, obj2, path="", axis=1, s=0, v=0, canCollide=true)

Return a `joint` that translates `obj2::`[`Object3D`](@ref) with respect to
`obj1::`[`Object3D`](@ref) along coordinate axis `axis` (`axis = 1,2,3,-1,-2,-3`)
of `obj1`. The initial position is `s` and the initial velocity is `v`. Negative
`axis` values describe axes in negative axis directions, i.e. the signs of the values
of `s` and `v` are reversed. If `canCollide=false`, no collision detection will occur
between `obj1` and `obj2` (and `Object3D`s that are directly or indirectly rigidly
fixed to `obj1` or `obj2`).

It is currently not supported that a `Prismatic` joint *closes a kinematic loop*.
"""
mutable struct Prismatic <: Modia3D.AbstractJoint
    path::String

    obj1::Object3D
    obj2::Object3D

    posAxis::Int             # = 1,2,3
    posMovement::Bool        # = true, if movement along posAxis, otherwise in negative posAxis
    eAxis::Frames.Vector3D   # Unit vector in direction of axis
    ndof::Int
    canCollide::Bool         # = false, if no collision between obj1 and obj2

    s::Float64
    v::Float64
    a::Float64
    f::Float64
    residue::Float64

    function Prismatic(; obj1::Object3D,
                         obj2::Object3D,
                         path::String="",
                         axis::Int = 1,
                         s::Real   = 0.0,
                         v::Real   = 0.0,
                         canCollide::Bool = true)

        (parent,obj,cutJoint) = attach(obj1, obj2)
        if cutJoint
            error("\nError from Prismatic joint connecting ", Modia3D.fullName(obj1), " with ", Modia3D.fullName(obj2), ":\n",
                "    This joint is a cut-joint which is currently not supported.!")
        end

        if !(1 <= abs(axis) <= 3)
            error("\nError from Prismatic joint connecting ", Modia3D.fullName(obj1), " with ", Modia3D.fullName(obj2), ":\n",
                "    axis = $axis, but must be 1,2,3 or -1,-2-3.!")
        end

        if !(typeof(s) <: AbstractFloat) || !(typeof(v) <: AbstractFloat)
            error("\nError from Prismatic joint connecting ", Modia3D.fullName(obj1), " with ", Modia3D.fullName(obj2), ":\n",
                "    s and/or v are not <: AbstractFloat!")
        end

        axis  = obj === obj2 ? axis : -axis
        eAxis = get_eAxis(axis)

        s = Modia3D.convertAndStripUnit(Float64, u"m"  , s)
        v = Modia3D.convertAndStripUnit(Float64, u"m/s", v)

        posAxis     = abs(axis)
        posMovement = axis > 0

        obj.joint      = new(path, parent, obj, posAxis, posMovement, eAxis, 1, canCollide, s, v, 0.0, 0.0, 0.0)
        obj.jointKind  = PrismaticKind
        obj.jointIndex = 0
        obj.ndof       = 1
        obj.canCollide = canCollide
        obj.r_rel      = eAxis*s
        obj.R_rel      = Modia3D.NullRotation

        parent.hasChildJoint = true
        return obj.joint
    end
end


text:
# second-order-cone constraints

MOI.supports_constraint(o::Optimizer, ::Type{VECTOR}, ::Type{SOC}) = true

function MOI.add_constraint(o::Optimizer, func::VECTOR, set::SOC)
    allow_modification(o)

    # check bounds on variable for right-hand side
    rhsvar = func.variables[1]
    rhslb = SCIPvarGetLbOriginal(var(o, rhsvar))
    if rhslb < 0.0
        error("SCIP needs non-negative variable on right-hand side of second-order cone constraint, have lower bound: $(rhslb)!")
    end

    varrefs = [VarRef(vi.value) for vi in func.variables]
    cr = add_second_order_cone_constraint(o.inner, varrefs)
    ci = CI{VECTOR, SOC}(cr.val)
    register!(o, ci)
    register!(o, cons(o, ci), cr)
    return ci
end

function MOI.delete(o::Optimizer, ci::CI{VECTOR, SOC})
    _throw_if_invalid(o, ci)
    allow_modification(o)
    delete!(o.constypes[VECTOR, SOC], ConsRef(ci.value))
    delete!(o.reference, cons(o, ci))
    delete(o.inner, ConsRef(ci.value))
    return nothing
end

function MOI.get(o::Optimizer, ::MOI.ConstraintFunction, ci::CI{VECTOR, SOC})
    _throw_if_invalid(o, ci)
    c = cons(o, ci)::Ptr{SCIP_CONS}
    nvars::Int = SCIPgetNLhsVarsSOC(o, c)
    vars = unsafe_wrap(Array{Ptr{SCIP_VAR}}, SCIPgetLhsVarsSOC(o, c), nvars)
    rhsvar = SCIPgetRhsVarSOC(o, c)

    ptr2index(p) = VI(ref(o, p).val)
    return VECTOR(vcat([ptr2index(rhsvar)], ptr2index.(vars)))
end

function MOI.get(o::Optimizer, ::MOI.ConstraintSet, ci::CI{VECTOR, SOC})
    _throw_if_invalid(o, ci)
    c = cons(o, ci)::Ptr{SCIP_CONS}
    nvars::Int = SCIPgetNLhsVarsSOC(o, c)
    return MOI.SecondOrderCone(nvars + 1)
end