text:
"""
Parameter struct that holds all parameters that define the default model setup.
With keywords such that default values can be changed at creation.
"""
@with_kw struct Parameters

    # NUMBER FORMATS
    NF::DataType=Float64    # number format

    # RESOLUTION
    trunc::Int=31                      # spectral truncation
    nlon::Int=roundup_fft(3*trunc+1)   # number of longitudes
    nlat::Int=nlon÷2                   # number of latitudes

    nlev::Int=8             # number of vertical levels
    ntracers::Int=1         # number of tracers (specific humidity is one)

    # PHYSICAL CONSTANTS
    R_earth::Real=6.371e6   # Radius of Earth [m]
    Ω::Real=7.292e-5        # Angular frequency of Earth's rotation [1/s]
    gravity::Real=9.81      # Gravitational acceleration [m/s^2]
    akap::Real=2/7          # Ratio of gas constant to specific heat of dry air
                            # at constant pressure = 1 - 1/γ where γ is the
                            # heat capacity ratio of a perfect diatomic gas (7/5)
    cp::Real=1004.0         # Specific heat at constant pressure [J/K/kg]
    R::Real=akap*cp         # Specific gas constant for dry air [J/kg/K]
    alhc::Real=2501.0       # Latent heat of condensation [J/g] for consistency with
                            # specific humidity [g/Kg]
    alhs::Real=2801.0       # Latent heat of sublimation [?]
    sbc::Real=5.67e-8       # Stefan-Boltzmann constant [W/m^2/K^4]

    # STANDARD ATMOSPHERE
    lapse_rate::Real=6.0    # Reference temperature lapse rate -dT/dz [K/km]
    temp_ref::Real=288      # Reference absolute temperature at surface z=0 [K]
    temp_top::Real=216      # Reference absolute temperature in stratosphere [K]
    scale_height::Real=7.5  # Reference scale height for pressure [km]
    pres_ref::Real=1013     # Reference surface pressure [hPa]
    scale_height_humid::Real=2.5    # Reference scale height for specific humidity [km]
    relhumid_ref::Real=0.7          # Reference relative humidity of near-surface air [1]
    water_pres_ref::Real=17         # Reference saturation water vapour pressure [Pa]

    # VERTICAL COORDINATES
    # of the nlev vertical levels, defined by a generalised logistic function,
    # interpolating ECMWF's L31 configuration
    GLcoefs::GenLogisticCoefs=GenLogisticCoefs()
    n_stratosphere_levels::Int=2    # number of vertical levels used for the stratosphere

    # DIFFUSION
    npowhd::Real=4          # Power of Laplacian in horizontal diffusion
    thd::Real=2.4           # Damping time [hrs] for diffusion (del^6) of temperature and vorticity
    thdd::Real=2.4          # Damping time [hrs] for diffusion (del^6) of divergence
    thds::Real=12           # Damping time [hrs] for extra diffusion (del^2) in the stratosphere
    tdrs::Real=24*30        # Damping time [hrs] for drag on zonal-mean wind in the stratosphere

    # PARAMETERIZATIONS
    seasonal_cycle::Bool=true   # Seasonal cycle?
    n_shortwave::Int=3          # Compute shortwave radiation every n steps
    sppt_on::Bool=false         # Turn on SPPT?

    # TIME STEPPING
    nstepsday::Int=36           # number of time steps in one day
    Δt::Real=86400/nstepsday    # time step in seconds
    ndays::Real=10              # number of days to integrate for

    # NUMERICS
    robert_filter::Real=0.05    # Robert (1966) time filter coefficeint for suppress comput. mode
    williams_filter::Real=0.53  # William's time filter (Amezcua 2011) coefficient for 3rd order acc
    α::Real=0.5                 # coefficient for semi-implicit computations
                                # 0 -> forward step for gravity wave terms,
                                # 1 -> backward implicit
                                # 0.5 -> centered implicit

    # BOUNDARY FILES
    boundary_path::String=""    # package location is default
    boundary_file::String="surface.nc"

    # INITIAL CONDITIONS
    initial_conditions::Symbol=:rest    # :rest or :restart

    # OUTPUT
    verbose::Bool=true          # print dialog for feedback
    output::Bool=false          # Store data in netCDF?
    output_dt::Real=6           # output time step [hours]
    output_startdate::DateTime=DateTime(2000,1,1)
    outpath::String=pwd()       # path to output folder
    output_vars::Vector{String}=["u","v","T","humid","logp0"]
    compression_level::Int=3    # 1=low but fast, 9=high but slow
    keepbits::Int=10            # mantissa bits to keep for every variable

    # TODO assert not allowed parameter values
    @assert α in [0,0.5,1] "Only semi-implicit α = 0, 0.5 or 1 allowed."
end


text:
using Base.Cartesian

##############################################
# Accelerated Convolutions for Efficient Multi-Scale Time to Contact Computation in Julia
# Alexander Amini, Alan Edelman, Berthold Horn
##############################################
@generated function fastconv(E::Array{T,N}, k::Array{T,N}) where {T,N}
    quote
        retsize = [size(E)...] + [size(k)...] .- 1
        retsize = tuple(retsize...)
        ret = zeros(T, retsize)

        convn!(ret,E,k)
        return ret

    end
end

# in place helper operation to speedup memory allocations
@generated function convn!(out::Array{T}, E::Array{T,N}, k::Array{T,N}) where {T,N}
    quote
        @inbounds begin
            @nloops $N x E begin
                @nloops $N i k begin
                    (@nref $N out d->(x_d + i_d - 1)) += (@nref $N E x) * (@nref $N k i)
                end
            end
        end
        return out
    end
end

ext: jl


text:
using Plots
using Random
Random.seed!(1)

include_ddp()

# Model
include_model("double_integrator")

function f(model::DoubleIntegratorContinuous, x, u, w)
    [x[2]; u[1]]
end

function fd(model::DoubleIntegratorContinuous{Midpoint, FixedTime}, x, u, w, h, t)
	if t == 1
		x1 = u[2:3]
    	return [x1 + h * f(model, x1 + 0.5 * h * f(model, x1, u, w), u, w); x1]
	else
		x1 = x[3:4]
		return [view(x, 1:2) + h * f(model, view(x, 1:2) + 0.5 * h * f(model, view(x, 1:2), u, w), u, w); x1]
	end
end

# Time
T = 11
Tm = 6
h = 0.1

model = DoubleIntegratorContinuous{Midpoint, FixedTime}(2, 1, 0)
n = [t == 1 ? model.n : 2 * model.n for t = 1:T]
m = [t == 1 ? model.m + model.n : model.m for t = 1:T]

# Initial conditions, controls, disturbances
x1 = [1.0; 0.0]
xM = [0.0; 0.0]
xT = [1.0; 0.0]
ū = [t == 1 ? [1.0 * randn(1); randn(2)] : 1.0 * randn(1) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)

# Objective
Q = [(t < T ? h : 1.0) * (t > 1 ? Diagonal([ones(2); zeros(2)]) : Diagonal(ones(2))) for t = 1:T]
q = [(t > 1 ? -2.0 * Q[t] * [xT; zeros(2)] : -2.0 * Q[t] * xT) for t = 1:T]
R = h * [t > 1 ? Diagonal(1.0 * ones(1)) : Diagonal([1.0; 1.0e-5 * ones(model.n)]) for t = 1:T-1]
r = [zeros(m[t]) for t = 1:T-1]
obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = o bj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

g(obj, x̄[1], ū[1], 1)
g(obj, x̄[2], ū[2], 2)
g(obj, x̄[T], nothing, T)
objective(obj, x̄, ū)

# Constraints
p = [t == 1 ? 2 * model.n + 2 * model.m : (t == T ? model.n : (t == Tm ? model.n + 2 * model.m : 2 * model.m)) for t = 1:T]
ul = [-5.0]
uu = [5.0]
info_1 = Dict(:x1l => [0.1; -1.0], :x1u => [10.0; 1.0],
	:ul => ul, :uu => uu,
	:inequality => (1:2 * model.m + 2 * model.n))
info_M = Dict(:xM => xM,
	:ul => ul, :uu => uu, :inequality => (1:2 * model.m))
info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * model.m))
info_T = Dict(:xT => xT)
con_set = [StageConstraint(p[t], t < T ? (t == 1 ? info_1 : (t == Tm ? info_M : info_t)) : info_T) for t = 1:T]

info_t
function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	if t == 1
		c[1:1] .= view(u, 1:1) - cons.con[t].info[:uu]
		c[2:2] .= cons.con[t].info[:ul] - view(u, 1:1)
		c[3:4] .= view(u, 2:3) - cons.con[t].info[:x1u]
		c[5:6] .= cons.con[t].info[:x1l] - view(u, 2:3)
	end
	if t > 1 && t < T
		c[1:1] .= view(u, 1:1) - cons.con[t].info[:uu]
		c[2:2] .= cons.con[t].info[:ul] - view(u, 1:1)
		if t == Tm
			c[3:4] .= view(x, 1:2) - cons.con[t].info[:xM]
		end
	end
	if t == T
		c[1:2] .= view(x, 1:2) - view(x, 3:4)
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T, n = n, m = m)

# Solve
@time constrained_ddp_solve!(prob,
    max_iter = 1000, max_al_iter = 10,
	ρ_init = 1.0, ρ_scale = 10.0)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

x[1] = u[1][2:3]
x̄[1] = ū[1][2:3]
norm(x[1] - x[T][1:2])

plt = plot(hcat([xt[1:2] for xt in x]...)',
	width = 2.0,
	color = [:cyan :orange],
	label = ["x" "ẋ"],
	xlabel = "time step",
	ylabel = "state")

savefig(plt,
	joinpath("/home/taylor/Research/parameter_optimization_manuscript/figures/di_gait_state.png"))

plt = plot(hcat([ut[1] for ut in u]..., u[end][1])',
	width = 2.0,
	color = :magenta,
	linetype = :steppost,
	xlabel = "time step",
	ylabel = "control",
	label = "")

savefig(plt,
	joinpath("/home/taylor/Research/parameter_optimization_manuscript/figures/di_gait_control.png"))

max_issues_repo_name: thowell/DirectMotionPlanning


text:
# constants
export σI, σx, σy, σz

# contructors
export qubit_rotation, random_unitary

"""
    σI :: Unitary{Int64}

Identity unitary:

```jldoctest
julia> σI
2×2 Unitary{Int64}:
 1  0
 0  1
```
"""
const σI = Unitary([1 0;0 1])

"""
    σx :: Unitary{Int64}

Pauli-X unitary:

```jldoctest
julia> σx
2×2 Unitary{Int64}:
 0  1
 1  0
```
"""
const σx = Unitary([0 1;1 0])

"""
    σy :: Unitary{Complex{Int64}}

Pauli-Y unitary:

```jldoctest
julia> σy
2×2 Unitary{Complex{Int64}}:
 0+0im  0-1im
 0+1im  0+0im
```
"""
const σy = Unitary([0 -im;im 0])

"""
    σz :: Unitary{Int64}

Pauli-Z unitary:

```
julia> σz
2×2 Unitary:
 1   0
 0  -1
```
"""
const σz = Unitary([1 0;0 -1])

"""
    qubit_rotation( θ :: Real; axis="x" :: String ) :: Unitary{Complex{Float64}}

Returns a unitary which performs a qubit rotation along bloch sphere. `θ` designates
the angle of rotation and `axis` (`"x"`, `"y"`, `"z"`) designates the cartesian
axis about which the qubit is rotated.
"""
function qubit_rotation(θ :: Real; axis="x" :: String)::Unitary{Complex{Float64}}

    pauli = [1 0; 0 1]
    if axis == "x"
        pauli = σx
    elseif axis == "y"
        pauli = σy
    elseif axis == "z"
        pauli = σz
    end

    Unitary(exp(-im * pauli * θ/2))
end

"""
    random_unitary( d :: Int64 ) :: Unitary{Complex{Float64}}

Constructs a `d x d` random unitary matrix according to the Haar measure.
"""
function random_unitary(d :: Int64) :: Unitary{Complex{Float64}}
    generator = Haar(2) # Specifies a complex hermitian unitary matrix
    Unitary(rand(generator, d))
end


text:
function sacsma(forcings::DataFrame, par::AbstractArray; prcpCol::Symbol=:precip, petCol::Symbol=:pet, initstate::AbstractArray=[0,0,500,500,500,0])

    uztwm  =  par[1]    # Upper zone tension water capacity [mm]
    uzfwm  =  par[2]    # Upper zone free water capacity [mm]
    lztwm  =  par[3]    # Lower zone tension water capacity [mm]
    lzfpm  =  par[4]    # Lower zone primary free water capacity [mm]
    lzfsm  =  par[5]    # Lower zone supplementary free water capacity [mm]
    uzk    =  par[6]    # Upper zone free water lateral depletion rate [1/day]
    lzpk   =  par[7]    # Lower zone primary free water depletion rate [1/day]
    lzsk   =  par[8]    # Lower zone supplementary free water depletion rate [1/day]
    zperc  =  par[9]    # Percolation demand scale parameter [-]
    rexp   =  par[10]   # Percolation demand shape parameter [-]
    pfree  =  par[11]   # Percolating water split parameter (decimal fraction)
    pctim  =  par[12]   # Impervious fraction of the watershed area (decimal fraction)
    adimp  =  par[13]   # Additional impervious areas (decimal fraction)
    riva   =  par[14]   # Riparian vegetation area (decimal fraction)
    side   =  par[15]   # The ratio of deep recharge to channel base flow [-]
    rserv  =  par[16]   # Fraction of lower zone free water not transferrable (decimal fraction)

    # Initial Storage States (SAC-SMA)
    uztwc = initstate[1] # Upper zone tension water storage
    uzfwc = initstate[2] # Upper zone free water storage
    lztwc = initstate[3] # Lower zone tension water storage
    lzfsc = initstate[4] # Lower zone supplementary free water storage
    lzfpc = initstate[5] # Upper zone primary free water storage
    adimc = initstate[6] # Additional impervious area storage

    pet = forcings[!,petCol]
    prcp = forcings[!,prcpCol]
    len = length(pet)

    # RESERVOIR STATE ARRAY INITIALIZATION
    simflow   = Array{Float64}(undef,len) # total outflow
    base_tot  = Array{Float64}(undef,len) # base flow
    surf_tot  = Array{Float64}(undef,len) # surface runoff

    thres_zero  = 0.00001 # Threshold to be considered as zero
    parea       = 1 - adimp - pctim

    for i in 1:len

        ### Set input precipitation and potential evapotranspiration
        pr = prcp[i]
        edmnd = pet[i]

        ## Compute for different compnents...
        # ET(1), ET from Upper zone tension water storage
        et1 = edmnd * uztwc/uztwm
        red = edmnd - et1  # residual ET demand
        uztwc = uztwc - et1

        # ET(2), ET from upper zone free water storage
        et2 = 0

        # in case et1 > uztws, no water in the upper tension water storage
        if (uztwc <= 0)
            et1 = et1 + uztwc #et1 = uztwc
            uztwc = 0
            red = edmnd - et1

            # when upper zone free water content is less than residual ET
            if (uzfwc < red)

                # all content at upper zone free water zone will be gone as ET
                et2 = uzfwc
                uzfwc = 0
                red = red - et2
                uztwc = (uztwc < thres_zero) ? 0 : uztwc
                uzfwc = (uzfwc < thres_zero) ? 0 : uzfwc

            # when upper zone free water content is more than residual ET
            else
                et2 = red  # all residual ET will be gone as ET
                uzfwc = uzfwc - et2
                red = 0
            end

        # in case et1 <= uztws, all maximum et (et1) are consumed at uztwc,
        # so no et from uzfwc (et2=0)
        else

            # There's possibility that upper zone free water ratio exceeds
            #upper zone tension water ratio. If so, free water is transferred to
            #tension water storage

            if ((uztwc / uztwm) < (uzfwc / uzfwm))
                uzrat = (uztwc + uzfwc) / (uztwm + uzfwm)
                uztwc = uztwm * uzrat
                uzfwc = uzfwm * uzrat
            end

            uztwc = (uztwc < thres_zero) ? 0 : uztwc
            uzfwc = (uzfwc < thres_zero) ? 0 : uzfwc

        end

        # ET(3), ET from Lower zone tension water storage when residual ET > 0
        et3 = red * lztwc / (uztwm + lztwm) #residual ET is always bigger than ET(3)
        lztwc = lztwc - et3

        # if lztwc is less than zero, et3 cannot exceed lztws
        if(lztwc < 0)
            et3   = et3 + lztwc  # et3 = lztwc
            lztwc = 0
        end

        # Water resupply from Lower free water storages to Lower tension water storage
        saved  = rserv * (lzfpm + lzfsm)
        ratlzt = lztwc / lztwm
        ratlz  = (lztwc + lzfpc + lzfsc - saved) / (lztwm + lzfpm + lzfsm - saved)

        # water is first taken from supplementary water storage for resupply
        if (ratlzt < ratlz)

            del = (ratlz - ratlzt) * lztwm
            lztwc = lztwc + del  # Transfer water from lzfss to lztws
            lzfsc = lzfsc - del

            # if tranfer exceeds lzfsc then remainder comes from lzfps
            if (lzfsc < 0)
                lzfpc = lzfpc + lzfsc
                lzfsc = 0
            end
        end

        lztwc = (lztwc < thres_zero) ? 0 : lztwc

        # ET(5), ET from additional impervious (ADIMP) area
        # ????? no idea where this come from, I think there's a possibility that et5 can be negative values
        et5   = et1 + (red + et2) * (adimc - et1 - uztwc) / (uztwm + lztwm)
        adimc = adimc - et5

        if (adimc < 0)
            #et5 cannot exceed adims
            et5 = et5 + adimc # et5 = adimc
            adimc = 0
        end

        et5 = et5 * adimp

        # Time interval available moisture in excess of uztw requirements
        twx = pr + uztwc - uztwm

        # all moisture held in uztw- no excess
        if (twx < 0)
            uztwc = uztwc + pr
            twx = 0
            # moisture available in excess of uztw storage
        else
            uztwc = uztwm
        end

        # for now twx is excess rainfall after filling the uztwc
        adimc = adimc + pr - twx

        # Compute Impervious Area Runoff
        roimp = pr * pctim

        # Initialize time interval sums
        sbf   = 0  # Sum of total baseflow(from primary and supplemental storages)
        ssur  = 0  # Sum of surface runoff
        sif   = 0  # Sum of interflow
        sperc = 0  # Time interval summation of percolation
        sdro  = 0  # Sum of direct runoff from the additional impervious area

        # Determine computational time increments for the basic time interval
        ninc = floor(1.0 + 0.2*(uzfwc+twx))  # Number of time increments that interval is divided into for further soil-moisture accountng
        dinc = 1.0 / ninc                    # Length of each increment in days
        pinc = twx / ninc                    # Amount of available moisture for each increment

        # Compute free water depletion fractions for the time increment
        #(basic depletions are for one day)
        duz   = 1 - (1 - uzk)^dinc
        dlzp  = 1 - (1 - lzpk)^dinc
        dlzs  = 1 - (1 - lzsk)^dinc

        # Start incremental for-loop for the time interval
        for n in 1:ninc

            adsur = 0 # Amount of surface runoff. This will be updated.

            # Compute direct runoff from adimp area
            ratio = (adimc - uztwc) / lztwm
            ratio = (ratio < 0) ? 0 : ratio

            # Amount of direct runoff from the additional impervious area
            addro = pinc*(ratio^2)

            # Compute baseflow and keep track of time interval sum
            # Baseflow from free water primary storage
            bf_p = lzfpc * dlzp
            lzfpc = lzfpc - bf_p

            if (lzfpc <= 0.0001)
                bf_p  = bf_p + lzfpc
                lzfpc = 0
            end

            sbf = sbf + bf_p

            # Baseflow from free water supplemental storage
            bf_s  = lzfsc * dlzs
            lzfsc = lzfsc - bf_s

            if (lzfsc <= 0.0001)
                bf_s = bf_s + lzfsc
                lzfsc = 0
            end

            # Total Baseflow from primary and supplemental storages
            sbf = sbf + bf_s

            # Compute PERCOLATION- if no water available then skip.
            if ((pinc + uzfwc) <= 0.01)
                uzfwc = uzfwc + pinc

            else

                # Limiting drainage rate from the combined saturated lower zone storages
                percm = lzfpm * dlzp + lzfsm * dlzs
                perc = percm * uzfwc / uzfwm

                # DEFR is the lower zone moisture deficiency ratio
                defr = 1.0 - (lztwc + lzfpc + lzfsc)/(lztwm + lzfpm + lzfsm)

                defr = (defr < 0) ? 0 : defr

                perc = perc * (1.0 + zperc * (defr^rexp))

                # Note. . . percolation occurs from uzfws before pav is added

                # Percolation rate exceeds uzfws
                perc = (perc >= uzfwc) ? uzfwc : perc

                uzfwc = uzfwc - perc    # Percolation rate is less than uzfws.

                # Check to see if percolation exceeds lower zone deficiency.
                check = lztwc + lzfpc + lzfsc + perc - lztwm - lzfpm - lzfsm
                if(check > 0)
                    perc = perc - check
                    uzfwc = uzfwc + check
                end

                # SPERC is the time interval summation of PERC
                sperc = sperc + perc

                # Compute interflow and keep track of time interval sum. Note that PINC has not yet been added.
                del = uzfwc * duz # The amount of interflow
                sif = sif + del
                uzfwc = uzfwc - del

                # Distribute percolated water into the lower zones. Tension water
                # must be filled first except for the PFREE area. PERCT is
                # percolation to tension water and PERCF is percolation going to
                # free water.

                perct = perc * (1.0 - pfree)  # Percolation going to the tension water storage
                if ((perct + lztwc) <= lztwm)

                    lztwc = lztwc + perct
                    percf = 0 # Pecolation going to th lower zone free water storages

                else
                    percf = lztwc + perct - lztwm
                    lztwc = lztwm
                end

                # Distribute percolation in excess of tension requirements among the free water storages.
                percf = percf + (perc * pfree)
                if(percf != 0)

                    # Relative size of the primary storage as compared with total lower zone free water storages.
                    hpl = lzfpm / (lzfpm + lzfsm)

                    # Relative fullness of each storage.
                    ratlp = lzfpc / lzfpm
                    ratls = lzfsc / lzfsm

                    # The fraction going to primary
                    fracp = hpl * 2 * (1 - ratlp) / (2 - ratlp - ratls)

                    fracp = (fracp > 1.0) ? 1.0 : fracp

                    percp = percf * fracp # Amount of the excess percolation going to primary
                    percs = percf - percp # Amount of the excess percolation going to supplemental
                    lzfsc = lzfsc + percs

                    if (lzfsc > lzfsm)
                        percs = percs - lzfsc + lzfsm
                        lzfsc = lzfsm
                    end

                    lzfpc = lzfpc + percf - percs

                    # Check to make sure lzfps does not exceed lzfpm
                    if (lzfpc >= lzfpm)
                        excess = lzfpc - lzfpm
                        lztwc = lztwc + excess
                        lzfpc = lzfpm
                    end
                end

                # Distribute PINC between uzfws and surface runoff
                if (pinc != 0)

                    # check if pinc exceeds uzfwm
                    if((pinc + uzfwc) <= uzfwm)

                        uzfwc = uzfwc + pinc  # no surface runoff
                    else
                        sur = pinc + uzfwc - uzfwm # Surface runoff
                        uzfwc = uzfwm

                        ssur = ssur + (sur * parea)

                        # ADSUR is the amount of surface runoff which comes from
                        # that portion of adimp which is not currently generating
                        # direct runoff. ADDRO/PINC is the fraction of adimp
                        # currently generating direct runoff.
                        adsur = sur * (1.0 - addro / pinc)
                        ssur = ssur + adsur * adimp

                    end
                end
            end

            adimc = adimc + pinc - addro - adsur
            if (adimc > (uztwm + lztwm))
                addro = addro + adimc - (uztwm + lztwm)
                adimc = uztwm + lztwm
            end

            # Direct runoff from the additional impervious area
            sdro  = sdro + (addro * adimp)

            adimc = (adimc < thres_zero) ? 0 : adimc

        end # END of incremental for loop

        # Compute sums and adjust runoff amounts by the area over which they are generated.

        # EUSED is the ET from PAREA which is 1.0 - adimp - pctim
        eused = et1 + et2 + et3
        sif = sif * parea

        # Separate channel component of baseflow from the non-channel component
        tbf = sbf * parea   # TBF is the total baseflow
        bfcc = tbf / (1 + side)    # BFCC is baseflow, channel component

        # Ground flow and Surface flow
        base = bfcc                       # Baseflow and Interflow are considered as Ground inflow to the channel
        surf = roimp + sdro + ssur + sif  # Surface flow consists of Direct runoff and Surface inflow to the channel

        # ET(4)- ET from riparian vegetation.
        et4 = (edmnd - eused) * riva  # no effect if riva is set to zero

        # Compute total evapotransporation - TET
        eused = eused * parea
        tet = eused + et4 + et5

        # Check that adims >= uztws
        adimc = (adimc < uztwc) ? uztwc : adimc

        # Total inflow to channel for a timestep
        tot_outflow = surf + base - et4

        ### ------- Adjustments to prevent negative flows -------------------------#

        # If total outflow <0 surface and baseflow needs to be updated
        if (tot_outflow < 0)

            tot_outflow = 0; surf = 0; base = 0;

        else

            surf_remainder = surf - et4
            surf = max(0,surf_remainder)

            if (surf_remainder < 0) # In this case, base is reduced
                base = base + surf_remainder
                base = max(base,0)
            end
        end

        # Total inflow to channel for a timestep
        simflow[i]  = tot_outflow
        surf_tot[i] = surf
        base_tot[i] = base

    end #close time-loop

    return simflow, surf_tot, base_tot

end # close function


text:
# you can
#     push!(LOAD_PATH, ".")
# then
#     using MyModule
# then
#     MyModule.gleep()

module MyModule

using Match

function foop(x)
  return @match x begin
    1 => "one"
    2 => "two"
    3 => "three"
    _ => "another"
  end
end

function gleep()
  for i in 1:8
    println(foop(i))
  end
end

end # module

max_stars_repo_name: catseye/Dipple

max_issues_repo_name: catseye/Dipple

max_forks_repo_name: catseye/Dipple


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/NodeNumbering.jl/blob/master/LICENSE

using Base.Test
using NodeNumbering

test_files = String[]
push!(test_files, "test_calc_bw.jl")
push!(test_files, "test_create_adjacency_graph.jl")
push!(test_files, "test_node_degrees.jl")
push!(test_files, "test_RCM.jl")
push!(test_files, "test_create_RCM_adjacency.jl")
push!(test_files, "test_renumbering.jl")
push!(test_files, "test_adjacency_visualization.jl")

using TimerOutputs
const to = TimerOutput()
@testset "NodeNumbering.jl" begin
    for fn in test_files
        timeit(to, fn) do
            include(fn)
        end
    end
end
println()
println("Test statistics:")
println(to)


text:
module VertexEliminationOrder

include(joinpath(pwd(), "src", "Heuristics.jl"))
include(joinpath(pwd(), "src", "DFS_BB.jl"))

export graph_from_gr, run

import BenchmarkTools as BT
import LightGraphs as LG
import GraphPlot as GP

function graph_from_gr(filename::String)::LG.SimpleGraph
    lines = readlines(filename)

    # Create a Graph with the correct number of vertices.
    num_vertices, num_edges = parse.(Int, split(lines[1], ' ')[3:end])
    G = LG.SimpleGraph(num_vertices)

    # Add an edge to the graph for every other line in the file.
    for line in lines[2:end]
        src, dst = parse.(Int32, split(line, ' '))
        LG.add_edge!(G, src, dst)
    end

    G
end

function run_alg()
    # Main
    # graph_file = "circuit_graphs/qflex_line_graph_files_decomposed_true_hyper_true/test.gr"
    # G = graph_from_gr(graph_file)
    G = LG.smallgraph("house")
    g = copy(G)
    GP.gplot(g)
    # draw(PNG("graph.png", 800, 600), gplot(g, edgestrokec = colorant"black"))

    @BT.time tf = min_fill(g)
    g = copy(G)
    @BT.time tw = min_width(g)
end

end

max_stars_repo_name: OleeL/treewidth-heuristics

max_issues_repo_name: OleeL/treewidth-heuristics

max_forks_repo_name: OleeL/treewidth-heuristics


text:
using CompileBot

botconfig = BotConfig(
  "MatLang";
  yml_path = "SnoopCompile.yml",
)

println("Benchmarking the inference time of `using MatLang`")
snoop_bench(
  botconfig,
  :(using MatLang),
)

println("Benchmarking the inference time of `using MatLang` and runtests")
snoop_bench(
  botconfig,
  "$(@__DIR__)/example_script.jl",
)

max_stars_repo_name: juliamatlab/MatLang

max_issues_repo_name: juliamatlab/MatLang

max_forks_repo_name: juliamatlab/MatLang


text:
module renderables

import ..measure: Measure
import ..segment: Segment
import Term: split_lines, reshape_text, do_by_line

export AbstractRenderable, Renderable, RenderableText

# ------------------------------- abstract type ------------------------------ #

"""
    AbstractRenderable
"""
abstract type AbstractRenderable end

Measure(renderable::AbstractRenderable) = renderable.measure

"""
 Base.:+(r1::Union{AbstractString, AstractRenderable}, r2::Union{AbstractString, AstractRenderable})

 Concatenates two abstract rendereables
 """
function Base.:+(
    r1::Union{AbstractString,AbstractRenderable},
    r2::Union{AbstractString,AbstractRenderable},
)
    r1 = Renderable(r1)
    r2 = Renderable(r2)
    segments = vcat(r1.segments, r2.segments)
    return Renderable(segments, Measure(segments))
end

"""
    Base.string(r::AbstractRenderable)::String

Creates a string representation of a renderable
"""
function Base.string(r::AbstractRenderable)::String
    lines = [seg.text for seg in r.segments]
    return join(lines, "\n")
end

function Base.show(io::IO, renderable::AbstractRenderable)
    if io == stdout
        for seg in renderable.segments
            println(io, seg)
        end
    else
        print(
            io,
            "$(typeof(renderable)) <: AbstractRenderable \e[2m(size: $(renderable.measure))\e[0m",
        )
    end
end

# ------------------------- generic renderable object ------------------------ #

"""
    Renderable

Generic `Renderable` object.
"""
mutable struct Renderable <: AbstractRenderable
    segments::Vector
    measure::Measure
end

Renderable() = Renderable([], Measure(0, 0))
function Renderable(
    str::Union{Vector,AbstractString}; width::Union{Nothing,Int,Symbol} = nothing
)
    return RenderableText(str; width = width)
end
Renderable(ren::AbstractRenderable; width::Union{Nothing,Int,Symbol} = nothing) = ren  # returns the renderable
function Renderable(segment::Segment; width::Union{Nothing,Int,Symbol} = nothing)
    return Renderable([segment], Measure([segment]))
end

# ---------------------------------------------------------------------------- #
#                                TEXT RENDERABLE                               #
# ---------------------------------------------------------------------------- #

"""
    RenderableText

`Renderable` represnting a text.

See also [`Renderable`](@ref), [`TextBox`](@ref)
"""
mutable struct RenderableText <: AbstractRenderable
    segments::Vector
    measure::Measure
    text::AbstractString
end

"""
    RenderableText(text::AbstractString; width::Union{Nothing, Int, Symbol}=nothing)

Construct a `RenderableText` out of a string.

If a `width` is passed the text is resized to match the width.
"""
function RenderableText(text::AbstractString; width::Union{Nothing,Int,Symbol} = nothing)
    # @info "creating RenderableText"  text width
    if width isa Number
        text = do_by_line((ln) -> reshape_text(ln, width), text)
    end

    segments = [Segment(line) for line in split_lines(text)]
    return RenderableText(segments, Measure(segments), text)
end

function RenderableText(
    text::Vector{AbstractString}; width::Union{Nothing,Int,Symbol} = nothing
)
    return RenderableText(join(text, "\n"); width = width)
end
function RenderableText(text::Vector; width::Union{Nothing,Int,Symbol} = nothing)
    return RenderableText(join(text, "\n"); width = width)
end

function RenderableText(
    text::AbstractString, style::AbstractString; width::Union{Nothing,Int,Symbol} = nothing
)
    if width isa Number
        text = do_by_line((ln) -> reshape_text(ln, width), text)
    end
    segments = [Segment(line, style) for line in split_lines(text)]
    return RenderableText(segments, Measure(segments), text)
end

# -------------------------------- union type -------------------------------- #
RenderablesUnion = Union{AbstractString,AbstractRenderable,RenderableText}

end


text:
using Profile
using PProf
using Traceur
using BenchmarkTools

import Moves
import Util
import Game
import Search
board = Util.startPositions()
colour = 'W'
# Game.play(board,colour,[])

Search.negamax(-Inf,Inf,4,colour,[],board,0)
@btime Search.negamax(-Inf,Inf,4,colour,[],board,0)

@profile Search.negamax(-Inf,Inf,4,colour,[],board,0)
pprof()

max_stars_repo_name: grahamclyne/chessEngine

max_issues_repo_name: grahamclyne/chessEngine

max_forks_repo_name: grahamclyne/chessEngine


text:
module gcSolver

using OrdinaryDiffEq
import LinearAlgebra: diag, norm, dot
import ForwardDiff
import ForwardDiff: value, Dual, partials, jacobian!, jacobian
using Optim
using Statistics
import ModelingToolkit: modelingtoolkitize
import CSV
using DataFrames
using StatsFuns

include("reaction.jl")
include("dataImport.jl")

" Check that we haven't been provided anything ridiculous by the user. "
function checkInputs(tps::Vector{Float64}, params::Vector)
    @assert all(tps .>= 0.0)
    @assert params[22] < 1.0
    @assert length(params) == Nparams
    @assert all(params .>= 0.0)
end

" This recompiles the ODE function into a symbolic Jacobian. "
function __init__()
    u0 = ones(Nspecies)
    params = 0.1ones(Nparams)

    prob = ODEProblem(fullDeriv, u0, (0.0, 1.0), params)
    deMT = modelingtoolkitize(prob)

    global modelFunc = ODEFunction(deMT; jac = true)
end

function runCkineSetup(tps::Vector{Float64}, p::Vector{T}) where {T}
    checkInputs(tps, p)
    u0 = solveAutocrine(p)

    return ODEProblem(modelFunc, u0, convert(T, maximum(tps)), p)
end

" Actually run the gc ODE model. "
function runCkine(tps::Vector{Float64}, params; pSTAT5 = false)
    if params isa Vector
        prob = runCkineSetup(tps, params)
    else
        prob = params
    end

    if pSTAT5
        sidx = pSTATidx
    else
        sidx = nothing
    end

    solAlg = TRBDF2()
    sol = solve(prob, solAlg; saveat = tps, save_idxs = sidx, reltol = 1.0e-3, maxiters = 1e7)

    if sol.retcode != :Success
        println("Solving failed with the following parameters.")
        println(sol.retcode)
        println(ForwardDiff.value.(params))
        @assert sol.retcode == :Success
    end

    sol = hcat(sol.u...)'
    @assert ndims(sol) == 2
    @assert size(sol, 1) == length(tps)

    if pSTAT5
        @assert size(sol) == (length(tps), 3)
        # Summation of active species
        return vec(sol[:, 1] + 2 * (sol[:, 2] + sol[:, 3]))
    end

    return sol
end

" Calculate the Jacobian of the model and perform variance propagation. "
function runCkineVarProp(tps::Vector, params::Vector, sigma, recAbund)::Vector
    checkInputs(tps, params)

    # Sigma is the covariance matrix of the input parameters
    function jacF(x)
        pp = vcat(params[1:24], x, params[28:end])
        return runCkine(tps, pp, pSTAT5 = true)
    end
    jac = zeros(3, length(tps))
    jacobian!(jac, jacF, params[25:27])
    for i in (1:3)
        jac[i, :] ./= (recAbund[i] / params[24 + i]) #maybe
    end

    # Just return the diagonal for the marginal variance
    return diag(transpose(jac) * sigma * jac)
end

include("fit.jl")

export runCkine, runCkineVarProp, runCkineCost

end # module


text:
using Base: PkgId
using CodeTracking: basepath

# A near-copy of the same method in `base/loading.jl`. However, this retains the full module path to the file.
function parse_cache_header(f::IO)
    modules = Vector{Pair{PkgId, UInt64}}()
    while true
        n = read(f, Int32)
        n == 0 && break
        sym = String(read(f, n)) # module name
        uuid = UUID((read(f, UInt64), read(f, UInt64))) # pkg UUID
        build_id = read(f, UInt64) # build UUID (mostly just a timestamp)
        push!(modules, PkgId(uuid, sym) => build_id)
    end
    totbytes = read(f, Int64) # total bytes for file dependencies
    # read the list of requirements
    # and split the list into include and requires statements
    includes = Tuple{Module, String, Float64}[]
    requires = Pair{Module, PkgId}[]
    while true
        n2 = read(f, Int32)
        n2 == 0 && break
        depname = String(read(f, n2))
        mtime = read(f, Float64)
        n1 = read(f, Int32)
        mod = (n1 == 0) ? Main : Base.root_module(modules[n1].first)
        if n1 != 0
            # determine the complete module path
            while true
                n1 = read(f, Int32)
                totbytes -= 4
                n1 == 0 && break
                submodname = String(read(f, n1))
                mod = getfield(mod, Symbol(submodname))
                totbytes -= n1
            end
        end
        if depname[1] != '\0'
            push!(includes, (mod, depname, mtime))
        end
        totbytes -= 4 + 4 + n2 + 8
    end
    @assert totbytes == 12 "header of cache file appears to be corrupt"
    # Determine which includes are included in the source-text cache.
    # These are the *.jl files that we need to track.
    srctextpos = read(f, Int64)
    if srctextpos == 0
        empty!(includes)
    else
        seek(f, srctextpos)
        keep = Set{String}()
        while !eof(f)
            filenamelen = read(f, Int32)
            filenamelen == 0 && break
            fn = String(read(f, filenamelen))
            len = read(f, UInt64)
            push!(keep, fn)
            seek(f, position(f) + len)
        end
        delids = Int[]
        for (i, inc) in enumerate(includes)
            inc[2] ∈ keep || push!(delids, i)
        end
        deleteat!(includes, delids)
    end
    return modules, (includes, requires)
end

function parse_cache_header(cachefile::String)
    io = open(cachefile, "r")
    try
        !Base.isvalid_cache_header(io) && throw(ArgumentError("Invalid header in cache file $cachefile."))
        return parse_cache_header(io)
    finally
        close(io)
    end
end

# This is an implementation of
#   filter(path->Base.stale_cachefile(sourcepath, path) !== true, paths)
# that's easier to precompile. (This is a hotspot in loading Revise.)
function filter_valid_cachefiles(sourcepath, paths)
    fpaths = String[]
    sourcepath === nothing && return fpaths
    for path in paths
        if Base.stale_cachefile(sourcepath, path) !== true
            push!(fpaths, path)
        end
    end
    return fpaths
end

function pkg_fileinfo(id::PkgId)
    uuid, name = id.uuid, id.name
    # Try to find the matching cache file
    paths = Base.find_all_in_cache_path(id)
    sourcepath = Base.locate_package(id)
    if length(paths) > 1
        fpaths = filter_valid_cachefiles(sourcepath, paths)
        paths = isempty(fpaths) ? paths : fpaths
        # Work-around for #371 (broken dependency prevents tracking):
        # find the most recent cache file. Presumably this is the one built
        # to load the package.
        sort!(paths; by=path->mtime(path), rev=true)
    end
    isempty(paths) && return nothing, nothing
    path = first(paths)
    provides, includes_requires = try
        parse_cache_header(path)
    catch
        return nothing, nothing
    end
    mods_files_mtimes, _ = includes_requires
    for (pkgid, buildid) in provides
        if pkgid.uuid === uuid && pkgid.name == name
            return path, mods_files_mtimes
        end
    end
end

"""
    parentfile, included_files = modulefiles(mod::Module)

Return the `parentfile` in which `mod` was defined, as well as a list of any
other files that were `include`d to define `mod`. If this operation is unsuccessful,
`(nothing, nothing)` is returned.

All files are returned as absolute paths.
"""
function modulefiles(mod::Module)
    function keypath(filename)
        filename = fixpath(filename)
        return get(src_file_key, filename, filename)
    end
    parentfile = String(first(methods(getfield(mod, :eval))).file)
    id = PkgId(mod)
    if id.name == "Base" || Symbol(id.name) ∈ stdlib_names
        parentfile = normpath(Base.find_source_file(parentfile))
        filedata = Base._included_files
    else
        use_compiled_modules() || return nothing, nothing   # FIXME: support non-precompiled packages
        _, filedata = pkg_fileinfo(id)
    end
    filedata === nothing && return nothing, nothing
    included_files = filter(mf->mf[1] == mod, filedata)
    return keypath(parentfile), [keypath(mf[2]) for mf in included_files]
end

"""
    parse_pkg_files(id::PkgId)

This function gets called by `watch_package` and runs when a package is first loaded.
Its job is to organize the files and expressions defining the module so that later we can
detect and process revisions.
"""
function parse_pkg_files(id::PkgId)
    if !haskey(pkgdatas, id)
        pkgdatas[id] = PkgData(id)
    end
    pkgdata = pkgdatas[id]
    modsym = Symbol(id.name)
    if use_compiled_modules()
        cachefile, mods_files_mtimes = pkg_fileinfo(id)
        if cachefile !== nothing
            for (mod, fname, _) in mods_files_mtimes
                if mod === Main && !isdefined(mod, modsym)  # issue #312
                    mod = Base.root_module(PkgId(pkgdata))
                end
                fname = relpath(fname, pkgdata)
                # For precompiled packages, we can read the source later (whenever we need it)
                # from the *.ji cachefile.
                push!(pkgdata, fname=>FileInfo(mod, cachefile))
            end
            CodeTracking._pkgfiles[id] = pkgdata.info
            return srcfiles(pkgdata)
        end
    end
    # Non-precompiled package(s). Here we rely on the `include` callbacks to have
    # already populated `included_files`; all we have to do is collect the relevant
    # files.
    queue_includes!(pkgdata, id)
    return srcfiles(pkgdata)
end

# The main trick here is that since `using` is recursive, `included_files`
# might contain files associated with many different packages. We have to figure
# out which correspond to a particular module `mod`, which we do by:
#   - checking the module in which each file is evaluated. This suffices to
#     detect "supporting" files, i.e., those `included` within the module
#     definition.
#   - checking the filename. Since the "top level" file is evaluated into Main,
#     we can't use the module-of-evaluation to find it. Here we hope that the
#     top-level filename follows convention and matches the module. TODO?: it's
#     possible that this needs to be supplemented with parsing.
function queue_includes!(pkgdata::PkgData, id::PkgId)
    modstring = id.name
    delids = Int[]
    for i = 1:length(included_files)
        mod, fname = included_files[i]
        if mod == Base.__toplevel__
            mod = Main
        end
        modname = String(Symbol(mod))
        if startswith(modname, modstring) || endswith(fname, modstring*".jl")
            modexsigs = parse_source(fname, mod)
            if modexsigs !== nothing
                instantiate_sigs!(modexsigs)
                fname = relpath(fname, pkgdata)
                push!(pkgdata, fname=>FileInfo(modexsigs))
            end
            push!(delids, i)
        end
    end
    deleteat!(included_files, delids)
    CodeTracking._pkgfiles[id] = pkgdata.info
    return srcfiles(pkgdata)
end

function queue_includes(mod::Module)
    id = PkgId(mod)
    if !haskey(pkgdatas, id)
        pkgdatas[id] = PkgData(id)
    end
    pkgdata = pkgdatas[id]
    files = queue_includes!(pkgdata, id)
    if has_writable_paths(pkgdata)
        init_watching(pkgdata, files)
    end
end

# A near-duplicate of some of the functionality of queue_includes!
# This gets called for silenced packages, to make sure they don't "contaminate"
# included_files
function remove_from_included_files(modsym::Symbol)
    i = 1
    modstring = string(modsym)
    while i <= length(included_files)
        mod, fname = included_files[i]
        modname = String(Symbol(mod))
        if startswith(modname, modstring) || endswith(fname, modstring*".jl")
            deleteat!(included_files, i)
        else
            i += 1
        end
    end
end

function read_from_cache(pkgdata::PkgData, file::AbstractString)
    fi = fileinfo(pkgdata, file)
    filep = joinpath(basedir(pkgdata), file)
    if fi.cachefile == basesrccache
        # Get the original path
        filec = get(cache_file_key, filep, filep)
        return open(basesrccache) do io
            Base._read_dependency_src(io, filec)
        end
    end
    Base.read_dependency_src(fi.cachefile, filep)
end

function maybe_parse_from_cache!(pkgdata::PkgData, file::AbstractString)
    if startswith(file, "REPL[")
        return add_definitions_from_repl(file)
    end
    fi = fileinfo(pkgdata, file)
    if isempty(fi.modexsigs) && !isempty(fi.cachefile)
        # Source was never parsed, get it from the precompile cache
        src = read_from_cache(pkgdata, file)
        filep = joinpath(basedir(pkgdata), file)
        filec = get(cache_file_key, filep, filep)
        topmod = first(keys(fi.modexsigs))
        if parse_source!(fi.modexsigs, src, filec, topmod) === nothing
            @error "failed to parse cache file source text for $file"
        end
        instantiate_sigs!(fi.modexsigs)
    end
    return fi
end

function maybe_add_includes_to_pkgdata!(pkgdata::PkgData, file, includes)
    for (mod, incs) in includes
        for inc in incs
            inc = joinpath(splitdir(file)[1], inc)
            hasfile = false
            for srcfile in srcfiles(pkgdata)
                if srcfile == inc
                    hasfile = true
                    break
                end
            end
            if !hasfile
                # Add the file to pkgdata
                push!(pkgdata.info.files, inc)
                fi = FileInfo(mod)
                push!(pkgdata.fileinfos, fi)
                # Parse the source of the new file
                fullfile = joinpath(basedir(pkgdata), inc)
                if isfile(fullfile)
                    parse_source!(fi.modexsigs, fullfile, mod)
                    instantiate_sigs!(fi.modexsigs; define=true)
                end
                # Add to watchlist
                init_watching(pkgdata, (inc,))
            end
        end
    end
end

function add_require(sourcefile, modcaller, idmod, modname, expr)
    expr isa Expr || return
    id = PkgId(modcaller)
    @async begin
        # If this fires when the module is first being loaded (because the dependency
        # was already loaded), Revise may not yet have the pkgdata for this package.
        while !haskey(pkgdatas, id)
            sleep(0.1)
        end
        # Get/create the FileInfo specifically for tracking @require blocks
        pkgdata = pkgdatas[id]
        filekey = relpath(sourcefile, pkgdata) * "__@require__"
        fileidx = fileindex(pkgdata, filekey)
        if fileidx === nothing
            files = srcfiles(pkgdata)
            fileidx = length(files) + 1
            push!(files, filekey)
            push!(pkgdata.fileinfos, FileInfo(modcaller))
        end
        fi = pkgdata.fileinfos[fileidx]
        # Tag the expr to ensure it is unique
        expr = Expr(:block, copy(expr))
        push!(expr.args, :(__pkguuid__ = $idmod))
        # Add the expression to the fileinfo
        exsnew = ExprsSigs()
        exsnew[expr] = nothing
        mexsnew = ModuleExprsSigs(modcaller=>exsnew)
        # Before executing the expression we need to set the load path appropriately
        prev = Base.source_path(nothing)
        tls = task_local_storage()
        tls[:SOURCE_PATH] = sourcefile
        # Now execute the expression
        mexsnew, includes = try
            eval_new!(mexsnew, fi.modexsigs)
        finally
            if prev === nothing
                delete!(tls, :SOURCE_PATH)
            else
                tls[:SOURCE_PATH] = prev
            end
        end
        # Add any new methods or `include`d files to tracked objects
        pkgdata.fileinfos[fileidx] = FileInfo(mexsnew, fi)
        maybe_add_includes_to_pkgdata!(pkgdata, filekey, includes)
    end
    return nothing
end

function watch_files_via_dir(dirname)
    try
        wait_changed(dirname)  # this will block until there is a modification
    catch e
        # issue #459
        (isa(e, InterruptException) && throwto_repl(e)) || throw(e)
    end
    latestfiles = Pair{String,PkgId}[]
    # Check to see if we're still watching this directory
    stillwatching = haskey(watched_files, dirname)
    if stillwatching
        wf = watched_files[dirname]
        for (file, id) in wf.trackedfiles
            fullpath = joinpath(dirname, file)
            if !file_exists(fullpath)
                # File may have been deleted. But be very sure.
                sleep(0.1)
                if !file_exists(fullpath)
                    push!(latestfiles, file=>id)
                    continue
                end
            end
            if newer(mtime(fullpath), wf.timestamp)
                push!(latestfiles, file=>id)
            end
        end
        isempty(latestfiles) || updatetime!(wf)  # ref issue #341
    end
    return latestfiles, stillwatching
end

"""
    watch_package(id::Base.PkgId)

Start watching a package for changes to the files that define it.
This function gets called via a callback registered with `Base.require`, at the completion
of module-loading by `using` or `import`.
"""
function watch_package(id::PkgId)
    # Because the callbacks are made with `invokelatest`, for reasons of performance
    # we need to make sure this function is fast to compile. By hiding the real
    # work behind a @async, we truncate the chain of dependency.
    @async _watch_package(id)
end

@noinline function _watch_package(id::PkgId)
    modsym = Symbol(id.name)
    if modsym ∈ dont_watch_pkgs
        if modsym ∉ silence_pkgs
            @warn "$modsym is excluded from watching by Revise. Use Revise.silence(\"$modsym\") to quiet this warning."
        end
        remove_from_included_files(modsym)
        return nothing
    end
    files = parse_pkg_files(id)
    pkgdata = pkgdatas[id]
    if has_writable_paths(pkgdata)
        init_watching(pkgdata, files)
    end
end

function has_writable_paths(pkgdata::PkgData)
    haswritable = false
    for file in srcfiles(pkgdata)
        haswritable |= iswritable(joinpath(basedir(pkgdata), file))
    end
    return haswritable
end

## Working with Pkg and code-loading

# Much of this is adapted from base/loading.jl

function manifest_file(project_file)
    if project_file isa String
        mfile = Base.project_file_manifest_path(project_file)
        if mfile isa String
            return mfile
        end
    end
    return nothing
end
manifest_file() = manifest_file(Base.active_project())

function manifest_paths!(pkgpaths::Dict, manifest_file::String)
    open(manifest_file) do io
        uuid = name = path = hash = id = nothing
        for line in eachline(io)
            if (m = match(Base.re_section_capture, line)) != nothing
                name = String(m.captures[1])
                path = hash = nothing
            elseif (m = match(Base.re_uuid_to_string, line)) != nothing
                uuid = UUID(m.captures[1])
                name === nothing && error("name not set for $uuid")
                id = PkgId(uuid, name)
                # UUID is last, so time to store
                if path !== nothing
                    pkgpaths[id] = path
                elseif hash !== nothing
                    path = find_from_hash(name, uuid, hash)
                    path === nothing && error("no path found for $id and hash $hash")
                    pkgpaths[id] = path
                end
                uuid = name = path = hash = id = nothing
            elseif (m = match(Base.re_path_to_string, line)) != nothing
                path = String(m.captures[1])
                path = normpath(abspath(dirname(manifest_file), path))
            elseif (m = match(Base.re_hash_to_string, line)) != nothing
                hash = Base.SHA1(m.captures[1])
            end
        end
    end
    return pkgpaths
end

manifest_paths(manifest_file::String) =
    manifest_paths!(Dict{PkgId,String}(), manifest_file)

function find_from_hash(name, uuid, hash)
    for slug in (Base.version_slug(uuid, hash, 4), Base.version_slug(uuid, hash))
        for depot in DEPOT_PATH
            path = abspath(depot, "packages", name, slug)
            if ispath(path)
                return path
            end
        end
    end
    return nothing
end

function watch_manifest(mfile)
    while true
        try
            wait_changed(mfile)
        catch e
            # issue #459
            (isa(e, InterruptException) && throwto_repl(e)) || throw(e)
        end
        try
            with_logger(_debug_logger) do
                @debug "Pkg" _group="manifest_update" manifest_file=mfile
                isfile(mfile) || return nothing
                pkgdirs = manifest_paths(mfile)
                for (id, pkgdir) in pkgdirs
                    if haskey(pkgdatas, id)
                        pkgdata = pkgdatas[id]
                        if pkgdir != basedir(pkgdata)
                            ## The package directory has changed
                            @debug "Pkg" _group="pathswitch" oldpath=basedir(pkgdata) newpath=pkgdir
                            # Stop all associated watching tasks
                            for dir in unique_dirs(srcfiles(pkgdata))
                                @debug "Pkg" _group="unwatch" dir=dir
                                delete!(watched_files, joinpath(basedir(pkgdata), dir))
                                # Note: if the file is revised, the task(s) will run one more time.
                                # However, because we've removed the directory from the watch list this will be a no-op,
                                # and then the tasks will be dropped.
                            end
                            # Revise code as needed
                            files = String[]
                            for file in srcfiles(pkgdata)
                                maybe_parse_from_cache!(pkgdata, file)
                                push!(revision_queue, (pkgdata, file))
                                push!(files, file)
                            end
                            # Update the directory
                            pkgdata.info.basedir = pkgdir
                            # Restart watching, if applicable
                            if has_writable_paths(pkgdata)
                                init_watching(pkgdata, files)
                            end
                        end
                    end
                end
            end
        catch err
            @static if VERSION >= v"1.2.0-DEV.253"
                put!(Base.active_repl_backend.response_channel, (Base.catch_stack(), true))
            else
                put!(Base.active_repl_backend.response_channel, (err, catch_backtrace()))
            end
        end
    end
end


text:
module DepParsing

include("token.jl")
include("state.jl")
include("beamsearch.jl")
include("model.jl")

end

max_stars_repo_name: hshindo/JukaiNLP

max_issues_repo_name: hshindo/JukaiNLP

max_forks_repo_name: hshindo/JukaiNLP


text:
"""
    Const <: Kernel

Constant kernel
```math
k(x,x') = σ²
```
with signal standard deviation ``σ``.
"""
mutable struct Const{T} <: Kernel where {T<:Real}
    "Signal variance"
    σ2::T
    "Priors for kernel parameters"
    priors::Array
end

"""
Constant kernel function

    Const(lσ::T)

# Arguments
  - `lσ::Real`: signal standard deviation (given on log scale)
"""
Const(lσ::T) where T = Const{T}(exp(2 * lσ), [])

function set_params!(cons::Const, hyp::AbstractVector)
    length(hyp) == 1 || throw(ArgumentError("Constant kernel has one parameter, received $(length(hyp))."))
    cons.σ2 = exp(2.0*hyp[1])
end

get_params(cons::Const{T}) where T = T[log(cons.σ2)/2.0]
get_param_names(cons::Const) = [:lσ]
num_params(cons::Const) = 1

cov(cons::Const) = cons.σ2
function cov(cons::Const, x::AbstractVector, y::AbstractVector)
    return cov(cons)
end

@inline dk_dlσ(cons::Const) = 2.0*cov(cons)
@inline function dKij_dθp(cons::Const, X1, X2, i::Int, j::Int, p::Int, dim::Int)
    if p == 1
        return dk_dlσ(cons)
    else
        return NaN
    end
end


text: # Curve evolution active contours


text:
module IntegrateBenchs

using Test
using InplaceArrays.Arrays
using InplaceArrays.Fields
using InplaceArrays.Fields: IntKernel
using InplaceArrays.Fields: OtherMockBasis, MockBasis, MockField
using FillArrays
using InplaceArrays.TensorValues

@noinline function loop(a,cache)
  for i in eachindex(a)
    ai = getindex!(cache,a,i)
  end
  nothing
end

@inline function repeat(n,f,args...)
  for i in 1:n
    f(args...)
  end
  nothing
end

function bench1(n)

  fi = 2.0
  wi = 3.0
  ji = TensorValue(1.0,0.0,0.0,2.0)
  np = 4
  f = fill(fi,np)
  w = fill(wi,np)
  j = fill(ji,np)
  k = IntKernel()
  ck = kernel_cache(k,f,w,j)
  @time repeat(n,apply_kernel!,ck,k,f,w,j)

end

function bench2(n)

  fi = 2.0
  wi = 3.0
  ji = TensorValue(1.0,0.0,0.0,2.0)
  np = 4
  f = fill(fi,np,10)
  w = fill(wi,np)
  j = fill(ji,np)
  k = IntKernel()
  ck = kernel_cache(k,f,w,j)
  @time repeat(n,apply_kernel!,ck,k,f,w,j)

end

function bench3(n)

  np = 4
  p = Point(2,2)
  x = fill(p,np)

  ndof = 8
  d = 2
  ri = 5.0
  r = MockBasis{d}(ri,ndof)

  v = 3.0
  ndof = 8
  w = fill(1/np,np)
  c = fill(1.5,ndof)
  f = OtherMockBasis{d}(ndof)

  l = n
  ac = fill(c,l)
  af = Fill(f,l)
  aϕ = lincomb(af,ac)
  aj = ∇(aϕ)
  aw = Fill(w,l)
  ax = Fill(x,l)
  ar = Fill(r,l)
  ab = attachmap(ar,aϕ)

  fmass = varinner(ab,ab)
  mmass = integrate(fmass,ax,aw,aj)

  cmmass = array_cache(mmass)
  @time loop(mmass,cmmass)

  fstif = varinner(∇(ab),∇(ab))
  mstif = integrate(fstif,ax,aw,aj)

  cmstif = array_cache(mstif)
  @time loop(mstif,cmstif)

end

function bench4(n)

  fi = 2.0
  wi = 3.0
  ji = TensorValue(1.0,0.0,0.0,2.0)
  np = 4
  f = fill(fi,np)
  w = fill(wi,np)
  j = fill(ji,np)
  l = n
  af = fill(f,l)
  aw = fill(w,l)
  aj = fill(j,l)

  k = IntKernel()
  s = apply(k,af,aw,aj)

  cs = array_cache(s)
  @time loop(s,cs)

end

function bench5(n)

  fi = 2.0
  wi = 3.0
  ji = TensorValue(1.0,0.0,0.0,2.0)
  np = 4
  ndof = 10
  f = fill(fi,np,ndof)
  w = fill(wi,np)
  j = fill(ji,np)
  l = n
  af = fill(f,l)
  aw = fill(w,l)
  aj = fill(j,l)

  k = IntKernel()
  s = apply(k,af,aw,aj)

  cs = array_cache(s)
  @time loop(s,cs)
  @time repeat(n,getindex!,cs,s,1)

end

for n in (1,1,10,1000,100000)
  @eval begin
    println("+++ runing suite for n = $($n) +++")
    bench1($n)
    bench2($n)
    bench3($n)
    bench4($n)
    bench5($n)
  end
end

end #module


text:
module BioProfiling
export  logtransform,
        logtransform!,
        normtransform,
        normtransform!,
        decorrelate,
        decorrelate!,
        decorrelate_by_mad!,
        mahalanobis,
        hellinger,
        Filter,
        CombinationFilter,
        MissingFilter,
        Experiment,
        filter_entries!,
        filter_entries,
        Selector,
        NameSelector,
        CombinationSelector,
        select_features,
        select_features!,
        select!,
        filter!,
        diagnostic,
        diagnostic_path,
        diagnostic_images,
        negation,
        getdata,
        shuffled_distance_robust_mahalanobis_median,
        distance_robust_mahalanobis_median,
        shuffled_distance_robust_hellinger,
        distance_robust_hellinger,
        shuffled_distance_mahalanobis_center,
        distance_mahalanobis_center,
        shuffled_distance_mahalanobis_median,
        distance_mahalanobis_median,
        robust_morphological_perturbation_value,
        most_variable_features,
        characteristic_features,
        most_correlated,
        freqtable,
        umap
using Statistics, StatsBase, DataFrames, Images, ImageMagick, UMAP, RCall, MultipleTesting
using Distributed, ParallelDataTransfer, FreqTables

using LinearAlgebra: det

include("struct.jl")
include("internal.jl")
include("transform.jl")
include("distances.jl")
include("diagnostic.jl")
include("visu.jl")
include("rmpv.jl")
include("interpret.jl")

end # module


text:
using SparseArrays
using LinearAlgebra
using SparsityDetection

struct ParaboloidStruct{T, Tm <: AbstractArray{T,2},
                           Tv <: AbstractArray{T}} <: Any where T<:Number
    mat::Tm
    vec::Tv
    xt::Tv
    alpha::T
end

function quad(x::Vector, param)
    mat = param.mat
    xt = x-param.vec
    return 0.5*dot(xt, mat*xt)
end

function _paraboloidproblem(N::Int;
                            mat::AbstractArray{T,2} = sparse(Diagonal(float(1:N))),
                            alpha::T=10.0,
                            x0::AbstractVector{T} = ones(N)) where T <: Number
    hessian_sparsity(quad,x0,ParaboloidStruct(mat, x0, similar(x0), alpha))
end

@test isdiag(_paraboloidproblem(10))


text:
#!/usr/bin/env bash
    #=
    exec julia --project="$(realpath $(dirname $0))" --color=yes --startup-file=no -e "include(popfirst!(ARGS))" \
    "${BASH_SOURCE[0]}" "$@"
    =#


text:
using Test

println("Hello")
include("structures.jl")


text:
@testset "Mixed" begin
    N = 10

    for (constructor, copier) in ((interpolate, x->x), (interpolate!, copy))
        A2 = rand(Float64, N, N) * 100
        for BC in (Flat,Line,Free,Periodic,Reflect,Natural), GT in (OnGrid, OnCell)
            itp_a = @inferred(constructor(copier(A2), (BSpline(Linear()), BSpline(Quadratic(BC(GT()))))))
            itp_b = @inferred(constructor(copier(A2), (BSpline(Quadratic(BC(GT()))), BSpline(Linear()))))
            isfullsize = constructor == interpolate || BC==Periodic
            if isfullsize
                @test @inferred(size(itp_a)) == size(A2)
                @test @inferred(size(itp_b)) == size(A2)
                @test @inferred(axes(itp_a)) == axes(A2)
                @test @inferred(axes(itp_b)) == axes(A2)
            else
                @test @inferred(size(itp_a)) == (N, N-2)
                @test @inferred(size(itp_b)) == (N-2, N)
                @test @inferred(axes(itp_a)) == (1:N, 2:N-1)
                @test @inferred(axes(itp_b)) == (2:N-1, 1:N)
            end
            @test_throws ArgumentError parent(itp_a)
            @test_throws ArgumentError parent(itp_b)

            for i in eachindex(itp_a)
                @test itp_a[i] ≈ A2[i] atol=sqrt(eps(A2[i]))
            end
            for i in eachindex(itp_b)
                @test itp_b[i] ≈ A2[i] atol=sqrt(eps(A2[i]))
            end

            for i = 1:10
                dx, dy = rand(), rand()
                @test itp_a(2 + dx,2) ≈ (1 - dx) * A2[2,2] + dx * A2[3,2]
                @test itp_b(2,2 + dy) ≈ (1 - dy) * A2[2,2] + dy * A2[2,3]
            end
        end
    end

    # AbstractArrays
    makesharedarray(::Type{T}, dims; kwargs...) where {T} = SharedArray{T}(dims; kwargs...)
    function copyshared(A)
        B = makesharedarray(eltype(A), size(A))
        copyto!(B, A)
    end

    for (constructor, copier) in ((interpolate, x->x), (interpolate!, copyshared))
        A2 = makesharedarray(Float64, (N,N), init=A->rand!(A))
        for i = 1:length(A2)
            A2[i] *= 100
        end
        for BC in (Flat,Line,Free,Periodic,Reflect,Natural), GT in (OnGrid, OnCell)
            itp_a = @inferred(constructor(copier(A2), (BSpline(Linear()), BSpline(Quadratic(BC(GT()))))))
            itp_b = @inferred(constructor(copier(A2), (BSpline(Quadratic(BC(GT()))), BSpline(Linear()))))
            if constructor == interpolate!
                @test isa(itp_a.coefs, SharedArray)
                @test isa(itp_b.coefs, SharedArray)
            end
            isfullsize = constructor == interpolate || BC==Periodic
            if isfullsize
                @test @inferred(size(itp_a)) == size(A2)
                @test @inferred(size(itp_b)) == size(A2)
                @test @inferred(axes(itp_a)) == axes(A2)
                @test @inferred(axes(itp_b)) == axes(A2)
            else
                @test @inferred(size(itp_a)) == (N, N-2)
                @test @inferred(size(itp_b)) == (N-2, N)
                @test @inferred(axes(itp_a)) == (1:N, 2:N-1)
                @test @inferred(axes(itp_b)) == (2:N-1, 1:N)
            end

            for j = 2:N-1, i = 2:N-1
                @test ≈(itp_a[i,j],A2[i,j],atol=sqrt(eps(A2[i,j])))
                @test ≈(itp_b[i,j],A2[i,j],atol=sqrt(eps(A2[i,j])))
            end

            for i = 1:10
                dx, dy = rand(), rand()
                @test itp_a(2 + dx,2) ≈ (1 - dx) * A2[2,2] + dx * A2[3,2]
                @test itp_b(2,2 + dy) ≈ (1 - dy) * A2[2,2] + dy * A2[2,3]
            end
        end
    end
end


text:
using MOCNeutronTransport
using Plots
pyplot()
T = Float64
Nmax = 100
Nset = 1:Nmax

p₁ = Point_3D(T, 0)
p₂ = Point_3D(T, 2)
p₃ = Point_3D(T, 2, 3)
p₄ = Point_3D(T, 0, 3)
p₅ = Point_3D(T, 3//2, 1//2)
p₆ = Point_3D(T, 5//2, 3//2)
p₇ = Point_3D(T, 3//2, 5//2)
p₈ = Point_3D(T, 0,    3//2)
p₈ = Point_3D(T, 0, 3//2, 2)
quad8 = Quadrilateral8_3D((p₁, p₂, p₃, p₄, p₅, p₆, p₇, p₈))
l = LineSegment_3D(Point_3D(T, 1, 0, 1//2),
                   Point_3D(T, 1, 3, 1//2))

p₁ = intersect(l, quad8; N=Nmax)[3]
p₂ = intersect(l, quad8; N=Nmax)[4]

p1_error = [norm(p₁ - intersect(l, quad8; N = n)[3]) for n = Nset]
p2_error = [norm(p₂ - intersect(l, quad8; N = n)[4]) for n = Nset]
p_error = p1_error + p2_error
plot(Nset[1:Nmax-1],
     p_error[1:Nmax-1],
     title = "Edge Subdivisions vs Quad8_3D sum of norm of point errors (Nsoln = $Nmax)",
     yaxis=:log)
#     xaxis=:log)
xlabel!("Edge Subdivisions")
ylabel!("sum of 2-norm of point error vectors")

max_stars_repo_name: KyleVaughn/MOCNeutronTransport

max_issues_repo_name: KyleVaughn/MOCNeutronTransport

max_forks_repo_name: KyleVaughn/MOCNeutronTransport


text:
function initial_plot(f; plot=true, redraw=true, newfig=true, kwargs...)
    global SURROGATE_PLOT

    if plot
        if newfig
            figure()
        end
        if redraw
            clf()
            SURROGATE_PLOT = nothing # imshow can now be called again (instead of using set_data)
            plot_objective_function(f)
        end
    end
end

global LIMITS = (X=[-15, 15], Y=[-15, 15]) # To be filled.
function set_limits(X, Y)
    global LIMITS
    LIMITS = (X=X, Y=Y)
end

function plot_objective_function(f, limits = LIMITS; bins=100, kwargs...)
    rx = range(limits.X[1], stop=limits.X[2], length=bins)
    ry = range(limits.Y[1], stop=limits.Y[2], length=bins)

    # note reverse and x, y switch (imshow has a weird mapping)
    imshow([f([x,y]; kwargs...) for y in reverse(ry), x in rx],
        extent=[limits.X[1], limits.X[2], limits.Y[1], limits.Y[2]],
        cmap="viridis_r")
end

function plot_cem(P, samples, order, m_elite, limits=LIMITS)
    gca().collections = [] # update figure instead of redrawing
    scatter(samples[1,:], samples[2,:], s=7, c="black", edgecolors="white", linewidths=1/2)
    scatter(samples[1,order[1:m_elite]], samples[2,order[1:m_elite]], s=7, c="red", edgecolors="white", linewidths=1/2)
    xlim(LIMITS.X)
    ylim(LIMITS.Y)

    # Plot covariance.
    X = range(LIMITS.X[1], stop=LIMITS.X[2], length=100)
    Y = range(LIMITS.Y[1], stop=LIMITS.Y[2], length=100)
    Z = [pdf(P, [x,y]) for y in Y, x in X] # Note: reverse in X, Y.
    contour(Z, extent=[LIMITS.X[1], LIMITS.X[2], LIMITS.Y[1], LIMITS.Y[2]], cmap="hot", alpha=0.3)
    gcf() # |> display # tell Juno to display figure
    sleep(1/2)
end

function plot_mixture(mm, limits=LIMITS; bins=100, newfig=true)
    if newfig
        figure()
    end
    rx = range(limits.X[1], stop=limits.X[2], length=bins)
    ry = range(limits.Y[1], stop=limits.Y[2], length=bins)
    # note reverse and x, y switch (imshow has a weird mapping)
    imshow([pdf(mm, [x,y]) for y in reverse(ry), x in rx], extent=[limits.X[1], limits.X[2], limits.Y[1], limits.Y[2]], cmap="viridis_r")
end

function plot_voronoi(𝒟, φ=x->x, limits=(X=[1,10], Y=[1,10]); newfig=false)
    # Voronoi diagram (Euclidean)
    f = (x1,x2)->nearest_neighbors(φ([x1,x2]), 𝒟, dist_euclidean)

    if newfig
        figure()
    end

    rx = range(limits.X[1], stop=limits.X[2], length=100)
    ry = range(limits.Y[1], stop=limits.Y[2], length=100)

    # note reverse and x, y switch (imshow has a weird mapping)
    imshow([f(x,y) for y in reverse(ry), x in rx], extent=[limits.X[1], limits.X[2], limits.Y[1], limits.Y[2]], cmap="viridis_r")

    xlim(limits.X)
    ylim(limits.Y)

end

global SURROGATE_PLOT = nothing
function plot_surrogate_model(f)
    global SURROGATE_PLOT

    # note reverse and x, y switch (imshow has a weird mapping)
    rx = range(LIMITS.X[1], stop=LIMITS.Y[2], length=100)
    ry = range(LIMITS.X[1], stop=LIMITS.Y[2], length=100)
    𝒟 = [f([x,y]) for y in reverse(ry), x in rx]
    if isnothing(SURROGATE_PLOT)
        SURROGATE_PLOT = imshow(𝒟, extent=[rx[1], rx[end], ry[1], ry[end]], cmap="viridis_r") #, alpha=0.33) # Greys
    else
        SURROGATE_PLOT.set_data(𝒟)
    end
end

function plot_covariance(𝐦, 𝐌; show_sub_layer_distributions=true)
    global LIMITS
    varX = range(LIMITS.X[1], stop=LIMITS.X[2], length=100)
    varY = range(LIMITS.Y[1], stop=LIMITS.Y[2], length=100)

    if show_sub_layer_distributions && !ismissing(𝐦)
        for 𝐦ᵢ in 𝐦
            Z = [pdf(𝐦ᵢ, [x,y]) for y in varY, x in varX] # Note: reverse in X, Y.
            contour(Z, extent=[LIMITS.X[1], LIMITS.X[2], LIMITS.Y[1], LIMITS.Y[2]], cmap="hot", alpha=0.3)
        end
    end

    Z = [pdf(𝐌, [x,y]) for y in varY, x in varX] # Note: reverse in X, Y.
    contour(Z, extent=[LIMITS.X[1], LIMITS.X[2], LIMITS.Y[1], LIMITS.Y[2]], cmap="hot", alpha=0.3)
end

"""
Handle settings for the `plotting` function
"""
@with_kw mutable struct PlotSettings
    show_non_elites=true
    highlight_sub_elites=true
    highlight_elites=true
    show_covariance_plots=true
    show_surrogate_model=false
    show_mixture=false
    show_voronoi=false
end

function plotting(k, 𝐦, 𝐌, 𝒮, 𝒟, m_elite, samples, sub_elites, elite, nextcycle; settings=PlotSettings())
    clear_scatter()

    viridis_r = ColorMaps.RGBArrayMap(ColorSchemes.viridis, interpolation_levels=m_elite, invert=true)

    # Color according to elite-class
    colors = fill(:black, m_elite) # [:green, :blue, :black, :cyan, :yellow, :gray]

    if settings.show_non_elites
        nonelite = columns(samples)
        for ne in nonelite
            scatter(ne[1], ne[2], s=7, c="black", edgecolors="white", linewidths=1/2)
        end
    end

    # Plot new sub-elites.
    if nextcycle
        if settings.highlight_sub_elites
            scatter(map(first, columns(sub_elites)), map(last, columns(sub_elites)), s=7, c="white", edgecolors="black", linewidths=1/2)
        else
            for (i,e) in enumerate(columns(sub_elites))
                scatter(e[1], e[2], s=7, c=colors[i], edgecolors="white", linewidths=1/2)
            end
        end
    end

    # Highlight elite samples, on top.
    if settings.highlight_elites
        scatter(map(first, columns(elite)), map(last, columns(elite)), s=7, c="red", edgecolors="white", linewidths=1/2)
    else
        for (i,e) in enumerate(columns(elite))
            scatter(e[1], e[2], s=7, c=colors[i], edgecolors="white", linewidths=1/2)
        end
    end

    title("CE-mixture method ($k)")
    xlabel(L"x_1")
    ylabel(L"x_2")

    if settings.show_covariance_plots
        plot_covariance(𝐦, 𝐌)
    end

    if settings.show_surrogate_model
        plot_surrogate_model(𝒮)
    end

    if settings.show_mixture
        plot_mixture(𝐌; newfig=false)
    end

    if settings.show_voronoi
        figure()
        plot_voronoi(𝒟, LIMITS)
    end

    xlim(LIMITS.X)
    ylim(LIMITS.Y)
    gcf() # |> display # tell Juno to display figure.
    sleep(1/20)
end

sample_components(𝐌) = map(rand, 𝐌.components)
component_means(𝐌) = map(m->m.μ, 𝐌.components)

plot_component_means(𝐌) = plot_samples(component_means(𝐌))
plot_sampled_components(𝐌) = plot_samples(sample_components(𝐌))

function plot_samples(𝐦)
    scatter(map(first, 𝐦), map(last, 𝐦); s=7, c="black", edgecolors="white", linewidths=1/2)
end

function clear_scatter()
    ax = gca()
    ax.collections = [] # update figure instead of redrawing
    plt.draw()
end


text:
## kepler.jl
##     Functions for orbital elements and Kepler motion

###################################################################
#                       Orbital elements
###################################################################

"""
# Orbital elements

- `a` : Semi-major axis
- `e` : Eccentricity
- `I` : Inclination
- `ω` : Argument of periapsis
- `Ω` : Longitude of ascending node
- `Φ` : Mean anomaly : tₚ = - Φ / n
- `tₚ` : Periapsis passage time

- `μ` : Standard gravitational parameter of the system (GM)
- `n` : Mean motion    : n = √(μ / a^3)
- `T` : Orbital period : T = 2π / n
"""
struct OrbitalElements{T}
    a::T
    e::T
    I::T
    ω::T
    Ω::T
    Φ::T
    tₚ::T

    μ::T
    n::T
    T::T
end

function OrbitalElements(params)

    a = params[:a] * AU      # semi-mojor axis [m]
    e = params[:e]           # eccentricity
    I = deg2rad(params[:I])  # inclination [rad]
    Ω = deg2rad(params[:Ω])  # longitude of the ascending node [rad]
    ω = deg2rad(params[:ω])  # argument of periapsis [rad]

    μ = params[:μ]  # Gravitational parameter
    n = √(μ / a^3)  # mean motion  [rad/sec]
    T = 2π / n      # orbital period [sec]

    if haskey(params, :Φ)
        Φ  = deg2rad(params[:Φ])  # mean anomaly at the epoch [rad]
        tₚ = - Φ / n              # periapsis passage time [sec]
    elseif haskey(params, :tₚ)
        tₚ = params[:tₚ]
        Φ  = - n * tp
    else
        println("Give [:Φ] or [:tp] in Dict.")
    end

    OrbitalElements(a, e, I, ω, Ω, Φ, tₚ, μ, n, T)
end

function OrbitalElements(a, e, I, μ)
    a *= AU
    I = deg2rad(I)
    ω = 2π * rand()
    Ω = 2π * rand()
    Φ = 2π * rand()

    n = √(μ / a^3)
    T = 2π / n
    tₚ = - Φ / n

    OrbitalElements(a, e, I, ω, Ω, Φ, tₚ, μ, n, T)
end

function Base.show(io::IO, elms::OrbitalElements)
    println(io, "Orbital elements")
    println(io, "----------------")

    println("Semi-mojor axis        : a  = ", elms.a / AU,     " [AU]")
    println("Eccentricity           : e  = ", elms.e,          " [-]")
    println("Lon. of ascending node : Ω  = ", rad2deg(elms.Ω), " [deg]")
    println("Argument of periapsis  : ω  = ", rad2deg(elms.ω), " [deg]")
    println("Inclination            : I  = ", rad2deg(elms.I), " [deg]")
    println("Periapsis passage time : tₚ = ", elms.tₚ,         " [sec]")
    println("Mean anomaly           : Φ  = ", rad2deg(elms.Φ), " [deg]")

    println()
    println("Other parameters")
    println("----------------")

    println("Gravitational parameter : μ = ", elms.μ,                      " [m^3/s^2]")
    println("Mean motion             : n = ", rad2deg(elms.n) * (3600*24), " [deg/day]")
    println("Orbital period          : T = ", elms.T / (3600*24),          " [day]")
end

###################################################################
#                   Coordinate transformation
#                   R (X, Y, Z) -> r (x, y, z)
###################################################################

ref_to_orb!(r, elms) = ref_to_orb!(r, elms.ω, elms.I, elms.Ω)
ref_to_orb(r, elms) = ref_to_orb(r, elms.ω, elms.I, elms.Ω)

"""
Transform from the reference coordinate system to the orbital plane system

# Parameters
- r : vector to be transformed
- ω : argument of periapsis
- I : inclination
- Ω : longitude of the ascneding node
"""
function ref_to_orb!(r, ω, I, Ω)
    rotateZ!(r, Ω)
    rotateX!(r, I)
    rotateZ!(r, ω)
end

function ref_to_orb(r, ω, I, Ω)
    r = rotateZ(r, Ω)
    r = rotateX(r, I)
    r = rotateZ(r, ω)
end

###################################################################
#                   Coordinate transformation
#                    r(x, y, z) -> R(X, Y, Z)
###################################################################

orb_to_ref!(r, elms) = orb_to_ref!(r, elms.ω, elms.I, elms.Ω)
orb_to_ref(r, elms) = orb_to_ref(r, elms.ω, elms.I, elms.Ω)

"""
Transform from the orbital plane system to the reference coordinate system

# Parameters
- r : vector to be transformed
- ω : argument of periapsis
- I : inclination
- Ω : longitude of the ascneding node
"""
function orb_to_ref!(r, ω, I, Ω)
    rotateZ!(r, -ω)
    rotateX!(r, -I)
    rotateZ!(r, -Ω)
end

function orb_to_ref(r, ω, I, Ω)
    r = rotateZ(r, -ω)
    r = rotateX(r, -I)
    r = rotateZ(r, -Ω)
end

###################################################################
#           R(X, Y, Z), V(Ẋ, Ẏ, Ż) -> Orbital elements
#             @Reference frame
###################################################################

OrbitalElements(R, V, M, m, t) = OrbitalElements(R, V, G*(M+m), t)

function OrbitalElements(R::AbstractVector, V::AbstractVector, μ, t)

    h = R × V

    a = get_a(R, V, μ)
    e = get_e(h, μ, a)
    I = get_I(h)
    Ω = get_Ω(h, I)

    ω = get_ω(R, V, a, e, I, Ω)
    f = get_f(R, V, a, e)

    E = f2E(f, e)   # eccentric anomaly (E or u)
    n = √(μ / a^3)  # mean motion
    T = 2π / n      # orbital period

    tₚ = get_tₚ(t%T, E, e, n)  # time of periapsis passage
    M = - n * tₚ              # mean anomaly (M or Φ)

    OrbitalElements(a, e, I, ω, Ω, M, tₚ, μ, n, T)
end

function OrbitalElements(p1, p2, t)
    ps = setParticles(p1, p2)
    r_G, v_G, a_G = getBaryCenter(ps)
    OrbitalElements(p2.r - r_G, p2.v - v_G, p1.m, p2.m, t)
end

get_a(R, V, M, m) = get_a(R, V, G*(M+m))
get_a(R, V, μ) = 1 / (2/norm(R) - norm(V)^2/μ)

get_e(R, V, μ, a) = get_e(R × V, μ, a)
get_e(h, μ, a) = √(1 - norm(h)^2/(μ*a))

get_I(R, V) = get_I(R × V)
get_I(h) = acos(h[3]/norm(h))

get_Ω(R, V, I) = get_Ω(R × V, I)

function get_Ω(h, I)
    # sinΩ =  sign(h[3]) * h[1] / (norm(h)*sin(I))
    # cosΩ = -sign(h[3]) * h[2] / (norm(h)*sin(I))
    # Ω = atan(sinΩ, cosΩ)

    Ω = atan(sign(h[3])*h[1], -sign(h[3])*h[2])
    Ω ≥ 0 ? Ω : Ω+2π
end

function get_ω(R, V, a, e, I, Ω)
    ω₊f = get_ω₊f(R, V, I, Ω)
    f = get_f(R, V, a, e)

    ω = ω₊f - f
    # ω ≥ 0 ? ω : ω+2π
end

function get_ω₊f(R, V, I, Ω)
    I == 0 && return asin((R[2]*cos(Ω) - R[1]*sin(Ω))/norm(R))

    sin_ω₊f = R[3] / (norm(R)*sin(I))
    cos_ω₊f = sec(Ω) * (R[1]/norm(R) + sin(Ω)*sin_ω₊f*cos(I))

    ω₊f = atan(sin_ω₊f, cos_ω₊f)
    # ω₊f ≥ 0 ? ω₊f : ω₊f + 2π
end

function get_f(R, V, a, e)
    h = R × V
    Ṙ = sign(R ⋅ V) * √(norm(V)^2 - norm(h)^2/norm(R)^2)

    sinf = a*(1-e*e) / (norm(h)*e) * Ṙ
    cosf = (a*(1-e*e) / norm(R) - 1) / e

    f = atan(sinf, cosf)
    # f ≥ 0 ? f : f+2π
end

function get_E(R, a, e)
    E = acos((a - norm(R)) / (a*e))  # eccentric anomaly (E or u)
    # E ≥ 0 ? E : E+2π
end

E2f(E, e) = 2 * atan(√((1+e)/(1-e)) * sin(E/2), cos(E/2))
f2E(f, e) = 2 * atan(√((1-e)/(1+e)) * sin(f/2), cos(f/2))

get_tₚ(t, E, e, M, m, a) = get_tₚ(t, E, G*(M+m), a)
get_tₚ(t, E, e, μ, a) = get_tₚ(t, E, e, √(μ/a^3))
get_tₚ(t, E, e, n) = t - (E - e*sin(E)) / n

# ****************************************************************
# ****************************************************************
# ****************************************************************

function getOrbitalElementsError(elms, elms_ref)
    Δa = elms.a - elms_ref.a
    Δe = elms.e - elms_ref.e
    Δω = elms.ω - elms_ref.ω
    ΔT = elms.T - elms_ref.T

    return Δa, Δe, Δω, ΔT
end

# ****************************************************************
#           Orbital elements -> r(x, y, z), v(ẋ, ẏ, ż)
#                                @Orbital plane frame
# ****************************************************************

solveKeplerEquation1(elms, ts::T) where T<:AbstractArray = [solveKeplerEquation1(elms, t) for t in ts]
solveKeplerEquation2(elms, ts::T) where T<:AbstractArray = [solveKeplerEquation2(elms, t) for t in ts]

solveKeplerEquation1(elms, t) = solveKeplerEquation1(elms.e, elms.n, elms.tₚ, t)
solveKeplerEquation2(elms, t) = solveKeplerEquation2(elms.e, elms.n, elms.tₚ, t)

get_rv(elms, u) = get_rv(elms.a, elms.e, elms.n, u)
get_r(elms, u) = get_r(elms.a, elms.e, u)
get_v(elms, u) = get_v(elms.a, elms.e, elms.n, u)

heliocentric_distance(elms, u) = norm(get_r(elms, u))
heliocentric_distance(elms, us::T) where T<:AbstractArray = [heliocentric_distance(elms, u) for u in us]

"""
    solveKeplerEquation1(e, n, tₚ, t) -> u_i

Solve Kepler's equation with the 1st order method

# Parameters
- `e`  : eccentricity
- `n`  : mean motion
- `tₚ` : periapsis passage time [sec]
- `t`  : time

# Return
- `u_i` : eccentric anomaly (i-th order solution)
"""
function solveKeplerEquation1(e, n, tₚ, t)

    Φ = n * (t - tₚ)  # mean anomaly

    u_i = Φ  # eccentirc anomaly at order 0

    for i in 1:20
        u_pri = u_i
        u_i = Φ + e * sin(u_i)
        err = abs(1 - u_pri / u_i)
        # println(i, " : ", u_i, " : ", err)
        err < 1e-10 && return u_i
    end
    u_i
end

"""
    solveKeplerEquation2(e, n, tₚ, t) -> u_i

Solve Kepler's equation with the 2nd order method (Newton's method)

# Parameters
- `e`  : eccentricity
- `n`  : mean motion
- `tₚ` : periapsis passage time [sec]
- `t`  : time

# Return
- `u_i` : eccentric anomaly (i-th order solution)
"""
function solveKeplerEquation2(e, n, tₚ, t)

    Φ = n * (t - tₚ)  # mean anomaly

    u_i = Φ  # eccentirc anomaly at order 0

    for i in 1:20
        u_pri = u_i
        u_i -= (u_i - e * sin(u_i) - Φ) / (1 - e * cos(u_i))
        err = abs(1 - u_pri / u_i)
        # println(i, " : ", u_i, " : ", err)
        err < 1e-10 && return u_i
    end
    u_i
end

"""
    u2ν(u, e) -> ν

Calculate true anomaly ν from eccentric anomaly u and eccentricity e

# Parameters
- `u` : eccentric anomaly
- `e` : eccentricity

# Return
- `ν` : true anomaly

"""
u2ν(u, e) = 2 * atan(√((1+e)/(1-e)) * sin(u*0.5), cos(u*0.5))
u2ν(u, elms::OrbitalElements) = u2ν(u, elms.e)

u2ν(us::T, e) where T<:AbstractArray = [u2ν(u, e) for u in us]
u2ν(us::T, elms::OrbitalElements) where T<:AbstractArray = u2ν(us, elms.e)

"""
    getrv(a, e, n, u) -> r, v

Get a body's positon and velocity based on Kepler's motion

# Parameters
- `a` : semi-major axis
- `e` : eccentricity
- `n` : mean motion
- `u` : eccentric anomaly

# Returns
- `r` : position
- `v` : velocity
(the coordinate system is in the orbital plane and its origin is the eclipse focus)
"""
function get_rv(a, e, n, u)
    sin_u = sin(u)
    cos_u = cos(u)

    x = a * (cos_u - e)
    y = a * √(1 - e*e) * sin_u
    z = 0.

    r = SA_F64[x, y, z]

    x = - a * n * sin_u / (1 - e * cos_u)
    y = a * n * √(1 - e*e) * cos_u / (1 - e * cos_u)
    z = 0.

    v = SA_F64[x, y, z]

    return r, v
end

"""
    get_r(a, e, u) -> r

Get a body's positon based on Kepler's motion

# Parameters
- `a` : semi-major axis
- `e` : eccentricity
- `u` : eccentric anomaly

# Return
- `r` : position
(the coordinate system is in the orbital plane and its origin is the eclipse focus)
"""
function get_r(a, e, u)
    x = a * (cos(u) - e)
    y = a * √(1 - e*e) * sin(u)
    z = 0.

    r = SA_F64[x, y, z]
end

"""
    get_v(a, e, n, u) -> v

Get a body's velocity based on Kepler's motion

# Parameters
- `a` : semi-major axis
- `e` : eccentricity
- `n` : mean motion
- `u` : eccentric anomaly

# Return
- `v` : velocity
(the coordinate system is in the orbital plane and its origin is the eclipse focus)
"""
function get_v(a, e, n, u)
    sin_u = sin(u)
    cos_u = cos(u)

    x = - a * n * sin_u / (1 - e * cos_u)
    y = a * n * √(1 - e*e) * cos_u / (1 - e * cos_u)
    z = 0.

    v = SA_F64[x, y, z]
end

################################################################
#                     Rotation matrix
################################################################

# Rx(θ) = [
#     1   0      0
#     0  cos(θ) sin(θ)
#     0 -sin(θ) cos(θ)
# ]

# Ry(θ) = [
#     cos(θ) 0 -sin(θ)
#      0     1   0
#     sin(θ) 0  cos(θ)
# ]

# Rz(θ) = [
#      cos(θ) sin(θ) 0
#     -sin(θ) cos(θ) 0
#       0      0     1
# ]

# orb2ref(r, elms) = orb2ref(r, elms.ω, elms.I, elms.Ω)
# ref2orb(r, elms) = ref2orb(r, elms.ω, elms.I, elms.Ω)

# orb2ref(r, ω, I, Ω) = Rz(-Ω) * Rx(-I) * Rz(-ω) * r
# ref2orb(r, ω, I, Ω) = Rz( ω) * Rx( I) * Rz( Ω) * r

max_stars_repo_name: MasanoriKanamaru/Astroshaper

max_issues_repo_name: MasanoriKanamaru/Astroshaper

max_forks_repo_name: MasanoriKanamaru/Astroshaper


text:
module TensorFactorizations
using TensorOperations
using LinearAlgebra
export tensorsvd, tensoreig, tensorsplit

"""
    tensorsplit(A, a, b; kwargs...)

Calls tensorsvd with the arguments given to it to decompose the given tensor
A with indices a on one side and indices b on the other.  It then splits
the diagonal matrix of singular values into two with a square root and
multiplies these weights into the isometric tensors.  Thus tensorsplit ends
up splitting A into two parts, which are then returned, possibly together
with auxiliary data such as a truncation error. If the keyword argument
hermitian=true, an eigenvalue decomposition is used in stead of an SVD. All
the keyword arguments are passed to either tensorsvd or tensoreig.

See tensorsvd and tensoreig for further documentation.
"""
function tensorsplit(args...; kwargs...)
    # Find the keyword argument hermitian.
    # TODO This is awful, why do I have to do this?
    hermitian = false
    for (key, value) in kwargs
        key == :hermitian && (hermitian = value)
    end

    if hermitian
        res = tensoreig(args...; kwargs...)
        S, U = res[1:2]
        Vt_perm = [ndims(U), (1:ndims(U)-1)...]
        Vt = conj!(tensorcopy(U, collect(1:ndims(U)), Vt_perm))
        S = Diagonal(S)
        if !isposdef(S)
            S = complex.(S)
        end
        auxdata = res[3:end]
    else
        res = tensorsvd(args...; kwargs...)
        U, S, Vt = res[1:3]
        S = Diagonal(S)
        auxdata = res[4:end]
    end
    S_sqrt = sqrt.(S)
    A1 = tensorcontract(U, (1:ndims(U)-1..., :a), S_sqrt, (:a, :b))
    A2 = tensorcontract(S_sqrt, (:b, :a), Vt, (:a, 1:ndims(Vt)-1...))
    return A1, A2, auxdata...
end

"""
    tensoreig(A, a, b; chis=nothing, eps=0,
              return_error=false, print_error=false,
              break_degenerate=false, degeneracy_eps=1e-6,
              norm_type=:frobenius, hermitian=false)

Finds the "right" eigenvectors and eigenvalues of A. The indices of A are
permuted so that the indices listed in the Array/Tuple a are on the "left"
side and indices listed in b are on the "right".  The resulting tensor is
then reshaped to a matrix, and eig is called on this matrix to get the vector
of eigenvalues E and matrix of eigenvectors U. Finally, U is reshaped to
a tensor that has as its last index the one that enumerates the eigenvectors
and the indices in a as its first indices.

Truncation and error printing work as with tensorsvd.

Note that no iterative techniques are used, which means that choosing to
truncate provides no performance benefits: All the eigenvalues are computed
in any case.

The keyword argument hermitian (false by default) tells the algorithm
whether the reshaped matrix is Hermitian or not. If hermitian=true, then A =
U*diagm(E)*U' up to the truncation error.

Output is E, U, and possibly error, if return_error=true. Here E is a
vector of eigenvalues values and U[:,...,:,k] is the kth eigenvector.
"""
function tensoreig(A, a, b; chis=nothing, eps=0, return_error=false,
                   print_error=false, break_degenerate=false,
                   degeneracy_eps=1e-6, norm_type=:frobenius,
                   hermitian=false)
    # Create the matrix and decompose it.
    A, shp_a, shp_b = to_matrix(A, a, b; return_tensor_shape=true)
    if hermitian
        A = Hermitian((A+A')/2)
    end
    fact = eigen(A)
    E, U = fact.values, fact.vectors
    # Sort the by largest magnitude eigenvalue first.
    perm = sortperm(abs.(E), rev=true)
    if perm != collect(1:length(E))
        E = E[perm]
        U = U[:,perm]
    end

    # Find the dimensions to truncate to and the error caused in doing so.
    chi, error = find_trunc_dim(E, chis, eps, break_degenerate, degeneracy_eps,
                                norm_type)

    # Truncate
    E = E[1:chi]
    U = U[:,1:chi]

    if print_error
        println("Relative truncation error ($norm_type norm) in eig: $error")
    end

    # Reshape U to a tensor with a shape matching the shape of A and
    # return.
    dim = size(E)[1]
    U_tens = reshape(U, shp_a..., dim)
    retval = (E, U_tens)
    if return_error
        retval = (retval..., error)
    end
    return retval
end

"""
    tensorsvd(A, a, b;
              chis=nothing, eps=0,
              return_error=false, print_error=false,
              break_degenerate=false, degeneracy_eps=1e-6,
              norm_type=:frobenius)

Singular valued decomposes a tensor A. The indices of A are
permuted so that the indices listed in the Array/Tuple a are on the "left"
side and indices listed in b are on the "right".  The resulting tensor is
then reshaped to a matrix, and this matrix is SVDed into U*diagm(S)*Vt.
Finally, the unitary matrices U and Vt are reshaped to tensors so that
they have a new index coming from the SVD, for U as the last index and for
Vt as the first, and U has indices a as its first indices and V has
indices b as its last indices.

If eps>0 then the SVD may be truncated if the relative error can be kept
below eps. For this purpose different dimensions to truncate to can be tried,
and these dimensions should be listed in chis. If chis is nothing (the
default) then the full range of possible dimensions is tried. If
break_degenerate=false (the default) then the truncation never cuts between
degenerate singular values. degeneracy_eps controls how close the values need
to be to be considered degenerate.

norm_type specifies the norm used to measure the error. This defaults to
:frobenius, which means that the error measured is the Frobenius norm of the
difference between A and the decomposition, divided by the Frobenius norm of
A.  This is the same thing as the 2-norm of the singular values that are
truncated out, divided by the 2-norm of all the singular values. The other
option is :trace, in which case a 1-norm is used instead.

If print_error=true the truncation error is printed. The default is false.

If return_error=true then the truncation error is also returned.
The default is false.

Note that no iterative techniques are used, which means choosing to truncate
provides no performance benefits: The full SVD is computed in any case.

Output is U, S, Vt, and possibly error. Here S is a vector of
singular values and U and Vt are isometric tensors (unitary if the matrix
that is SVDed is square and there is no truncation) such that  U*diag(S)*Vt =
A, up to truncation errors.
"""
function tensorsvd(A, a, b; chis=nothing, eps=0, return_error=false,
                   print_error=false, break_degenerate=false,
                   degeneracy_eps=1e-6, norm_type=:frobenius)
    # Create the matrix and SVD it.
    A, shp_a, shp_b = to_matrix(A, a, b; return_tensor_shape=true)
    fact = svd(A)
    U, S, Vt = fact.U, fact.S, fact.Vt

    # Find the dimensions to truncate to and the error caused in doing so.
    chi, error = find_trunc_dim(S, chis, eps, break_degenerate, degeneracy_eps,
                                norm_type)

    # Truncate
    S = S[1:chi]
    U = U[:,1:chi]
    Vt = Vt[1:chi,:]

    if print_error
        println("Relative truncation error ($norm_type norm) in SVD: $error")
    end

    # Reshape U and V to tensors with shapes matching the shape of A and
    # return.
    dim = size(S)[1]
    U_tens = reshape(U, shp_a..., dim)
    Vt_tens = reshape(Vt, dim, shp_b...)
    retval = (U_tens, S, Vt_tens)
    if return_error
        retval = (retval..., error)
    end
    return retval
end

"""
Format the bond dimensions listed in chis to a standard format.
"""
function format_trunc_chis(v, chis, eps)
    max_dim = length(v)
    if chis == nothing
        if eps > 0
            # Try all possible chis.
            chis = collect(1:max_dim)
        else
            # No truncation.
            chis = [max_dim]
        end
    else
        if isa(chis, Number)
            # Wrap an individual number in an Array.
            chis = [chis]
        else
            # Make sure chis is an Array, and not, say, a tuple.
            chis = collect(chis)
        end
        if eps == 0
            chis = [maximum(chis)]
        else
            sort!(chis)
        end
    end
    # If some of the chis are larger than max_dim, get rid of them.
    for (i, chi) in enumerate(chis)
        if chi >= max_dim
            chis[i] = max_dim
            chis = chis[1:i]
            break
        end
    end
    return chis
end

"""
Transpose A so that the indices listed in a are on the left and the indices
listed in b on the right, and reshape A into a matrix.

a and b should be arrays of Integers that together include the numbers from 1
to ndims(A). Alternatively they can be just individual Integers, if only one
index is left on one side.

If return_tensor_shape is true (by default it's not) return, in addition to the
matrix, the shape of the tensor after the transpose but before the reshape.
"""
function to_matrix(A, a, b; return_tensor_shape=false)
    # Make sure a and b are Arrays.
    if isa(a, Number)
        a = [a]
    elseif !isa(a, Array)
        a = collect(a)
    end
    if isa(b, Number)
        b = [b]
    elseif !isa(b, Array)
        b = collect(b)
    end

    # Permute the indices of A to the right order
    perm = vcat(a, b)
    A = tensorcopy(A, collect(1:ndims(A)), perm)
    # The lists shp_a and shp_b list the dimensions of the bonds in a and b
    shp = size(A)
    shp_a = shp[1:length(a)]
    shp_b = shp[end+1-length(b):end]

    # Compute the dimensions of the the matrix that will be formed when
    # indices of a and b are joined together.
    dim_a = prod(shp_a)
    dim_b = prod(shp_b)

    A = reshape(A, dim_a, dim_b)

    if return_tensor_shape
        return A, shp_a, shp_b
    else
        return A
    end
end

"""
Finds the dimension to which v should be truncated, and the error caused in
this truncation. See documentation for tensorsvd for the meaning of the
different arguments.
"""
function find_trunc_dim(v, chis, eps,
                        break_degenerate=false, degeneracy_eps=1e-6,
                        norm_type=:frobenius)
    # Put chis in a standard format.
    chis = format_trunc_chis(v, chis, eps)
    v = abs.(v)
    if !issorted(v, rev=true)
        sort!(v, rev=true)
    end

    if norm_type==:frobenius
        v = v.^2
        eps = eps^2
    elseif norm_type == :trace
        # Nothing to be done
    else
        throw(ArgumentError("Unknown norm_type $norm_type."))
    end
    sum_all = sum(v)
    # Find the smallest chi for which the error is small enough.
    # If none is found, use the largest chi.
    if sum_all != 0
        for i in 1:length(chis)
            chi = chis[i]
            if !break_degenerate
                # Make sure that we don't break degenerate singular values by
                # including one but not the other by decreasing chi if
                # necessary.
                while 0 < chi < length(v)
                    last_in = v[chi]
                    last_out = v[chi+1]
                    rel_diff = abs(last_in - last_out)/last_in
                    if rel_diff < degeneracy_eps
                        chi -= 1
                    else
                        break
                    end
                end
            end
            sum_disc = sum(v[chi+1:end])
            error = sum_disc/sum_all
            if error <= eps
                break
            end
        end
        if norm_type==:frobenius
            error = sqrt(error)
        end
    else
        error = 0
        chi = minimum(chis)
    end
    return chi, error
end

end