text:
function make_batch(o, samples, vocab)
    images = make_images_batch(o, map(s->s[1], samples))
    x, y, batchsizes = make_captions_batch(o, map(s->s[2], samples), vocab)
    return images, x, y, batchsizes
end

function make_images_batch(o, filenames)
    images = h5open(o[:images], "r") do f
        reduce(
            (x...)->cat(o[:finetune]?4:2, x...),
            map(x->read(f,x), filenames))
    end

    batch = nothing
    if o[:finetune]
        batch = zeros(typeof(images[1]), size(images,1:3...)..., o[:batchsize])
        batch[:,:,:,1:length(filenames)] = images
    else
        # batch = zeros(typeof(images[1]), size(images,1), o[:batchsize])
        # batch[:,1:length(filenames)] = images
        batch = images
    end

    return batch
end

function make_captions_batch(o, words, vocab)
    tokens = map(wi->sen2vec(vocab, wi), words)
    longest = mapreduce(length, max, tokens)
    batchsizes = zeros(Int, longest-1)
    x = Int32[]; y = Int32[]
    for t = 1:longest-1
        for i = 1:length(tokens)
            length(tokens[i])-1 < t && break
            # @show i, t, size(tokens[i])
            push!(x, tokens[i][t])
            push!(y, tokens[i][t+1])
            batchsizes[t] += 1
        end
    end
    return x, y, batchsizes
end

max_stars_repo_name: ilkerkesen/kusrp

max_issues_repo_name: ilkerkesen/kusrp

max_forks_repo_name: ilkerkesen/kusrp


text:
#!/usr/bin/julia

using DataStructures: CircularBuffer

function initialize!(cb::CircularBuffer, value)
    push!(cb, -1)
    push!(cb, value)
    return cb
end

function count_trees(nums::AbstractArray{T, 1}, iter::Int=2020) where {T<:Integer}
    mem = Dict{T, CircularBuffer{Int}}(
        n => initialize!(CircularBuffer{Int}(2), i)
        for (i, n) in enumerate(nums)
    )
    # sizehint!(mem, iter÷5)
    num_nums = length(nums)
    last = nums[end]
    for i in num_nums+1:iter
        if (mem[last][1] < 0)
            last = 0
        else
            last = mem[last][2] - mem[last][1]
        end
        if !haskey(mem, last)
            mem[last] = initialize!(CircularBuffer{Int}(2), i)
            continue
        end
        push!(mem[last], i)
    end
    return last
end

function main()
    test_cases = [
        [1,3,2],
        [2,1,3],
        [1,2,3],
        [2,3,1],
        [3,2,1],
        [3,1,2],
        [1,0,15,2,10,13],
    ]

    iter_num = 2020
    for nums in test_cases
        @info iter_num, nums, count_trees(nums, iter_num)
        @time count_trees(nums, iter_num)
    end

    iter_num = 30000000
    for nums in test_cases
        @info iter_num, nums, count_trees(nums, iter_num)
        @time count_trees(nums, iter_num)
    end

    nothing
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end

ext: jl


text:
using POMDPs
using POMDPModels
using POMDPModelTools
using StaticArrays
using JLD2, FileIO
using Random
using Logging
using Distributions
using GridInterpolations
using LocalFunctionApproximation
using LocalApproximationValueIteration
using JSON
using CMSSPs

global_logger(SimpleLogger(stderr, Logging.Error))

rng = MersenneTwister(3456)

param_files = ["./paramsets/scale-1.toml","./paramsets/simtime-1.toml","./paramsets/cost-1.toml"]

scale_file = param_files[1]
simtime_file = param_files[2]
cost_file = param_files[3]

inhor_fn = ARGS[1]
outhor_fn = ARGS[2]
infhor_fn = "./policies/dreamr-uf-params111.jld2"

ep_file_prefix = ARGS[3]
num_eps = parse(Int64, ARGS[4])
outfn = ARGS[5]

# Parse parameters
params = parse_params(scale_file=scale_file, simtime_file=simtime_file, cost_file=cost_file)

@info "Loading policies"
flight_policy = load_localapproxvi_policy_from_jld2(infhor_fn)
hopon_policy = load_modal_horizon_policy_localapprox(inhor_fn, outhor_fn)
ride_policy = DREAMRDeterministicPolicy(params)

# Create modal policies
modal_policies = Dict(FLIGHT => ModalFinInfHorPolicy(hopon_policy, flight_policy),
                      RIDE => ModalFinInfHorPolicy(ride_policy,flight_policy))

deltaT = convert(Int64, (params.time_params.MAX_REPLAN_TIMESTEP/params.time_params.MDP_TIMESTEP))

costs = Vector{Float64}(undef, 0)
steps = Vector{Int64}(undef, 0)
mode_switches = Vector{Int64}(undef, 0)
unsuccessful = Vector{Int64}(undef, 0)

for iter = 1:num_eps

    episode_dict = Dict

    ep_file = string(ep_file_prefix,"-",iter,".json")

    @show ep_file

    open(ep_file,"r") do f
        episode_dict = JSON.parse(f)
    end

    context_set = get_dreamr_episode_context(episode_dict["epochs"])

    start_pos = Point(episode_dict["start_pos"][1], episode_dict["start_pos"][2])
    goal_pos = Point(episode_dict["goal_pos"][1], episode_dict["goal_pos"][2])

    start_state = DREAMRStateType{MultiRotorUAVState}(FLIGHT, get_state_at_rest(MultiRotorUAVState, start_pos))

    # Create CMSSP
    cmssp = create_dreamr_cmssp(MultiRotorUAVState, MultiRotorUAVAction, context_set, goal_pos, params)

    # Create solver
    dreamr_solver = DREAMRSolverType{MultiRotorUAVState,typeof(rng)}(0, modal_policies,
                                    deltaT, [FLIGHT], start_state, DREAMRBookkeeping(), rng)

    (cost, timesteps, successful, num_mode_switches) = solve(dreamr_solver, cmssp)

    if successful
        energy_cost = cost - params.cost_params.TIME_COEFFICIENT*params.time_params.MDP_TIMESTEP*timesteps
        push!(costs, energy_cost)
        push!(steps, timesteps)
        push!(mode_switches, num_mode_switches)
    else
        push!(unsuccessful, iter)
    end
end

results_dict = Dict("costs"=>costs, "steps"=>steps, "mode_switches"=>mode_switches, "failed_iters"=>unsuccessful)

open(outfn, "w") do f
    JSON.print(f, results_dict, 2)
end


text:
using StaticArrays

"""
    butcher_tableau_radau_2stages()

Returns (A,b,c) corresponding to the Butcher tableau for the 2 stage Radau IIA scheme.
"""
function butcher_tableau_radau_2stages()
	A = @SMatrix [5/12   -1/12
	              3/4    1/4];
	b = @SVector [3/4, 1/4];
	c = @SVector [1/3, 1.0];
	return (A, b, c);
end

"""
    butcher_tableau_radau_3stages()

Returns (A,b,c) corresponding to the Butcher tableau for the 3 stage Radau IIA scheme.
"""
function butcher_tableau_radau_3stages()
	A = @SMatrix [(88.0-7.0*sqrt(6.0))/360.0       (296.0-169.0*sqrt(6.0))/1800.0   (-2.0+3.0*sqrt(6.0))/225.0
	              (296.0+169.0*sqrt(6.0))/1800.0   (88.0+7.0*sqrt(6.0))/360.0       (-2.0-3.0*sqrt(6.0))/225.0
	              (16.0-sqrt(6.0))/36.0            (16.0+sqrt(6.0))/36.0             1/9];
	b = @SVector [(16.0-sqrt(6.0))/36.0,    (16.0+sqrt(6.0))/36.0,    1/9];
	c = @SVector [(4.0-sqrt(6.0))/10.0,     (4.0+sqrt(6.0))/10.0,     1.0];
	return (A, b, c);
end


text:
using BenchmarkTools

using Memoize
using Memoization
using MemoizedMethods

Memoize.@memoize fib_Memoize(n::Int)::Int = n <= 1 ? 1 : fib_Memoize(n - 1) + fib_Memoize(n - 2)
Memoization.@memoize fib_Memoization(n::Int)::Int = n <= 1 ? 1 : fib_Memoization(n - 1) + fib_Memoization(n - 2)
MemoizedMethods.@memoize fib_MemoizedMethods(n::Int)::Int = n <= 1 ? 1 : fib_MemoizedMethods(n - 1) + fib_MemoizedMethods(n - 2)

const N = 10000

println("Default IdDict (initialize):")

println("fib_Memoize($N):")
@time fib_Memoize(N)

println("fib_Memoization($N):")
@time fib_Memoization(N)

println("fib_MemoizedMethods($N):")
@time fib_MemoizedMethods(N)

println()
println("Default IdDict (refill):")

println("fib_Memoize($N):")
@btime fib_Memoize($N) setup = empty!(Memoize.memoize_cache(fib_Memoize))

println("fib_Memoization($N):")
@btime fib_Memoization($N) setup = Memoization.empty_cache!(fib_Memoization)

println("fib_MemoizedMethods($N):")
@btime fib_MemoizedMethods($N) setup = map(forget!, methods(fib_MemoizedMethods))

# Note that for Memoize and Memoization, we cannot overwrite the old cache without supplying a new function name.
Memoize.@memoize Dict{Tuple{Int}}{Int} fib_Memoize_2(n::Int)::Int = n <= 1 ? 1 : fib_Memoize_2(n - 1) + fib_Memoize_2(n - 2)
Memoization.@memoize Dict{Tuple{Tuple{Int},NamedTuple{(),Tuple{}}}}{Int} fib_Memoization_2(n::Int)::Int = n <= 1 ? 1 : fib_Memoization_2(n - 1) + fib_Memoization_2(n - 2)
MemoizedMethods.@memoize Dict{__Key__}{__Value__}() fib_MemoizedMethods(n::Int)::Int = n <= 1 ? 1 : fib_MemoizedMethods(n - 1) + fib_MemoizedMethods(n - 2)

println()
println("Typed Dict (initialize):")

println("fib_Memoize($N):")
@time fib_Memoize_2(N)

println("fib_Memoization($N):")
@time fib_Memoization_2(N)

println("fib_MemoizedMethods($N):")
@time fib_MemoizedMethods(N)

println()
println("Typed Dict (refill):")

println("fib_Memoize($N):")
@btime fib_Memoize_2($N) setup = empty!(Memoize.memoize_cache(fib_Memoize_2))

println("fib_Memoization($N):")
@btime fib_Memoization_2($N) setup = Memoization.empty_cache!(fib_Memoization_2)

println("fib_MemoizedMethods($N):")
@btime fib_MemoizedMethods($N) setup = map(forget!, methods(fib_MemoizedMethods))

Memoize.@memoize Dict{Tuple{Int, Int}}{Int} kwfib_Memoize(n::Int; b::Int = 1)::Int = n <= 1 ? b : kwfib_Memoize(n - 1; b=b) + kwfib_Memoize(n - 2; b=b)
Memoization.@memoize Dict{Tuple{Tuple{Int},NamedTuple{(:b,),Tuple{Int}}}}{Int} kwfib_Memoization(n::Int; b::Int = 1)::Int = n <= 1 ? b : kwfib_Memoization(n - 1; b=b) + kwfib_Memoization(n - 2; b=b)
MemoizedMethods.@memoize Dict{__Key__}{__Value__}() kwfib_MemoizedMethods(n::Int; b::Int = 1)::Int = n <= 1 ? b : kwfib_MemoizedMethods(n - 1; b=b) + kwfib_MemoizedMethods(n - 2; b=b)

const N = 10000

println("Typed Dict With Kwargs (initialize):")

println("kwfib_Memoize($N):")
@time kwfib_Memoize(N, b=1)

println("kwfib_Memoization($N):")
@time kwfib_Memoization(N, b=1)

println("kwfib_MemoizedMethods($N):")
@time kwfib_MemoizedMethods(N, b=1)

println()
println("Typed Dict With Kwargs (refill):")

println("kwfib_Memoize($N):")
@btime kwfib_Memoize($N, b=1) setup = empty!(Memoize.memoize_cache(kwfib_Memoize))

println("kwfib_Memoization($N):")
@btime kwfib_Memoization($N, b=1) setup = Memoization.empty_cache!(kwfib_Memoization)

println("kwfib_MemoizedMethods($N):")
@btime kwfib_MemoizedMethods($N, b=1) setup = map(forget!, methods(kwfib_MemoizedMethods))


text:
using Hypatia
using Empirikos
using Test

# We compare the chi^2 F-localization intervals against the intervals reported
# by Lord and Cressie.

lord_cressie = LordCressie.load_table()
Zs = LordCressie.ebayes_samples()

idx_test = [1, 13, 20]
z_cressie = lord_cressie.x[idx_test]
lower_test = lord_cressie.Lower1[idx_test]
upper_test = lord_cressie.Upper1[idx_test]

gcal = DiscretePriorClass(range(0.0,stop=1.0,length=300));

postmean_targets = Empirikos.PosteriorMean.(BinomialSample.(z_cressie,20));
chisq_floc = Empirikos.ChiSquaredFLocalization(0.05)

floc_method_chisq = FLocalizationInterval(flocalization = chisq_floc,
                                       convexclass= gcal, solver=Hypatia.Optimizer)

chisq_cis = confint.(floc_method_chisq, postmean_targets, Zs)

lower_chisq_ci = getproperty.(chisq_cis, :lower)
upper_chisq_ci = getproperty.(chisq_cis, :upper)

@test lower_chisq_ci ≈ lower_test atol=1e-3
@test upper_chisq_ci ≈ upper_test atol=0.005

Zs_collapse = begin
	Zs_collapse = deepcopy(Zs)
	n0 = pop!(Zs_collapse.store, BinomialSample(0, 20))
	n1 = pop!(Zs_collapse.store, BinomialSample(1, 20))
	updated_keys =  [BinomialSample(Interval(0,1), 20); collect(keys(Zs_collapse))]
	updated_values = [n0+n1; collect(values(Zs_collapse))]
	Empirikos.MultinomialSummary(updated_keys, updated_values)
end

discr = integer_discretizer(1:20)
Zs_collapse2 = discr(Zs)
discr(Zs_collapse) == Zs_collapse2

@test nobs(Zs_collapse) == nobs(Zs_collapse2)

floc_method_dkw = FLocalizationInterval(
							flocalization = DvoretzkyKieferWolfowitz(;α=0.05),
                            convexclass= gcal, solver=Hypatia.Optimizer)

fitted_dkw = fit(DvoretzkyKieferWolfowitz(;α=0.05), Zs_collapse)
fitted_dkw2 = fit(DvoretzkyKieferWolfowitz(;α=0.05), Zs_collapse2)

postmean_target = postmean_targets[2]
dkw_ci = confint(floc_method_dkw, postmean_target, Zs_collapse);
dkw_ci2 = confint(floc_method_dkw, postmean_target, Zs_collapse2);

@test dkw_ci == dkw_ci2

lam_chisq_withF = Empirikos.AMARI(
							convexclass=gcal,
                            flocalization = Empirikos.ChiSquaredFLocalization(0.01),
                            solver=Hypatia.Optimizer, discretizer=discr,
                            modulus_model = Empirikos.ModulusModelWithF
                           )

postmean_ci_lam_withF = confint(lam_chisq_withF, postmean_target, Zs_collapse2)

lam_chisq_withoutF = Empirikos.AMARI(
							convexclass=gcal,
                            flocalization = Empirikos.ChiSquaredFLocalization(0.01),
                            solver=Hypatia.Optimizer, discretizer=discr,
                            modulus_model = Empirikos.ModulusModelWithoutF
                           )

postmean_ci_lam_withoutF = confint(lam_chisq_withoutF, postmean_target, Zs_collapse2)

@test postmean_ci_lam_withoutF.lower ≈ postmean_ci_lam_withF.lower
@test postmean_ci_lam_withoutF.upper ≈ postmean_ci_lam_withF.upper


text:
using Random, Distributions, StatsPlots; pyplot()
Random.seed!(0)

mu1, sigma1 = 10, 5
mu2, sigma2 = 40, 12
dist1, dist2 = Normal(mu1,sigma1), Normal(mu2,sigma2)
p = 0.3
mixRv() = (rand() <= p) ? rand(dist1) : rand(dist2)

n = 2000
data = [mixRv() for _ in 1:n]

density(data, c=:blue, label="Density via StatsPlots",
        xlims=(-20,80), ylims=(0,0.035))
stephist!(data, bins=50, c=:black, norm=true,
    label="Histogram", xlabel="x", ylabel = "Density")

max_stars_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_issues_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_forks_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook


text:
#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modeling language for Julia
# See https://github.com/jump-dev/JuMP.jl
#############################################################################

"""
    JuMP

An algebraic modeling language for Julia.

For more information, go to https://jump.dev.
"""
module JuMP

using LinearAlgebra
using SparseArrays

import MutableArithmetics
const _MA = MutableArithmetics

import MathOptInterface

"""
    MOI

Shorthand for the MathOptInterface package.
"""
const MOI = MathOptInterface

"""
    MOIU

Shorthand for the MathOptInterface.Utilities package.
"""
const MOIU = MOI.Utilities

"""
    MOIB

Shorthand for the MathOptInterface.Bridges package.
"""
const MOIB = MOI.Bridges

import Calculus
import DataStructures.OrderedDict
import ForwardDiff
include("_Derivatives/_Derivatives.jl")
using ._Derivatives

include("Containers/Containers.jl")

# Exports are at the end of the file.

# Deprecations for JuMP v0.18 -> JuMP v0.19 transition
Base.@deprecate(getobjectivevalue, JuMP.objective_value)
Base.@deprecate(getobjectivebound, JuMP.objective_bound)
Base.@deprecate(getvalue, JuMP.value)
Base.@deprecate(getdual, JuMP.dual)
Base.@deprecate(numvar, JuMP.num_variables)
Base.@deprecate(numnlconstr, JuMP.num_nl_constraints)
Base.@deprecate(setlowerbound, JuMP.set_lower_bound)
Base.@deprecate(setupperbound, JuMP.set_upper_bound)
Base.@deprecate(linearterms, JuMP.linear_terms)

function writeLP(args...; kargs...)
    return error(
        "writeLP has been removed from JuMP. Use `write_to_file` instead.",
    )
end
function writeMPS(args...; kargs...)
    return error(
        "writeMPS has been removed from JuMP. Use `write_to_file` instead.",
    )
end

include("utils.jl")

const _MOIVAR = MOI.VariableIndex
const _MOICON{F,S} = MOI.ConstraintIndex{F,S}

"""
    optimizer_with_attributes(optimizer_constructor, attrs::Pair...)

Groups an optimizer constructor with the list of attributes `attrs`. Note that
it is equivalent to `MOI.OptimizerWithAttributes`.

When provided to the `Model` constructor or to [`set_optimizer`](@ref), it
creates an optimizer by calling `optimizer_constructor()`, and then sets the
attributes using [`set_optimizer_attribute`](@ref).

## Example

```julia
model = Model(
    optimizer_with_attributes(
        Gurobi.Optimizer, "Presolve" => 0, "OutputFlag" => 1
    )
)
```
is equivalent to:
```julia
model = Model(Gurobi.Optimizer)
set_optimizer_attribute(model, "Presolve", 0)
set_optimizer_attribute(model, "OutputFlag", 1)
```

## Note

The string names of the attributes are specific to each solver. One should
consult the solver's documentation to find the attributes of interest.

See also: [`set_optimizer_attribute`](@ref), [`set_optimizer_attributes`](@ref),
[`get_optimizer_attribute`](@ref).
"""
function optimizer_with_attributes(optimizer_constructor, args::Pair...)
    return MOI.OptimizerWithAttributes(optimizer_constructor, args...)
end

function with_optimizer(constructor; kwargs...)
    if isempty(kwargs)
        deprecation_message = """
`with_optimizer` is deprecated. Adapt the following example to update your code:
`with_optimizer(Ipopt.Optimizer)` becomes `Ipopt.Optimizer`.
"""
        Base.depwarn(deprecation_message, :with_optimizer)
        return constructor
    else
        deprecation_message = """
`with_optimizer` is deprecated. Adapt the following example to update your code:
`with_optimizer(Ipopt.Optimizer, max_cpu_time=60.0)` becomes `optimizer_with_attributes(Ipopt.Optimizer, "max_cpu_time" => 60.0)`.
"""
        Base.depwarn(deprecation_message, :with_optimizer_kw)
        params =
            [MOI.RawParameter(string(kw.first)) => kw.second for kw in kwargs]
        return MOI.OptimizerWithAttributes(constructor, params)
    end
end
function with_optimizer(constructor, args...; kwargs...)
    if isempty(kwargs)
        deprecation_message = """
`with_optimizer` is deprecated. Adapt the following example to update your code:
`with_optimizer(Gurobi.Optimizer, env)` becomes `() -> Gurobi.Optimizer(env)`.
"""
        Base.depwarn(deprecation_message, :with_optimizer_args)
        if !applicable(constructor, args...)
            error(
                "$constructor does not have any method with arguments $args.",
                " The first argument of `with_optimizer` should be callable with",
                " the other argument of `with_optimizer`.",
            )
        end
        return with_optimizer(() -> constructor(args...); kwargs...)
    else
        deprecation_message = """
`with_optimizer` is deprecated. Adapt the following example to update your code:
`with_optimizer(Gurobi.Optimizer, env, Presolve=0)` becomes `optimizer_with_attributes(() -> Gurobi.Optimizer(env), "Presolve" => 0)`.
"""
        Base.depwarn(deprecation_message, :with_optimizer_args_kw)
        if !applicable(constructor, args...)
            error(
                "$constructor does not have any method with arguments $args.",
                " The first argument of `with_optimizer` should be callable with",
                " the other argument of `with_optimizer`.",
            )
        end
        params =
            [MOI.RawParameter(string(kw.first)) => kw.second for kw in kwargs]
        return MOI.OptimizerWithAttributes(() -> constructor(args...), params)
    end
end

include("shapes.jl")

# Model

"""
    ModelMode

An enum to describe the state of the CachingOptimizer inside a JuMP model.
"""
@enum(ModelMode, AUTOMATIC, MANUAL, DIRECT)
@doc(
    "`moi_backend` field holds a CachingOptimizer in AUTOMATIC mode.",
    AUTOMATIC
)
@doc("`moi_backend` field holds a CachingOptimizer in MANUAL mode.", MANUAL)
@doc(
    "`moi_backend` field holds an AbstractOptimizer. No extra copy of the " *
    "model is stored. The `moi_backend` must support `add_constraint` etc.",
    DIRECT,
)

"""
    AbstractModel

An abstract type that should be subtyped for users creating JuMP extensions.
"""
abstract type AbstractModel end
# All `AbstractModel`s must define methods for these functions:
# num_variables, object_dictionary

"""
    Model

A mathematical model of an optimization problem.
"""
mutable struct Model <: AbstractModel
    # In MANUAL and AUTOMATIC modes, CachingOptimizer.
    # In DIRECT mode, will hold an AbstractOptimizer.
    moi_backend::MOI.AbstractOptimizer
    # List of shapes of constraints that are not `ScalarShape` or `VectorShape`.
    shapes::Dict{_MOICON,AbstractShape}
    # List of bridges to add in addition to the ones added in
    # `MOI.Bridges.full_bridge_optimizer`. With `BridgeableConstraint`, the
    # same bridge may be added many times so we store them in a `Set` instead
    # of, e.g., a `Vector`.
    bridge_types::Set{Any}
    # Hook into a solve call...function of the form f(m::Model; kwargs...),
    # where kwargs get passed along to subsequent solve calls.
    optimize_hook::Any
    # TODO: Document.
    nlp_data::Any
    # Dictionary from variable and constraint names to objects.
    obj_dict::Dict{Symbol,Any}
    # Number of times we add large expressions. Incremented and checked by
    # the `operator_warn` method.
    operator_counter::Int
    # Enable extensions to attach arbitrary information to a JuMP model by
    # using an extension-specific symbol as a key.
    ext::Dict{Symbol,Any}
end

"""
    Model(; caching_mode::MOIU.CachingOptimizerMode=MOIU.AUTOMATIC)

Return a new JuMP model without any optimizer; the model is stored the model in
a cache. The mode of the `CachingOptimizer` storing this cache is
`caching_mode`. Use [`set_optimizer`](@ref) to set the optimizer before
calling [`optimize!`](@ref).
"""
function Model(;
    caching_mode::MOIU.CachingOptimizerMode = MOIU.AUTOMATIC,
    solver = nothing,
)
    if solver !== nothing
        error(
            "The solver= keyword is no longer available in JuMP 0.19 and " *
            "later. See the JuMP documentation " *
            "(https://jump.dev/JuMP.jl/latest/) for latest syntax.",
        )
    end
    universal_fallback = MOIU.UniversalFallback(MOIU.Model{Float64}())
    caching_opt = MOIU.CachingOptimizer(universal_fallback, caching_mode)
    return direct_model(caching_opt)
end

"""
    Model(optimizer_factory;
          caching_mode::MOIU.CachingOptimizerMode=MOIU.AUTOMATIC,
          bridge_constraints::Bool=true)

Return a new JuMP model with the provided optimizer and bridge settings. This
function is equivalent to:
```julia
    model = Model()
    set_optimizer(model, optimizer_factory,
                  bridge_constraints=bridge_constraints)
    return model
```
See [`set_optimizer`](@ref) for the description of the `optimizer_factory` and
`bridge_constraints` arguments.

## Examples

The following creates a model with the optimizer set to `Ipopt`:
```julia
model = Model(Ipopt.Optimizer)
```
"""
function Model(optimizer_factory; bridge_constraints::Bool = true, kwargs...)
    model = Model(; kwargs...)
    set_optimizer(
        model,
        optimizer_factory,
        bridge_constraints = bridge_constraints,
    )
    return model
end

"""
    direct_model(backend::MOI.ModelLike)

Return a new JuMP model using [`backend`](@ref) to store the model and solve it.

As opposed to the [`Model`](@ref) constructor, no cache of the model is stored
outside of [`backend`](@ref) and no bridges are automatically applied to
[`backend`](@ref).

## Notes

The absence of a cache reduces the memory footprint but, it is important to bear
in mind the following implications of creating models using this *direct* mode:

* When [`backend`](@ref) does not support an operation, such as modifying
  constraints or adding variables/constraints after solving, an error is
  thrown. For models created using the [`Model`](@ref) constructor, such
  situations can be dealt with by storing the modifications in a cache and
  loading them into the optimizer when `optimize!` is called.
* No constraint bridging is supported by default.
* The optimizer used cannot be changed the model is constructed.
* The model created cannot be copied.
"""
function direct_model(backend::MOI.ModelLike)
    @assert MOI.is_empty(backend)
    return Model(
        backend,
        Dict{_MOICON,AbstractShape}(),
        Set{Any}(),
        nothing,
        nothing,
        Dict{Symbol,Any}(),
        0,
        Dict{Symbol,Any}(),
    )
end

Base.broadcastable(model::Model) = Ref(model)

"""
    backend(model::Model)

Return the lower-level MathOptInterface model that sits underneath JuMP. This
model depends on which operating mode JuMP is in (see [`mode`](@ref)), and
whether there are any bridges in the model.

If JuMP is in `DIRECT` mode (i.e., the model was created using
[`direct_model`](@ref)), the backend will be the optimizer passed to
[`direct_model`](@ref).

If JuMP is in `MANUAL` or `AUTOMATIC` mode, the backend is a
`MOI.Utilities.CachingOptimizer`.

This function should only be used by advanced users looking to access low-level
MathOptInterface or solver-specific functionality.
"""
backend(model::Model) = model.moi_backend

moi_mode(model::MOI.ModelLike) = DIRECT
function moi_mode(model::MOIU.CachingOptimizer)
    if model.mode == MOIU.AUTOMATIC
        return AUTOMATIC
    else
        return MANUAL
    end
end

"""
    mode(model::Model)

Return the [`ModelMode`](@ref) ([`DIRECT`](@ref), [`AUTOMATIC`](@ref), or
[`MANUAL`](@ref)) of `model`.
"""
function mode(model::Model)
    # The type of `backend(model)` is not type-stable, so we use a function
    # barrier (`moi_mode`) to improve performance.
    return moi_mode(backend(model))
end

# Direct mode
moi_bridge_constraints(model::MOI.ModelLike) = false
function moi_bridge_constraints(model::MOIU.CachingOptimizer)
    return model.optimizer isa MOI.Bridges.LazyBridgeOptimizer
end

# Internal function.
function _try_get_solver_name(model_like)
    try
        return MOI.get(model_like, MOI.SolverName())::String
    catch ex
        if isa(ex, ArgumentError)
            return "SolverName() attribute not implemented by the optimizer."
        else
            rethrow(ex)
        end
    end
end

"""
    solver_name(model::Model)

If available, returns the `SolverName` property of the underlying optimizer.

Returns `"No optimizer attached"` in `AUTOMATIC` or `MANUAL` modes when no
optimizer is attached.

Returns `"SolverName() attribute not implemented by the optimizer."` if the
attribute is not implemented.
"""
function solver_name(model::Model)
    if mode(model) != DIRECT && MOIU.state(backend(model)) == MOIU.NO_OPTIMIZER
        return "No optimizer attached."
    else
        return _try_get_solver_name(backend(model))
    end
end

"""
    bridge_constraints(model::Model)

When in direct mode, return `false`.
When in manual or automatic mode, return a `Bool` indicating whether the
optimizer is set and unsupported constraints are automatically bridged
to equivalent supported constraints when an appropriate transformation is
available.
"""
function bridge_constraints(model::Model)
    # The type of `backend(model)` is not type-stable, so we use a function
    # barrier (`moi_bridge_constraints`) to improve performance.
    return moi_bridge_constraints(backend(model))
end

function _moi_add_bridge(
    model::Nothing,
    BridgeType::Type{<:MOI.Bridges.AbstractBridge},
)
    # No optimizer is attached, the bridge will be added when one is attached
    return
end
function _moi_add_bridge(
    model::MOI.ModelLike,
    BridgeType::Type{<:MOI.Bridges.AbstractBridge},
)
    return error(
        "Cannot add bridge if `bridge_constraints` was set to `false` in the",
        " `Model` constructor.",
    )
end
function _moi_add_bridge(
    bridge_opt::MOI.Bridges.LazyBridgeOptimizer,
    BridgeType::Type{<:MOI.Bridges.AbstractBridge},
)
    MOI.Bridges.add_bridge(bridge_opt, BridgeType{Float64})
    return
end
function _moi_add_bridge(
    caching_opt::MOIU.CachingOptimizer,
    BridgeType::Type{<:MOI.Bridges.AbstractBridge},
)
    _moi_add_bridge(caching_opt.optimizer, BridgeType)
    return
end

"""
     add_bridge(model::Model,
                BridgeType::Type{<:MOI.Bridges.AbstractBridge})

Add `BridgeType` to the list of bridges that can be used to transform
unsupported constraints into an equivalent formulation using only constraints
supported by the optimizer.
"""
function add_bridge(
    model::Model,
    BridgeType::Type{<:MOI.Bridges.AbstractBridge},
)
    push!(model.bridge_types, BridgeType)
    # The type of `backend(model)` is not type-stable, so we use a function
    # barrier (`_moi_add_bridge`) to improve performance.
    _moi_add_bridge(JuMP.backend(model), BridgeType)
    return
end

"""
     print_bridge_graph([io::IO,] model::Model)

Print the hyper-graph containing all variable, constraint, and objective types
that could be obtained by bridging the variables, constraints, and objectives
that are present in the model.

Each node in the hyper-graph corresponds to a variable, constraint, or objective
type.
  * Variable nodes are indicated by `[ ]`
  * Constraint nodes are indicated by `( )`
  * Objective nodes are indicated by `| |`
The number inside each pair of brackets is an index of the node in the
hyper-graph.

Note that this hyper-graph is the full list of possible transformations. When
the bridged model is created, we select the shortest hyper-path(s) from this
graph, so many nodes may be un-used.

For more information, see Legat, B., Dowson, O., Garcia, J., and Lubin, M.
(2020).  "MathOptInterface: a data structure for mathematical optimization
problems." URL: [https://arxiv.org/abs/2002.03447](https://arxiv.org/abs/2002.03447)
"""
print_bridge_graph(model::Model) = print_bridge_graph(Base.stdout, model)

function print_bridge_graph(io::IO, model::Model)
    # The type of `backend(model)` is not type-stable, so we use a function
    # barrier (`_moi_print_bridge_graph`) to improve performance.
    return _moi_print_bridge_graph(io, backend(model))
end

function _moi_print_bridge_graph(io::IO, model::MOI.Bridges.LazyBridgeOptimizer)
    return MOI.Bridges.print_graph(io, model)
end

function _moi_print_bridge_graph(io::IO, model::MOIU.CachingOptimizer)
    return _moi_print_bridge_graph(io, model.optimizer)
end

function _moi_print_bridge_graph(::IO, ::MOI.ModelLike)
    return error(
        "Cannot print bridge graph if `bridge_constraints` was set to " *
        "`false` in the `Model` constructor.",
    )
end

"""
    empty!(model::Model)::Model

Empty the model, that is, remove all variables, constraints and model
attributes but not optimizer attributes. Always return the argument.

Note: removes extensions data.
"""
function Base.empty!(model::Model)::Model
    # The method changes the Model object to, basically, the state it was when
    # created (if the optimizer was already pre-configured). The exceptions
    # are:
    # * optimize_hook: it is basically an optimizer attribute and we promise
    #   to leave them alone (as do MOI.empty!).
    # * bridge_types: for consistency with MOI.empty! for
    #   MOI.Bridges.LazyBridgeOptimizer.
    # * operator_counter: it is just a counter for a single-time warning
    #   message (so keeping it helps to discover inneficiencies).
    MOI.empty!(model.moi_backend)
    empty!(model.shapes)
    model.nlp_data = nothing
    empty!(model.obj_dict)
    empty!(model.ext)
    return model
end

"""
    num_variables(model::Model)::Int64

Returns number of variables in `model`.
"""
num_variables(model::Model)::Int64 = MOI.get(model, MOI.NumberOfVariables())

"""
    num_nl_constraints(model::Model)

Returns the number of nonlinear constraints associated with the `model`.
"""
function num_nl_constraints(model::Model)
    return model.nlp_data !== nothing ? length(model.nlp_data.nlconstr) : 0
end

"""
    object_dictionary(model::Model)

Return the dictionary that maps the symbol name of a variable, constraint, or
expression to the corresponding object.

Objects are registered to a specific symbol in the macros.
For example, `@variable(model, x[1:2, 1:2])` registers the array of variables
`x` to the symbol `:x`.

This method should be defined for any subtype of `AbstractModel`.
"""
object_dictionary(model::Model) = model.obj_dict

"""
    unregister(model::Model, key::Symbol)

Unregister the name `key` from `model` so that a new variable, constraint, or
expression can be created with the same key.

Note that this will not delete the object `model[key]`; it will just remove the
reference at `model[key]`. To delete the object, use
```julia
delete(model, model[key])
unregister(model, key)
```

See also: [`object_dictionary`](@ref).

## Examples

```jldoctest; setup=:(model = Model())
julia> @variable(model, x)
x

julia> @variable(model, x)
ERROR: An object of name x is already attached to this model. If
this is intended, consider using the anonymous construction syntax,
e.g., `x = @variable(model, [1:N], ...)` where the name of the object
does not appear inside the macro.

Alternatively, use `unregister(model, :x)` to first unregister the
existing name from the model. Note that this will not delete the object;
it will just remove the reference at `model[:x]`.
[...]

julia> num_variables(model)
1

julia> unregister(model, :x)

julia> @variable(model, x)
x

julia> num_variables(model)
2
```
"""
function unregister(model::AbstractModel, key::Symbol)
    delete!(object_dictionary(model), key)
    return
end

"""
    termination_status(model::Model)

Return the reason why the solver stopped (i.e., the MathOptInterface model
attribute `TerminationStatus`).
"""
function termination_status(model::Model)
    return MOI.get(model, MOI.TerminationStatus())::MOI.TerminationStatusCode
end

"""
    raw_status(model::Model)

Return the reason why the solver stopped in its own words (i.e., the
MathOptInterface model attribute `RawStatusString`).
"""
function raw_status(model::Model)
    return MOI.get(model, MOI.RawStatusString())
end

"""
    primal_status(model::Model; result::Int = 1)

Return the status of the most recent primal solution of the solver (i.e., the
MathOptInterface model attribute `PrimalStatus`) associated with the result
index `result`.

See also: [`result_count`](@ref).
"""
function primal_status(model::Model; result::Int = 1)
    return MOI.get(model, MOI.PrimalStatus(result))::MOI.ResultStatusCode
end

"""
    dual_status(model::Model; result::Int = 1)

Return the status of the most recent dual solution of the solver (i.e., the
MathOptInterface model attribute `DualStatus`) associated with the result
index `result`.

See also: [`result_count`](@ref).
"""
function dual_status(model::Model; result::Int = 1)
    return MOI.get(model, MOI.DualStatus(result))::MOI.ResultStatusCode
end

"""
    set_optimize_hook(model::Model, f::Union{Function,Nothing})

Set the function `f` as the optimize hook for `model`.

`f` should have a signature `f(model::Model; kwargs...)`, where the `kwargs` are
those passed to [`optimize!`](@ref).

## Notes

 * The optimize hook should generally modify the model, or some external state
   in some way, and then call `optimize!(model; ignore_optimize_hook = true)` to
   optimize the problem, bypassing the hook.
 * Use `set_optimize_hook(model, nothing)` to unset an optimize hook.

## Examples

```julia
model = Model()
function my_hook(model::Model; kwargs...)
    print(kwargs)
    return optimize!(model; ignore_optimize_hook = true)
end
set_optimize_hook(model, my_hook)
optimize!(model; test_arg = true)
```
"""
set_optimize_hook(model::Model, f) = (model.optimize_hook = f)

"""
    solve_time(model::Model)

If available, returns the solve time reported by the solver.
Returns "ArgumentError: ModelLike of type `Solver.Optimizer` does not support accessing
the attribute MathOptInterface.SolveTime()" if the attribute is
not implemented.
"""
function solve_time(model::Model)
    return MOI.get(model, MOI.SolveTime())
end

"""
    set_optimizer_attribute(model::Model, name::String, value)

Sets solver-specific attribute identified by `name` to `value`.

Note that this is equivalent to
`set_optimizer_attribute(model, MOI.RawParameter(name), value)`.

## Example

```julia
set_optimizer_attribute(model, "SolverSpecificAttributeName", true)
```

See also: [`set_optimizer_attributes`](@ref), [`get_optimizer_attribute`](@ref).
"""
function set_optimizer_attribute(model::Model, name::String, value)
    return set_optimizer_attribute(model, MOI.RawParameter(name), value)
end

"""
    set_optimizer_attribute(
        model::Model, attr::MOI.AbstractOptimizerAttribute, value
    )

Set the solver-specific attribute `attr` in `model` to `value`.

## Example

```julia
set_optimizer_attribute(model, MOI.Silent(), true)
```

See also: [`set_optimizer_attributes`](@ref), [`get_optimizer_attribute`](@ref).
"""
function set_optimizer_attribute(
    model::Model,
    attr::MOI.AbstractOptimizerAttribute,
    value,
)
    return MOI.set(model, attr, value)
end

@deprecate set_parameter set_optimizer_attribute

"""
    set_optimizer_attributes(model::Model, pairs::Pair...)

Given a list of `attribute => value` pairs, calls
`set_optimizer_attribute(model, attribute, value)` for each pair.

## Example

```julia
model = Model(Ipopt.Optimizer)
set_optimizer_attributes(model, "tol" => 1e-4, "max_iter" => 100)
```
is equivalent to:
```julia
model = Model(Ipopt.Optimizer)
set_optimizer_attribute(model, "tol", 1e-4)
set_optimizer_attribute(model, "max_iter", 100)
```

See also: [`set_optimizer_attribute`](@ref), [`get_optimizer_attribute`](@ref).
"""
function set_optimizer_attributes(model::Model, pairs::Pair...)
    for (name, value) in pairs
        set_optimizer_attribute(model, name, value)
    end
end

@deprecate set_parameters set_optimizer_attributes

"""
    get_optimizer_attribute(model, name::String)

Return the value associated with the solver-specific attribute named `name`.

Note that this is equivalent to
`get_optimizer_attribute(model, MOI.RawParameter(name))`.

## Example

```julia
get_optimizer_attribute(model, "SolverSpecificAttributeName")
```

See also: [`set_optimizer_attribute`](@ref), [`set_optimizer_attributes`](@ref).
"""
function get_optimizer_attribute(model::Model, name::String)
    return get_optimizer_attribute(model, MOI.RawParameter(name))
end

"""
    get_optimizer_attribute(
        model::Model, attr::MOI.AbstractOptimizerAttribute
    )

Return the value of the solver-specific attribute `attr` in `model`.

## Example

```julia
get_optimizer_attribute(model, MOI.Silent())
```

See also: [`set_optimizer_attribute`](@ref), [`set_optimizer_attributes`](@ref).
"""
function get_optimizer_attribute(
    model::Model,
    attr::MOI.AbstractOptimizerAttribute,
)
    return MOI.get(model, attr)
end

"""
    set_silent(model::Model)

Takes precedence over any other attribute controlling verbosity and requires the
solver to produce no output.

See also: [`unset_silent`](@ref).
"""
function set_silent(model::Model)
    return MOI.set(model, MOI.Silent(), true)
end

"""
    unset_silent(model::Model)

Neutralize the effect of the `set_silent` function and let the solver attributes
control the verbosity.

See also: [`set_silent`](@ref).
"""
function unset_silent(model::Model)
    return MOI.set(model, MOI.Silent(), false)
end

"""
    set_time_limit_sec(model::Model, limit)

Set the time limit (in seconds) of the solver.

Can be unset using [`unset_time_limit_sec`](@ref) or with `limit` set to
`nothing`.

See also: [`unset_time_limit_sec`](@ref), [`time_limit_sec`](@ref).
"""
function set_time_limit_sec(model::Model, limit)
    return MOI.set(model, MOI.TimeLimitSec(), limit)
end

"""
    unset_time_limit_sec(model::Model)

Unset the time limit of the solver.

See also: [`set_time_limit_sec`](@ref), [`time_limit_sec`](@ref).
"""
function unset_time_limit_sec(model::Model)
    return MOI.set(model, MOI.TimeLimitSec(), nothing)
end

"""
    time_limit_sec(model::Model)

Return the time limit (in seconds) of the `model`.

Returns `nothing` if unset.

See also: [`set_time_limit_sec`](@ref), [`unset_time_limit_sec`](@ref).
"""
function time_limit_sec(model::Model)
    return MOI.get(model, MOI.TimeLimitSec())
end

"""
    simplex_iterations(model::Model)

Gets the cumulative number of simplex iterations during the most-recent optimization.

Solvers must implement `MOI.SimplexIterations()` to use this function.
"""
function simplex_iterations(model::Model)
    return MOI.get(model, MOI.SimplexIterations())
end

"""
    barrier_iterations(model::Model)

Gets the cumulative number of barrier iterations during the most recent optimization.

Solvers must implement `MOI.BarrierIterations()` to use this function.
"""
function barrier_iterations(model::Model)
    return MOI.get(model, MOI.BarrierIterations())
end

"""
    node_count(model::Model)

Gets the total number of branch-and-bound nodes explored during the most recent
optimization in a Mixed Integer Program.

Solvers must implement `MOI.NodeCount()` to use this function.
"""
function node_count(model::Model)
    return MOI.get(model, MOI.NodeCount())
end

# Abstract base type for all scalar types
# The subtyping of `AbstractMutable` will allow calls of some `Base` functions
# to be redirected to a method in MA that handles type promotion more carefuly
# (e.g. the promotion in sparse matrix products in SparseArrays usually does not
# work for JuMP types) and exploits the mutability of `AffExpr` and `QuadExpr`.
abstract type AbstractJuMPScalar <: _MA.AbstractMutable end
Base.ndims(::Type{<:AbstractJuMPScalar}) = 0

# These are required to create symmetric containers of AbstractJuMPScalars.
LinearAlgebra.symmetric_type(::Type{T}) where {T<:AbstractJuMPScalar} = T
LinearAlgebra.symmetric(scalar::AbstractJuMPScalar, ::Symbol) = scalar
# This is required for linear algebra operations involving transposes.
LinearAlgebra.adjoint(scalar::AbstractJuMPScalar) = scalar

"""
    owner_model(s::AbstractJuMPScalar)

Return the model owning the scalar `s`.
"""
function owner_model end

Base.iterate(x::AbstractJuMPScalar) = (x, true)
Base.iterate(::AbstractJuMPScalar, state) = nothing
Base.isempty(::AbstractJuMPScalar) = false

# Check if two arrays of AbstractJuMPScalars are equal. Useful for testing.
function isequal_canonical(
    x::AbstractArray{<:JuMP.AbstractJuMPScalar},
    y::AbstractArray{<:JuMP.AbstractJuMPScalar},
)
    return size(x) == size(y) && all(JuMP.isequal_canonical.(x, y))
end

include("constraints.jl")
include("variables.jl")
include("objective.jl")

function Base.zero(::Type{V}) where {V<:AbstractVariableRef}
    return zero(GenericAffExpr{Float64,V})
end
Base.zero(v::AbstractVariableRef) = zero(typeof(v))
function Base.one(::Type{V}) where {V<:AbstractVariableRef}
    return one(GenericAffExpr{Float64,V})
end
Base.one(v::AbstractVariableRef) = one(typeof(v))

mutable struct VariableNotOwnedError <: Exception
    context::String
end
function Base.showerror(io::IO, ex::VariableNotOwnedError)
    return print(
        io,
        "VariableNotOwnedError: Variable not owned by model present in $(ex.context)",
    )
end

_moi_optimizer_index(model::MOI.AbstractOptimizer, index::MOI.Index) = index
function _moi_optimizer_index(model::MOIU.CachingOptimizer, index::MOI.Index)
    if MOIU.state(model) == MOIU.NO_OPTIMIZER
        throw(NoOptimizer())
    elseif MOIU.state(model) == MOIU.EMPTY_OPTIMIZER
        error(
            "There is no `optimizer_index` as the optimizer is not ",
            "synchronized with the cached model. Call ",
            "`MOIU.attach_optimizer(model)` to synchronize it.",
        )
    else
        @assert MOIU.state(model) == MOIU.ATTACHED_OPTIMIZER
        return _moi_optimizer_index(
            model.optimizer,
            model.model_to_optimizer_map[index],
        )
    end
end
function _moi_optimizer_index(
    model::MOI.Bridges.LazyBridgeOptimizer,
    index::MOI.Index,
)
    if index isa MOI.ConstraintIndex && MOI.Bridges.is_bridged(model, index)
        error(
            "There is no `optimizer_index` for $(typeof(index)) constraints",
            " because they are bridged.",
        )
    else
        return _moi_optimizer_index(model.model, index)
    end
end

"""
    optimizer_index(v::VariableRef)::MOI.VariableIndex

Return the index of the variable that corresponds to `v` in the optimizer model.
It throws [`NoOptimizer`](@ref) if no optimizer is set and throws an
`ErrorException` if the optimizer is set but is not attached.
"""
function optimizer_index(v::VariableRef)
    model = owner_model(v)
    if mode(model) == DIRECT
        return index(v)
    else
        return _moi_optimizer_index(backend(model), index(v))
    end
end

"""
    optimizer_index(cr::ConstraintRef{Model})::MOI.ConstraintIndex

Return the index of the constraint that corresponds to `cr` in the optimizer
model. It throws [`NoOptimizer`](@ref) if no optimizer is set and throws an
`ErrorException` if the optimizer is set but is not attached or if the
constraint is bridged.
"""
function optimizer_index(cr::ConstraintRef{Model})
    if mode(cr.model) == DIRECT
        return index(cr)
    else
        return _moi_optimizer_index(backend(cr.model), index(cr))
    end
end

"""
    index(cr::ConstraintRef)::MOI.ConstraintIndex

Return the index of the constraint that corresponds to `cr` in the MOI backend.
"""
index(cr::ConstraintRef) = cr.index

"""
    struct OptimizeNotCalled <: Exception end

A result attribute cannot be queried before [`optimize!`](@ref) is called.
"""
struct OptimizeNotCalled <: Exception end

"""
    struct NoOptimizer <: Exception end

No optimizer is set. The optimizer can be provided to the [`Model`](@ref)
constructor or by calling [`set_optimizer`](@ref).
"""
struct NoOptimizer <: Exception end

# Throws an error if `optimize!` has not been called, i.e., if there is no
# optimizer attached or if the termination status is `MOI.OPTIMIZE_NOT_CALLED`.
function _moi_get_result(model::MOI.ModelLike, args...)
    if MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED
        throw(OptimizeNotCalled())
    end
    return MOI.get(model, args...)
end
function _moi_get_result(model::MOIU.CachingOptimizer, args...)
    if MOIU.state(model) == MOIU.NO_OPTIMIZER
        throw(NoOptimizer())
    elseif MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED
        throw(OptimizeNotCalled())
    end
    return MOI.get(model, args...)
end

"""
    get(model::Model, attr::MathOptInterface.AbstractModelAttribute)

Return the value of the attribute `attr` from the model's MOI backend.
"""
function MOI.get(model::Model, attr::MOI.AbstractModelAttribute)
    if MOI.is_set_by_optimize(attr) &&
       !(attr isa MOI.TerminationStatus) && # Before `optimize!` is called, the
       !(attr isa MOI.PrimalStatus) &&      # statuses are `OPTIMIZE_NOT_CALLED`
       !(attr isa MOI.DualStatus)           # and `NO_SOLUTION`
        _moi_get_result(backend(model), attr)
    else
        MOI.get(backend(model), attr)
    end
end
"""
    get(model::Model, attr::MathOptInterface.AbstractOptimizerAttribute)

Return the value of the attribute `attr` from the model's MOI backend.
"""
function MOI.get(model::Model, attr::MOI.AbstractOptimizerAttribute)
    return MOI.get(backend(model), attr)
end
function MOI.get(
    model::Model,
    attr::MOI.AbstractVariableAttribute,
    v::VariableRef,
)
    check_belongs_to_model(v, model)
    if MOI.is_set_by_optimize(attr)
        return _moi_get_result(backend(model), attr, index(v))
    else
        return MOI.get(backend(model), attr, index(v))
    end
end
function MOI.get(
    model::Model,
    attr::MOI.AbstractConstraintAttribute,
    cr::ConstraintRef,
)
    check_belongs_to_model(cr, model)
    if MOI.is_set_by_optimize(attr)
        return _moi_get_result(backend(model), attr, index(cr))
    else
        return MOI.get(backend(model), attr, index(cr))
    end
end

function MOI.set(m::Model, attr::MOI.AbstractOptimizerAttribute, value)
    return MOI.set(backend(m), attr, value)
end
function MOI.set(m::Model, attr::MOI.AbstractModelAttribute, value)
    return MOI.set(backend(m), attr, value)
end
function MOI.set(
    model::Model,
    attr::MOI.AbstractVariableAttribute,
    v::VariableRef,
    value,
)
    check_belongs_to_model(v, model)
    return MOI.set(backend(model), attr, index(v), value)
end
function MOI.set(
    model::Model,
    attr::MOI.AbstractConstraintAttribute,
    cr::ConstraintRef,
    value,
)
    check_belongs_to_model(cr, model)
    return MOI.set(backend(model), attr, index(cr), value)
end

const _Constant = Union{Number,UniformScaling}
_constant_to_number(x::Number) = x
_constant_to_number(J::UniformScaling) = J.λ

# GenericAffineExpression, AffExpr, AffExprConstraint
include("aff_expr.jl")

# GenericQuadExpr, QuadExpr
# GenericQuadConstraint, QuadConstraint
include("quad_expr.jl")

include("mutable_arithmetics.jl")

include("sets.jl")

# Indicator constraint
include("indicator.jl")
# Complementarity constraint
include("complement.jl")
# SDConstraint
include("sd.jl")

"""
    Base.getindex(m::JuMP.AbstractModel, name::Symbol)

To allow easy accessing of JuMP Variables and Constraints via `[]` syntax.
Returns the variable, or group of variables, or constraint, or group of constraints, of the given name which were added to the model. This errors if multiple variables or constraints share the same name.
"""
function Base.getindex(m::JuMP.AbstractModel, name::Symbol)
    obj_dict = object_dictionary(m)
    if !haskey(obj_dict, name)
        throw(KeyError(name))
    elseif obj_dict[name] === nothing
        error(
            "There are multiple variables and/or constraints named $name that are already attached to this model. If creating variables programmatically, use the anonymous variable syntax x = @variable(m, [1:N], ...). If creating constraints programmatically, use the anonymous constraint syntax con = @constraint(m, ...).",
        )
    else
        return obj_dict[name]
    end
end

"""
    Base.setindex!(m::JuMP.AbstractModel, value, name::Symbol)

stores the object `value` in the model `m` using so that it can be accessed via `getindex`.  Can be called with `[]` syntax.
"""
function Base.setindex!(model::AbstractModel, value, name::Symbol)
    # if haskey(object_dictionary(model), name)
    #     warn("Overwriting the object $name stored in the model. Consider using anonymous variables and constraints instead")
    # end
    return object_dictionary(model)[name] = value
end

"""
    haskey(model::AbstractModel, name::Symbol)

Determine whether the model has a mapping for a given name.
"""
function Base.haskey(model::AbstractModel, name::Symbol)
    return haskey(object_dictionary(model), name)
end

"""
    operator_warn(model::AbstractModel)
    operator_warn(model::Model)

This function is called on the model whenever two affine expressions are added
together without using `destructive_add!`, and at least one of the two
expressions has more than 50 terms.

For the case of `Model`, if this function is called more than 20,000 times then
a warning is generated once.
"""
function operator_warn(::AbstractModel) end
function operator_warn(model::Model)
    model.operator_counter += 1
    if model.operator_counter > 20000
        @warn(
            "The addition operator has been used on JuMP expressions a large " *
            "number of times. This warning is safe to ignore but may " *
            "indicate that model generation is slower than necessary. For " *
            "performance reasons, you should not add expressions in a loop. " *
            "Instead of x += y, use add_to_expression!(x,y) to modify x in " *
            "place. If y is a single variable, you may also use " *
            "add_to_expression!(x, coef, y) for x += coef*y.",
            maxlog = 1
        )
    end
end

# TODO: rename "m" field to "model" for style compliance
"""
    NonlinearExpression

A struct to represent a nonlinear expression.

Create an expression using [`@NLexpression`](@ref).
"""
struct NonlinearExpression
    m::Model
    index::Int
end

"""
    NonlinearParameter

A struct to represent a nonlinear parameter.

Create a parameter using [`@NLparameter`](@ref).
"""
struct NonlinearParameter <: AbstractJuMPScalar
    m::Model
    index::Int
end

include("copy.jl")
include("operators.jl")
include("macros.jl")
include("optimizer_interface.jl")
include("nlp.jl")
include("print.jl")
include("lp_sensitivity.jl")
include("lp_sensitivity2.jl")
include("callbacks.jl")
include("file_formats.jl")
include("feasibility_checker.jl")

# JuMP exports everything except internal symbols, which are defined as those
# whose name starts with an underscore. Macros whose names start with
# underscores are internal as well. If you don't want all of these symbols
# in your environment, then use `import JuMP` instead of `using JuMP`.

# Do not add JuMP-defined symbols to this exclude list. Instead, rename them
# with an underscore.
const _EXCLUDE_SYMBOLS = [Symbol(@__MODULE__), :eval, :include]

for sym in names(@__MODULE__, all = true)
    sym_string = string(sym)
    if sym in _EXCLUDE_SYMBOLS ||
       startswith(sym_string, "_") ||
       startswith(sym_string, "@_")
        continue
    end
    if !(
        Base.isidentifier(sym) ||
        (startswith(sym_string, "@") && Base.isidentifier(sym_string[2:end]))
    )
        continue
    end
    @eval export $sym
end

if Base.VERSION >= v"1.4.2"
    include("precompile.jl")
    _precompile_()
end

end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
"""
    This creates the overall material parameters structure
"""

module MaterialParameters
using Unitful: Energy
using Unitful
using Parameters
using ..Units

import Base.show, Base.convert
using GeoParams: AbstractMaterialParam, AbstractMaterialParamsStruct

export
    MaterialParams, SetMaterialParams

# Link the modules with various definitions:
include("./PhaseDiagrams/PhaseDiagrams.jl")
include("./CreepLaw/CreepLaw.jl")
include("./Density/Density.jl")
include("./GravitationalAcceleration/GravitationalAcceleration.jl")
include("./Energy/HeatCapacity.jl")
include("./Energy/Conductivity.jl")
include("./Energy/LatentHeat.jl")
include("./Energy/RadioactiveHeat.jl")
include("./Energy/Shearheating.jl")

"""
    MaterialParams

Structure that holds all material parameters for a given phase

"""
 @with_kw_noshow mutable struct MaterialParams <: AbstractMaterialParamsStruct
    #
    Name::String         =   ""                  #       Description/name of the phase
    Phase::Int64         =   1;                  #       Number of the phase (optional)
    Nondimensional::Bool =   false;              #       Are all fields non-dimensionalized or not?
    Density              =   nothing             #       Density equation of state
    Gravity              =   nothing             #       Gravitational acceleration (set automatically)
    CreepLaws            =   nothing             #       Creep laws
    Elasticity           =   nothing             #       Elastic parameters
    Plasticity           =   nothing             #       Plasticity
    Conductivity         =   nothing             #       Parameters related to the energy equation
    HeatCapacity         =   nothing             #       Heat capacity
    EnergySourceTerms    =   nothing             #       Source terms in energy conservation equation (such as radioactive heat)
    Melting              =   nothing             #       Melting model
end

"""
    SetMaterialParams(; Name::String="", Phase::Int64=1,
                        Density             =   nothing,
                        Gravity             =   nothing,
                        CreepLaws           =   nothing,
                        Elasticity          =   nothing,
                        Plasticity          =   nothing,
                        Conductivity        =   nothing,
                        HeatCapacity        =   nothing,
                        EnergySourceTerms   =   nothing,
                        Melting             =   nothing,
                        CharDim::GeoUnits   =   nothing)

Sets material parameters for a given phase.

If `CharDim` is specified the input parameters are non-dimensionalized.
Note that if `Density` is specified, we also set `Gravity` even if not explicitly listed

# Examples

Define two viscous creep laws & constant density:
```julia-repl
julia> Phase = SetMaterialParams(Name="Viscous Matrix",
                       Density   = ConstantDensity(),
                       CreepLaws = (PowerlawViscous(), LinearViscous(η=1e21Pa*s)))
Phase 1 : Viscous Matrix
        | [dimensional units]
        |
        |-- Density           : Constant density: ρ=2900 kg m⁻³
        |-- Gravity           : Gravitational acceleration: g=9.81 m s⁻²
        |-- CreepLaws         : Powerlaw viscosity: η0=1.0e18 Pa s, n=2.0, ε0=1.0e-15 s⁻¹
        |                       Linear viscosity: η=1.0e21 Pa s
```

Define two viscous creep laws & P/T dependent density and nondimensionalize
```julia-repl
julia> CharUnits_GEO   =   GEO_units(viscosity=1e19, length=1000km);
julia> Phase = SetMaterialParams(Name="Viscous Matrix", Phase=33,
                              Density   = PT_Density(),
                              CreepLaws = (PowerlawViscous(n=3), LinearViscous(η=1e23Pa*s)),
                              CharDim   = CharUnits_GEO)
Phase 33: Viscous Matrix
        | [non-dimensional units]
        |
        |-- Density           : P/T-dependent density: ρ0=2.9e-16, α=0.038194500000000006, β=0.01, T0=0.21454659702313156, P0=0.0
        |-- Gravity           : Gravitational acceleration: g=9.810000000000002e18
        |-- CreepLaws         : Powerlaw viscosity: η0=0.1, n=3, ε0=0.001
        |                       Linear viscosity: η=10000.0
```

You can also create an array that holds several parameters:
```julia-repl
julia> MatParam        =   Array{MaterialParams, 1}(undef, 2);
julia> Phase           =   1;
julia> MatParam[Phase] =   SetMaterialParams(Name="Upper Crust", Phase=Phase,
                            CreepLaws= (PowerlawViscous(), LinearViscous(η=1e23Pa*s)),
                            Density  = ConstantDensity(ρ=2900kg/m^3));
julia> Phase           =   2;
julia> MatParam[Phase] =   SetMaterialParams(Name="Lower Crust", Phase=Phase,
                            CreepLaws= (PowerlawViscous(n=5), LinearViscous(η=1e21Pa*s)),
                            Density  = PT_Density(ρ0=3000kg/m^3));
julia> MatParam
2-element Vector{MaterialParams}:
 Phase 1 : Upper Crust
    | [dimensional units]
    |
    |-- Density           : Constant density: ρ=2900 kg m⁻³
    |-- Gravity           : Gravitational acceleration: g=9.81 m s⁻²
    |-- CreepLaws         : Powerlaw viscosity: η0=1.0e18 Pa s, n=2.0, ε0=1.0e-15 s⁻¹
    |                       Linear viscosity: η=1.0e23 Pa s

 Phase 2 : Lower Crust
    | [dimensional units]
    |
    |-- Density           : P/T-dependent density: ρ0=3000 kg m⁻³, α=3.0e-5 K⁻¹, β=1.0e-9 Pa⁻¹, T0=0 °C, P0=0 MPa
    |-- Gravity           : Gravitational acceleration: g=9.81 m s⁻²
    |-- CreepLaws         : Powerlaw viscosity: η0=1.0e18 Pa s, n=5, ε0=1.0e-15 s⁻¹
    |                       Linear viscosity: η=1.0e21 Pa s
```

"""
function SetMaterialParams(; Name::String="", Phase=1,
            Density             =   nothing,
            Gravity             =   nothing,
            CreepLaws           =   nothing,
            Elasticity          =   nothing,
            Plasticity          =   nothing,
            Conductivity        =   nothing,
            HeatCapacity        =   nothing,
            EnergySourceTerms   =   nothing,
            Melting             =   nothing,
            CharDim             =   nothing)

        # In case density is defined and gravity not, set gravity to default value
        if ~isnothing(Density) & isnothing(Gravity)
            Gravity = GravitationalAcceleration.ConstantGravity();
        end

        # define struct for phase, while also specifying the maximum number of definitions for every field
        phase = MaterialParams(Name, Phase, false,
                                     ConvField(Density,             :Density,           maxAllowedFields=1),
                                     ConvField(Gravity,             :Gravity,           maxAllowedFields=1),
                                     ConvField(CreepLaws,           :Creeplaws),
                                     ConvField(Elasticity,          :Elasticity,        maxAllowedFields=1),
                                     ConvField(Plasticity,          :Plasticity),
                                     ConvField(Conductivity,        :Conductivity,      maxAllowedFields=1),
                                     ConvField(HeatCapacity,        :HeatCapacity,      maxAllowedFields=1),
                                     ConvField(EnergySourceTerms,   :EnergySourceTerms),
                                     ConvField(Melting,             :Melting,           maxAllowedFields=1))

        # [optionally] non-dimensionalize the struct
        if ~isnothing(CharDim)
            if typeof(CharDim) <: GeoUnits
                Nondimensionalize!(phase, CharDim)
            else
                error("CharDim should be of type GeoUnits")
            end

        end

        return phase
end

# Helper function that converts a field to a Tuple, provided it is not nothing
# This also checks for the maximum allowed number of definitions
# (some rheological phases may allow for an arbitrary combination per phase; others like density EoS not)
function ConvField(field, fieldname::Symbol; maxAllowedFields=1e6)
    if ~isnothing(field)
        if typeof(field) <: AbstractMaterialParam
            field = (field, )       # transform to tuple
        end
        if typeof(field[1]) <: AbstractMaterialParam
            if length(field)>maxAllowedFields
                error("Maximum $(maxAllowedFields) field allowed for: $fieldname")
            end
        end
    end
    return field
end

# Helper that prints info about each of the material parameters
#  for this to look nice, you need to define a Base.show
function Print_MaterialParam(io::IO, name::Symbol, Data)
    if ~isnothing(Data)
        if typeof(Data[1]) <: AbstractMaterialParam
            print(io, "        |-- $(rpad(name,18)):")
            for i=1:length(Data)
                if i==1
                    print(io, " $(Data[1]) \n")
                else
                    print(io, "        |  $(rpad("     ",18))   $(Data[i]) \n")
                end
            end
        end
    end
end

# Specify how the printing of the MaterialParam structure is done
function Base.show(io::IO, phase::MaterialParams)
    println(io, "Phase $(rpad(phase.Phase,2)): $(phase.Name)")
    if phase.Nondimensional
        println(io,"        | [non-dimensional units]")
    else
        println(io,"        | [dimensional units]")

    end
    println(io,"        | ")

    for param in fieldnames(typeof(phase))
        Print_MaterialParam(io, param, getfield(phase, param))
    end

end

end


text:
"""
Returns the variance of data
"""
function Statistics.var(data1::VNamedD)
	σ=0.0
	μ=Statistics.mean(data1)
	n=0
	for iss in 1:length(data1)
		for ifield in names(data1[iss].d)[1]
			for ir = 1:data1[iss].n, it = 1:length(data1[iss].grid)
				n += 1
				σ += (data1[iss].d[ifield][it, ir]-μ)^2
			end
		end
	end
	return σ*inv(n)
end

function Statistics.mean(data1::VNamedD)
	n=0
	μ=0.0
	for iss in 1:length(data1)
		for ifield in names(data1[iss].d)[1]
			for ir = 1:data1[iss].n, it = 1:length(data1[iss].grid)
				n += 1
				μ += data1[iss].d[ifield][it, ir]
			end
		end
	end
	return μ*inv(n)
end

function addnoise!(dataN::VNamedD, data::VNamedD, SNR)

	σx=Statistics.var(data)

	σxN=sqrt(σx^2*inv(10^(SNR/10.)))

	# factor to be multiplied to each scalar
	α=sqrt(σxN)
	for iss in 1:length(data)
		for ifield in names(data[iss].d)[1]
			for ir = 1:data[iss].n, it = 1:length(data[iss].grid)
				dataN[iss].d[ifield][it, ir] = dataN[iss].d[ifield][it, ir] + α*Random.randn()
			end
		end
	end
end


text:
using Jive
runtests(@__DIR__, skip=["revise.jl"])


text:
import Celeste.SDSSIO: RunCamcolField, PlainFITSStrategy

function parse_rcfs_nsrcs(rcf_nsrcs_file)
    all_rcfs = Vector{RunCamcolField}()
    all_rcf_nsrcs = Vector{Int16}()
    open(rcf_nsrcs_file) do f
        for ln in eachline(f)
            lp = split(ln, '\t')
            run = parse(Int16, lp[1])
            camcol = parse(UInt8, lp[2])
            field = parse(Int16, lp[3])
            nsrc = parse(Int16, lp[4])
            push!(all_rcfs, RunCamcolField(run, camcol, field))
            push!(all_rcf_nsrcs, nsrc)
        end
    end
    all_rcfs, all_rcf_nsrcs
end

function parse_boxfile(boxfile)
    boxes = Vector{BoundingBox}()
    boxes_rcf_idxs = Vector{Vector{Int32}}()
    open(boxfile) do f
        for ln in eachline(f)
            lp = split(ln, '\t')
            if length(lp) != 5
                Log.one_message("ERROR: malformed line in box file:\n> $ln ")
                return boxes, boxes_rcf_idxs
            end

            ss = split(lp[4], ' ')
            ramin = parse(Float64, ss[1])
            ramax = parse(Float64, ss[2])
            decmin = parse(Float64, ss[3])
            decmax = parse(Float64, ss[4])
            bb = BoundingBox(ramin, ramax, decmin, decmax)
            push!(boxes, bb)

            ris = split(lp[5], ',')
            rcf_idxs = [parse(Int32, ri) for ri in ris]
            push!(boxes_rcf_idxs, rcf_idxs)
        end
    end
    if length(boxes) < 1
        Log.one_message("$boxfile is empty?")
    end
    return boxes, boxes_rcf_idxs
end


text:
# ---
# title: 300. Longest Increasing Subsequence
# id: problem300
# author: Indigo
# date: 2021-03-04
# difficulty: Medium
# categories: Binary Search, Dynamic Programming
# link: <https://leetcode.com/problems/longest-increasing-subsequence/description/>
# hidden: true
# ---
#
# Given an integer array `nums`, return the length of the longest strictly
# increasing subsequence.
#
# A **subsequence** is a sequence that can be derived from an array by deleting
# some or no elements without changing the order of the remaining elements. For
# example, `[3,6,2,7]` is a subsequence of the array `[0,3,1,6,2,2,7]`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [10,9,2,5,3,7,101,18]
#     Output: 4
#     Explanation: The longest increasing subsequence is [2,3,7,101], therefore the length is 4.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [0,1,0,3,2,3]
#     Output: 4
#
#
# **Example 3:**
#
#
#
#     Input: nums = [7,7,7,7,7,7,7]
#     Output: 1
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 2500`
#   * `-104 <= nums[i] <= 104`
#
#
#
# **Follow up:**
#
#   * Could you come up with the `O(n2)` solution?
#   * Could you improve it to `O(n log(n))` time complexity?
#
#
## @lc code=start
using LeetCode

function length_of_LIS(nums::Vector{Int})
    f = [first(nums)]
    for num in @view(nums[2:end])
        if num > f[end]
            push!(f, num)
        else
            f[searchsortedfirst(f, num)] = num
        end
    end
    length(f)
end
## @lc code=end


text:
# =================================================================================================
# =                                         Hyland83.jl                                           =
# =================================================================================================

# Appendix of ref. [1]. First table
# Specific volume of saturated ice:
@test 1.0768e-3 ≈ Psychro.volumeice(173.15) atol=0.0001e-3
@test 1.0829e-3 ≈ Psychro.volumeice(223.15) atol=0.0001e-3
@test 1.0909e-3 ≈ Psychro.volumeice(273.16) atol=0.0001e-3

# Specific volume of saturated liquid water
@test 1.00021e-3 ≈ Psychro.volumewater(273.16) atol=0.00001e-3
@test 1.01215e-3 ≈ Psychro.volumewater(323.15) atol=0.00001e-3
@test 1.04346e-3 ≈ Psychro.volumewater(373.15) atol=0.00001e-3 # In the paper there is a typo in the power of 10: 1.04346e10
@test 1.09050e-3 ≈ Psychro.volumewater(423.15) atol=0.00001e-3
@test 1.15653e-3 ≈ Psychro.volumewater(473.15) atol=0.00001e-3

# Saturation enthalpy of saturated ice
@test -507.215e3 ≈ Psychro.enthalpyice(173.15) atol=0.001e3
@test -429.413e3 ≈ Psychro.enthalpyice(223.15) atol=0.001e3
@test -333.429e3 ≈ Psychro.enthalpyice(273.15) atol=0.001e3  # The table presents -333.409. Typo???

# Saturation enthalpy of saturated liquid water
@test 0.0 ≈ Psychro.enthalpywater(273.16) atol=1.0
@test 209.330e3 ≈ Psychro.enthalpywater(323.15) atol=5
@test 419.158e3 ≈ Psychro.enthalpywater(373.15) atol=5
@test 632.210e3 ≈ Psychro.enthalpywater(423.15) atol=9
@test 852.329e3 ≈ Psychro.enthalpywater(473.15) atol=15

# Table 2 of reference [2]
# Enhancement factor
P = 0.1e6
@test 1.0105 ≈ Psychro.efactor(173.15, P) atol=0.0001
@test 1.0039 ≈ Psychro.efactor(273.15, P) atol=0.0001
@test 1.0039 ≈ Psychro.efactor(363.15, P) atol=0.0001
P = 0.5e6
@test 1.054 ≈ Psychro.efactor(173.15, P) atol=0.001
@test 1.0177 ≈ Psychro.efactor(273.15, P) atol=0.0006  # I think the table has a typo!!?
@test 1.0180 ≈ Psychro.efactor(363.15, P) atol=0.0001
@test 1.0188 ≈ Psychro.efactor(373.15, P) atol=0.0001
@test 1.0022 ≈ Psychro.efactor(423.15, P) atol=0.0001
P = 1.0e6
@test 1.113 ≈ Psychro.efactor(173.15, P) atol=0.001
@test 1.0353 ≈ Psychro.efactor(273.15, P) atol=0.0015  # I think the table has a typo!!?
@test 1.0284 ≈ Psychro.efactor(363.15, P) atol=0.0001
@test 1.0295 ≈ Psychro.efactor(373.15, P) atol=0.0001
@test 1.0288 ≈ Psychro.efactor(423.15, P) atol=0.0001
@test 1.0235 ≈ Psychro.efactor(433.15, P) atol=0.0001
@test 1.0142 ≈ Psychro.efactor(443.15, P) atol=0.0001

P = 5.0e6
@test 1.820 ≈ Psychro.efactor(173.15, P) atol=0.001
@test 1.191 ≈ Psychro.efactor(273.15, P) atol=0.008  # I think the table has a typo!!?
@test 1.1102 ≈ Psychro.efactor(363.15, P) atol=0.0001
@test 1.1082 ≈ Psychro.efactor(373.15, P) atol=0.0001
@test 1.111 ≈ Psychro.efactor(423.15, P) atol=0.001
@test 1.114 ≈ Psychro.efactor(433.15, P) atol=0.001
@test 1.116 ≈ Psychro.efactor(443.15, P) atol=0.001
@test 1.117 ≈ Psychro.efactor(453.15, P) atol=0.001
@test 1.116 ≈ Psychro.efactor(473.15, P) atol=0.001

# Saturation pressure of water vapor over ice
# Second table in the appendix of reference [1]
@test 1.40510e-3 ≈ Psychro.Pws_s(173.15) atol=0.00001e-3
@test 6.11153e2 ≈ Psychro.Pws_s(273.15) atol=0.00001e2

# Saturation pressure of water vapor over liquid water
# Second table in the appendix of reference [1]
@test 6.11213e2 ≈ Psychro.Pws_l(273.15) atol=0.00001e2
@test 1.01419e5 ≈ Psychro.Pws_l(373.15) atol=0.00001e5
@test 1.55507e6 ≈ Psychro.Pws_l(473.15) atol=0.00001e6

# Testing saturation temperature function:
@test Psychro.Tws(Psychro.Pws_s(173.15)) ≈ 173.15 atol=1e-5
@test Psychro.Tws(Psychro.Pws_s(223.15)) ≈ 223.15 atol=1e-5
@test Psychro.Tws(Psychro.Pws_s(273.15)) ≈ 273.15 atol=1e-5
@test Psychro.Tws(Psychro.Pws_s(273.16)) ≈ 273.16 atol=1e-5
@test Psychro.Tws(Psychro.Pws_l(273.16)) ≈ 273.16 atol=1e-5
@test Psychro.Tws(Psychro.Pws_l(323.15)) ≈ 323.15 atol=1e-5
@test Psychro.Tws(Psychro.Pws_l(373.15)) ≈ 373.15 atol=1e-5
@test Psychro.Tws(Psychro.Pws_l(473.15)) ≈ 473.15 atol=1e-5

# First virial coefficient of saturated vapor B'
# Second table in the appendix of reference [1]
@test -3.2939e-5 ≈ Psychro.Blin(173.15) atol=0.0001e-5
@test -8.3497e-7 ≈ Psychro.Blin(273.15) atol=0.0001e-7
@test -1.4658e-7 ≈ Psychro.Blin(373.15) atol=0.0001e-5
@test -5.0508e-8 ≈ Psychro.Blin(473.15) atol=0.0001e-5

# Second virial coefficient of saturated vapor C'
# Second table in the appendix of reference [1]
@test -4.6563e-9 ≈ Psychro.Clin(173.15) atol=0.0001e-9
@test -2.0928e-12 ≈ Psychro.Clin(273.15) atol=0.0001e-12
@test -5.7548e-14 ≈ Psychro.Clin(373.15) atol=0.0001e-14
@test -6.3933e-15 ≈ Psychro.Clin(473.15) atol=0.0001e-15

# Specific enthalpy of saturated water vapor
# First table of the appendix of ref. [1]
@test 2315.87 ≈ Psychro.enthalpyvapor(173.15)/1000 atol=0.01
@test 2408.41 ≈ Psychro.enthalpyvapor(223.15)/1000 atol=0.01
@test 2500.81 ≈ Psychro.enthalpyvapor(273.16)/1000 atol=0.01
@test 2591.29 ≈ Psychro.enthalpyvapor(323.15)/1000 atol=0.01
@test 2675.46 ≈ Psychro.enthalpyvapor(373.15)/1000 atol=0.01
@test 2746.15 ≈ Psychro.enthalpyvapor(423.15)/1000 atol=0.01
@test 2793.11 ≈ Psychro.enthalpyvapor(473.15)/1000 atol=0.01

# Specific volume of saturated water vapor
# First table of the appendix of ref. [1]
@test 5.6873e7 ≈ Psychro.volumevapor(173.15) atol=0.0001e7
@test 2.6146e4 ≈ Psychro.volumevapor(223.15) atol=0.0001e4
@test 2.0601e2 ≈ Psychro.volumevapor(273.16) atol=0.0001e2
@test 1.2030e1 ≈ Psychro.volumevapor(323.15) atol=0.0001e1
@test 1.6718e0 ≈ Psychro.volumevapor(373.15) atol=0.0001e0
@test 3.9253e-1 ≈ Psychro.volumevapor(423.15) atol=0.0001e-1
@test 1.2722e-1 ≈ Psychro.volumevapor(473.15) atol=0.0001e0

# Specific entropy of saturated water vapor
# First table of the appendix of ref. [1]
@test 14.30387 ≈ Psychro.entropyvapor(173.15)/1000 atol=0.00005
@test 11.10965 ≈ Psychro.entropyvapor(223.15)/1000 atol=0.00005
@test  9.15510 ≈ Psychro.entropyvapor(273.16)/1000 atol=0.00005
@test  8.07477 ≈ Psychro.entropyvapor(323.15)/1000 atol=0.00005
@test  7.35365 ≈ Psychro.entropyvapor(373.15)/1000 atol=0.00005
@test  6.83731 ≈ Psychro.entropyvapor(423.15)/1000 atol=0.00005
@test  6.43218 ≈ Psychro.entropyvapor(473.15)/1000 atol=0.00005


text:
# prints solution of both parts
function x()
	io = open("2.txt", "r")

	n = 0
	m = 0

	for i=1:1000
		a = parse(Int, readuntil(io, "-"))
		b = parse(Int, readuntil(io, " "))
		c = read(io, Char)
		read(io, Char)
		read(io, Char)
		d = readline(io)

		if a <= count(x->x==c, d) <= b
			m += 1
		end

		if d[a] == c && d[b] != c || d[a] != c && d[b] == c
			n += 1
		end
	end
	println(m)
	println(n)
end


text:
using TypeClasses

"""
considerations: because `Base.typename(Vector).wrapper == Array`, we dispatch everything on AbstractArray to support AbstractVector.

If later on the AbstractArray should be separated from AbstractVector, someone needs to think about a neat way to circumvent this problem.
As of now this seems to be a great solution.
"""

# MonoidAlternative
# =================

TypeClasses.neutral(T::Type{<:AbstractArray}) = convert(T, [])
TypeClasses.combine(v1::AbstractArray, v2::AbstractArray) = [v1; v2]  # general syntax which is overloaded by concrete AbstractVector types

# FunctorApplicativeMonad
# =======================

TypeClasses.pure(T::Type{<:AbstractArray}, a) = convert(Base.typename(T).wrapper, [a])
TypeClasses.ap(fs::AbstractArray, v::AbstractArray) = vcat((map(f, v) for f in fs)...)

# for flattening we solve type-safety by converting to Vector elementwise
# this also gives well-understandable error messages if something goes wrong
function TypeClasses.flatmap(f, v::AbstractArray)
    type = Base.typename(typeof(v)).wrapper
    vcat((convert(type, f(x)) for x in v)...)
end

# FlipTypes
# =========

# we define flip_types for all Vector despite it only works if the underlying element defines `ap`
# as there is no other sensible definition for Iterable, an error that the element does not implement `ap`
# is actually the correct error
flip_types(v::AbstractArray) = default_flip_types_having_pure_combine_apEltype(v)


text:
module Repulsion

import NBodyIPs: NBodyFunction, descriptor, evaluate_I, evaluate_I_ed,
                 inv_transform, fcut, fast

import JuLIP: AbstractCalculator, energy, forces, virial, decode_dict
import JuLIP.Potentials: @pot, evaluate, evaluate_d, PairPotential, @D, cutoff,
                         @analytic
import Base: Dict, convert

struct RepulsiveCore{TV1, TV2, DT} <: NBodyFunction{2, DT}
   Vout::TV1         # the outer pair potential
   Vin::TV2          # the inner (repulsive) pair potential
   ri::Float64       # the interface between the two
   e0::Float64
   D::DT             # shouldn't be needed, but can't get it to work otherwise...
end

@pot RepulsiveCore

cutoff(V::RepulsiveCore) = cutoff(V.Vout)
descriptor(V::RepulsiveCore) = V.D

fast(V::RepulsiveCore) = RepulsiveCore(fast(V.Vout), V.Vin, V.ri, V.e0, V.D)

function evaluate_I(V::RepulsiveCore, II)
   I1, I2 = II
   @assert length(I1) == 1
   r = inv_transform(V.D.transform, I1[1])
   if r > V.ri
      return evaluate_I(V.Vout, II)
   end
   return V.Vin(r) / fcut(V.Vout.D, r)
end

function evaluate_I_ed(V::RepulsiveCore, II)
   I1, I2, dI1, dI2 = II
   @assert length(I1) == 1
   r = inv_transform(V.D.transform, I1[1])
   if r > V.ri
      return evaluate_I_ed(V.Vout, II)
   end
   eV = V.Vin(r)
   dV = @D V.Vin(r)
   return eV, dV
end

function RepulsiveCore(Vout::NBodyFunction{2}, ri, e0=0.0)
   v = Vout(ri)
   dv = @D Vout(ri)
   if dv >= 0.0
      @error("The slope `Vout'(ri)` is required to be negative")
   end
   if dv > -1.0
      @warn("""The slope `Vout'(ri) = $dv` may not be steep enough to attach a
               repulsive core. Proceed at your own risk.""")
   end
   if v-e0 <= 0.0
      @error("it is required that `Vout(ri) > 0`.")
   end
   if v-e0 <= 1.0
      @warn("""Ideally the repulsive core should not be attached at small
               values of `Vout(ri) = $v`. Proceed at your own risk.""")
   end
   # e0 + B e^{-A (r/ri-1)} * ri/r
   #    => e0 + B = Vout(ri) => = Vout(ri) - e0 = v - e0
   # dv = - A*B/ri e^{-A (r/ri-1)} * ri/r - B*ri*e^{...} / r^2
   #    = - A/ri * (v - 1/ri * (v = - (1+A)/ri * (v-e0)
   #    => -(1+A)/ri * (v-e0) = dv
   #    => 1+A = - ri dv / (v-e0)
   #    => A = -1 - ri dv / (v-e0)
   Vin = let A = -1 - ri * dv / (v-e0), B = v-e0, e0=e0, ri = ri
      @show A
      @show B
      @analytic r -> e0 + B * exp( - A * (r/ri-1) ) * ri/r
   end
   # @show ri
   # @show Vout(ri), (@D Vout(ri))
   # @show Vin(ri), (@D Vin(ri))
   # construct the piecewise potential
   return RepulsiveCore(Vout, Vin, ri, e0, descriptor(Vout))
end

# ----------------------------------------------------
#  File IO
# ----------------------------------------------------

Dict(V::RepulsiveCore) = Dict("__id__" => "NBodyIPs_RepulsiveCore",
                              "Vout" => Dict(V.Vout),
                              "e0" => V.e0,
                              "ri" => V.ri)

function RepulsiveCore(D::Dict)
   if haskey(D, "e0")
      return RepulsiveCore(decode_dict(D["Vout"]), D["ri"], D["e0"])
   else
      return RepulsiveCore(decode_dict(D["Vout"]), D["ri"])
   end
end

convert(::Val{:NBodyIPs_RepulsiveCore}, D::Dict) = RepulsiveCore(D)

end

# EXAMPLE CODE => TURN THiS INTO A TEST!
# ## try out the repulsive potential
# Vfit = IPall.components[2]
# Vrep = NBodyIPs.Repulsion.RepulsiveCore(Vfit, 2.1)
#
# rp = range(0.3*r0, 9.0, length=500)
# plot(rp, Vfit.(rp), lw=2, label ="fit")
# plot!(rp, Vrep.Vin.(rp), lw=2, label ="inner", ylims = [-2.0, 10.0] )
# plot!(rp, Vrep.(rp), label = "combined")

max_stars_repo_name: cortner/ManyBodyIPs

max_issues_repo_name: cortner/ManyBodyIPs

max_forks_repo_name: cortner/ManyBodyIPs


text:
@static if VERSION < v"0.7.0-DEV.481"
    const _mul! = isdefined(:mul!) ? mul! : Base.A_mul_B!
    const _mul2! = isdefined(:mul!) ? mul! : Base.scale!
else
    const _mul! = (@isdefined mul!) ? mul! : Base.A_mul_B!
    const _mul2! = (@isdefined mul!) ? mul! : scale.A_mul_B!
end

# Consider permutedims as an alternative to direct multiplication.
# Multiplication is an O(m*N) cost compared to an O(N) cost for tridiagonal algorithms.
# However, when the multiplication is only a small fraction of the total time
# (for example, when m is small), then these can be convenient and avoid the need for calls to permutedims.

# Multiplication using permutedims
"""
`A_mul_B_perm!(dest, M, src, dim)` computes `M*x` for slices `x` of `src` along dimension `dim`,
storing the result in `dest`. `M` must be an `AbstractMatrix`. This uses `permutedims` to make dimension
`dim` into the first dimension, performs a standard matrix multiplication, and restores the original
dimension ordering. In many cases, this algorithm exhibits the best cache behavior.
"""
function A_mul_B_perm!(dest, M::AbstractMatrix, src, dim::Integer)
    check_matmul_sizes(dest, M, src, dim)
    order = [dim; setdiff(1:ndims(src), dim)]
    srcp = permutedims(src, order)
    tmp = Array{eltype(dest), 2}(undef, size(dest, dim), div(length(dest), size(dest, dim)))
    _mul!(tmp, M, reshape(srcp, (size(src,dim), div(length(srcp), size(src,dim)))))
    iorder = [2:dim; 1; dim+1:ndims(src)]
    permutedims!(dest, reshape(tmp, size(dest)[order]), iorder)
    dest
end

# Direct (temporary-free) multiplication
"""
`A_mul_B_md!(dest, M, src, dim)` computes `M*x` for slices `x` of `src` along dimension `dim`,
storing the result in `dest`. `M` must be an `AbstractMatrix`. This uses an in-place naive algorithm.
"""
function A_mul_B_md!(dest, M::AbstractMatrix, src, dim::Integer)
    check_matmul_sizes(dest, M, src, dim)
    if size(M,1) == size(M,2) == 1
        return _mul2!(dest, src, M[1,1])
    end
    R2 = CartesianIndices(size(dest)[dim+1:end])
    if dim > 1
        R1 = CartesianIndices(size(dest)[1:dim-1])
        _A_mul_B_md!(dest, M, src, R1, R2)
    else
        _A_mul_B_md!(dest, M, src, R2)
    end
end

# Multiplication along the first dimension
# Here we expect that M will typically be small and fit in cache, whereas src and dest do not
function _A_mul_B_md!(dest, M::AbstractMatrix, src, R2::CartesianIndices)
    m, n = size(M, 1), size(M, 2)
    if m == n == 2
        return _A_mul_B_md_2x2!(dest, M, src, R2)
    end
    for I2 in R2
        @inbounds for i = 1:m
            dest[i,I2] = zero(eltype(dest))
        end
        @inbounds for j = 1:n
            b = src[j,I2]
            @simd for i = 1:m
                dest[i,I2] += M[i,j]*b
            end
        end
    end
    dest
end
_A_mul_B_md(M::AbstractMatrix, src, R2::CartesianIndices) = _A_mul_B_md!(alloc_matmul(M, src, 1), M, src, R2)

function _A_mul_B_md_2x2!(dest, M::AbstractMatrix, src, R2::CartesianIndices)
    a, b, c, d = M[1,1], M[1,2], M[2,1], M[2,2]
    @simd for I2 in R2
        @inbounds begin
            s1, s2 = src[1,I2], src[2,I2]
            dest[1,I2] = a*s1 + b*s2
            dest[2,I2] = c*s1 + d*s2
        end
    end
    dest
end

function _A_mul_B_md!(dest, M::SparseMatrixCSC, src, R2::CartesianIndices)
    m, n = size(M,1), size(M,2)
    nzv = M.nzval
    rv = M.rowval
    cp = M.colptr
    for I2 in R2
        @inbounds for i = 1:m
            dest[i,I2] = zero(eltype(dest))
        end
        for j = 1:n
            b = src[j,I2]
            @inbounds for k = cp[j]:(cp[j+1]-1)
                dest[rv[k],I2] += nzv[k]*b
            end
        end
    end
    dest
end

# Multiplication along any other dimension
function _A_mul_B_md!(dest, M::AbstractMatrix, src,  R1::CartesianIndices, R2::CartesianIndices)
    m, n = size(M, 1), size(M, 2)
    if m == n == 2
        return _A_mul_B_md_2x2!(dest, M, src, R1, R2)
    end
    fill!(dest, zero(eltype(dest)))
    for I2 in R2
        for j = 1:n
            @inbounds for i = 1:m
                Mij = M[i,j]
                @simd for I1 in R1
                    dest[I1,i,I2] += Mij*src[I1,j,I2]
                end
            end
        end
    end
    dest
end
_A_mul_B_md(M::AbstractMatrix, src,  R1::CartesianIndices, R2::CartesianIndices) = _A_mul_B_md!(alloc_matmul(M, src, ndims(R1)+1), M, src, R1, R2)

function _A_mul_B_md_2x2!(dest, M::AbstractMatrix, src,  R1::CartesianIndices, R2::CartesianIndices)
    a, b, c, d = M[1,1], M[1,2], M[2,1], M[2,2]
    for I2 in R2
        @simd for I1 in R1
            @inbounds begin
                s1, s2 = src[I1,1,I2], src[I1,2,I2]
                dest[I1,1,I2] = a*s1 + b*s2
                dest[I1,2,I2] = c*s1 + d*s2
            end
        end
    end
    dest
end

function _A_mul_B_md!(dest, M::SparseMatrixCSC, src, R1::CartesianIndices, R2::CartesianIndices)
    m, n = size(M,1), size(M,2)
    nzv = M.nzval
    rv = M.rowval
    cp = M.colptr
    fill!(dest, zero(eltype(dest)))
    for I2 in R2
        for j = 1:n
            @inbounds for k = cp[j]:(cp[j+1]-1)
                i, Mij = rv[k], nzv[k]
                @simd for I1 in R1
                    dest[I1,i,I2] += Mij*src[I1,j,I2]
                end
            end
        end
    end
    dest
end

function check_matmul_sizes(dest, M::AbstractMatrix, src, dim)
    1 <= dim <= max(ndims(dest),ndims(src)) || throw(DimensionMismatch("The chosen dimension $dim is larger than $(ndims(src)) and $(ndims(dest))"))
    m, n = size(M, 1), size(M, 2)
    n == size(src, dim) && m == size(dest, dim) || throw(DimensionMismatch("Sizes $m, $n, $(size(src,dim)), and $(size(dest,dim)) do not match"))
    for i = 1:max(ndims(src), ndims(dest))
        i == dim && continue
        if size(src,i) != size(dest,i)
            throw(DimensionMismatch("Sizes $(size(dest)), $(size(src)) do not match"))
        end
    end
    nothing
end


text:
function rms(s::Vector{T} where {T<:Number})
  # iteratively mask the signal until the RMS
  # does not change
  return 0
end


text:
# ******************************************************************************
# Response
# ******************************************************************************
struct Response{T, F, P}
    ne :: AbstractArray{T}
    Pnl :: AbstractArray{Complex{T}}
    J :: AbstractArray{Complex{T}}
    func! :: F
    p :: P
end

function calculate_response!(resp::Response, E, z)
    resp.func!(resp.ne, resp.Pnl, resp.J, E, resp.p, z)
    return nothing
end

# ******************************************************************************
# Model
# ******************************************************************************
struct Model{T, TG, TF, TR}
    grid :: TG
    field :: TF
    response :: TR
    guard :: Guard{T}
    KK :: AbstractArray{Complex{T}}
    QQ :: AbstractArray{Complex{T}}
    QJ :: AbstractVector{Complex{T}}
    zu :: T
    z :: T
    zmax :: T
    nonlinearity :: Bool
end

function Model(
    grid::Grid{T}, field, response;
    zu=1, z=0, zmax, kparaxial=true, qparaxial=true, nonlinearity=true,
) where T
    (; arch, Nt, w, wu) = grid

    guard = Guard(grid, field)

    # Linear propagator:
    KK = KKpropagator(grid, field, zu; paraxial=kparaxial)
    KK = adapt_array(arch, KK)
    apply_guard_spec_domain!(KK, guard)

    # Nonlinear propagator:
    QQ = QQpropagator(grid, field, zu; paraxial=qparaxial)
    QQ = adapt_array(arch, QQ)
    apply_guard_spec_domain!(QQ, guard)

    QJ = zeros(ComplexF64, Nt)   # Float64 for higher precision
    for it=1:Nt
        QJ[it] = QJfunc(w[it] * wu)
    end
    QJ = Array{Complex{T}}(QJ)
    QJ = adapt_array(arch, QJ)
    @. QJ = QJ * guard.W

    TG = typeof(grid)
    TF = typeof(field)
    TR = typeof(response)
    return Model{T, TG, TF, TR}(
        grid, field, response, guard, KK, QQ, QJ, zu, z, zmax, nonlinearity,
    )
end

# ******************************************************************************
# K propagator
# ******************************************************************************
function KKfunc(k, w; paraxial=true)
    beta = beta_func(w)
    KK = 0
    if paraxial
        if beta != 0
            KK = beta - k^2 / (2 * beta)
        end
    else
        KK = sqrt(beta^2 - k^2 + 0im)
    end
    return KK
end

function KKpropagator(grid::GridT{T}, field, zu; paraxial=true) where T
    (; Nt, w, wu) = grid
    (; w0) = field
    vf = group_velocity(w0)   # frame velocity
    KK = zeros(ComplexF64, Nt)   # Float64 for higher precision
    for it=1:Nt
        if w[it] > 0
            KK[it] = KKfunc(0, w[it] * wu; paraxial) * zu
            KK[it] = KK[it] - (w[it] * wu) / vf * zu
        end
    end
    return Array{Complex{T}}(KK)
end

function KKpropagator(grid::GridRT{T}, field, zu; paraxial=true) where T
    (; Nr, Nt, k, ku, w, wu) = grid
    (; w0) = field
    vf = group_velocity(w0)   # frame velocity
    KK = zeros(ComplexF64, (Nr, Nt))   # Float64 for higher precision
    for it=1:Nt
        if w[it] > 0
            for ir=1:Nr
                KK[ir,it] = KKfunc(k[ir]*ku, w[it]*wu; paraxial)
                KK[ir,it] = (KK[ir,it] - (w[it] * wu) / vf) * zu
            end
        end
    end
    return Array{Complex{T}}(KK)
end

# ******************************************************************************
# Q propagator
# ******************************************************************************
function QQfunc(k, w; paraxial=true)
    mu = permeability(w)
    beta = beta_func(w)
    QQ = 0
    if paraxial
        if beta != 0
            QQ = MU0 * mu * w^2 / (2 * beta)
        end
    else
        KK = sqrt(beta^2 - k^2 + 0im)
        if KK != 0
            QQ = MU0 * mu * w^2 / (2 * KK)
        end
    end
    return QQ
end

function QQpropagator(grid::GridT{T}, field, zu; paraxial=true) where T
    (; Nt, w, wu) = grid
    (; Eu) = field
    QQ = zeros(ComplexF64, Nt)   # Float64 for higher precision
    for it=1:Nt
        if w[it] > 0
            QQ[it] = QQfunc(0, w[it] * wu; paraxial) * zu / Eu
        end
    end
    return Array{Complex{T}}(QQ)
end

function QQpropagator(grid::GridRT{T}, field, zu; paraxial=true) where T
    (; Nr, Nt, k, ku, w, wu) = grid
    (; Eu) = field
    QQ = zeros(ComplexF64, (Nr, Nt))   # Float64 for higher precision
    for it=1:Nt
        if w[it] > 0
            for ir=1:Nr
                QQ[ir,it] = QQfunc(k[ir]*ku, w[it]*wu; paraxial) * zu / Eu
            end
        end
    end
    return Array{Complex{T}}(QQ)
end

function QJfunc(w)
    QJ = 0
    if w != 0
        QJ = 1im / w
    end
    return QJ
end

# ******************************************************************************
# Q func
# ******************************************************************************
function q_func!(dE, E, p, z)
    model, = p
    (; QQ, QJ, field, response, guard) = model
    (; FFT) = field
    (; Pnl, J) = response

    frequency2time!(E, FFT)

    calculate_response!(response, E, z)

    # nonlinear polarization:
    apply_guard_real_domain!(Pnl, guard)
    rsig2aspec!(Pnl, FFT)   # real signal -> analytic spectrum
    apply_guard_spec_domain!(Pnl, guard)

    # current:
    apply_guard_real_domain!(J, guard)
    rsig2aspec!(J, FFT)   # real signal -> analytic spectrum
    apply_guard_spec_domain!(J, guard)

    dE_update!(dE, Pnl, J, QQ, QJ)   # @. dE = 1im * QQ * (Pnl + QJ * J)

    time2frequency!(E, FFT)

    return nothing
end

function dE_update!(
    dE::TA, Pnl::TA, J::TA, QQ::TA, QJ::TA,
) where TA<:AbstractVector
    @. dE = 1im * QQ * (Pnl + QJ * J)
    return nothing
end

function dE_update!(
    dE::TA, Pnl::TA, J::TA, QQ::TA, QJ::TB,
) where {TA<:AbstractMatrix, TB<:AbstractVector}
    Nr, Nt = size(dE)
    for it=1:Nt
    for ir=1:Nr
        dE[ir,it] = 1im * QQ[ir,it] * (Pnl[ir,it] + QJ[it] * J[ir,it])
    end
    end
    return nothing
end

function dE_update!(
    dE::TA, Pnl::TA, J::TA, QQ::TA, QJ::TB,
) where {TA<:CuMatrix, TB<:CuVector}
    N = length(dE)
    ckernel = @cuda launch=false dE_update_kernel(dE, Pnl, J, QQ, QJ)
    config = launch_configuration(ckernel.fun)
    threads = min(N, config.threads)
    blocks = cld(N, threads)
    ckernel(dE, Pnl, J, QQ, QJ; threads=threads, blocks=blocks)
    return nothing
end

function dE_update_kernel(dE, Pnl, J, QQ, QJ)
    id = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride = blockDim().x * gridDim().x
    ci = CartesianIndices(size(dE))
    for ici=id:stride:length(ci)
        it = ci[ici][2]
        dE[ici] = 1im * QQ[ici] * (Pnl[ici] + QJ[it] * J[ici])
    end
    return nothing
end