text:
using DICE
using JuMP
using Ipopt
using Test

DICE.@base mutable struct TestType{T}
    one
    two::Int
    three::T
    four::Vector{T}
end

DICE.@extend mutable struct ExtendTest <: TestType
    five::T
end

@testset "Abstractions" begin
    @test ExtendTest <: TestType

    etest = ExtendTest(10,10,10,[10],5)
    @test isa(etest, ExtendTest{Int})
    @test fieldnames(typeof(etest)) == (:one,:two,:three,:four,:five)
    @test isa(etest.one, Int)
    @test isa(etest.two, Int)
    @test isa(etest.three, Int)
    @test isa(etest.four, Vector{Int})
    @test isa(etest.five, Int)

    @test_throws LoadError @macroexpand DICE.@extend mutable struct missingTest <: Missing end
end

# The tests shouldn't need to converge for long times,
# so dump early and fail rather than wasting resouces on failure.
const linear_solver = DICE.selectLinearSolver();
const isMumps = linear_solver == "mumps";
optimizer = optimizer_with_attributes(Ipopt.Optimizer, "max_iter" => 1000, "print_frequency_iter" => 500, "sb" => "yes", "linear_solver" => linear_solver);
model = Model(optimizer);
modelrr = Model(optimizer);
model2016 = Model(optimizer);
model2016r2 = Model(optimizer);
vanilla_opt = options(v2013R());
rr_opt = options(v2013R(RockyRoad));
v2016_opt = options(v2016R());
v2016r2_opt = options(v2016R2());
vanilla_params = DICE.generate_parameters(vanilla_opt);
rr_params = DICE.generate_parameters(rr_opt, modelrr);
v2016_params = DICE.generate_parameters(v2016_opt, model2016);
v2016r2_params = DICE.generate_parameters(v2016r2_opt, model2016r2);

@testset "Types Display" begin
    @testset "Flavour" begin
        @test sprint(show, Vanilla) == "Vanilla flavour"
        @test sprint(show, RockyRoad) == "Rocky Road flavour"
    end
    @testset "Version" begin
        @test sprint(show, v2013R()) == "v2013R (Vanilla flavour)"
        @test sprint(show, v2013R(RockyRoad)) == "v2013R (Rocky Road flavour)"
        @test sprint(show, v2016R()) == "v2016R beta"
        @test sprint(show, v2016R2()) == "v2016R2"
    end
    @testset "Scenario" begin
        @test sprint(show, BasePrice) == "Base (current policy) carbon price"
        @test sprint(show, OptimalPrice) == "Optimal carbon price"
        @test sprint(show, Limit2Degrees) == "Limit warming to 2° with damages"
        @test sprint(show, Stern) == "Stern"
        @test sprint(show, SternCalibrated) == "Calibrated Stern"
        @test sprint(show, Copenhagen) == "Copenhagen participation"
    end
    @testset "Default Options and Parameters" begin
        @test sprint(show, "text/plain", vanilla_opt) == "Options for Vanilla 2013R version\nTime step\nN: 60, tstep: 5\nPreferences\nα: 1.45, ρ: 0.015\nPopulation and Technology\nγₑ: 0.3, pop₀: 6838, popadj: 0.134, popasym: 10500, δk: 0.1\nq₀: 63.69, k₀: 135.0, a₀: 3.8, ga₀: 0.079, δₐ: 0.006\nEmissions Parameters\ngσ₁: -0.01, δσ: -0.001, eland₀: 3.3\ndeland: 0.2, e₀: 33.61, μ₀: 0.039\nCarbon Cycle\nmat₀: 830.4, mu₀: 1527.0, ml₀: 10010.0\nmateq: 588.0, mueq: 1350.0, mleq: 10000.0\nFlow Parameters\nϕ₁₂: 0.088, ϕ₂₃: 0.0025\nClimate Model Parameters\nt2xco2: 2.9, fₑₓ0: 0.25, fₑₓ1: 0.7\ntocean₀: 0.0068, tatm₀: 0.8, ξ₁: 0.098\nξ₃: 0.088, ξ₄: 0.025, η: 3.8\nClimate Damage Parameters\nψ₁: 0.0, ψ₂: 0.00267, ψ₃: 2.0\nAbatement Cost\nθ₂: 2.8, pback: 344.0, gback: 0.025, limμ: 1.2\ntnopol: 45.0, cprice₀: 1.0, gcprice: 0.02\nParticipation Parameters\nperiodfullpart: 21.0, partfract2010: 1.0, partfractfull: 1.0\nFossil Fuel Availability\nfosslim: 6000.0\nScaling Parameters\nscale1: 0.016408662, scale2: -3855.106895"
        @test sprint(show, "text/plain", rr_opt) == "Options for Rocky Road 2013R version\nTime step\nN: 60, tstep: 5\nPreferences\nα: 1.45, ρ: 0.015\nPopulation and Technology\nγₑ: 0.3, pop₀: 6838, popadj: 0.134, popasym: 10500, δk: 0.1\nq₀: 63.69, k₀: 135.0, a₀: 3.8, ga₀: 0.079, δₐ: 0.006\nEmissions Parameters\ngσ₁: -0.01, δσ: -0.001, eland₀: 3.3\ndeland: 0.2, e₀: 33.61, μ₀: 0.039\nCarbon Cycle\nmat₀: 830.4, mu₀: 1527.0, ml₀: 10010.0\nmateq: 588.0, mueq: 1350.0, mleq: 10000.0\nFlow Parameters\nϕ₁₂: 0.088, ϕ₂₃: 0.0025\nClimate Model Parameters\nt2xco2: 2.9, fₑₓ0: 0.25, fₑₓ1: 0.7\ntocean₀: 0.0068, tatm₀: 0.8, ξ₁₀: 0.098, ξ₁β: 0.01243\nξ₁: 0.098, ξ₃: 0.088, ξ₄: 0.025, η: 3.8\nClimate Damage Parameters\nψ₁₀: 0.0, ψ₂₀: 0.00267, ψ₁: 0.0, ψ₂: 0.0, ψ₃: 2.0\nAbatement Cost\nθ₂: 2.8, pback: 344.0, gback: 0.025, limμ: 1.2\ntnopol: 45.0, cprice₀: 1.0, gcprice: 0.02\nParticipation Parameters\nperiodfullpart: 21.0, partfract2010: 1.0, partfractfull: 1.0\nFossil Fuel Availability\nfosslim: 6000.0\nScaling Parameters\nscale1: 0.016408662, scale2: -3855.106895"
        @test sprint(show, "text/plain", v2016_opt) == "Options for 2016R beta version\nTime step\nN: 100, tstep: 5\nPreferences\nα: 1.45, ρ: 0.015\nPopulation and Technology\nγₑ: 0.3, pop₀: 7403, popadj: 0.134, popasym: 11500, δk: 0.1\nq₀: 105.5, k₀: 223.0, a₀: 5.115, ga₀: 0.076, δₐ: 0.005\nEmissions Parameters\ngσ₁: -0.0152, δσ: -0.001, eland₀: 2.6\ndeland: 0.115, e₀: 35.85, μ₀: 0.03\nCarbon Cycle\nmat₀: 851.0, mu₀: 460.0, ml₀: 1740.0\nmateq: 588.0, mueq: 360.0, mleq: 1720.0\nFlow Parameters\nϕ₁₂: 0.12, ϕ₂₃: 0.007\nClimate Model Parameters\nt2xco2: 3.1, fₑₓ0: 0.5, fₑₓ1: 1.0\ntocean₀: 0.0068, tatm₀: 0.85, ξ₁: 0.1005\nξ₃: 0.088, ξ₄: 0.025, η: 3.6813\nClimate Damage Parameters\nψ₁₀: 0.0, ψ₁: 0.0, ψ₂: 0.00236, ψ₃: 2.0\nAbatement Cost\nθ₂: 2.6, pback: 550.0, gback: 0.025, limμ: 1.2\ntnopol: 45.0, cprice₀: 2.0, gcprice: 0.02\nFossil Fuel Availability\nfosslim: 6000.0\nScaling Parameters\nscale1: 0.0302455265681763, scale2: -10993.704"
        @test sprint(show, "text/plain", v2016r2_opt) == "Options for 2016R2 version\nTime step\nN: 100, tstep: 5\nPreferences\nα: 1.45, ρ: 0.015\nPopulation and Technology\nγₑ: 0.3, pop₀: 7403, popadj: 0.134, popasym: 11500, δk: 0.1\nq₀: 105.5, k₀: 223.0, a₀: 5.115, ga₀: 0.076, δₐ: 0.005\nEmissions Parameters\ngσ₁: -0.0152, δσ: -0.001, eland₀: 2.6\ndeland: 0.115, e₀: 35.85, μ₀: 0.03\nCarbon Cycle\nmat₀: 851.0, mu₀: 460.0, ml₀: 1740.0\nmateq: 588.0, mueq: 360.0, mleq: 1720.0\nFlow Parameters\nϕ₁₂: 0.12, ϕ₂₃: 0.007\nClimate Model Parameters\nt2xco2: 3.1, fₑₓ0: 0.5, fₑₓ1: 1.0\ntocean₀: 0.0068, tatm₀: 0.85, ξ₁: 0.1005\nξ₃: 0.088, ξ₄: 0.025, η: 3.6813\nClimate Damage Parameters\nψ₁₀: 0.0, ψ₁: 0.0, ψ₂: 0.00236, ψ₃: 2.0\nAbatement Cost\nθ₂: 2.6, pback: 550.0, gback: 0.025, limμ: 1.0\ntnopol: 40.0, cprice₀: 2.0, gcprice: 0.02\nFossil Fuel Availability\nfosslim: 6000.0\nScaling Parameters\nscale1: 0.0302455265681763, scale2: -10993.704"
        @test sprint(show, "text/plain", vanilla_params) == "Calculated Parameters for Vanilla 2013R\nOptimal savings rate: 0.2582781456953642\nCarbon cycle transition matrix coefficients\nϕ₁₁: 0.912, ϕ₂₁: 0.03832888888888889, ϕ₂₂: 0.9591711111111112, ϕ₃₂: 0.0003375, ϕ₃₃: 0.9996625\n2010 Carbon intensity: 0.5491283628802297\nClimate model parameter: 1.3103448275862069\nBackstop price: \n[344.0, 335.4, 327.015, 318.84, 310.869, 303.097, 295.519, 288.132, 280.928, 273.905  …  97.0039, 94.5788, 92.2143, 89.909, 87.6613, 85.4697, 83.333, 81.2497, 79.2184, 77.238]\nGrowth rate of productivity: [0.079, 0.0766652, 0.0743994, 0.0722006, 0.0700667, 0.0679959, 0.0659863, 0.0640362, 0.0621436, 0.060307  …  0.0176273, 0.0171063, 0.0166007, 0.0161101, 0.015634, 0.0151719, 0.0147235, 0.0142884, 0.0138661, 0.0134563]\nEmissions from deforestation: [3.3, 2.64, 2.112, 1.6896, 1.35168, 1.08134, 0.865075, 0.69206, 0.553648, 0.442919  …  4.70992e-5, 3.76793e-5, 3.01435e-5, 2.41148e-5, 1.92918e-5, 1.54335e-5, 1.23468e-5, 9.87741e-6, 7.90193e-6, 6.32154e-6]\nAvg utility social discout rate: [1.0, 0.92826, 0.861667, 0.799852, 0.74247, 0.689206, 0.639762, 0.593866, 0.551262, 0.511715  …  0.0241818, 0.022447, 0.0208367, 0.0193419, 0.0179543, 0.0166663, 0.0154706, 0.0143608, 0.0133305, 0.0123742]\nBase case carbon price: [1.0, 1.10408, 1.21899, 1.34587, 1.48595, 1.64061, 1.81136, 1.99989, 2.20804, 2.43785  …  141.268, 155.971, 172.205, 190.128, 209.916, 231.765, 255.887, 282.52, 311.925, 344.39]\nPopulation and labour: [6838.0, 7242.49, 7612.06, 7947.31, 8249.55, 8520.56, 8762.43, 8977.44, 9167.89, 9336.08  …  10496.6, 10497.1, 10497.5, 10497.8, 10498.1, 10498.4, 10498.6, 10498.8, 10498.9, 10499.1]\nTotal factor productivity: [3.8, 4.12595, 4.46853, 4.82771, 5.2034, 5.59545, 6.00368, 6.42783, 6.8676, 7.32266  …  31.9603, 32.5337, 33.1, 33.6587, 34.2098, 34.7532, 35.2886, 35.8159, 36.3351, 36.846]\nΔσ: [-0.01, -0.0099501, -0.00990045, -0.00985105, -0.00980189, -0.00975298, -0.00970431, -0.00965589, -0.0096077, -0.00955976  …  -0.00778703, -0.00774818, -0.00770951, -0.00767104, -0.00763276, -0.00759468, -0.00755678, -0.00751907, -0.00748155, -0.00744422]\nσ: [0.549128, 0.522347, 0.496996, 0.472992, 0.45026, 0.428725, 0.408319, 0.38898, 0.370647, 0.353262  …  0.0597958, 0.0575124, 0.0553269, 0.0532348, 0.0512316, 0.0493133, 0.0474758, 0.0457154, 0.0440286, 0.0424121]\nθ₁: [0.0674643, 0.0625697, 0.0580447, 0.0538603, 0.0499898, 0.046409, 0.0430951, 0.0400277, 0.0371875, 0.0345572  …  0.00207158, 0.00194266, 0.00182212, 0.00170939, 0.00160394, 0.00150528, 0.00141296, 0.00132656, 0.00124567, 0.00116994]\nExogenious forcing: [0.25, 0.275, 0.3, 0.325, 0.35, 0.375, 0.4, 0.425, 0.45, 0.475  …  0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7]\nFraction of emissions in control regieme: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0  …  1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]"
        @test sprint(show, "text/plain", rr_params) == "Calculated Parameters for Rocky Road 2013R\nOptimal savings rate: 0.2582781456953642\nCarbon cycle transition matrix coefficients\nϕ₁₁: 0.912, ϕ₂₁: 0.03832888888888889, ϕ₂₂: 0.9591711111111112, ϕ₃₂: 0.0003375, ϕ₃₃: 0.9996625\n2010 Carbon intensity: 0.5491283628802297\nClimate model parameter: 1.3103448275862069\nTransient TSC Correction: 0.098\nBackstop price: \n[344.0, 335.4, 327.015, 318.84, 310.869, 303.097, 295.519, 288.132, 280.928, 273.905  …  97.0039, 94.5788, 92.2143, 89.909, 87.6613, 85.4697, 83.333, 81.2497, 79.2184, 77.238]\nGrowth rate of productivity: [0.079, 0.0766652, 0.0743994, 0.0722006, 0.0700667, 0.0679959, 0.0659863, 0.0640362, 0.0621436, 0.060307  …  0.0176273, 0.0171063, 0.0166007, 0.0161101, 0.015634, 0.0151719, 0.0147235, 0.0142884, 0.0138661, 0.0134563]\nEmissions from deforestation: [3.3, 2.64, 2.112, 1.6896, 1.35168, 1.08134, 0.865075, 0.69206, 0.553648, 0.442919  …  4.70992e-5, 3.76793e-5, 3.01435e-5, 2.41148e-5, 1.92918e-5, 1.54335e-5, 1.23468e-5, 9.87741e-6, 7.90193e-6, 6.32154e-6]\nAvg utility social discout rate: [1.0, 0.92826, 0.861667, 0.799852, 0.74247, 0.689206, 0.639762, 0.593866, 0.551262, 0.511715  …  0.0241818, 0.022447, 0.0208367, 0.0193419, 0.0179543, 0.0166663, 0.0154706, 0.0143608, 0.0133305, 0.0123742]\nBase case carbon price: [1.0, 1.10408, 1.21899, 1.34587, 1.48595, 1.64061, 1.81136, 1.99989, 2.20804, 2.43785  …  141.268, 155.971, 172.205, 190.128, 209.916, 231.765, 255.887, 282.52, 311.925, 344.39]\nPopulation and labour: [6838.0, 7242.49, 7612.06, 7947.31, 8249.55, 8520.56, 8762.43, 8977.44, 9167.89, 9336.08  …  10496.6, 10497.1, 10497.5, 10497.8, 10498.1, 10498.4, 10498.6, 10498.8, 10498.9, 10499.1]\nTotal factor productivity: [3.8, 4.12595, 4.46853, 4.82771, 5.2034, 5.59545, 6.00368, 6.42783, 6.8676, 7.32266  …  31.9603, 32.5337, 33.1, 33.6587, 34.2098, 34.7532, 35.2886, 35.8159, 36.3351, 36.846]\nΔσ: [-0.01, -0.0099501, -0.00990045, -0.00985105, -0.00980189, -0.00975298, -0.00970431, -0.00965589, -0.0096077, -0.00955976  …  -0.00778703, -0.00774818, -0.00770951, -0.00767104, -0.00763276, -0.00759468, -0.00755678, -0.00751907, -0.00748155, -0.00744422]\nσ: [0.549128, 0.522347, 0.496996, 0.472992, 0.45026, 0.428725, 0.408319, 0.38898, 0.370647, 0.353262  …  0.0597958, 0.0575124, 0.0553269, 0.0532348, 0.0512316, 0.0493133, 0.0474758, 0.0457154, 0.0440286, 0.0424121]\nθ₁: [0.0674643, 0.0625697, 0.0580447, 0.0538603, 0.0499898, 0.046409, 0.0430951, 0.0400277, 0.0371875, 0.0345572  …  0.00207158, 0.00194266, 0.00182212, 0.00170939, 0.00160394, 0.00150528, 0.00141296, 0.00132656, 0.00124567, 0.00116994]\nExogenious forcing: [0.25, 0.275, 0.3, 0.325, 0.35, 0.375, 0.4, 0.425, 0.45, 0.475  …  0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7, 0.7]\nFraction of emissions in control regieme: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0  …  1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]"
        @test sprint(show, "text/plain", v2016r2_params) == "Calculated Parameters for 2016 versions.\nOptimal savings rate: 0.2582781456953642\nCarbon cycle transition matrix coefficients\nϕ₁₁: 0.88, ϕ₂₁: 0.196, ϕ₂₂: 0.797, ϕ₃₂: 0.0014651162790697675, ϕ₃₃: 0.9985348837209302\n2015 Carbon intensity: 0.35032002736111795\nClimate model parameter: 1.187516129032258\nBackstop price: \n[550.0, 536.25, 522.844, 509.773, 497.028, 484.603, 472.488, 460.675, 449.158, 437.93  …  56.335, 54.9266, 53.5534, 52.2146, 50.9092, 49.6365, 48.3956, 47.1857, 46.0061, 44.8559]\nGrowth rate of productivity: [0.076, 0.0741236, 0.0722934, 0.0705085, 0.0687676, 0.0670698, 0.0654138, 0.0637987, 0.0622235, 0.0606872  …  0.00801034, 0.00781257, 0.00761967, 0.00743154, 0.00724806, 0.0070691, 0.00689456, 0.00672434, 0.00655831, 0.00639639]\nEmissions from deforestation: [2.6, 2.301, 2.03639, 1.8022, 1.59495, 1.41153, 1.2492, 1.10554, 0.978407, 0.86589  …  4.36383e-5, 3.86199e-5, 3.41786e-5, 3.02481e-5, 2.67695e-5, 2.3691e-5, 2.09666e-5, 1.85554e-5, 1.64215e-5, 1.45331e-5]\nAvg utility social discout rate: [1.0, 0.92826, 0.861667, 0.799852, 0.74247, 0.689206, 0.639762, 0.593866, 0.551262, 0.511715  …  0.00123107, 0.00114275, 0.00106077, 0.000984669, 0.000914029, 0.000848457, 0.000787589, 0.000731088, 0.00067864, 0.000629954]\nBase case carbon price: [2.0, 2.20816, 2.43799, 2.69174, 2.97189, 3.28121, 3.62272, 3.99978, 4.41608, 4.87571  …  14828.8, 16372.2, 18076.3, 19957.7, 22034.9, 24328.3, 26860.4, 29656.0, 32742.7, 36150.6]\nPopulation and labour: [7403.0, 7853.09, 8264.92, 8638.97, 8976.56, 9279.54, 9550.18, 9790.92, 10004.3, 10192.8  …  11500.0, 11500.0, 11500.0, 11500.0, 11500.0, 11500.0, 11500.0, 11500.0, 11500.0, 11500.0]\nTotal factor productivity: [5.115, 5.53571, 5.97889, 6.44481, 6.93369, 7.44572, 7.981, 8.53961, 9.12155, 9.72679  …  85.334, 86.0231, 86.7005, 87.3662, 88.0203, 88.6629, 89.2941, 89.9141, 90.5228, 91.1204]\nΔσ: [-0.0152, -0.0151242, -0.0150487, -0.0149736, -0.0148989, -0.0148245, -0.0147506, -0.0146769, -0.0146037, -0.0145308  …  -0.00968977, -0.00964141, -0.0095933, -0.00954543, -0.0094978, -0.00945041, -0.00940325, -0.00935633, -0.00930964, -0.00926318]\nσ: [0.35032, 0.324682, 0.301035, 0.279215, 0.259074, 0.240476, 0.223296, 0.20742, 0.192744, 0.179171  …  0.00140155, 0.00133527, 0.00127243, 0.00121283, 0.00115631, 0.00110268, 0.00105179, 0.00100348, 0.000957617, 0.000914064]\ncumtree: [100.0, 103.546, 106.684, 109.462, 111.92, 114.095, 116.02, 117.724, 119.232, 120.566  …  130.835, 130.835, 130.835, 130.835, 130.835, 130.835, 130.835, 130.835, 130.835, 130.835]\nθ₁: [0.0741062, 0.0669657, 0.0605362, 0.0547447, 0.0495259, 0.0448213, 0.0405787, 0.0367512, 0.0332971, 0.0301786  …  3.03679e-5, 2.82084e-5, 2.62088e-5, 2.43568e-5, 2.26411e-5, 2.10512e-5, 1.95776e-5, 1.82115e-5, 1.69447e-5, 1.57697e-5]\nExogenious forcing: [0.5, 0.529412, 0.558824, 0.588235, 0.617647, 0.647059, 0.676471, 0.705882, 0.735294, 0.764706  …  1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]"
    end
end

#Use same limits for RR case for simplicity
μ_ubound = [if t < 30 1.0 else vanilla_opt.limμ*vanilla_params.partfract[t] end for t in 1:vanilla_opt.N];
vanilla_vars = DICE.model_vars(v2013R(), model, vanilla_opt.N, vanilla_opt.fosslim, μ_ubound, vanilla_params.cpricebase);
rr_vars = DICE.model_vars(v2013R(RockyRoad), modelrr, vanilla_opt.N, vanilla_opt.fosslim, μ_ubound, fill(Inf, vanilla_opt.N));

vanilla_eqs = DICE.model_eqs(OptimalPrice, model, vanilla_opt, vanilla_params, vanilla_vars, isMumps);
rr_eqs = DICE.model_eqs(OptimalPrice, modelrr, rr_opt, rr_params, rr_vars, isMumps);

v2016_vars = DICE.model_vars(v2016R(), model2016, v2016_opt.N, v2016_opt.fosslim, [if t < 30 1.0 else v2016_opt.limμ end for t in 1:v2016_opt.N], fill(Inf, v2016_opt.N));
v2016_eqs = DICE.model_eqs(OptimalPrice, model2016, v2016_opt, v2016_params, v2016_vars, isMumps);
v2016r2_vars = DICE.model_vars(v2016R2(), model2016r2, v2016r2_opt.N, v2016r2_opt.fosslim, [if t < 30 1.0 else v2016r2_opt.limμ end for t in 1:v2016r2_opt.N], fill(Inf, v2016r2_opt.N));
v2016r2_eqs = DICE.model_eqs(OptimalPrice, model2016r2, v2016r2_opt, v2016r2_params, v2016r2_vars, isMumps);
@testset "Model Construction" begin
    @testset "Variables" begin
        @test typeof(vanilla_vars) <: DICE.VariablesV2013
        @test typeof(rr_vars) <: DICE.VariablesV2013
        @test supertype(typeof(vanilla_vars)) == supertype(typeof(rr_vars))
        @test typeof(v2016_vars) <: DICE.VariablesV2016
        @test all(x->supertype(typeof(x)) <: DICE.Variables, [vanilla_vars,rr_vars,v2016_vars])
    end
    @testset "Equations" begin
        @test typeof(vanilla_eqs) <: DICE.VanillaEquations
        @test typeof(rr_eqs) <: DICE.RockyRoadEquations
        @test supertype(typeof(vanilla_eqs)) == supertype(typeof(rr_eqs))
        @test typeof(v2016_eqs) <: DICE.EquationsV2016
        @test all(x->supertype(typeof(x)) <: DICE.Equations, [vanilla_eqs,rr_eqs,v2016_eqs])
    end
    @testset "Scenarios" begin
        @testset "2013R (Vanilla)" begin
            @test DICE.assign_scenario(BasePrice, vanilla_opt, vanilla_params) == vanilla_params.cpricebase
            @test DICE.assign_scenario(OptimalPrice, vanilla_opt, vanilla_params) ==fill(Inf, vanilla_opt.N)
        end
        @testset "2013R (Rocky Road)" begin
            DICE.assign_scenario(BasePrice, modelrr, rr_opt, rr_params, rr_vars);
            @test all(x->isfinite.(JuMP.upper_bound.(rr_vars.CPRICE)[x]), 1:Int(rr_opt.tnopol))
            DICE.assign_scenario(OptimalPrice, modelrr, rr_opt, rr_params, rr_vars);
            @test JuMP.upper_bound(rr_vars.μ[1]) == rr_opt.μ₀
            DICE.assign_scenario(Limit2Degrees, modelrr, rr_opt, rr_params, rr_vars);
            @test JuMP.upper_bound(rr_vars.Tₐₜ[2]) ≈ 2.0
            DICE.scenario_alterations(Stern, rr_opt, rr_params);
            @test rr_opt.α ≈ 1.01
            DICE.scenario_alterations(SternCalibrated, rr_opt, rr_params);
            @test rr_opt.α ≈ 2.1
            DICE.scenario_alterations(Copenhagen, rr_opt, rr_params);
            @test rr_params.partfract[2] ≈ 0.390423082
            # These two are empty functions
            @test DICE.scenario_alterations(Limit2Degrees, rr_opt, rr_params) === nothing
            @test DICE.assign_scenario(Stern, modelrr, rr_opt, rr_params, rr_vars) === nothing
        end
        @testset "2016R beta" begin
            DICE.assign_scenario(BasePrice, model2016, v2016_opt, v2016_params, v2016_vars);
            @test all(x->JuMP.has_upper_bound.(v2016_vars.CPRICE)[x], 1:Int(v2016_opt.tnopol))
            DICE.assign_scenario(OptimalPrice, model2016, v2016_opt, v2016_params, v2016_vars);
            @test JuMP.fix_value(v2016_vars.μ[1]) == v2016_opt.μ₀
        end
        @testset "2016R2" begin
            DICE.assign_scenario(BasePrice, model2016r2, v2016r2_opt, v2016r2_params, v2016r2_vars);
            @test all(x->JuMP.has_upper_bound.(v2016r2_vars.CPRICE)[x], 1:Int(v2016r2_opt.tnopol))
            DICE.assign_scenario(OptimalPrice, model2016r2, v2016r2_opt, v2016r2_params, v2016r2_vars);
            @test JuMP.fix_value(v2016r2_vars.μ[1]) == v2016r2_opt.μ₀
        end
    end
    @testset "Invalid Scenarios" begin
        @test_throws ErrorException DICE.solve(Limit2Degrees, v2013R())
        @test_throws ErrorException DICE.solve(Stern, v2013R())
        @test_throws ErrorException DICE.solve(SternCalibrated, v2013R())
        @test_throws ErrorException DICE.solve(Copenhagen, v2013R())
        @test_throws ErrorException DICE.solve(Limit2Degrees, v2016R())
        @test_throws ErrorException DICE.solve(Stern, v2016R())
        @test_throws ErrorException DICE.solve(SternCalibrated, v2016R())
        @test_throws ErrorException DICE.solve(Copenhagen, v2016R())
    end
end

# Optimisation tests.
# Utility values here are generated by GAMS output from original source files, solved via CONOPT unlsee otherwise stated.
@info "Base Price Scenario with v2013R(Vanilla)"
base = DICE.solve(BasePrice, v2013R(), optimizer = optimizer);
#Tracking issues with some results in Libbum/DICE.jl#21
@testset "Utility" begin
    @testset "2013R (Vanilla)" begin
        @test base.results.UTILITY ≈ 2668.2118865871 atol=1e-4
        @info "Optimal Price Scenario with v2013R(Vanilla)"
        result = DICE.solve(OptimalPrice, v2013R(), optimizer = optimizer);
        @test result.results.UTILITY ≈ 2689.1761542629 atol=1e-4
    end
    @testset "2013R (RockyRoad)" begin
        # Slightly better utility than CONOPT found here.
        @info "Base Price Scenario with v2013R(RockyRoad)"
        result = DICE.solve(BasePrice, v2013R(RockyRoad), optimizer = optimizer);
        @test result.results.UTILITY ≈ 2669.8017192359343 atol=1e-4
        @info "Optimal Price Scenario with v2013R(RockyRoad)"
        result = DICE.solve(OptimalPrice, v2013R(RockyRoad), optimizer = optimizer);
        @test result.results.UTILITY ≈ 2741.2318406380 atol=1e-4
        # Slightly better utility than CONOPT found here.
        @info "Limit 2 Degrees Scenario with v2013R(RockyRoad)"
        result = DICE.solve(Limit2Degrees, v2013R(RockyRoad), optimizer = optimizer);
        @test result.results.UTILITY ≈ 2696.8586446996173 atol=1e-4
        @info "Stern Scenario with v2013R(RockyRoad)"
        result = DICE.solve(Stern, v2013R(RockyRoad), optimizer = optimizer);
        @test result.results.UTILITY ≈ 124305.6025535197 atol=1e-4
        # Tolerance is adjusted here due to small differences between output of Ipopt 3.12.10 and 3.12.13
        @info "Stern Calibrated Scenario with v2013R(RockyRoad)"
        result = DICE.solve(SternCalibrated, v2013R(RockyRoad), optimizer = optimizer);
        @test result.results.UTILITY ≈ -8469.005859988243 atol=1e-3
        @info "Copenhagen Scenario with v2013R(RockyRoad)"
        result = DICE.solve(Copenhagen, v2013R(RockyRoad), optimizer = optimizer);
        # The value is lower here because GAMS results truncate a lot,
        # giving us a slightly lower utility due to the inverse relationship of cprice and mu
        # NOTE: This may not be the only thing. MA97 gets closer to the CONOPT result but not exact.
        if linear_solver == "mumps"
            @test result.results.UTILITY ≈ 2724.1458144691937 atol=1e-4
        else
            @test result.results.UTILITY ≈ 2725.4158209963616 atol=1e-4
        end
    end
    @testset "2016R beta" begin
        @info "Base Price Scenario with v2016R"
        result = DICE.solve(BasePrice, v2016R(), optimizer = optimizer);
        @test result.results.UTILITY ≈ 4485.7434565264 atol=1e-4
        @info "Optimal Price Scenario with v2016R"
        result = DICE.solve(OptimalPrice, v2016R(), optimizer = optimizer);
        @test result.results.UTILITY ≈ 4517.3146811528 atol=1e-4
    end
    @testset "2016R2" begin
        @info "Base Price Scenario with v2016R2"
        result = DICE.solve(BasePrice, v2016R2(), optimizer = optimizer);
        @test result.results.UTILITY ≈ 4485.6767040826 atol=1e-4
        @info "Optimal Price Scenario with v2016R2"
        result = DICE.solve(OptimalPrice, v2016R2(), optimizer = optimizer);
        @test result.results.UTILITY ≈ 4515.8311036492 atol=1e-4
    end
end

#Show model
@testset "Display" begin
    if linear_solver == "mumps"
        @test sprint(show, "text/plain", base) == "Base (current policy) carbon price scenario using v2013R (Vanilla flavour).\nA JuMP Model\nMaximization problem with:\nVariables: 1561\nObjective function type: VariableRef\n`GenericAffExpr{Float64,VariableRef}`-in-`MathOptInterface.EqualTo{Float64}`: 656 constraints\n`GenericQuadExpr{Float64,VariableRef}`-in-`MathOptInterface.EqualTo{Float64}`: 240 constraints\n`VariableRef`-in-`MathOptInterface.EqualTo{Float64}`: 15 constraints\n`VariableRef`-in-`MathOptInterface.GreaterThan{Float64}`: 716 constraints\n`VariableRef`-in-`MathOptInterface.LessThan{Float64}`: 298 constraints\nNonlinear: 539 constraints\nModel mode: AUTOMATIC\nCachingOptimizer state: ATTACHED_OPTIMIZER\nSolver name: Ipopt\nNames registered in the model: C, CCA, CEMUTOTPER, CPC, CPRICE, DAMAGES, E, Eind, FORC, I, K, MCABATE, Mᵤₚ, Mₐₜ, Mₗₒ, PERIODU, RI, S, Tₐₜ, Tₗₒ, UTILITY, Y, YGROSS, YNET, Λ, Ω, μ"
    else
        @test sprint(show, "text/plain", base) == "Base (current policy) carbon price scenario using v2013R (Vanilla flavour).\nA JuMP Model\nMaximization problem with:\nVariables: 1561\nObjective function type: VariableRef\n`GenericAffExpr{Float64,VariableRef}`-in-`MathOptInterface.EqualTo{Float64}`: 655 constraints\n`GenericQuadExpr{Float64,VariableRef}`-in-`MathOptInterface.EqualTo{Float64}`: 240 constraints\n`VariableRef`-in-`MathOptInterface.EqualTo{Float64}`: 17 constraints\n`VariableRef`-in-`MathOptInterface.GreaterThan{Float64}`: 714 constraints\n`VariableRef`-in-`MathOptInterface.LessThan{Float64}`: 297 constraints\nNonlinear: 538 constraints\nModel mode: AUTOMATIC\nCachingOptimizer state: ATTACHED_OPTIMIZER\nSolver name: Ipopt\nNames registered in the model: C, CCA, CEMUTOTPER, CPC, CPRICE, DAMAGES, E, Eind, FORC, I, K, MCABATE, Mᵤₚ, Mₐₜ, Mₗₒ, PERIODU, RI, S, Tₐₜ, Tₗₒ, UTILITY, Y, YGROSS, YNET, Λ, Ω, μ"
    end
end


text:
for func in (:daily, :weekly, :monthly)
    x = "digital_currency_$(func)"
    fname = Symbol(x)
    @eval begin
        function ($fname)(symbol::String, market::String="USD"; client = GLOBAL[], outputsize::String="compact", datatype::String="json", parser = "default")
            @argcheck in(outputsize, ["compact", "full"])
            @argcheck in(datatype, ["json", "csv"])
            uri = _form_uri_head(client, uppercase($x)) * "&symbol=$symbol&market=$market" * _form_uri_tail(client, outputsize, datatype)
            data = retry(_get_request, delays=Base.ExponentialBackOff(n=3, first_delay=5, max_delay=1000))(uri)
            p = _parser(parser, datatype)
            return p(data)
        end
        export $fname
    end
end

function crypto_rating(symbol::String; client = GLOBAL[], parser = "default")
    uri = _form_uri_head(client, "CRYPTO_RATING") * "&symbol=$symbol" * _form_uri_tail(client, nothing, nothing)
    data = retry(_get_request, delays=Base.ExponentialBackOff(n=3, first_delay=5, max_delay=1000))(uri)

    p = _parser(parser, "json")
    return p(data)
end

ext: jl


text:
# degree

degree(model::Union{UBCM, UBCMCompact}, i::Int) = sum(a(i,j,model.x) for j in eachindex(model.x) if j≠i)
degree(model::Union{UBCM, UBCMCompact}, v::Vector{Int}=collect(1:length(model.x))) = [degree(model,i) for i in v]

indegree(model::Union{UBCM, UBCMCompact}, i::Int) = degree(model, i)
outdegree(model::Union{UBCM, UBCMCompact}, v::Vector{Int}=collect(1:length(model.x))) = degree(model, v)

#⇽(x) = x + oneunit(x)

#f⟨a⟩(x) = x + oneunit(x)


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

# Sparse assembly test, 1d poisson equation:
# $$u'' = 0$$,
# $$u(0) = 0$$,
# $$u(1) = 1$$,
# discretized to $$N$$ elements.

using Test, BenchmarkTools, LinearAlgebra, SparseArrays
using Revise
using FEMSparse

function fill_dense(N)
    h = 1.0/N
    Ke = h*[1.0 -1.0; -1.0 1.0]
    K = zeros(N+1, N+1)
    for i in 1:N
        K[i:i+1,i:i+1] .+= Ke
    end
    return K
end

function fill_sparse_csc(N)
    h = 1.0/N
    Ke = h*[1.0 -1.0; -1.0 1.0]
    K = spzeros(N+1, N+1)
    for i in 1:N
        K[i:i+1,i:i+1] += Ke
    end
    return K
end

function fill_sparse_coo(N)
    h = 1.0/N
    Ke = h*[1.0 -1.0; -1.0 1.0]
    K = SparseMatrixCOO()
    for i in 1:N
        add!(K, i:i+1, i:i+1, Ke)
    end
    return SparseMatrixCSC(K)
end

function test_assembly(K)
    N = size(K,1)-1
    # set boundary conditions: u(0) = 0, u(1) = 1
    u = zeros(N+1)
    u[end] = 1.0
    f = -K*u
    K[1,:] .= 0.0
    K[:,1] .= 0.0
    K[1,1] = 1.0
    f[end] = 0.0
    K[end,:] .= 0.0
    K[:,end] .= 0.0
    K[end,end] = 1.0
    f[end] = 1.0
    u = K\f
    u_acc = collect(range(0.0, stop=1.0, length=N+1))
    return isapprox(u, u_acc; atol=1.0e-12)
end

@info("Warm-up")

@test test_assembly(fill_dense(10))
@test test_assembly(fill_sparse_csc(10))
@test test_assembly(fill_sparse_coo(10))

@info("Test")

@info("Dense matrix:")
@btime fill_dense(30_000)

@info("Sparse matrix (CSC format):")
@btime fill_sparse_csc(30_000)

@info("Sparse matrix (COO format):")
@btime fill_sparse_coo(30_000)


text:
function rmsd_simd(a, b)
    N = length(a)
    N === length(b) || throw(DimensionMismatch())
    err = zero(promote_type(eltype(a), eltype(b)))
    @turbo for i ∈ eachindex(a,b)
        err += abs2(a[i] - b[i])
    end
    return sqrt(err/size(a,1))
end

logrange(x1, x2, n) = (exp10(y) for y in range(log10(x1), log10(x2), length=n))

"""
    calcresiduals(data, fc; fs)

Calculate the residuals of the data for a range of cutoff frequencies.

Following the method outlined by Winter (1990), the residual for a given cutoff freqency is defined as:

```math
R(f_c) = \\sqrt{\\frac{1}{N}\\sum_{i=1}^{N}(X_i - X̂_i)^2}
```
where ``f_c`` is the cutoff frequency, ``X_i`` is raw data at the *i* th sample, and ``X̂_i``
is filtered data at the *i* th sample [^1].

[1] D.A. Winter, Choice of Cutoff Frequency - Residual Analysis, in: Biomechanics and Motor
Control of Human Movement, 4th ed., John Wiley & Sons, Inc., Hoboken, New Jersey, 2009: pp.
70–73.

"""
function calcresiduals(data, fc=[range(1, fs/2/1.247, length=200);]; fs, db=false)
    sort!(fc)
    last(fc) ≤ fs/2/1.247 || throw(ArgumentError("cutoff frequency cannot be greater than Nyquist, corrected for dual-pass, at $(fs/2/1.247)"))
    numFc = length(fc)
    R = zeros(Float64,numFc)
    designmethod = Butterworth(2)
    filtdata = similar(data)

    for j in 1:numFc
        responsetype = Lowpass(fc[j]*1.247; fs=fs)
        thefilter = digitalfilter(responsetype, designmethod)
        filtdata[:] = filtfilt(thefilter, data)
        R[j] = rmsd_simd(data, filtdata)
    end

    if db
        baseP = rmsd_simd(data, zero(data))
        R[:] = amp2db.(R./baseP)
    end

    return R
end

"""
    optfc(data, fc; fs, [tol=0.9, rmsnoise])

Find the optimal cutoff frequency.

The optimal cutoff frequency is the point at which the signal distortion is equal to the
baseline RMS noise. The RMS noise is calculated as the intercept of the residual noise line,
which is the longest linear section of residuals, ending at the residual for the highest
frequency evaluated, with a correlation coefficient ≥`tol` [^1].

[1] D.A. Winter, Choice of Cutoff Frequency - Residual Analysis, in: Biomechanics and Motor
Control of Human Movement, 4th ed., John Wiley & Sons, Inc., Hoboken, New Jersey, 2009: pp.
70–73.
"""
function optfc(data, fc=[range(1, fs/2/1.247, length=200);]; fs, tol=0.9, rmsnoise=nothing)
    Rs = calcresiduals(data, fc; fs)
    if isnothing(rmsnoise)
        linregs = [ @views(((; a,b = _linreg(fc[n:end], Rs[n:end]))...,
                            r2 = cor((fc[n:end], Rs[n:end]))))
            for n in 1:length(fc)-1 ]
        d = findlast(x -> x.r2 < tol && x.b < 0, linregs)
        (isnothing(d) || d === lastindex(linregs)) &&
            throw(error("a linear section of residuals was not found with an r² ≥ $tol; try changing the tolerance"))
        rmsnoise = linregs[d+1].a
    end

    i = findlast(>(rmsnoise), Rs)
    @assert !isnothing(i) # At least one residual must be less than the intercept

    return fc[i], rmsnoise
end


text:
using ReinforcementLearningZoo
using Test
using ReinforcementLearningBase
using ReinforcementLearningCore
using ReinforcementLearningEnvironments
using Flux
using Statistics
using Random
using OpenSpiel
using StableRNGs
import GridWorlds

function get_optimal_kuhn_policy(α = 0.2)
    TabularRandomPolicy(
        table = Dict(
            "0" => [1 - α, α],
            "0pb" => [1.0, 0.0],
            "1" => [1.0, 0.0],
            "1pb" => [2.0 / 3.0 - α, 1.0 / 3.0 + α],
            "2" => [1 - 3 * α, 3 * α],
            "2pb" => [0.0, 1.0],
            "0p" => [2.0 / 3.0, 1.0 / 3.0],
            "0b" => [1.0, 0.0],
            "1p" => [1.0, 0.0],
            "1b" => [2.0 / 3.0, 1.0 / 3.0],
            "2p" => [0.0, 1.0],
            "2b" => [0.0, 1.0],
        ),
    )
end

@testset "ReinforcementLearningZoo.jl" begin

    @testset "training" begin
        mktempdir() do dir
            for method in (:BasicDQN, :BC, :DQN, :PrioritizedDQN, :Rainbow, :IQN, :VPG)
                res = run(
                    Experiment(
                        Val(:JuliaRL),
                        Val(method),
                        Val(:CartPole),
                        nothing;
                        save_dir = joinpath(dir, "CartPole", string(method)),
                    ),
                )
                @info "stats for $method" avg_reward = mean(res.hook[1].rewards) avg_fps =
                    1 / mean(res.hook[2].times)
            end

            for method in (:BasicDQN, :DQN)
                res = run(
                    Experiment(
                        Val(:JuliaRL),
                        Val(method),
                        Val(:MountainCar),
                        nothing;
                        save_dir = joinpath(dir, "MountainCar", string(method)),
                    ),
                )
                @info "stats for $method" avg_reward = mean(res.hook[1].rewards) avg_fps =
                    1 / mean(res.hook[2].times)
            end

            for method in (:A2C, :A2CGAE, :PPO, :MAC)
                res = run(
                    Experiment(
                        Val(:JuliaRL),
                        Val(method),
                        Val(:CartPole),
                        nothing;
                        save_dir = joinpath(dir, "CartPole", string(method)),
                    ),
                )
                @info "stats for $method" avg_reward =
                    mean(Iterators.flatten(res.hook[1].rewards))
            end

            for method in (:DDPG, :SAC, :TD3)
                res = run(
                    Experiment(
                        Val(:JuliaRL),
                        Val(method),
                        Val(:Pendulum),
                        nothing;
                        save_dir = joinpath(dir, "Pendulum", string(method)),
                    ),
                )
                @info "stats for $method" avg_reward =
                    mean(Iterators.flatten(res.hook[1].rewards))
            end
        end
    end

    @testset "minimax" begin
        e = E`JuliaRL_Minimax_OpenSpiel(tic_tac_toe)`
        run(e)
        @test e.hook[1][] == e.hook[0][] == [0.0]
    end

    @testset "GridWorlds" begin
        mktempdir() do dir
            for method in (:BasicDQN,)
                res = run(
                    Experiment(
                        Val(:JuliaRL),
                        Val(method),
                        Val(:EmptyRoom),
                        nothing;
                        save_dir = joinpath(dir, "EmptyRoom", string(method)),
                    ),
                )
                @info "stats for $method" avg_reward = mean(res.hook[1].rewards) avg_fps =
                    1 / mean(res.hook[2].times)
            end
        end
    end

    @testset "TabularCFR" begin
        e = E`JuliaRL_TabularCFR_OpenSpiel(kuhn_poker)`
        run(e)

        reset!(e.env)
        expected_values = expected_policy_values(e.policy, e.env)
        @test isapprox(expected_values[1], -1 / 18; atol = 0.001)
        @test isapprox(expected_values[2], 1 / 18; atol = 0.001)
    end

    include("cfr/cfr.jl")
end


text:
abstract type Trace end
"""
    FactTrace

See [`Trace`](@ref).
"""
struct FactTrace{FT,T,S,S2,R} <: Trace
    F::FT
    t0::T
    x0::S
    θ0::S2
    events::R
end

"""
    FactTrace

See [`Trace`](@ref).
"""
struct PDMPTrace{FT,T,S,S2,S3,R} <: Trace
    F::FT
    t0::T
    x0::S
    θ0::S2
    f::S3
    events::R
end

"""
    Trace(t0::T, x0, θ0, F::Union{ZigZag,FactBoomerang})

Trace object for exact trajectory of pdmp samplers. Returns an iterable `FactTrace` object.
Note that iteration iterates pairs `t => x` where the vector `x` is modified
inplace, so copies have to be made if the `x` is to be saved.
`collect` applied to a trace object automatically copies `x`.
[`discretize`](@ref) returns a discretized version.
"""
Trace(t0::T, x0, θ0, F::Union{ZigZag,FactBoomerang,JointFlow}) where {T} = FactTrace(F, t0, x0, θ0, Tuple{T,Int,eltype(x0),eltype(θ0)}[])
Trace(t0::T, x0::U, θ0::U2, F::Union{BouncyParticle,Boomerang}) where {T, U, U2} = PDMPTrace(F, t0, x0, θ0, ones(Bool, length(x0)), Tuple{T,U,U2,Nothing}[])
Trace(t0::T, x0::U, θ0::U2, f::U3, F::Union{BouncyParticle,Boomerang}) where {T, U, U2, U3} = PDMPTrace(F, t0, x0, θ0, f, Tuple{T,U,U2,U3}[])

Base.length(FT::Trace) = 1 + length(FT.events)

function Base.iterate(FT::FactTrace)
    t, x, θ = FT.t0, copy(FT.x0), copy(FT.θ0)
    t => x, (t, x, θ, 1)
end

function Base.iterate(FT::PDMPTrace)
    t, x, θ = FT.t0, FT.x0, FT.θ0
    t => x, (t, x, θ, 1)
end

function Base.iterate(FT::FactTrace, (t, x, θ, k))
    k > length(FT.events) && return nothing
    k == length(FT.events) && return t => x, (t, x, θ, k + 1)
    t2, i, xi, θi = FT.events[k]
    t, x, θ = move_forward!(t2 - t, t, x, θ, FT.F)
    t = t2
    x[i] = xi
    θ[i] = θi
    return t => x, (t, x, θ, k + 1)
end

function Base.iterate(FT::PDMPTrace, (t, x, θ, k))
    k > length(FT.events) && return nothing
    k == length(FT.events) && return t => x, (t, x, θ, k + 1)
    t, x, θ, _ = FT.events[k]
    return t => x, (t, x, θ, k + 1)
end

Base.push!(FT::Trace, ev) = push!(FT.events, ev)
Base.collect(FT::FactTrace) = collect(t=>copy(x) for (t, x) in FT)
Base.collect(FT::PDMPTrace) = collect(t=>x for (t, x) in FT)

struct Discretize{T, S}
    FT::T
    dt::S
end
Base.IteratorSize(::Discretize) = Iterators.SizeUnknown()

"""
    discretize(trace::Trace, dt)

Discretize `trace` with step-size dt. Returns iterable object
iterating pairs `t => x`.

Iteration changes the vector `x` inplace,
`collect` creates necessary copies.
"""
discretize(FT, dt) = Discretize(FT, dt)

function Base.iterate(D::Discretize{<:FactTrace})
    FT = D.FT
    t, x, θ = FT.t0, copy(FT.x0), copy(FT.θ0)
    t => x, (t, x, θ, 1)
end

function Base.iterate(D::Discretize{<:PDMPTrace})
    FT = D.FT
    t, x, θ, f = FT.t0, copy(FT.x0), copy(FT.θ0), copy(FT.f)
    t => x, (t, x, θ, f, 1)
end

function Base.iterate(D::Discretize{<:FactTrace{T}}, (t, x, θ, k)) where {T<:Union{FactBoomerang,ZigZag,JointFlow}}
    dt = D.dt
    FT = D.FT
    while true
        k > length(FT.events) && return nothing
        ti, i, xi, θi = FT.events[k]
        if t + dt < ti
            t, x, θ = move_forward!(dt, t, x, θ, FT.F)
            return t => x, (t, x, θ, k)
        else # move not more than to ti to change direction
            Δt = ti - t
            dt = dt - Δt
            t, x, θ = move_forward!(Δt, t, x, θ, FT.F)
            t = ti
            x[i] = xi
            θ[i] = θi
            k = k + 1
        end
    end
end

function Base.iterate(D::Discretize{<:PDMPTrace{T}}, (t, x, θ, f, k)) where {T<:Union{Boomerang,BouncyParticle}}
    dt = D.dt
    FT = D.FT
    while true
        k > length(FT.events) && return nothing
        tn = first(FT.events[k])
        if t + dt < tn
            t, x, θ = smove_forward!(dt, t, x, θ, f, FT.F)
            return t => x, (t, x, θ, f, k)
        else # move not more than to ti to change direction
            Δt = tn - t
            dt = dt - Δt
            t, x, θ, f_ = FT.events[k]
            if !isnothing(f_)
                f = f_
            end
            k = k + 1
        end
    end
end

function Base.collect(D::Discretize{<:FactTrace})
    collect(t=>copy(x) for (t, x) in D)
end
function Base.collect(D::Discretize{<:PDMPTrace})
    collect(t=>copy(x) for (t, x) in D)
end

function Statistics.mean(trace::ZigZagBoomerang.Trace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    T = trace.events[end][1]
    t2 = trace.t0
    t = fill(t2, length(x))
    k = 1
    scale = 1/(2T)
    while k <= length(trace.events)
        t2, i, xi, θi = trace.events[k]
        k += 1
        y[i] += (x[i]+xi)*(t2-t[i])*scale
        t[i] = t2
        x[i] = xi
        θ[i] = θi
    end
    y
end

function cummean(trace::FactTrace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    T = trace.events[end][1]
    t2 = trace.t0
    ys = [(t=[t2], y=[xi]) for xi in x]

    t = fill(t2, length(x))
    k = 1
    while k <= length(trace.events)
        t2, i, xi, θi = trace.events[k]
        k += 1
        y[i] += (x[i]+xi)*(t2-t[i])
        t[i] = t2
        x[i] = xi
        θ[i] = θi
        push!(ys[i].t, t[i])
        push!(ys[i].y, y[i]/(2t[i]))
    end
    ys
end
export cummean

function Statistics.mean(trace::PDMPTrace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    ys = fill(y, 0)
    T = trace.events[end][1]
    t2 = trace.t0
    t = t2
    k = 1
    while k <= length(trace.events)
        t2, x2, _ =trace.events[k]
        k += 1
        y  += (x + x2)*(t2-t)
        t = t2
        x  = x2
    end
    y/T
end

function cummean(trace::PDMPTrace)
    x = copy(trace.x0)
    θ = copy(trace.θ0)
    y = 0*x
    ys = fill(y, 0)
    T = trace.events[end][1]
    t2 = trace.t0
    t = t2
    k = 1
    while k <= length(trace.events)
        t2, x2, _ =trace.events[k]
        k += 1
        y  += (x + x2)*(t2-t)
        t = t2
        x  = x2
        push!(ys, y/(2t))
    end
    ys
end

"""
    subtrace(tr, J)

Compute the trace of a subvector `x[J]`,
returns a trace object.
"""
function subtrace(tr, J)
    @assert issorted(J)
    F = tr.F
    t0 = tr.t0
    x0 = tr.x0[J]
    θ0 = tr.θ0[J]

    str = Trace(t0, x0, θ0, F)
    events = str.events
    for ev in tr.events
        r = searchsorted(J, ev[2])
        isempty(r) && continue
        push!(events, (ev[1], r[1], ev[3:end]...))
    end
    str
end


text:
# Before running this, please make sure to activate and instantiate the environment
# corresponding to [this `Project.toml`](https://raw.githubusercontent.com/alan-turing-institute/MLJTutorials/master/Project.toml) and [this `Manifest.toml`](https://raw.githubusercontent.com/alan-turing-institute/MLJTutorials/master/Manifest.toml)
# so that you get an environment which matches the one used to generate the tutorials:
#
# ```julia
# cd("MLJTutorials") # cd to folder with the *.toml
# using Pkg; Pkg.activate("."); Pkg.instantiate()
# ```

using MLJ, RDatasets, PrettyPrinting

@load DecisionTreeClassifier pkg=DecisionTree

carseats = dataset("ISLR", "Carseats")

first(carseats, 3) |> pretty

High = ifelse.(carseats.Sales .<= 8, "No", "Yes") |> categorical;
carseats[!, :High] = High;

X = select(carseats, Not([:Sales, :High]))
y = carseats.High;

@pipeline HotTreeClf(hot = OneHotEncoder(),
                     tree = DecisionTreeClassifier()) is_probabilistic=true

mdl = HotTreeClf()
mach = machine(mdl, X, y)
fit!(mach);

ypred = predict_mode(mach, X)
misclassification_rate(ypred, y)

train, test = partition(eachindex(y), 0.5, shuffle=true, rng=333)
fit!(mach, rows=train)
ypred = predict_mode(mach, rows=test)
misclassification_rate(ypred, y[test])

r_mpi = range(mdl, :(tree.max_depth), lower=1, upper=10)
r_msl = range(mdl, :(tree.min_samples_leaf), lower=1, upper=50)

tm = TunedModel(model=mdl, ranges=[r_mpi, r_msl], tuning=Grid(resolution=8),
                resampling=CV(nfolds=5, rng=112),
                operation=predict_mode, measure=misclassification_rate)
mtm = machine(tm, X, y)
fit!(mtm, rows=train)

ypred = predict_mode(mtm, rows=test)
misclassification_rate(ypred, y[test])

fitted_params(mtm).best_model.tree

@load DecisionTreeRegressor pkg=DecisionTree

boston = dataset("MASS", "Boston")

y, X = unpack(boston, ==(:MedV), col -> true)

train, test = partition(eachindex(y), 0.5, shuffle=true, rng=551);

scitype(X)

X = coerce(X, autotype(X, rules=(:discrete_to_continuous,)))

dtr_model = DecisionTreeRegressor()
dtr = machine(dtr_model, X, y)

fit!(dtr, rows=train)

ypred = predict(dtr, rows=test)
round(rms(ypred, y[test]), sigdigits=3)

r_depth = range(dtr_model, :max_depth, lower=2, upper=20)

tm = TunedModel(model=dtr_model, ranges=[r_depth], tuning=Grid(resolution=10),
                resampling=CV(nfolds=5, rng=1254), measure=rms)
mtm = machine(tm, X, y)

fit!(mtm, rows=train)

ypred = predict(mtm, rows=test)
round(rms(ypred, y[test]), sigdigits=3)

@load RandomForestRegressor pkg=ScikitLearn

rf_mdl = RandomForestRegressor()
rf = machine(rf_mdl, X, y)
fit!(rf, rows=train)

ypred = predict(rf, rows=test)
round(rms(ypred, y[test]), sigdigits=3)

@load XGBoostRegressor

xgb_mdl = XGBoostRegressor(num_round=10, max_depth=10)
xgb = machine(xgb_mdl, X, y)
fit!(xgb, rows=train)

ypred = predict(xgb, rows=test)
round(rms(ypred, y[test]), sigdigits=3)

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

max_stars_repo_name: ven-k/MLJTutorials

max_issues_repo_name: ven-k/MLJTutorials

max_forks_repo_name: ven-k/MLJTutorials


text:
x, y = generators(MPoly{Float64}, :x, :y)

@test vars(x) == [:x, :y]
@test nvars(x) == 2
@test eltype(x) == Float64


text:
leftgroupjoin(left, right) = leftgroupjoin(identity, identity, left, right)
leftgroupjoin(lkey, rkey, left, right) = leftgroupjoin(lkey, rkey, merge, left, right)
leftgroupjoin(lkey, rkey, f, left, right) = leftgroupjoin(lkey, rkey, f, isequal, left, right)

"""
    leftgroupjoin(lkey, rkey, f, comparison, left, right)

Creates a collection if groups labelled by `lkey(l)` where each group contains elements
`f(l, r)` which satisfy `comparison(lkey(l), rkey(r))`. If there are no matches, the group
is still created (with an empty collection).

This operation shares some similarities with an SQL left outer join.

### Example

```jldoctest
julia> leftgroupjoin(iseven, iseven, tuple, ==, [1,2,3,4], [0,1,2])
Dict{Bool,Array{Tuple{Int64,Int64},1}} with 2 entries:
  false => Tuple{Int64,Int64}[(1, 1), (3, 1)]
  true  => Tuple{Int64,Int64}[(2, 0), (2, 2), (4, 0), (4, 2)]
```
"""
function leftgroupjoin(lkey, rkey, f, comparison, left, right)
    # The O(length(left)*length(right)) generic method when nothing about `comparison` is known

    # TODO Do this inference-free, like comprehensions...
    T = promote_op(f, eltype(left), eltype(right))
    K = promote_op(lkey, eltype(left))
    V = Vector{T}
    out = Dict{K, V}()
    for a ∈ left
        key = lkey(a)
        group = get!(() -> T[], out, key)
        for b ∈ right
            if comparison(lkey(a), rkey(b))
                push!(group, f(a, b))
            end
        end
    end
    return out
end

function leftgroupjoin(lkey, rkey, f, ::typeof(isequal), left, right)
    # isequal heralds a hash-based approach, roughly O(length(left) * log(length(right)))

    # TODO Do this inference-free, like comprehensions...
    T = promote_op(f, eltype(left), eltype(right))
    K = promote_op(rkey, eltype(right))
    V = eltype(right)
    dict = Dict{K,Vector{V}}() # maybe a different stategy if right is unique
    for b ∈ right
        key = rkey(b)
        push!(get!(()->Vector{V}(), dict, key), b)
    end

    K2 = promote_op(lkey, eltype(left))
    out = Dict{K2, Vector{T}}()
    for a ∈ left
        key = lkey(a)
        group = get!(() -> T[], out, key)
        dict_index = Base.ht_keyindex(dict, key)
        if dict_index > 0 # -1 if key not found
            for b ∈ dict.vals[dict_index]
                push!(group, f(a, b))
            end
        end
    end
    return out
end

function leftgroupjoin(lkey, rkey, f, ::typeof(isequal), left::AbstractArray, right::AbstractArray)
    # isequal heralds a hash-based approach, roughly O(length(left) * log(length(right)))

    # TODO Do this inference-free, like comprehensions...
    T = promote_op(f, eltype(left), eltype(right))
    K = promote_op(rkey, eltype(right))
    V = eltype(right)
    dict = Dict{K,Vector{V}}() # maybe a different stategy if right is unique
    for b ∈ right
        key = rkey(b)
        push!(get!(()->Vector{V}(), dict, key), b)
    end

    K2 = promote_op(lkey, eltype(left))
    out = Dict{K2, typeof(empty(left, T))}()
    for a ∈ left
        key = lkey(a)
        group = get!(() -> empty(left, T), out, key)
        dict_index = Base.ht_keyindex(dict, key)
        if dict_index > 0 # -1 if key not found
            for b ∈ dict.vals[dict_index]
                push!(group, f(a, b))
            end
        end
    end
    return out
end


text:
#=

  Combinatorial auction in ConstraintSolver.jl

  http://en.wikipedia.org/wiki/Combinatorial_auction
  """
  A combinatorial auction is an auction in which bidders can place
  bids on combinations of items, or "packages," rather than
  just individual items. Simple combinatorial auctions have been
  used for many years in estate auctions, where a common procedure
  is to auction the individual items and then at the end to accept
  bids for packages of items.
  """

  This simple example is from the lecture slides
  Constraint Satisfaction Problems, Constraint Optimization
  by Bernhard Nebel and Stefan Wölfl
  http://www.informatik.uni-freiburg.de/~ki/teaching/ws0910/csp/csp10-handout4.pdf
  """
  In combinatorial auctions, bidders can give bids for set of items.
  The auctioneer [then] has to generate an optimial selection, e.g.
  one that maximizes revenue.

  Definition
  The combinatorial auction problem  is specified as follows:
    Given: A set of items Q = {q1,...,qn} and a set of bids
           B = {b1,...,bm} such that each bid is bi = (Qi, ri),
           where Qi (= Q and ri is a strictly positive real number.
    Task: Find a subset of bids B'(= B such that any two bids in B'
          do not share an item maximizing Sum(Qi,ri) (= Biri.

  ...

  Example Auction

  Consider the following auction:
    b1 = {1,2,3,4}, r1 = 8
    b2 = {2,3,6},   r2 = 6
    b3 = {1,4,5},   r3 = 5
    b4 = {2,8},     r4 = 2
    b5 = {5,6},     r5 = 2

  What is the optimal assignment?
  """

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Picat page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function combinatorial_auction(problem, print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    num_items = problem[:num_items]
    num_bids  = problem[:num_bids]
    packages  = problem[:packages]
    bids      = problem[:bids]

    @variable(model, x[1:num_bids], Bin)
    @variable(model, 0 <= total <= sum(bids), Int)

    @constraint(model, total == sum(x.*bids))

    # ensure that each item is selected atmost once
    for j in 1:num_items
        @constraint(model, sum([x[i] for i in 1:num_bids if j in packages[i]]) <= 1)
    end

    @objective(model,Max,total)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                total_val = convert.(Integer,JuMP.value.(total; result=sol))
                println("x:$x_val")
                println("total:$total_val")
                for b in 1:num_bids
                    if x_val[b] == 1
                        println("bid:$b items:$(packages[b])")
                    end
                end
                println()

            end
        end
    else
        println("status:$status")
    end

    return status, num_sols
end

#
# Problem instances
#
problems = Dict(
    #
    # The example cited above
    #
    :1 => Dict(
        :num_items => 7,
        :num_bids => 5,
        :packages => [[1,2,3,4],
                      [2,3,6],
                      [1,4,5],
                      [2,7],
                      [5,6]],
        :bids => [8,6,5,2,2]
    ),

    #
    # From Numberjack Tutorial, page 24 (slide 51/175)
    #
    :2 => Dict(
            :num_items => 4,
            :num_bids => 5,
            :packages => [[1,2],
                          [1,3],
                          [2,4],
                          [2,3,4],
                          [1]],
            :bids => [8,6,5,2,2]
        ),

    #
    # From "A Faster Core Constraint Generation Algorithm for Combinatorial Auctions"
    # Benedikt Bunz, Sven Seuken, Benjamin Lubin
    # page 4
    #
    :3  => Dict(
        :num_items => 5,
        :num_bids => 7,
        :packages => [[1],
                      [2],
                      [3],
                      [4],
                      [1,2,3,4,5],
                      [1,2,5],
                      [3,4,5]],
        :bids => [10,10,10,10,12,8,8]
    ),

    #
    # From "Combinatorial Auctions:  Complexity and Algorithms"
    # Martin Bichler
    # page 4 (table 1)
    #
    # Line  Bids                       B1    B2   B3   B4
    #    1  1000t grain in Berlin      1     0    1    1
    #    2  800t grain in Munich       0     1    1    1
    #    3  800t grain in Vienna       1     1    1    0
    #    4  Bid price (in thousands) 150   125   300 125
    #
    :4 => Dict(
            :num_items => 3,
            :num_bids => 4,
            :packages => [[1,3],
                          [2,3],
                          [1,2,3],
                          [1,2]],
            :bids => [150,125,300,125],
    ),

)

# @time combinatorial_auction(problems[:1])
for p in sort(collect(keys(problems)))
    println("\nproblem:$p")
    @time combinatorial_auction(problems[p])
end

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
function madlibs(template)
    println("The story template is:\n", template)
    fields = Set(getfield.(collect(eachmatch(r"<[^>]+>", template)), :match))
    print("\nInput a comma-separated list of words to replace the following items:",
          join(fields, ", "), "\n -> ")
    values = split(readline(STDIN), ",")
    for (m, v) in zip(fields, values)
        template = replace(template, m, v)
    end
    println("\nThe story becomes:\n", template)
end

const template = """
<name> went for a walk in the park. <he or she>
found a <noun>. <name> decided to take it home.
"""

madlibs(template)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
###

# type to throw on succesful convergence
mutable struct StateConverged
    x0::Real
end

# type to throw on failure
mutable struct ConvergenceFailed
    reason::AbstractString
end

bracketing_error = """The interval [a,b] is not a bracketing interval.
You need f(a) and f(b) to have different signs (f(a) * f(b) < 0).
Consider a different bracket or try fzero(f, c) with an initial guess c.

"""

## Methods for root finding which use a bracket

## Bisection for FLoat64 values.
##
## From Jason Merrill https://gist.github.com/jwmerrill/9012954
## cf. http://squishythinking.com/2014/02/22/bisecting-floats/
# Alternative "mean" definition that operates on the binary representation
# of a float. Using this definition, bisection will never take more than
# 64 steps.
const FloatNN = Union{Float64, Float32, Float16}
_float_int_pairs = Dict(Float64 => UInt64, Float32 => UInt32, Float16 => UInt16)

function _middle(x::T, y::T) where {T <: FloatNN}
    # Use the usual float rules for combining non-finite numbers
    if !isfinite(x) || !isfinite(y)
        return x + y
    end
    # Always return 0.0 when inputs have opposite sign
    if sign(x) != sign(y) && !iszero(x) && ! iszero(y)
        return zero(T)
    end

    negate = sign(x) < 0 || sign(y) < 0

    # do division over unsigned integers with bit shift
    xint = reinterpret(_float_int_pairs[T], abs(x))
    yint = reinterpret(_float_int_pairs[T], abs(y))
    mid = (xint + yint) >> 1

    # reinterpret in original floating point
    unsigned = reinterpret(T, mid)

    negate ? -unsigned : unsigned
end

## fall back or non Floats
function _middle(x::Real, y::Real)
    x + (y-x)/2
end

"""

    Roots.bisection64(f, a, b)

(unexported)

* `f`: a callable object, like a function

* `a`, `b`: Real values specifying a *bracketing* interval (one with
`f(a) * f(b) < 0`). These will be converted to `Float64` values.

Runs the bisection method using midpoints determined by a trick
leveraging 64-bit floating point numbers. After ensuring the
intermediate bracketing interval does not straddle 0, the "midpoint"
is half way between the two values onces converted to unsigned 64-bit
integers. This means no more than 64 steps will be taken, fewer if `a`
and `b` already share some bits.

The process is guaranteed to return a value `c` with `f(c)` one of
`0`, `Inf`, or `NaN`; *or* one of `f(prevfloat(c))*f(c) < 0` or
`f(c)*f(nextfloat(c)) > 0` holding.

This function is a bit faster than the slightly more general
`find_zero(f, [a,b], Bisection())` call.

Due to Jason Merrill.

"""
function bisection64(@nospecialize(f), a0::FloatNN, b0::FloatNN)

    a,b = promote(a0, b0)

    if a > b
        b,a = a, b
    end

    m = _middle(a,b)

    fa, fb = sign(f(a)), sign(f(b))

    fa * fb > 0 && throw(ArgumentError(bracketing_error))
    (iszero(fa) || isnan(fa) || isinf(fa)) && return a
    (iszero(fb) || isnan(fb) || isinf(fb)) && return b

    while a < m < b

        fm = sign(f(m))

        if iszero(fm) || isnan(fm) || isinf(fm)
            return m
        elseif fa * fm < 0
            b,fb=m,fm
        else
            a,fa=m,fm
        end
        m = _middle(a,b)
    end
    return m
end

####
## find_zero interface.
"""

    Bisection()

Use the bisection method over `Float64` values. The bisection method
starts with a bracketing interval `[a,b]` and splits it into two
intervals `[a,c]` and `[c,b]`, If `c` is not a zero, then one of these
two will be a bracketing interval and the process continues. The
computation of `c` is done by `_middle`, which reinterprets floating
point values as unsigned 64-bit integers and splits there. This method
avoids floating point issues and guarantees a "best" solution (one
where a zero is found or the bracketing interval is of the type `[a,
nextfloat(a)]`).

"""
mutable struct Bisection <: AbstractBisection end

"""
    Roots.A42()

Bracketing method which finds the root of a continuous function within
a provided interval [a, b], without requiring derivatives. It is based
on algorithm 4.2 described in: 1. G. E. Alefeld, F. A. Potra, and
Y. Shi, "Algorithm 748: enclosing zeros of continuous functions," ACM
Trans. Math. Softw. 21, 327–344 (1995).
"""
mutable struct A42 <: AbstractBisection end

function find_zero(f, x0::Vector{T}, method::AbstractBisection; kwargs...) where {T <: Real}

    find_zero(f, (x0[1], x0[2]), method; kwargs...)
end

## For speed, bypass find_zero setup for bisection+Float64
function find_zero(f, x0::Tuple{T,S}, method::Bisection; verbose=false, kwargs...) where {T <: FloatNN, S <: FloatNN}
    x = adjust_bracket(x0)
    if verbose
        prob, options = derivative_free_setup(method, DerivativeFree(f), x; verbose=verbose, kwargs...)
        find_zero(prob, method, options)
    else
        bisection64(f, x[1], x[2])::eltype(x)
    end
end

## This is a bit clunky, as we use `a42` for bisection when we don't have `Float64` values.
## As we don't have the `A42` algorithm implemented through `find_zero`, we adjust here.
function find_zero(f, x0::Tuple{T,S}, method::M;  kwargs...) where {M <: Union{Bisection, A42}, T <: Real, S <: Real}

    x = adjust_bracket(x0)
    ## round about way to get options and state
    prob, options = derivative_free_setup(method, DerivativeFree(f), x[1];  kwargs...)
    o = init_state(method, prob.fs, prob.x0, prob.bracket)

    o.xn1 = a42(f, x[1], x[2]; xtol=options.xreltol, maxeval=options.maxevals, verbose=options.verbose)
    o.message = "Used Alefeld-Potra-Shi method, `Roots.a42`, to find the zero. Iterations and function evaluations are not counted properly."
    o.stopped = o.x_converged = true

    options.verbose && show_trace(prob.fs, o, [o.xn1], [o.fxn1], method)
    o.xn1
end

# this is hacky.
# we reach here from calling Bisection midway through the Order0() routine with "big" values.
# call a42 in this case
const BigSomething = Union{BigFloat, BigInt}
function find_zero(method::Bisection, fs, o::UnivariateZeroState{T, S}, options::UnivariateZeroOptions) where {T<:BigSomething, S}
    xn0, xn1 = o.xn0 < o.xn1 ? (o.xn0, o.xn1) : (o.xn1, o.xn0)
    o.xn1 = a42(fs, o.xn0, o.xn1; xtol=options.xreltol, maxeval=options.maxevals, verbose=options.verbose)
    o.message = "Used Alefeld-Potra-Shi method, `Roots.a42`, to find the zero. Iterations and function evaluations are not counted properly."
    o.stopped = o.x_converged = true

    nothing
end

function init_state(method::AbstractBisection, fs, x::Union{Tuple{T,T}, Vector{T}}, bracket) where {T <: Real}
    x0, x2 = adjust_bracket(x)
    y0 = fs(x0)
    y2 = fs(x2)

    sign(y0) * sign(y2) > 0 && throw(ArgumentError(bracketing_error))

    state = UnivariateZeroStateBase(x0, x2,
                                    y0, y2,
                                    ismissing(bracket) ? bracket : float.(bracket),
                                    0, 2,
                                    false, false, false, false,
                                    "")
    state
end

## This uses _middle bisection
## Find zero using modified bisection method for Float64 arguments.
## This is guaranteed to take no more than 64 steps. The `a42` alternative usually has
## fewer iterations, but this seems to find the value with fewer function evaluations.
##
## Terminates with `x1` when the bracket length of `[x0,x2]` is `<= max(xtol, xtolrel*abs(x1))` where `x1` is the midpoint .
## The tolerances can be set to 0, in which case, the termination occurs when `nextfloat(x0) = x2`.
## The bracket `[a,b]` must be bounded.

function update_state(method::Bisection, fs, o::Roots.UnivariateZeroState{T,S}, options::UnivariateZeroOptions) where {T<:FloatNN,S}
    x0, x2 = o.xn0, o.xn1
    y0, y2 = o.fxn0, o.fxn1

    x1 = _middle(x0, x2)

    y1 = fs(x1)
    incfn(o)

    if sign(y0) * sign(y1) > 0
        x0, x2 = x1, x2
        y0, y2 = y1, y2
    else
        x0, x2 = x0, x1
        y0, y2 = y0, y1
    end

    o.xn0, o.xn1 = x0, x2
    o.fxn0, o.fxn1 = y0, y2
    incsteps(o)
    nothing
end

## convergence is much differen there, as we bound between x0, nextfloat(x0) is not measured by eps(), but eps(x0)
function assess_convergence(method::Bisection, fs, state, options)
    x0, x2 = state.xn0, state.xn1
    if x0 > x2
        x0, x2 = x2, x0
    end

    x1 = _middle(x0, x2)
    if iszero(state.fxn1)
        state.message = ""
        state.stopped = state.x_converged = true
        return true
    end
    x0 < x1 && x1 < x2 && return false

    state.message = ""
    state.stopped = state.x_converged = true
    true
end

##################################################

"""

    Roots.a42(f, a, b; kwargs...)

(not exported)

Finds the root of a continuous function within a provided
interval [a, b], without requiring derivatives. It is based on algorithm 4.2
described in: 1. G. E. Alefeld, F. A. Potra, and Y. Shi, "Algorithm 748:
enclosing zeros of continuous functions," ACM Trans. Math. Softw. 21,
327–344 (1995).

input:
    `f`: function to find the root of
    `a`, `b`: the initial bracket, with: a < b, f(a)*f(b) < 0
    `xtol`: acceptable error (it's safe to set zero for machine precision)
    `maxeval`:  maximum number of iterations

output:
    an estimate of the zero of f

By John Travers

"""
function a42(f, a, b;
      xtol=zero(float(a)),
      maxeval::Int=15,
      verbose::Bool=false)

    if a > b
        a,b = b,a
    end
    fa, fb = f(a), f(b)

    if a >= b || sign(fa)*sign(fb) >= 0
        error("on input a < b and f(a)f(b) < 0 must both hold")
    end
    if xtol < 0.0
        error("tolerance must be >= 0.0")
    end

    c, fc = secant(f, a, fa, b, fb)
    a42a(f, a, fa, b, fb, c, fc,
         xtol=xtol, maxeval=maxeval, verbose=verbose)
end

"""

Split Alefeld, F. A. Potra, and Y. Shi algorithm 4.2 into a function
where `c` is passed in.

Solve f(x) = 0 over bracketing interval [a,b] starting at c, with a < c < b

"""
a42a(f, a, b, c=(a+b)/2; args...) = a42a(f, a, f(a), b, f(b), c, f(c); args...)

function a42a(f, a, fa, b, fb, c, fc;
       xtol=zero(float(a)),
       maxeval::Int=15,
       verbose::Bool=false)

    try
        # re-bracket and check termination
        a, fa, b, fb, d, fd = bracket(f, a, fa, b, fb, c, fc, xtol)
        ee, fee = d, fd
        for n = 2:maxeval
            # use either a cubic (if possible) or quadratic interpolation
            if n > 2 && distinct(a, fa, b, fb, d, fd, ee, fee)
                c, fc = ipzero(f, a, fa, b, fb, d, fd, ee, fee)
            else
                c, fc = newton_quadratic(f, a, fa, b, fb, d, fd, 2)
            end
            # re-bracket and check termination
            ab, fab, bb, fbb, db, fdb = bracket(f, a, fa, b, fb, c, fc, xtol)
            eb, feb = d, fd
            # use another cubic (if possible) or quadratic interpolation
            if distinct(ab, fab, bb, fbb, db, fdb, eb, feb)
                cb, fcb = ipzero(f, ab, fab, bb, fbb, db, fdb, eb, feb)
            else
                cb, fcb = newton_quadratic(f, ab, fab, bb, fbb, db, fdb, 3)
            end
            # re-bracket and check termination
            ab, fab, bb, fbb, db, fdb = bracket(f, ab, fab, bb, fbb, cb, fcb, xtol)
            # double length secant step; if we fail, use bisection
            if abs(fab) < abs(fbb)
                u, fu = ab, fab
            else
                u, fu = bb, fbb
            end
            # u = abs(fab) < abs(fbb) ? ab : bb
            cb = u - 2*fu/(fbb - fab)*(bb - ab)
            fcb = f(cb)
            if abs(cb - u) > (bb - ab)/2
                ch, fch = ab+(bb-ab)/2, f(ab+(bb-ab)/2)
            else
                ch, fch = cb, fcb
            end
            # ch = abs(cb - u) > (bb - ab)/2 ? ab + (bb - ab)/2 : cb
            # re-bracket and check termination
            ah, fah, bh, fbh, dh, fdh = bracket(f, ab, fab, bb, fbb, ch, fch, xtol)
            # if not converging fast enough bracket again on a bisection
            if bh - ah < 0.5*(b - a)
                a, fa = ah, fah
                b, fb = bh, fbh
                d, fd = dh, fdh
                ee, fee = db, fdb
            else
                ee, fee = dh, fdh
                a, fa, b, fb, d, fd = bracket(f, ah, fah, bh, fbh, ah + (bh - ah)/2, f(ah+(bh-ah)/2), xtol)
            end

            verbose && println(@sprintf("a=%18.15f, n=%s", float(a), n))

            if nextfloat(ch) * prevfloat(ch) <= 0
                throw(StateConverged(ch))
            end
            if nextfloat(a) >= b
                throw(StateConverged(a))
            end
        end
        throw(ConvergenceFailed("More than $maxeval iterations before convergence"))
    catch ex
        if isa(ex, StateConverged)
            return ex.x0
        else
            rethrow(ex)
        end
    end
    throw(ConvergenceFailed("More than $maxeval iterations before convergence"))
end

# calculate a scaled tolerance
# based on algorithm on page 340 of [1]
function tole(a::S, b::R, fa, fb, tol) where {S,R}
    T = promote_type(S,R)
    u = abs(fa) < abs(fb) ? abs(a) : abs(b)
    2u*eps(T) + tol
end

# bracket the root
# inputs:
#     - f: the function
#     - a, b: the current bracket with a < b, f(a)f(b) < 0
#     - c within (a,b): current best guess of the root
#     - tol: desired accuracy
#
# if root is not yet found, return
#     ab, bb, d
# with:
#     - [ab, bb] a new interval within [a, b] with f(ab)f(bb) < 0
#     - d a point not inside [ab, bb]; if d < ab, then f(ab)f(d) > 0,
#       and f(d)f(bb) > 0 otherwise
#
# if the root is found, throws a StateConverged instance with x0 set to the
# root.
#
# based on algorithm on page 341 of [1]
function bracket(f, a, fa, b, fb, c, fc, tol)

    if !(a <= c <= b)
        error("c must be in (a,b)")
    end
    delta = 0.7*tole(a, b, fa, fb, tol)
    if b - a <= 4delta
        c = (a + b)/2
        fc = f(c)
    elseif c <= a + 2delta
        c = a + 2delta
        fc = f(c)
    elseif c >= b - 2delta
        c = b - 2delta
        fc = f(c)
    end
    if fc == 0
        throw(Roots.StateConverged(c))
    elseif sign(fa)*sign(fc) < 0
        aa, faa = a, fa
        bb, fbb = c, fc
        db, fdb = b, fb
    else
        aa, faa = c, fc
        bb, fbb = b, fb
        db, fdb = a, fa
    end
    if bb - aa < 2*tole(aa, bb, faa, fbb, tol)
        x0 = abs(faa) < abs(fbb) ? aa : bb
        throw(Roots.StateConverged(x0))
    end
    aa, faa, bb, fbb, db, fdb
end

# take a secant step, if the resulting guess is very close to a or b, then
# use bisection instead
function secant(f, a::T, fa, b, fb) where {T}

    c = a - fa/(fb - fa)*(b - a)
    tol = eps(T)*5
    if isnan(c) || c <= a + abs(a)*tol || c >= b - abs(b)*tol
        return a + (b - a)/2, f(a+(b-a)/2)
    end
    return c, f(c)
end

# approximate zero of f using quadratic interpolation
# if the new guess is outside [a, b] we use a secant step instead
# based on algorithm on page 330 of [1]
function newton_quadratic(f, a, fa, b, fb, d, fd, k::Int)
    B = (fb - fa)/(b - a)
    A = ((fd - fb)/(d - b) - B)/(d - a)
    if A == 0
        return secant(f, a, fa, b, fb)
    end
    r = A*fa > 0 ? a : b
    for i = 1:k
        r -= (fa + (B + A*(r - b))*(r - a))/(B + A*(2*r - a - b))
    end
    if isnan(r) || (r <= a || r >= b)
        r, fr = secant(f, a, fa, b, fb)
    else
        fr = f(r)
    end
    return r, fr
end

# approximate zero of f using inverse cubic interpolation
# if the new guess is outside [a, b] we use a quadratic step instead
# based on algorithm on page 333 of [1]
function ipzero(f, a, fa, b, fb, c, fc, d, fd)

    Q11 = (c - d)*fc/(fd - fc)
    Q21 = (b - c)*fb/(fc - fb)
    Q31 = (a - b)*fa/(fb - fa)
    D21 = (b - c)*fc/(fc - fb)
    D31 = (a - b)*fb/(fb - fa)
    Q22 = (D21 - Q11)*fb/(fd - fb)
    Q32 = (D31 - Q21)*fa/(fc - fa)
    D32 = (D31 - Q21)*fc/(fc - fa)
    Q33 = (D32 - Q22)*fa/(fd - fa)
    c = a + (Q31 + Q32 + Q33)
    if (c <= a) || (c >= b)
        return newton_quadratic(f, a, fa, b, fb, d, fd, 3)
    end
    return c, f(c)
end

# floating point comparison function
function almost_equal(x::T, y::T) where {T}
    min_diff = realmin(T)*32
    abs(x - y) < min_diff
end

# check that all interpolation values are distinct
function distinct(a, f1, b, f2, d, f3, e, f4)
    !(almost_equal(f1, f2) || almost_equal(f1, f3) || almost_equal(f1, f4) ||
      almost_equal(f2, f3) || almost_equal(f2, f4) || almost_equal(f3, f4))
end

## ----------------------------

"""

    FalsePosition()

Use the [false
position](https://en.wikipedia.org/wiki/False_position_method) method
to find a zero for the function `f` within the bracketing interval
`[a,b]`.

The false position method is a modified bisection method, where the
midpoint between `[a_k, b_k]` is chosen to be the intersection point
of the secant line with the x axis, and not the average between the
two values.

To speed up convergence for concave functions, this algorithm
implements the 12 reduction factors of Galdino (*A family of regula
falsi root-finding methods*). These are specified by number, as in
`FalsePosition(2)` or by one of three names `FalsePosition(:pegasus)`,
`FalsePosition(:illinois)`, or `FalsePosition(:anderson_bjork)` (the
default). The default choice has generally better performance than the
others, though there are exceptions.

For some problems, the number of function calls can be greater than
for the `bisection64` method, but generally this algorithm will make
fewer function calls.

Examples
```
find_zero(x -> x^5 - x - 1, [-2, 2], FalsePosition())
```
"""
# mutable struct FalsePosition <: AbstractBisection
#     reduction_factor::Union{Int, Symbol}
#     FalsePosition(x=:anderson_bjork) = new(x)
# end
struct FalsePosition{R} <: AbstractBisection
    FalsePosition(x=:anderson_bjork) = new{:anderson_bjork}()
end

function update_state(method::FalsePosition, fs, o, options)

    fs
    a, b =  o.xn0, o.xn1

    fa, fb = o.fxn0, o.fxn1

    lambda = fb / (fb - fa)
    tau = 1e-10                   # some engineering to avoid short moves
    if !(tau < norm(lambda) < 1-tau)
        lambda = 1/2
    end
    x = b - lambda * (b-a)
    fx = fs(x)
    incfn(o)
    incsteps(o)

    if iszero(fx)
        o.xn1 = x
        o.fxn1 = fx
        return
    end

    if sign(fx)*sign(fb) < 0
        a, fa = b, fb
    else
        fa = reduction(method, fa, fb, fx)
    end
    b, fb = x, fx

    o.xn0, o.xn1 = a, b
    o.fxn0, o.fxn1 = fa, fb

    nothing
end

# the 12 reduction factors offered by Galadino
const galdino = Dict{Union{Int,Symbol},Function}(:1 => (fa, fb, fx) -> fa*fb/(fb+fx),
                                           :2 => (fa, fb, fx) -> (fa - fb)/2,
                                           :3 => (fa, fb, fx) -> (fa - fx)/(2 + fx/fb),
                                           :4 => (fa, fb, fx) -> (fa - fx)/(1 + fx/fb)^2,
                                           :5 => (fa, fb, fx) -> (fa -fx)/(1.5 + fx/fb)^2,
                                           :6 => (fa, fb, fx) -> (fa - fx)/(2 + fx/fb)^2,
                                           :7 => (fa, fb, fx) -> (fa + fx)/(2 + fx/fb)^2,
                                           :8 => (fa, fb, fx) -> fa/2,
                                           :9 => (fa, fb, fx) -> fa/(1 + fx/fb)^2,
                                           :10 => (fa, fb, fx) -> (fa-fx)/4,
                                           :11 => (fa, fb, fx) -> fx*fa/(fb+fx),
                                           :12 => (fa, fb, fx) -> (fa * (1-fx/fb > 0 ? 1-fx/fb : 1/2))
)
# give common names
for (nm, i) in [(:pegasus, 1), (:illinois, 8), (:anderson_bjork, 12)]
    galdino[nm] = galdino[i]
end
@generated function reduction(method::FalsePosition{R}, fa, fb, fx) where R
    f = galdino[R]
    :( $f(fa, fb, fx) )
end

function find_zero(f, x0::Tuple{T,S}, method::FalsePosition; kwargs...) where {T<:Real,S<:Real}
    x = adjust_bracket(x0)
    prob, options = derivative_free_setup(method, DerivativeFree(f), x; kwargs...)
    find_zero(prob, method, options)
end
## helper function
function adjust_bracket(x0)
    u, v = float.(x0)
    if u > v
        u, v = v, u
    end

    if isinf(u)
        u = nextfloat(u)
    end
    if isinf(v)
        v = prevfloat(v)
    end
    u, v
end

## --------------------------------------

"""

Searches for zeros  of `f` in an interval [a, b].

Basic algorithm used:

* split interval [a,b] into `no_pts` subintervals.
* For each bracketing interval finds a bracketed zero.
* For other subintervals does a quick search with a derivative-free method.

If there are many zeros relative to the number of points, the process
is repeated with more points, in hopes of finding more zeros for
oscillating functions.

Called by `fzeros` or `Roots.find_zeros`.

"""
function find_zeros(f, a::Real, b::Real, args...;
                    no_pts::Int=100,
                    abstol::Real=10*eps(), reltol::Real=10*eps(), ## should be abstol, reltol as used.
                    kwargs...)

    a, b = a < b ? (a,b) : (b,a)
    rts = eltype(promote(float(a),b))[]
    xs = vcat(a, a .+ (b-a) .* sort(rand(no_pts)), b)

    ## Look in [ai, bi)
    for i in 1:(no_pts+1)
        ai,bi=xs[i:i+1]
        if isapprox(f(ai), 0.0, rtol=reltol, atol=abstol)
            push!(rts, ai)
        elseif sign(f(ai)) * sign(f(bi)) < 0
            push!(rts, find_zero(f, [ai, bi], Bisection()))
        else
            try
                x = find_zero(f, ai + (0.5)* (bi-ai), Order8(); maxevals=10, abstol=abstol, reltol=reltol)
                if ai < x < bi
                    push!(rts, x)
                end
            catch e
            end
        end
    end
    ## finally, b?
    isapprox(f(b), 0.0, rtol=reltol, atol=abstol) && push!(rts, b)

    ## redo if it appears function oscillates alot in this interval...
    if length(rts) > (1/4) * no_pts
        return(find_zeros(f, a, b, args...; no_pts = 10*no_pts, abstol=abstol, reltol=reltol, kwargs...))
    else
        return(sort(rts))
    end
end


text:
using PolyChaosODE, BenchmarkTools, DifferentialEquations, Distributions, Parameters

### Performance vs dimensionality example ###

struct LinearODEComp{T1,T2}
    Λ::T1
    error_vec::T2
end

function ode_func!(du, u, Λ, error_vec, x, t)
    du .= (Λ .+ error_vec .* x) .* u
end

function ode_func_closure!(du, u, comp::LinearODEComp, x, t)
    @unpack Λ, error_vec = comp
    ode_func!(du, u, Λ, error_vec, x, t)
end

max_dim = 5
μ = 0.0
σ = sqrt(2.0) / 2

# run benchmark
for prob_dim = 1:max_dim
    vars = Tuple(Normal(μ, σ) for i = 1:prob_dim)
    Λ = rand(Uniform(-0.200, -0.100), prob_dim)
    error_vec = error_vec = rand(Uniform(0.050, 0.100), prob_dim)
    # ODE params
    u0 = 10.0 .* randn(prob_dim)
    tspan = (0.0, 5.0)
    compartment = LinearODEComp(Λ, error_vec)
    # interval for plotting
    step_t = 0.01
    interval_t = tspan[1]:step_t:tspan[2]
    stoch_galerkin_ode = StochGalerkinODE(ode_func_closure!, prob_dim, vars)
    println("~~~~~~~~~~~~~")
    @show prob_dim
    # get the solution for plotting
    sol = stoch_galerkin_ode(u0, tspan, compartment)
    # after that, time it
    @time stoch_galerkin_ode(u0, tspan, compartment)
    @btime (
        x -> sqrt.(x)
    ).(
        compute_expectation_and_diag_variance(
            stoch_galerkin_ode,
            interval_t,
            [$sol(t) for t in interval_t],
        )[2],
    )
end

max_stars_repo_name: d-netto/PolyChaosODE

max_issues_repo_name: d-netto/PolyChaosODE

max_forks_repo_name: d-netto/PolyChaosODE


text:
# Fixed point classification in 2D
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

alleigs = []
fig, axs = subplots(2,3; figsize = (figx, 2figy))
ax = gca()
xgrid = -5:0.05:5
ygrid = xgrid
ux = zeros(length(xgrid), length(ygrid))
uy = copy(ux)

Ms = [
     [1.0  0.8;
      0.4 1.0],
     [-1.0  -0.8;
      1.8  -1.0],
     [-1.0  0;
      1  1],
     [0 -1;
      1  0],
]

titles = ["repulsive node", "attractive spiral", "hyperbolic/saddle", "center"]

using LinearAlgebra
function stream_eigs!(ax, M, c = "C0")
    for (i, x) in enumerate(xgrid)
        for (j, y) in enumerate(ygrid)
            ux[i, j], uy[i, j] = M * [x, y]
        end
    end

    ax.streamplot(Vector(xgrid), Vector(ygrid), ux', uy';
        linewidth = 1.5, density = 0.5, color = c, arrowsize = 2,
    )
    ev = eigen(M)
    push!(alleigs, ev.values)
    if eltype(ev.values) <: Float64
        e1 = ev.vectors[:, 1]
        e2 = ev.vectors[:, 2]
        for e in (e1, e2)
            ax.plot(e[1] .* 2xgrid, e[2] .* 2ygrid; color = "C2", ls = "dashed")
        end
    end
    ax.set_xlim(xgrid[1], xgrid[end])
    ax.set_ylim(ygrid[1], ygrid[end])
    ax.set_xticks([])
    ax.set_yticks([])
    ax.plot(0, 0; marker = "o", mec = "C2", mew = 1,
        markersize = 12, mfc = :white, zorder = 99
    )
end

function tovec(e)
    if eltype(e) <: Real
        e1 = e
        e2 = zeros(2)
    else
        e1 = [e[1].re, e[2].re]
        e2 = [e[1].im, e[2].im]
    end
    return e1, e2
end

for (i, M) in enumerate(Ms)
    ax = axs[i]
    color = i < 3 ? "C$(i-1)" : "C$(i)"
    stream_eigs!(ax, M, color)
    ax.set_title(titles[i]; color = color)
    # Plot eigenvalues
    e1, e2 = tovec(alleigs[i])
    axs[5].scatter(e1, e2; color = color, s = 200, zorder = 99)
end

# Set axis of eigenval plot
axs[5].spines["left"].set_position("center")
axs[5].spines["bottom"].set_position("center")
axs[5].grid(false)
axs[5].spines["right"].set_color("none")
axs[5].spines["top"].set_color("none")
axs[5].set_title("eigenvalues"; color = "k")
# axs[5].set_yticklabels([])
# axs[5].set_xticklabels([])
axs[5].set_xlim(-2, 2)
axs[5].set_ylim(-1.8, 1.8)

# Plot limit cycle
function vanderpoll(u, p, t)
    x, y = u;
    xdot = p*(x - x^3/3 - y)
    ydot = x/p
    return SVector(xdot, ydot)
end

ds = ContinuousDynamicalSystem(vanderpoll, [0.1, 0.2], 0.5)
tr = trajectory(ds, 7.0; Ttr = 100.0)
axs[6].clear()
axs[6].plot(columns(tr)...; color = "C2", linestyle = "--")
for (i, x) in enumerate(xgrid)
    for (j, y) in enumerate(ygrid)
        ux[i, j], uy[i, j] = ds.f(SVector(x, y), ds.p, 0)
    end
end
axs[6].streamplot(Vector(xgrid), Vector(ygrid), ux', uy';
    linewidth = 1.5, density = 0.5, color = "C5", arrowsize = 2
)
ax = axs[6]
ax.set_xlim(xgrid[1], xgrid[end])
ax.set_ylim(ygrid[1], ygrid[end])
ax.set_xticks([])
ax.set_yticks([])
ax.set_title("attractive limit cycle"; color = "C5")

fig.subplots_adjust(bottom = 0.02, left = 0.02, top = 0.92, right = 0.97, hspace = 0.2)
wsave(plotsdir("2", "2ddynamics"), fig)

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook


text:
using SubMatrixSelectionSVD
@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

@testset "Basic Tests" begin
	X = randn(10,7)

	U,Σ,V,ps,signalDimensions,selectedVariables = smssvd(X,4,10 .^ range(-2,stop=0,length=5),nbrIter=2)
	@test size(U)==(10,4)
	@test size(Σ)==(4,)
	@test size(V)==(7,4)
	@test size(ps)==(4,5)
	@test length(signalDimensions)<=4
	@test sum(signalDimensions)==4
	@test all(x->size(x)==(10,), selectedVariables)

	U,Σ,V,ps,signalDimensions,selectedVariables = smssvd(X,7,10 .^ range(-2,stop=0,length=5),nbrIter=2)
	@test size(U)==(10,7)
	@test size(Σ)==(7,)
	@test size(V)==(7,7)
	@test size(ps)==(7,5)
	@test length(signalDimensions)<=7
	@test sum(signalDimensions)==7
	@test all(x->size(x)==(10,), selectedVariables)

	X = randn(7,10)

	U,Σ,V,ps,signalDimensions,selectedVariables = smssvd(X,4,10 .^ range(-2,stop=0,length=5),nbrIter=2)
	@test size(U)==(7,4)
	@test size(Σ)==(4,)
	@test size(V)==(10,4)
	@test size(ps)==(4,5)
	@test length(signalDimensions)<=4
	@test sum(signalDimensions)==4
	@test all(x->size(x)==(7,), selectedVariables)

	U,Σ,V,ps,signalDimensions,selectedVariables = smssvd(X,7,10 .^ range(-2,stop=0,length=5),nbrIter=2)
	@test size(U)==(7,7)
	@test size(Σ)==(7,)
	@test size(V)==(10,7)
	@test size(ps)==(7,5)
	@test length(signalDimensions)<=7
	@test sum(signalDimensions)==7
	@test all(x->size(x)==(7,), selectedVariables)
end


text:
# Testing:
#
# - MC estimates of mean, variance
#
# TODO: add variance testing for asymmetric bounds

using Statistics, DiffModels
using Test

nanmean(x) = mean(x[.!isnan.(x)])
nanvar(x) = var(x[.!isnan.(x)])

function estimatemoments(d::AbstractDrift, b::AbstractBounds, n)
    t = fill(NaN, n)
    c = fill(NaN, n)
    s = sampler(d, b)
    for i = 1:n
        t[i], ci = rand(s)
        c[i] = ci ? 1.0 : 0.0
    end
    return [nanmean(t), nanvar(t), mean(c),
            nanmean(t[c .== 1.0]), nanvar(t[c .== 1.0]),
            nanmean(t[c .== 0.0]), nanvar(t[c .== 0.0])]
end

# symmetric bounds

dmsymmeant(mu::Real, bound::Real) =
    isapprox(mu, zero(mu)) ? abs2(bound) : bound / mu * tanh(bound * mu)
function dmsymvart(mu::Real, bound::Real)
    if isapprox(mu, zero(mu))
        return 2abs2(abs2(bound)) / 3
    end
    a = mu * bound
    return bound * (tanh(a) - a * abs2(sech(a))) / mu^3
end
dmsympubound(mu::Real, bound::Real) = 1 / (1 + exp(-2mu * bound))

n = 10^5
dt = 0.001

# normal/exponential sampler for small mu
bound = 0.9
mu = 0.5
m = estimatemoments(ConstDrift(mu, dt), ConstSymBounds(bound, dt), n)
@test m[1] ≈ dmsymmeant(mu, bound) atol=0.1
@test m[2] ≈ dmsymvart(mu, bound) atol=0.1
@test m[3] ≈ dmsympubound(mu, bound) atol=0.05

# inverse-normal sampler for large mu
bound = 0.9
mu = 2
m = estimatemoments(ConstDrift(mu, dt), ConstSymBounds(bound, dt), n)
@test m[1] ≈ dmsymmeant(mu, bound) atol=0.1
@test m[2] ≈ dmsymvart(mu, bound) atol=0.1
@test m[3] ≈ dmsympubound(mu, bound) atol=0.05

# generic sampler
bound = 0.9
mu = 1.1
m = estimatemoments(VarDrift([mu], dt, 15.0), ConstSymBounds(bound, dt), n)
@test m[1] ≈ dmsymmeant(mu, bound) atol=0.1
@test m[2] ≈ dmsymvart(mu, bound) atol=0.1
@test m[3] ≈ dmsympubound(mu, bound) atol=0.05

# asymmetric bounds

dmasymmeantup(mu::Real, blo::Real, bup::Real) = isapprox(mu, zero(mu)) ?
    (abs2(bup) - 2bup * blo) / 3 :
    (bup - blo) / mu * coth((bup - blo) * mu) + blo / mu * coth(- blo * mu)
dmasymmeantlo(mu::Real, blo::Real, bup::Real) = isapprox(mu, zero(mu)) ?
    (abs2(blo) - 2bup * blo) / 3 :
    (bup - blo) / mu * coth((bup - blo) * mu) - bup / mu * coth(bup * mu)
dmasympubound(mu::Real, blo::Real, bup::Real) = isapprox(mu, zero(mu)) ?
    blo / (blo + bup) : (exp(-2mu * blo) - 1) / (exp(-2mu * blo) - exp(-2mu * bup))

# fast sampler for constant drift/bounds
blo = -1.5
bup = 1.0
mu = 1.1
m = estimatemoments(ConstDrift(mu, dt), ConstAsymBounds(bup, blo, dt), n)
@test m[3] ≈ dmasympubound(mu, blo, bup) atol=0.05
@test m[4] ≈ dmasymmeantup(mu, blo, bup) atol=0.1
@test m[6] ≈ dmasymmeantlo(mu, blo, bup) atol=0.1

# generic sampler
blo = -1.5
bup = 1.0
mu = 1.1
m = estimatemoments(VarDrift([mu], dt, 15.0), ConstAsymBounds(bup, blo, dt), n)
@test m[3] ≈ dmasympubound(mu, blo, bup) atol=0.05
@test m[4] ≈ dmasymmeantup(mu, blo, bup) atol=0.1
@test m[6] ≈ dmasymmeantlo(mu, blo, bup) atol=0.1


text:
push!(LOAD_PATH, abspath("src"))
using Yao, Yao.Blocks
using LinearAlgebra, Statistics
using QMPS
using Test, Random

# make it cluster state
@testset "convert wave function check" begin
    chem = model(:su2; nbit=9, nlayer=2, B=10, V=5, pairs=pair_ring(5))
    c = random_circuit(1, 4, 2, 5, pair_ring(5))
    circuit = chem2circuit(chem)
    @test zero_state(nqubits(circuit)) |> circuit |> statevec |> length == 2^10
end

@testset "measure check" begin
    Random.seed!(5)
    chem = model(:su2; nbit=9, nlayer=2, B=10000, V=5, pairs=pair_ring(5))
    circuit = chem2circuit(chem)

    for (i, j) in [(3,5), (5,3), (3,7), (7,3), (6,8), (8,6)]
        @show (i,j)
        mean35 = expect(heisenberg_ij(nqubits(circuit), i, j), zero_state(nqubits(circuit)) |> circuit) |> real
        eng = sum(g->measure_corr(chem, i=>g, j=>g), [X, Y, Z])
        @test isapprox(mean35, eng, rtol=0.3)
    end
end

@testset "j1j2" begin
    j1j2 = J1J2(4; periodic=false, J2=0.5)
    @test get_bonds(j1j2) == [(1, 2, 1.0),(2, 3, 1.0), (3, 4, 1.0), (1,3, 0.5), (2,4, 0.5)]
    j1j2 = J1J2(4; periodic=true, J2=0.5)
    @test get_bonds(j1j2) == [(1, 2, 1.0), (2, 3, 1.0), (3, 4, 1.0), (4, 1, 1.0), (1,3, 0.5), (2,4, 0.5), (3,1, 0.5), (4,2, 0.5)]

    j1j2 = J1J2(3, 3; periodic=false, J2=0.5)
    prs = [1=>2, 1=>4, 2=>3, 2=>5, 2=>4, 3=>6, 3=>5, 4=>5, 4=>7, 5=>6, 5=>8, 5=>1, 5=>7, 6=>9, 6=>2, 6=>8, 7=>8, 8=>9, 8=>4, 9=>5]
    vs =  [1.0,   1,    1,    1,   0.5,  1.0,   0.5,  1.0,  1.0,  1.0, 1.0,  0.5,   0.5,  1.0,  0.5, 0.5,   1.0, 1.0,  0.5,  0.5]
    tps = [(i,j,v) for ((i,j), v) in zip(prs, vs)]
    @test sort(get_bonds(j1j2)) == sort(tps)
    j1j2 = J1J2(3, 3; periodic=true, J2=0.5)
    @test sort(get_bonds(j1j2)) == sort([(1,2,1.0), (1,4,1.0), (1,9,0.5), (1,6,0.5), (2,3,1.0), (2,5,1.0), (2,7,0.5), (2,4,0.5), (3,1,1.0), (3,6,1.0), (3,8,0.5), (3,5,0.5),
        (4,5,1.0), (4,7,1.0), (4,3,0.5), (4,9,0.5), (5,6,1.0), (5,8,1.0), (5,1,0.5), (5,7,0.5), (6,4,1.0), (6,9,1.0), (6,2,0.5), (6,8,0.5), (7,8,1.0), (7,1,1.0), (7,6,0.5), (7,3,0.5),
        (8,9,1.0), (8,2,1.0), (8,4,0.5), (8,1,0.5), (9,7,1.0), (9,3,1.0), (9,5,0.5), (9,2,0.5)])
end

@testset "energy-goodqn" begin
    Random.seed!(5)
    for hei in [Heisenberg(10; periodic=false), Heisenberg(3, 3; periodic=false), J1J2(3,3; J2=0.5, periodic=false)]
        nbit = nspin(hei)
        for xmodel in [:u1, :su2]
            pairs = pair_ring(xmodel==:su2 ? 4 : 5)
            chem = model(:general; nbit=nbit, B=10000, V=4, pairs=pairs)
            println("Number of parameters is ", chem.circuit |> nparameters)
            circuit = chem2circuit(chem)
            eng = energy(chem, hei)
            hami = hamiltonian(hei)
            eng_exact = expect(hami, product_state(nbit, chem.input_state |> Yao.Intrinsics.packbits) |> circuit) |> real
            @test isapprox(eng, eng_exact, rtol=0.3)
        end
    end
end

max_stars_repo_name: Roger-luo/QuantumMPS

max_issues_repo_name: Roger-luo/QuantumMPS

max_forks_repo_name: Roger-luo/QuantumMPS