text:
using Test
using ClimateMachine.Thermodynamics
using ClimateMachine.TemperatureProfiles
using UnPack
using NCDatasets
using Random
using RootSolvers
const TD = Thermodynamics
using LinearAlgebra

using CLIMAParameters
using CLIMAParameters.Planet

struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

# Tolerances for tested quantities:
atol_temperature = 5e-1
atol_pressure = MSLP(param_set) * 2e-2
atol_energy = cv_d(param_set) * atol_temperature
rtol_temperature = 1e-1
rtol_density = rtol_temperature
rtol_pressure = 1e-1
rtol_energy = 1e-1

array_types = [Array{Float32}, Array{Float64}]

include("profiles.jl")
include("data_tests.jl")

compare_moisture(a::ThermodynamicState, b::ThermodynamicState) =
    compare_moisture(a, PhasePartition(b))

compare_moisture(ts::PhaseEquil, q_pt::PhasePartition) =
    getproperty(PhasePartition(ts), :tot) ≈ getproperty(q_pt, :tot)

compare_moisture(ts::PhaseNonEquil, q_pt::PhasePartition) = all((
    getproperty(PhasePartition(ts), :tot) ≈ getproperty(q_pt, :tot),
    getproperty(PhasePartition(ts), :liq) ≈ getproperty(q_pt, :liq),
    getproperty(PhasePartition(ts), :ice) ≈ getproperty(q_pt, :ice),
))

@testset "Thermodynamics - isentropic processes" begin
    for ArrayType in array_types
        FT = eltype(ArrayType)

        _R_d = FT(R_d(param_set))
        _molmass_ratio = FT(molmass_ratio(param_set))
        _cp_d = FT(cp_d(param_set))
        _cp_v = FT(cp_v(param_set))
        _cp_l = FT(cp_l(param_set))
        _cv_d = FT(cv_d(param_set))
        _cv_v = FT(cv_v(param_set))
        _cv_l = FT(cv_l(param_set))
        _cv_i = FT(cv_i(param_set))
        _T_0 = FT(T_0(param_set))
        _e_int_v0 = FT(e_int_v0(param_set))
        _e_int_i0 = FT(e_int_i0(param_set))
        _LH_v0 = FT(LH_v0(param_set))
        _LH_s0 = FT(LH_s0(param_set))
        _cp_i = FT(cp_i(param_set))
        _LH_f0 = FT(LH_f0(param_set))
        _press_triple = FT(press_triple(param_set))
        _R_v = FT(R_v(param_set))
        _T_triple = FT(T_triple(param_set))
        _T_freeze = FT(T_freeze(param_set))
        _T_min = FT(T_min(param_set))
        _MSLP = FT(MSLP(param_set))
        _T_max = FT(T_max(param_set))
        _kappa_d = FT(kappa_d(param_set))

        profiles = PhaseEquilProfiles(param_set, ArrayType)
        @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
        @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

        # Test state for thermodynamic consistency (with ideal gas law)
        T_idgl = TD.air_temperature_from_ideal_gas_law.(param_set, p, ρ, q_pt)
        @test all(T .≈ T_idgl)

        Φ = FT(1)
        Random.seed!(15)
        perturbation = FT(0.1) * rand(FT, length(T))

        # TODO: Use reasonable values for ambient temperature/pressure
        T∞, p∞ = T .* perturbation, p .* perturbation
        @test air_temperature.(param_set, p, θ_liq_ice, DryAdiabaticProcess()) ≈
              (p ./ _MSLP) .^ (_R_d / _cp_d) .* θ_liq_ice
        @test TD.air_pressure_given_θ.(
            param_set,
            θ_liq_ice,
            Φ,
            DryAdiabaticProcess(),
        ) ≈ _MSLP .* (1 .- Φ ./ (θ_liq_ice .* _cp_d)) .^ (_cp_d / _R_d)
        @test air_pressure.(param_set, T, T∞, p∞, DryAdiabaticProcess()) ≈
              p∞ .* (T ./ T∞) .^ (FT(1) / _kappa_d)
    end
end

@testset "Thermodynamics - correctness" begin
    FT = Float64
    _R_d = FT(R_d(param_set))
    _molmass_ratio = FT(molmass_ratio(param_set))
    _cp_d = FT(cp_d(param_set))
    _cp_v = FT(cp_v(param_set))
    _cp_l = FT(cp_l(param_set))
    _cv_d = FT(cv_d(param_set))
    _cv_v = FT(cv_v(param_set))
    _cv_l = FT(cv_l(param_set))
    _cv_i = FT(cv_i(param_set))
    _T_0 = FT(T_0(param_set))
    _e_int_v0 = FT(e_int_v0(param_set))
    _e_int_i0 = FT(e_int_i0(param_set))
    _LH_v0 = FT(LH_v0(param_set))
    _LH_s0 = FT(LH_s0(param_set))
    _cp_i = FT(cp_i(param_set))
    _LH_f0 = FT(LH_f0(param_set))
    _press_triple = FT(press_triple(param_set))
    _R_v = FT(R_v(param_set))
    _T_triple = FT(T_triple(param_set))
    _T_freeze = FT(T_freeze(param_set))
    _T_min = FT(T_min(param_set))
    _MSLP = FT(MSLP(param_set))
    _T_max = FT(T_max(param_set))
    _kappa_d = FT(kappa_d(param_set))
    _T_icenuc = FT(T_icenuc(param_set))

    # ideal gas law
    @test air_pressure(param_set, FT(1), FT(1), PhasePartition(FT(1))) === _R_v
    @test air_pressure(
        param_set,
        FT(1),
        FT(2),
        PhasePartition(FT(1), FT(0.5), FT(0)),
    ) === _R_v
    @test air_pressure(param_set, FT(1), FT(1)) === _R_d
    @test air_pressure(param_set, FT(1), FT(2)) === 2 * _R_d
    @test air_density(param_set, FT(1), FT(1)) === 1 / _R_d
    @test air_density(param_set, FT(1), FT(2)) === 2 / _R_d

    # gas constants and heat capacities
    @test gas_constant_air(param_set, PhasePartition(FT(0))) === _R_d
    @test gas_constant_air(param_set, PhasePartition(FT(1))) === _R_v
    @test gas_constant_air(param_set, PhasePartition(FT(0.5), FT(0.5))) ≈
          _R_d / 2
    @test gas_constant_air(param_set, FT) == _R_d

    @test cp_m(param_set, PhasePartition(FT(0))) === _cp_d
    @test cp_m(param_set, PhasePartition(FT(1))) === _cp_v
    @test cp_m(param_set, PhasePartition(FT(1), FT(1))) === _cp_l
    @test cp_m(param_set, PhasePartition(FT(1), FT(0), FT(1))) === _cp_i
    @test cp_m(param_set, FT) == _cp_d

    @test cv_m(param_set, PhasePartition(FT(0))) === _cp_d - _R_d
    @test cv_m(param_set, PhasePartition(FT(1))) === _cp_v - _R_v
    @test cv_m(param_set, PhasePartition(FT(1), FT(1))) === _cv_l
    @test cv_m(param_set, PhasePartition(FT(1), FT(0), FT(1))) === _cv_i
    @test cv_m(param_set, FT) == _cv_d

    # speed of sound
    @test soundspeed_air(param_set, _T_0 + 20, PhasePartition(FT(0))) ==
          sqrt(_cp_d / _cv_d * _R_d * (_T_0 + 20))
    @test soundspeed_air(param_set, _T_0 + 100, PhasePartition(FT(1))) ==
          sqrt(_cp_v / _cv_v * _R_v * (_T_0 + 100))

    # specific latent heats
    @test latent_heat_vapor(param_set, _T_0) ≈ _LH_v0
    @test latent_heat_fusion(param_set, _T_0) ≈ _LH_f0
    @test latent_heat_sublim(param_set, _T_0) ≈ _LH_s0

    # saturation vapor pressure and specific humidity
    p = FT(1.e5)
    q_tot = FT(0.23)
    ρ = FT(1.0)
    ρ_v_triple = _press_triple / _R_v / _T_triple
    @test saturation_vapor_pressure(param_set, _T_triple, Liquid()) ≈
          _press_triple
    @test saturation_vapor_pressure(param_set, _T_triple, Ice()) ≈ _press_triple

    phase_type = PhaseDry
    @test q_vap_saturation(
        param_set,
        _T_triple,
        ρ,
        phase_type,
        PhasePartition(FT(0)),
    ) == ρ_v_triple / ρ
    phase_type = PhaseNonEquil
    @test q_vap_saturation(
        param_set,
        _T_triple,
        ρ,
        phase_type,
        PhasePartition(q_tot, q_tot),
    ) == ρ_v_triple / ρ

    @test q_vap_saturation_generic(param_set, _T_triple, ρ, Liquid()) ==
          ρ_v_triple / ρ
    @test q_vap_saturation_generic(param_set, _T_triple, ρ, Ice()) ==
          ρ_v_triple / ρ
    @test q_vap_saturation_generic(param_set, _T_triple - 20, ρ, Liquid()) >=
          q_vap_saturation_generic(param_set, _T_triple - 20, ρ, Ice())

    # test the wrapper for q_vap_saturation over liquid water and ice
    ρ = FT(1)
    ρu = FT[1, 2, 3]
    ρe = FT(1100)
    e_pot = FT(93)
    e_int = internal_energy(ρ, ρe, ρu, e_pot)
    q_pt = PhasePartition(FT(0.02), FT(0.002), FT(0.002))
    ts = PhaseNonEquil(param_set, e_int, ρ, q_pt)
    @test q_vap_saturation_generic(
        param_set,
        air_temperature(ts),
        ρ,
        Liquid(),
    ) ≈ q_vap_saturation_liquid(ts)
    @test q_vap_saturation_generic(param_set, air_temperature(ts), ρ, Ice()) ≈
          q_vap_saturation_ice(ts)
    @test q_vap_saturation_ice(ts) <= q_vap_saturation_liquid(ts)

    phase_type = PhaseDry
    @test saturation_excess(
        param_set,
        _T_triple,
        ρ,
        phase_type,
        PhasePartition(q_tot),
    ) == q_tot - ρ_v_triple / ρ
    @test saturation_excess(
        param_set,
        _T_triple,
        ρ,
        phase_type,
        PhasePartition(q_tot / 1000),
    ) == 0.0

    @test supersaturation(
        param_set,
        PhasePartition(q_tot, 1e-3 * q_tot, 1e-3 * q_tot),
        ρ,
        _T_triple,
        Liquid(),
    ) ≈ 0.998 * q_tot / ρ_v_triple / ρ - 1

    @test supersaturation(
        param_set,
        PhasePartition(q_tot, 1e-3 * q_tot, 1e-3 * q_tot),
        ρ,
        _T_triple,
        Ice(),
    ) ≈ 0.998 * q_tot / ρ_v_triple / ρ - 1

    # energy functions and inverse (temperature)
    T = FT(300)
    e_kin = FT(11)
    e_pot = FT(13)
    @test air_temperature(param_set, _cv_d * (T - _T_0)) === FT(T)
    @test air_temperature(
        param_set,
        _cv_d * (T - _T_0),
        PhasePartition(FT(0)),
    ) === FT(T)

    @test air_temperature(
        param_set,
        cv_m(param_set, PhasePartition(FT(0))) * (T - _T_0),
        PhasePartition(FT(0)),
    ) === FT(T)
    @test air_temperature(
        param_set,
        cv_m(param_set, PhasePartition(FT(q_tot))) * (T - _T_0) +
        q_tot * _e_int_v0,
        PhasePartition(q_tot),
    ) ≈ FT(T)

    @test total_energy(param_set, FT(e_kin), FT(e_pot), _T_0) ===
          FT(e_kin) + FT(e_pot)
    @test total_energy(param_set, FT(e_kin), FT(e_pot), FT(T)) ≈
          FT(e_kin) + FT(e_pot) + _cv_d * (T - _T_0)
    @test total_energy(param_set, FT(0), FT(0), _T_0, PhasePartition(q_tot)) ≈
          q_tot * _e_int_v0

    # phase partitioning in equilibrium
    q_liq = FT(0.1)
    T = FT(_T_icenuc - 10)
    ρ = FT(1.0)
    q_tot = FT(0.21)
    phase_type = PhaseDry
    @test liquid_fraction(param_set, T, phase_type) === FT(0)
    phase_type = PhaseNonEquil
    @test liquid_fraction(
        param_set,
        T,
        phase_type,
        PhasePartition(q_tot, q_liq, q_liq),
    ) === FT(0.5)
    phase_type = PhaseDry
    q = PhasePartition_equil(param_set, T, ρ, q_tot, phase_type)
    @test q.liq ≈ FT(0)
    @test 0 < q.ice <= q_tot

    T = FT(_T_freeze + 10)
    ρ = FT(0.1)
    q_tot = FT(0.60)
    @test liquid_fraction(param_set, T, phase_type) === FT(1)
    phase_type = PhaseNonEquil
    @test liquid_fraction(
        param_set,
        T,
        phase_type,
        PhasePartition(q_tot, q_liq, q_liq / 2),
    ) === FT(2 / 3)
    phase_type = PhaseDry
    q = PhasePartition_equil(param_set, T, ρ, q_tot, phase_type)
    @test 0 < q.liq <= q_tot
    @test q.ice ≈ 0

    # saturation adjustment in equilibrium (i.e., given the thermodynamic
    # variables E_int, p, q_tot, compute the temperature and partitioning of the phases
    q_tot = FT(0)
    ρ = FT(1)
    phase_type = PhaseEquil
    @test TD.saturation_adjustment_SecantMethod(
        param_set,
        internal_energy_sat(param_set, 300.0, ρ, q_tot, phase_type),
        ρ,
        q_tot,
        phase_type,
        10,
        rtol_temperature,
    ) ≈ 300.0
    @test abs(
        TD.saturation_adjustment(
            param_set,
            internal_energy_sat(param_set, 300.0, ρ, q_tot, phase_type),
            ρ,
            q_tot,
            phase_type,
            10,
            rtol_temperature,
        ) - 300.0,
    ) < rtol_temperature

    q_tot = FT(0.21)
    ρ = FT(0.1)
    @test isapprox(
        TD.saturation_adjustment_SecantMethod(
            param_set,
            internal_energy_sat(param_set, 200.0, ρ, q_tot, phase_type),
            ρ,
            q_tot,
            phase_type,
            10,
            rtol_temperature,
        ),
        200.0,
        rtol = rtol_temperature,
    )
    @test abs(
        TD.saturation_adjustment(
            param_set,
            internal_energy_sat(param_set, 200.0, ρ, q_tot, phase_type),
            ρ,
            q_tot,
            phase_type,
            10,
            rtol_temperature,
        ) - 200.0,
    ) < rtol_temperature
    q = PhasePartition_equil(param_set, T, ρ, q_tot, phase_type)
    @test q.tot - q.liq - q.ice ≈
          vapor_specific_humidity(q) ≈
          q_vap_saturation(param_set, T, ρ, phase_type)

    ρ = FT(1)
    ρu = FT[1, 2, 3]
    ρe = FT(1100)
    e_pot = FT(93)
    @test internal_energy(ρ, ρe, ρu, e_pot) ≈ 1000.0

    # internal energies for dry, vapor, liquid and ice
    T = FT(300)
    q = PhasePartition(FT(20 * 1e-3), FT(5 * 1e-3), FT(2 * 1e-3))
    q_vap = vapor_specific_humidity(q)
    Id = internal_energy_dry(param_set, T)
    Iv = internal_energy_vapor(param_set, T)
    Il = internal_energy_liquid(param_set, T)
    Ii = internal_energy_ice(param_set, T)
    @test internal_energy(param_set, T, q) ≈
          (1 - q.tot) * Id + q_vap * Iv + q.liq * Il + q.ice * Ii
    @test internal_energy(param_set, T) ≈ Id

    # potential temperatures
    T = FT(300)
    @test TD.liquid_ice_pottemp_given_pressure(param_set, T, _MSLP) === T
    @test TD.liquid_ice_pottemp_given_pressure(param_set, T, _MSLP / 10) ≈
          T * 10^(_R_d / _cp_d)
    @test TD.liquid_ice_pottemp_given_pressure(
        param_set,
        T,
        _MSLP / 10,
        PhasePartition(FT(1)),
    ) ≈ T * 10^(_R_v / _cp_v)

    # dry potential temperatures. FIXME: add correctness tests
    T = FT(300)
    p = FT(1.e5)
    q_tot = FT(0.23)
    @test TD.dry_pottemp_given_pressure(
        param_set,
        T,
        p,
        PhasePartition(q_tot),
    ) isa typeof(p)
    @test TD.air_temperature_given_θpq(
        param_set,
        TD.dry_pottemp_given_pressure(param_set, T, p, PhasePartition(q_tot)),
        p,
        PhasePartition(q_tot),
    ) ≈ T

    # Exner function. FIXME: add correctness tests
    p = FT(1.e5)
    q_tot = FT(0.23)
    @test TD.exner_given_pressure(param_set, p, PhasePartition(q_tot)) isa
          typeof(p)
end

@testset "Thermodynamics - default behavior accuracy" begin
    # Input arguments should be accurate within machine precision
    # Temperature is approximated via saturation adjustment, and should be within a physical tolerance

    or(a, b) = a || b
    for ArrayType in array_types
        FT = eltype(ArrayType)
        profiles = PhaseEquilProfiles(param_set, ArrayType)
        @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
        @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

        RH_sat_mask = or.(RH .> 1, RH .≈ 1)
        RH_unsat_mask = .!or.(RH .> 1, RH .≈ 1)
        ts = PhaseEquil.(param_set, e_int, ρ, q_tot)
        @test all(saturated.(ts[RH_sat_mask]))
        @test !any(saturated.(ts[RH_unsat_mask]))

        # PhaseEquil (freezing)
        _T_freeze = FT(T_freeze(param_set))
        e_int_upper =
            internal_energy_sat.(
                param_set,
                Ref(_T_freeze + sqrt(eps(FT))),
                ρ,
                q_tot,
                phase_type,
            )
        e_int_lower =
            internal_energy_sat.(
                param_set,
                Ref(_T_freeze - sqrt(eps(FT))),
                ρ,
                q_tot,
                phase_type,
            )
        _e_int = (e_int_upper .+ e_int_lower) / 2
        ts = PhaseEquil.(param_set, _e_int, ρ, q_tot)
        @test all(air_temperature.(ts) .== Ref(_T_freeze))

        # Args needs to be in sync with PhaseEquil:
        ts =
            PhaseEquil.(
                param_set,
                _e_int,
                ρ,
                q_tot,
                8,
                FT(1e-1),
                TD.saturation_adjustment_SecantMethod,
            )
        @test all(air_temperature.(ts) .== Ref(_T_freeze))

        # PhaseEquil
        ts_exact = PhaseEquil.(param_set, e_int, ρ, q_tot, 100, FT(1e-3))
        ts = PhaseEquil.(param_set, e_int, ρ, q_tot)
        @test all(isapprox.(T, air_temperature.(ts), rtol = rtol_temperature))

        # Should be machine accurate (because ts contains `e_int`,`ρ`,`q_tot`):
        @test all(compare_moisture.(ts, ts_exact))
        @test all(internal_energy.(ts) .≈ internal_energy.(ts_exact))
        @test all(air_density.(ts) .≈ air_density.(ts_exact))
        # Approximate (temperature must be computed via saturation adjustment):
        @test all(isapprox.(
            air_pressure.(ts),
            air_pressure.(ts_exact),
            rtol = rtol_pressure,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol_temperature,
        ))

        dry_mask = abs.(q_tot .- 0) .< eps(FT)
        q_dry = q_pt[dry_mask]
        @test all(
            condensate.(q_pt) .==
            getproperty.(q_pt, :liq) .+ getproperty.(q_pt, :ice),
        )
        @test all(has_condensate.(q_dry) .== false)

        e_tot = total_energy.(e_kin, e_pot, ts)
        @test all(
            specific_enthalpy.(ts) .≈
            e_int .+ gas_constant_air.(ts) .* air_temperature.(ts),
        )
        @test all(
            total_specific_enthalpy.(ts, e_tot) .≈
            specific_enthalpy.(ts) .+ e_kin .+ e_pot,
        )
        @test all(
            moist_static_energy.(ts, e_pot) .≈ specific_enthalpy.(ts) .+ e_pot,
        )
        @test all(
            moist_static_energy.(ts, e_pot) .≈
            total_specific_enthalpy.(ts, e_tot) .- e_kin,
        )

        # PhaseEquil
        ts_exact =
            PhaseEquil.(
                param_set,
                e_int,
                ρ,
                q_tot,
                100,
                FT(1e-3),
                TD.saturation_adjustment_SecantMethod,
            )
        ts =
            PhaseEquil.(
                param_set,
                e_int,
                ρ,
                q_tot,
                35,
                FT(1e-1),
                TD.saturation_adjustment_SecantMethod,
            ) # Needs to be in sync with default
        # Should be machine accurate (because ts contains `e_int`,`ρ`,`q_tot`):
        @test all(compare_moisture.(ts, ts_exact))
        @test all(internal_energy.(ts) .≈ internal_energy.(ts_exact))
        @test all(air_density.(ts) .≈ air_density.(ts_exact))
        # Approximate (temperature must be computed via saturation adjustment):
        @test all(isapprox.(
            air_pressure.(ts),
            air_pressure.(ts_exact),
            rtol = rtol_pressure,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol_temperature,
        ))

        # PhaseEquil_ρθq
        ts_exact = PhaseEquil_ρθq.(param_set, ρ, θ_liq_ice, q_tot, 45, FT(1e-3))
        ts = PhaseEquil_ρθq.(param_set, ρ, θ_liq_ice, q_tot)
        # Should be machine accurate:
        @test all(air_density.(ts) .≈ air_density.(ts_exact))
        @test all(compare_moisture.(ts, ts_exact))
        # Approximate (temperature must be computed via saturation adjustment):
        @test all(isapprox.(
            internal_energy.(ts),
            internal_energy.(ts_exact),
            atol = atol_energy,
        ))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            liquid_ice_pottemp.(ts_exact),
            rtol = rtol_temperature,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol_temperature,
        ))

        # PhaseEquil_pθq
        ts_exact = PhaseEquil_pθq.(param_set, p, θ_liq_ice, q_tot, 40, FT(1e-3))
        ts = PhaseEquil_pθq.(param_set, p, θ_liq_ice, q_tot)
        # Should be machine accurate:
        @test all(compare_moisture.(ts, ts_exact))
        # Approximate (temperature must be computed via saturation adjustment):
        @test all(isapprox.(
            air_density.(ts),
            air_density.(ts_exact),
            rtol = rtol_density,
        ))
        @test all(isapprox.(
            internal_energy.(ts),
            internal_energy.(ts_exact),
            atol = atol_energy,
        ))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            liquid_ice_pottemp.(ts_exact),
            rtol = rtol_temperature,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol_temperature,
        ))

        # PhaseNonEquil_ρθq
        ts_exact =
            PhaseNonEquil_ρθq.(param_set, ρ, θ_liq_ice, q_pt, 40, FT(1e-3))
        ts = PhaseNonEquil_ρθq.(param_set, ρ, θ_liq_ice, q_pt)
        # Should be machine accurate:
        @test all(compare_moisture.(ts, ts_exact))
        @test all(air_density.(ts) .≈ air_density.(ts_exact))
        # Approximate (temperature must be computed via non-linear solve):
        @test all(isapprox.(
            internal_energy.(ts),
            internal_energy.(ts_exact),
            atol = atol_energy,
        ))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            liquid_ice_pottemp.(ts_exact),
            rtol = rtol_temperature,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol_temperature,
        ))

    end

end

@testset "Thermodynamics - exceptions on failed convergence" begin

    ArrayType = Array{Float64}
    FT = eltype(ArrayType)
    profiles = PhaseEquilProfiles(param_set, ArrayType)
    @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
    @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

    @test_throws ErrorException TD.saturation_adjustment.(
        param_set,
        e_int,
        ρ,
        q_tot,
        Ref(phase_type),
        2,
        FT(1e-10),
    )

    @test_throws ErrorException TD.saturation_adjustment_SecantMethod.(
        param_set,
        e_int,
        ρ,
        q_tot,
        Ref(phase_type),
        2,
        FT(1e-10),
    )

    T_virt = T # should not matter: testing for non-convergence
    @test_throws ErrorException TD.temperature_and_humidity_given_TᵥρRH.(
        param_set,
        T_virt,
        ρ,
        RH,
        Ref(phase_type),
        2,
        ResidualTolerance(FT(1e-10)),
    )

    @test_throws ErrorException TD.air_temperature_given_θρq_nonlinear.(
        param_set,
        θ_liq_ice,
        ρ,
        2,
        ResidualTolerance(FT(1e-10)),
        q_pt,
    )

    @test_throws ErrorException TD.saturation_adjustment_given_ρθq.(
        param_set,
        ρ,
        θ_liq_ice,
        q_tot,
        Ref(phase_type),
        2,
        ResidualTolerance(FT(1e-10)),
    )

    @test_throws ErrorException TD.saturation_adjustment_given_pθq.(
        param_set,
        p,
        θ_liq_ice,
        q_tot,
        Ref(phase_type),
        2,
        ResidualTolerance(FT(1e-10)),
    )

    @test_throws ErrorException TD.saturation_adjustment_ρpq.(
        param_set,
        ρ,
        p,
        q_tot,
        Ref(phase_type),
        2,
        FT(1e-10),
    )

end

@testset "Thermodynamics - constructor consistency" begin

    # Make sure `ThermodynamicState` arguments are returned unchanged

    for ArrayType in array_types
        FT = eltype(ArrayType)
        _MSLP = FT(MSLP(param_set))

        profiles = PhaseDryProfiles(param_set, ArrayType)
        @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
        @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

        # PhaseDry
        ts = PhaseDry.(param_set, e_int, ρ)
        @test all(internal_energy.(ts) .≈ e_int)
        @test all(air_density.(ts) .≈ ρ)

        ts_pT = PhaseDry_pT.(param_set, p, T)
        @test all(internal_energy.(ts_pT) .≈ internal_energy.(param_set, T))
        @test all(air_density.(ts_pT) .≈ ρ)

        θ_dry = dry_pottemp.(param_set, T, ρ)
        ts_pθ = PhaseDry_pθ.(param_set, p, θ_dry)
        @test all(internal_energy.(ts_pθ) .≈ internal_energy.(param_set, T))
        @test all(air_density.(ts_pθ) .≈ ρ)

        ts_ρθ = PhaseDry_ρθ.(param_set, ρ, θ_dry)
        @test all(internal_energy.(ts_ρθ) .≈ internal_energy.(param_set, T))
        @test all(air_density.(ts_ρθ) .≈ ρ)

        ts_ρT = PhaseDry_ρT.(param_set, ρ, T)
        @test all(air_density.(ts_ρT) .≈ air_density.(ts))
        @test all(internal_energy.(ts_ρT) .≈ internal_energy.(ts))

        ts = PhaseDry_ρp.(param_set, ρ, p)
        @test all(air_density.(ts) .≈ ρ)
        @test all(air_pressure.(ts) .≈ p)
        e_tot_proposed =
            TD.total_energy_given_ρp.(param_set, ρ, p, e_kin, e_pot)
        @test all(total_energy.(e_kin, e_pot, ts) .≈ e_tot_proposed)

        profiles = PhaseEquilProfiles(param_set, ArrayType)
        @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
        @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

        # PhaseEquil
        ts =
            PhaseEquil.(
                param_set,
                e_int,
                ρ,
                q_tot,
                40,
                FT(1e-1),
                Ref(TD.saturation_adjustment_SecantMethod),
            )
        @test all(internal_energy.(ts) .≈ e_int)
        @test all(getproperty.(PhasePartition.(ts), :tot) .≈ q_tot)
        @test all(air_density.(ts) .≈ ρ)

        ts = PhaseEquil.(param_set, e_int, ρ, q_tot)
        @test all(internal_energy.(ts) .≈ e_int)
        @test all(getproperty.(PhasePartition.(ts), :tot) .≈ q_tot)
        @test all(air_density.(ts) .≈ ρ)

        ts = PhaseEquil_ρpq.(param_set, ρ, p, q_tot, true)
        @test all(air_density.(ts) .≈ ρ)
        @test all(air_pressure.(ts) .≈ p)
        @test all(getproperty.(PhasePartition.(ts), :tot) .≈ q_tot)

        # Test against total_energy_given_ρp when not iterating
        ts = PhaseEquil_ρpq.(param_set, ρ, p, q_tot, false)
        e_tot_proposed =
            TD.total_energy_given_ρp.(
                param_set,
                ρ,
                p,
                e_kin,
                e_pot,
                PhasePartition.(q_tot),
            )
        @test all(total_energy.(e_kin, e_pot, ts) .≈ e_tot_proposed)

        # PhaseNonEquil
        ts = PhaseNonEquil.(param_set, e_int, ρ, q_pt)
        @test all(internal_energy.(ts) .≈ e_int)
        @test all(compare_moisture.(ts, q_pt))
        @test all(air_density.(ts) .≈ ρ)

        # TD.air_temperature_given_θpq-liquid_ice_pottemp inverse
        θ_liq_ice_ =
            TD.liquid_ice_pottemp_given_pressure.(param_set, T, p, q_pt)
        @test all(
            TD.air_temperature_given_θpq.(param_set, θ_liq_ice_, p, q_pt) .≈ T,
        )

        # liquid_ice_pottemp-TD.air_temperature_given_θpq inverse
        T = TD.air_temperature_given_θpq.(param_set, θ_liq_ice, p, q_pt)
        @test all(
            TD.liquid_ice_pottemp_given_pressure.(param_set, T, p, q_pt) .≈
            θ_liq_ice,
        )

        # Accurate but expensive `PhaseNonEquil_ρθq` constructor (Non-linear temperature from θ_liq_ice)
        T_non_linear =
            TD.air_temperature_given_θρq_nonlinear.(
                param_set,
                θ_liq_ice,
                ρ,
                20,
                ResidualTolerance(FT(5e-5)),
                q_pt,
            )
        T_expansion =
            TD.air_temperature_given_θρq.(param_set, θ_liq_ice, ρ, q_pt)
        @test all(isapprox.(T_non_linear, T_expansion, rtol = rtol_temperature))
        e_int_ = internal_energy.(param_set, T_non_linear, q_pt)
        ts = PhaseNonEquil.(param_set, e_int_, ρ, q_pt)
        @test all(T_non_linear .≈ air_temperature.(ts))
        @test all(isapprox(
            θ_liq_ice,
            liquid_ice_pottemp.(ts),
            rtol = rtol_temperature,
        ))

        # PhaseEquil_ρθq
        ts = PhaseEquil_ρθq.(param_set, ρ, θ_liq_ice, q_tot, 45, FT(1e-3))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            θ_liq_ice,
            rtol = rtol_temperature,
        ))
        @test all(isapprox.(air_density.(ts), ρ, rtol = rtol_density))
        @test all(getproperty.(PhasePartition.(ts), :tot) .≈ q_tot)

        # The PhaseEquil_pθq constructor
        # passes the consistency test within sufficient physical precision,
        # however, it fails to satisfy the consistency test within machine
        # precision for the input pressure.

        # PhaseEquil_pθq
        ts = PhaseEquil_pθq.(param_set, p, θ_liq_ice, q_tot, 35, FT(1e-3))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            θ_liq_ice,
            rtol = rtol_temperature,
        ))
        @test all(compare_moisture.(ts, q_pt))
        @test all(isapprox.(air_pressure.(ts), p, atol = atol_pressure))

        # PhaseNonEquil_pθq
        ts = PhaseNonEquil_pθq.(param_set, p, θ_liq_ice, q_pt)
        @test all(liquid_ice_pottemp.(ts) .≈ θ_liq_ice)
        @test all(air_pressure.(ts) .≈ p)
        @test all(compare_moisture.(ts, q_pt))

        ts = PhaseNonEquil_ρpq.(param_set, ρ, p, q_pt)
        @test all(air_density.(ts) .≈ ρ)
        @test all(air_pressure.(ts) .≈ p)
        @test all(
            getproperty.(PhasePartition.(ts), :tot) .≈ getproperty.(q_pt, :tot),
        )
        @test all(
            getproperty.(PhasePartition.(ts), :liq) .≈ getproperty.(q_pt, :liq),
        )
        @test all(
            getproperty.(PhasePartition.(ts), :ice) .≈ getproperty.(q_pt, :ice),
        )
        e_tot_proposed =
            TD.total_energy_given_ρp.(param_set, ρ, p, e_kin, e_pot, q_pt)
        @test all(total_energy.(e_kin, e_pot, ts) .≈ e_tot_proposed)

        # PhaseNonEquil_ρθq
        ts = PhaseNonEquil_ρθq.(param_set, ρ, θ_liq_ice, q_pt, 5, FT(1e-3))
        @test all(isapprox.(
            θ_liq_ice,
            liquid_ice_pottemp.(ts),
            rtol = rtol_temperature,
        ))
        @test all(air_density.(ts) .≈ ρ)
        @test all(compare_moisture.(ts, q_pt))

        profiles = PhaseEquilProfiles(param_set, ArrayType)
        @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
        @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

        # Test that relative humidity is 1 for saturated conditions
        q_sat = q_vap_saturation.(param_set, T, ρ, Ref(phase_type))
        q_pt_sat = PhasePartition.(q_sat)
        q_vap = vapor_specific_humidity.(q_pt_sat)
        @test all(getproperty.(q_pt_sat, :liq) .≈ 0)
        @test all(getproperty.(q_pt_sat, :ice) .≈ 0)
        @test all(q_vap .≈ q_sat)

        # Compute thermodynamic consistent pressure
        p_sat = air_pressure.(param_set, T, ρ, q_pt_sat)

        # Test that density remains consistent
        ρ_rec = air_density.(param_set, T, p_sat, q_pt_sat)
        @test all.(ρ_rec ≈ ρ)

        RH_sat =
            relative_humidity.(param_set, T, p_sat, Ref(phase_type), q_pt_sat)

        # TODO: Add this test back in
        @test all(RH_sat .≈ 1)

        # Test that RH is zero for dry conditions
        q_pt_dry = PhasePartition.(zeros(FT, length(T)))
        p_dry = air_pressure.(param_set, T, ρ, q_pt_dry)
        RH_dry =
            relative_humidity.(param_set, T, p_dry, Ref(phase_type), q_pt_dry)
        @test all(RH_dry .≈ 0)

        # Test virtual temperature and inverse functions:
        _R_d = FT(R_d(param_set))
        T_virt = virtual_temperature.(param_set, T, ρ, q_pt)
        @test all(T_virt ≈ gas_constant_air.(param_set, q_pt) ./ _R_d .* T)

        T_rec_qpt_rec =
            TD.temperature_and_humidity_given_TᵥρRH.(
                param_set,
                T_virt,
                ρ,
                RH,
                Ref(phase_type),
            )

        T_rec = first.(T_rec_qpt_rec)
        q_pt_rec = last.(T_rec_qpt_rec)

        # Test convergence of virtual temperature iterations
        @test all(isapprox.(
            T_virt,
            virtual_temperature.(param_set, T_rec, ρ, q_pt_rec),
            atol = sqrt(eps(FT)),
        ))

        # Test that reconstructed specific humidity is close
        # to original specific humidity
        q_tot_rec = getproperty.(q_pt_rec, :tot)
        RH_moist = q_tot .> eps(FT)
        @test all(isapprox.(q_tot[RH_moist], q_tot_rec[RH_moist], rtol = 5e-2))

        # Update temperature to be exactly consistent with
        # p, ρ, q_pt_rec; test that this is equal to T_rec
        T_local =
            TD.air_temperature_from_ideal_gas_law.(param_set, p, ρ, q_pt_rec)
        @test all(isapprox.(T_local, T_rec, atol = sqrt(eps(FT))))
    end

end

@testset "Thermodynamics - type-stability" begin

    # NOTE: `Float32` saturation adjustment tends to have more difficulty
    # with converging to the same tolerances as `Float64`, so they're relaxed here.
    ArrayType = Array{Float32}
    FT = eltype(ArrayType)
    profiles = PhaseEquilProfiles(param_set, ArrayType)
    @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
    @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

    ρu = FT[1.0, 2.0, 3.0]
    @test typeof.(internal_energy.(ρ, ρ .* e_int, Ref(ρu), e_pot)) ==
          typeof.(e_int)

    θ_dry = dry_pottemp.(param_set, T, ρ)
    ts_dry = PhaseDry.(param_set, e_int, ρ)
    ts_dry_ρp = PhaseDry_ρp.(param_set, ρ, p)
    ts_dry_pT = PhaseDry_pT.(param_set, p, T)
    ts_dry_ρθ = PhaseDry_ρθ.(param_set, ρ, θ_dry)
    ts_dry_pθ = PhaseDry_pθ.(param_set, p, θ_dry)
    ts_eq = PhaseEquil.(param_set, e_int, ρ, q_tot, 15, FT(1e-1))
    e_tot = total_energy.(e_kin, e_pot, ts_eq)

    ts_T =
        PhaseEquil_ρTq.(
            param_set,
            air_density.(ts_dry),
            air_temperature.(ts_dry),
            q_tot,
        )
    ts_Tp =
        PhaseEquil_pTq.(
            param_set,
            air_pressure.(ts_dry),
            air_temperature.(ts_dry),
            q_tot,
        )

    ts_ρp =
        PhaseEquil_ρpq.(
            param_set,
            air_density.(ts_dry),
            air_pressure.(ts_dry),
            q_tot,
        )

    @test all(air_temperature.(ts_T) .≈ air_temperature.(ts_Tp))
    # @test all(isapprox.(air_pressure.(ts_T), air_pressure.(ts_Tp), atol = _MSLP * 2e-2)) # TODO: Fails, needs fixing / better test
    @test all(total_specific_humidity.(ts_T) .≈ total_specific_humidity.(ts_Tp))

    ts_neq = PhaseNonEquil.(param_set, e_int, ρ, q_pt)
    ts_T_neq = PhaseNonEquil_ρTq.(param_set, ρ, T, q_pt)

    ts_θ_liq_ice_eq =
        PhaseEquil_ρθq.(param_set, ρ, θ_liq_ice, q_tot, 45, FT(1e-3))
    ts_θ_liq_ice_eq_p =
        PhaseEquil_pθq.(param_set, p, θ_liq_ice, q_tot, 40, FT(1e-3))
    ts_θ_liq_ice_neq = PhaseNonEquil_ρθq.(param_set, ρ, θ_liq_ice, q_pt)
    ts_θ_liq_ice_neq_p = PhaseNonEquil_pθq.(param_set, p, θ_liq_ice, q_pt)

    for ts in (
        ts_dry,
        ts_dry_ρp,
        ts_dry_pT,
        ts_dry_ρθ,
        ts_dry_pθ,
        ts_eq,
        ts_T,
        ts_Tp,
        ts_ρp,
        ts_neq,
        ts_T_neq,
        ts_θ_liq_ice_eq,
        ts_θ_liq_ice_eq_p,
        ts_θ_liq_ice_neq,
        ts_θ_liq_ice_neq_p,
    )
        @test typeof.(soundspeed_air.(ts)) == typeof.(e_int)
        @test typeof.(gas_constant_air.(ts)) == typeof.(e_int)
        @test typeof.(specific_enthalpy.(ts)) == typeof.(e_int)
        @test typeof.(vapor_specific_humidity.(ts)) == typeof.(e_int)
        @test typeof.(relative_humidity.(ts)) == typeof.(e_int)
        @test typeof.(air_pressure.(ts)) == typeof.(e_int)
        @test typeof.(air_density.(ts)) == typeof.(e_int)
        @test typeof.(total_specific_humidity.(ts)) == typeof.(e_int)
        @test typeof.(liquid_specific_humidity.(ts)) == typeof.(e_int)
        @test typeof.(ice_specific_humidity.(ts)) == typeof.(e_int)
        @test typeof.(cp_m.(ts)) == typeof.(e_int)
        @test typeof.(cv_m.(ts)) == typeof.(e_int)
        @test typeof.(air_temperature.(ts)) == typeof.(e_int)
        @test typeof.(internal_energy_sat.(ts)) == typeof.(e_int)
        @test typeof.(internal_energy.(ts)) == typeof.(e_int)
        @test typeof.(internal_energy_dry.(ts)) == typeof.(e_int)
        @test typeof.(internal_energy_vapor.(ts)) == typeof.(e_int)
        @test typeof.(internal_energy_liquid.(ts)) == typeof.(e_int)
        @test typeof.(internal_energy_ice.(ts)) == typeof.(e_int)
        @test typeof.(latent_heat_vapor.(ts)) == typeof.(e_int)
        @test typeof.(latent_heat_sublim.(ts)) == typeof.(e_int)
        @test typeof.(latent_heat_fusion.(ts)) == typeof.(e_int)
        @test typeof.(q_vap_saturation.(ts)) == typeof.(e_int)
        @test typeof.(q_vap_saturation_liquid.(ts)) == typeof.(e_int)
        @test typeof.(q_vap_saturation_ice.(ts)) == typeof.(e_int)
        @test typeof.(saturation_excess.(ts)) == typeof.(e_int)
        @test typeof.(liquid_fraction.(ts)) == typeof.(e_int)
        @test typeof.(liquid_ice_pottemp.(ts)) == typeof.(e_int)
        @test typeof.(dry_pottemp.(ts)) == typeof.(e_int)
        @test typeof.(exner.(ts)) == typeof.(e_int)
        @test typeof.(liquid_ice_pottemp_sat.(ts)) == typeof.(e_int)
        @test typeof.(specific_volume.(ts)) == typeof.(e_int)
        @test typeof.(supersaturation.(ts, Ice())) == typeof.(e_int)
        @test typeof.(supersaturation.(ts, Liquid())) == typeof.(e_int)
        @test typeof.(virtual_pottemp.(ts)) == typeof.(e_int)
        @test eltype.(gas_constants.(ts)) == typeof.(e_int)

        @test typeof.(total_specific_enthalpy.(ts, e_tot)) == typeof.(e_int)
        @test typeof.(moist_static_energy.(ts, e_pot)) == typeof.(e_int)
        @test typeof.(getproperty.(PhasePartition.(ts), :tot)) == typeof.(e_int)
    end

end

@testset "Thermodynamics - dry limit" begin

    ArrayType = Array{Float64}
    FT = eltype(ArrayType)
    profiles = PhaseEquilProfiles(param_set, ArrayType)
    @unpack T, p, RS, e_int, ρ, θ_liq_ice, phase_type = profiles
    @unpack q_tot, q_liq, q_ice, q_pt, RH, e_kin, e_pot = profiles

    # PhasePartition test is noisy, so do this only once:
    ts_dry = PhaseDry(param_set, first(e_int), first(ρ))
    ts_eq = PhaseEquil(param_set, first(e_int), first(ρ), typeof(first(ρ))(0))
    @test PhasePartition(ts_eq).tot ≈ PhasePartition(ts_dry).tot
    @test PhasePartition(ts_eq).liq ≈ PhasePartition(ts_dry).liq
    @test PhasePartition(ts_eq).ice ≈ PhasePartition(ts_dry).ice

    ts_dry = PhaseDry.(param_set, e_int, ρ)
    ts_eq = PhaseEquil.(param_set, e_int, ρ, q_tot .* 0)

    @test all(gas_constant_air.(ts_eq) .≈ gas_constant_air.(ts_dry))
    @test all(relative_humidity.(ts_eq) .≈ relative_humidity.(ts_dry))
    @test all(air_pressure.(ts_eq) .≈ air_pressure.(ts_dry))
    @test all(air_density.(ts_eq) .≈ air_density.(ts_dry))
    @test all(specific_volume.(ts_eq) .≈ specific_volume.(ts_dry))
    @test all(
        total_specific_humidity.(ts_eq) .≈ total_specific_humidity.(ts_dry),
    )
    @test all(
        liquid_specific_humidity.(ts_eq) .≈ liquid_specific_humidity.(ts_dry),
    )
    @test all(ice_specific_humidity.(ts_eq) .≈ ice_specific_humidity.(ts_dry))
    @test all(cp_m.(ts_eq) .≈ cp_m.(ts_dry))
    @test all(cv_m.(ts_eq) .≈ cv_m.(ts_dry))
    @test all(air_temperature.(ts_eq) .≈ air_temperature.(ts_dry))
    @test all(internal_energy.(ts_eq) .≈ internal_energy.(ts_dry))
    @test all(internal_energy_sat.(ts_eq) .≈ internal_energy_sat.(ts_dry))
    @test all(internal_energy_dry.(ts_eq) .≈ internal_energy_dry.(ts_dry))
    @test all(internal_energy_vapor.(ts_eq) .≈ internal_energy_vapor.(ts_dry))
    @test all(internal_energy_liquid.(ts_eq) .≈ internal_energy_liquid.(ts_dry))
    @test all(internal_energy_ice.(ts_eq) .≈ internal_energy_ice.(ts_dry))
    @test all(soundspeed_air.(ts_eq) .≈ soundspeed_air.(ts_dry))
    @test all(supersaturation.(ts_eq, Ice()) .≈ supersaturation.(ts_dry, Ice()))
    @test all(
        supersaturation.(ts_eq, Liquid()) .≈ supersaturation.(ts_dry, Liquid()),
    )
    @test all(latent_heat_vapor.(ts_eq) .≈ latent_heat_vapor.(ts_dry))
    @test all(latent_heat_sublim.(ts_eq) .≈ latent_heat_sublim.(ts_dry))
    @test all(latent_heat_fusion.(ts_eq) .≈ latent_heat_fusion.(ts_dry))
    @test all(q_vap_saturation.(ts_eq) .≈ q_vap_saturation.(ts_dry))
    @test all(
        q_vap_saturation_liquid.(ts_eq) .≈ q_vap_saturation_liquid.(ts_dry),
    )
    @test all(q_vap_saturation_ice.(ts_eq) .≈ q_vap_saturation_ice.(ts_dry))
    @test all(saturation_excess.(ts_eq) .≈ saturation_excess.(ts_dry))
    @test all(liquid_fraction.(ts_eq) .≈ liquid_fraction.(ts_dry))
    @test all(liquid_ice_pottemp.(ts_eq) .≈ liquid_ice_pottemp.(ts_dry))
    @test all(dry_pottemp.(ts_eq) .≈ dry_pottemp.(ts_dry))
    @test all(virtual_pottemp.(ts_eq) .≈ virtual_pottemp.(ts_dry))
    @test all(liquid_ice_pottemp_sat.(ts_eq) .≈ liquid_ice_pottemp_sat.(ts_dry))
    @test all(exner.(ts_eq) .≈ exner.(ts_dry))

    @test all(
        saturation_vapor_pressure.(ts_eq, Ice()) .≈
        saturation_vapor_pressure.(ts_dry, Ice()),
    )
    @test all(
        saturation_vapor_pressure.(ts_eq, Liquid()) .≈
        saturation_vapor_pressure.(ts_dry, Liquid()),
    )
    @test all(first.(gas_constants.(ts_eq)) ≈ first.(gas_constants.(ts_dry)))
    @test all(last.(gas_constants.(ts_eq)) ≈ last.(gas_constants.(ts_dry)))

end

@testset "Thermodynamics - ProfileSet Iterator" begin
    ArrayType = Array{Float64}
    FT = eltype(ArrayType)
    profiles = PhaseEquilProfiles(param_set, ArrayType)
    @unpack T, q_pt, z, phase_type = profiles
    @test all(z .≈ (nt.z for nt in profiles))
    @test all(T .≈ (nt.T for nt in profiles))
    @test all(getproperty.(q_pt, :tot) .≈ (nt.q_pt.tot for nt in profiles))
    @test all(phase_type .== (nt.phase_type for nt in profiles))
end


text:
module EqualityConstrainedMinimization
export Solve_ECM
using Calculus
using LinearAlgebra
include("algorithms.jl")

#######################################################################
#######################################################################
#=

This is a non-linear minimization solver that works with all convex problems
with linear equality constraints. It is intended to solve problems in the form
"minimize f(x) such that Ax = b" and "minimize f(x)" for the unconstrained case.
I am taking the algorithms used directly from "Convex Optimization" by Stephen Boyd
and Lieven Vandenberghe. This package is intended for learning purposes and while
it may not be the fastest solver, it is certainly one of the easiest to understand.

Algorithms used for unconstrained problem:
    -Gradient descent
    -Newton's method
    -Unless force_gradient = true, the solver will use gradient descent method until
     stopping criterion η is met. Then, it will switch to Newton's Method until ϵ is met
     Note that η > ϵ.

Algorithms used for equality constrained problem:
    -Feasible Start Newton Method
    -Infeasible Start Newton Method

# Arguments
    - 'f': The convex objective. Must be convex or the solution may converge to a local solution.
    - 'x': The starting point. May be feasible or infeasible though starting feasible may converge faster.
    - 'A': The coefficient matrix on the Ax = b constraint. This system needs to be underdetermined. Not including this will imply an unconstrained problem.
    - 'b': The constant vector on the Ax = b constraint. Not including this will imply feasible start.
    - 'α': A parameter used in the backtracking line search. Note that α ∈ (0, 0.5)
    - 'β': Another parameter used in the backtracking line search. β ∈ (0, 1)
    - 'η': A stopping criterion used to determine when to switch from gradient descent to Newton's method.
    - 'ϵ': Overall stopping criterion.
    - 'maxIt': Max iteration stopping criterion. Meant as a failsafe.
    - 'verbose': Setting this to true will show data from each iteration.
    - 'force_gradient': Setting this to true will force only the gradient descent method for unconstrained problems. Do this when ϵ is relatively large.

x^2 + y^2 -> x[1]^2 + x[2]^2

Main solver function:
Solve_ECM(
        f,
        x::Array{Float64,1};
        A::Array{Float64,2} = zeros(2, 2),
        b::Array{Float64,1} = zeros(2),
        α = 0.3,
        β = 0.8,
        η = 1e-4,
        ϵ = 1e-8,
        maxIt = 1e3,
        verbose = true,
        force_gradient = false
    )

See 'UsageExamples.jl' for some examples.

=#
#######################################################################
#######################################################################

#####################################
#####################################
#
# Main Solver
#
#####################################
#####################################
function Solve_ECM(
        f,
        x::Array{Float64,1};
        A::Array{Float64,2} = zeros(2, 2),
        b::Array{Float64,1} = zeros(2),
        α = 0.3,
        β = 0.8,
        η = 1e-4,
        ϵ = 1e-8,
        maxIt = 1e3,
        verbose = true,
        force_gradient = false
    )

    n = length(x)
    m = size(A,1)

    uncon = true
    feasible = true

    if η >= ϵ && n >= m

        if A != zeros(2, 2)
            if b != zeros(2)
                feasible = false
            end
            uncon = false
        end

        sol = Tuple{Array{Float64,1},Float64}

        #Determine which method to use
        if uncon

            print("\n\nUnconstrained objective detected")

            if force_gradient
                print("\nSolving using gradient descent\n")
                sol = Unconstrained_Grad_Descent(f, x, n, α, β, ϵ, maxIt, verbose)
            else
                print("\nSolving using gradient descent and Newton's Method\n")
                sol = Solve_UMP(f, x, n, α, β, η, ϵ, maxIt, verbose)
            end
        else
            if feasible
                print("\n\nVector b not given. Therefore, x is assumed feasible")
                print("\nSolving using Feasible Start Newton Method")
                sol = Newton_FS(f, x, A, n, m, α, β, ϵ, maxIt, verbose)
            else
                print("\n\nVector b given. Therefore, x is assumed infeasible")
                print("\nSolving using Infeasible Start Newton Method")
                sol = Newton_IFS(f, x, A, b, n, m, α, β, ϵ, maxIt, verbose)
            end
        end

        print("\n\n------------------------")
        print("\nSolution found at point:")
        print("\n\nx = ", sol[1],"\n")
        print("\nObjective Value = ", sol[2], "\n\n")
    else
        print("\nParameters not correct. Either η < ϵ or Ax = b is an overdetermined system")
    end

    return sol[2]
end

end

ext: jl

max_stars_repo_name: ihawn/EqualityConstrainedMinimization

max_issues_repo_name: ihawn/EqualityConstrainedMinimization

max_forks_repo_name: ihawn/EqualityConstrainedMinimization


text:
using libnode_jll: node as node_cmd_func, artifact_dir

const NPM_DIR = joinpath(artifact_dir, "lib", "node_modules", "npm")

function cmd_gen(args...)
    cmd = node_cmd_func()
    for arg in args
        cmd = `$cmd $arg`
    end
    cmd
end

macro node_cmd(args)
    :(let cmd = @cmd $args
        cmd_gen(cmd)
    end)
end

macro npm_cmd(args)
    :(let cmd = @cmd($args), jsfile = joinpath(NPM_DIR, "bin", "npm-cli.js")
        cmd_gen(jsfile, cmd)
    end)
end

macro npx_cmd(args)
    :(let cmd = @cmd($args), jsfile = joinpath(NPM_DIR, "bin", "npx-cli.js")
        cmd_gen(jsfile, cmd)
    end)
end

function npx(args...)
    cmd = npx``
    append!(cmd.exec, args)
    Base.run(cmd)
end


text:
# About PML setup
#
# - Let's define the computational domain as the region where we perform computation, and
# the physical domain where actual physical phenomena are simulated.  The physical domain is
# the computational domain outside PML, because PML is an artificial region that is
# introduced to absorb physical waves.
#
# Currently, we specify the size of the computational domain first, and then specify PML
# such that it "eats in" the computational domain to leave the physical domain.  This means
# you always need to define the computational domain larger than the physical domain by the
# PML thickness.
#
# This seems inconvenient.  It may seem more convenient to specify the physical domain first,
# specify the PML thickness, and then construct the computational domain by attaching PML
# to the physical domain.  That way, we would be able to turn on and off the absorbing
# boundary without changing the size of the physical domain of interest.  Also, increasing
# the PML thickness would be done simply by increasing the PML thickness, wherease with the
# current method it requires increasing the computational domain as well to keep the size of
# the physical domain the same.
#
# However, there is a reason for the current method.  One of the most popular use cases of
# PML is the simulation of an infinite waveguide, where we need to define the waveguide
# inside PML.  Therefore, we cannot "hide" PML from the users: we have to let the users put
# objects inside PML.  In that case, it is better to let the users specify the size of the
# region including PML, because that way the users would know how deep they should place the
# waveguide.
#
# You way think that you hide PML from the users by automatically constructing the material
# parameters inside PML, by taking the material parameters on the boundary of the physical
# domain and extending it homogeneously into PML.  However, such a homogeneous extension is
# not always what we want.  For example, to construct PML for a phothonic crystal waveguide,
# we have to specify the photonic crystal structure inside PML.  In that case the material
# composition inside PML is not simply the homogeneous extension of their values at the PML
# boundary
#
# - Another question that could be raised about the current method is the way to set the
# PML thickness.  Currently we specify the number of grid cells inside PML and multiply the
# cell size ∆l to figure out the PML thickness, rather than specifying the PML thickness and
# divide it by ∆l to figure out the number of grid cells inside PML.
#
# The latter look like a better method, because the reflectance of PML is a function of the
# PML loss parameter and thickness.  Therefore, for a given PML loss parameter, the
# reflectance does not change with the spatial resolution of the grid if the PML thickness
# is maintained.  Conversely, if we fix the number of grid cells inside PML, the PML
# thickness changes with the spacial resolution, so the PML loss parameter needs to change
# as well to achieve the same PML reflectance.
#
# However, in my experience, when using iterative methods to solve Maxwell's equations, the
# number of grid cells inside PML affects the convergence of iterative methods significantly,
# regardless of the spatial resolution.  For example, when I doubled the spatial resolution
# while keeping the PML thickness and therefore doubled the number of grid cells inside PML,
# the iterative methods no longer converged.  On the other hand, even if I doubled the
# spatial resolution, if I keep the number of grid cells inside PML to halve the PML
# thickness, I still got convergence.
#
# Of course, I have to use a stronger PML loss parameter to achieve the same reflectance
# with a thinner PML, and if we use a too strong loss parameter, the discretized PML model
# deviates from the continuous model and we get stronger reflectance than the target value.
# Still, I found that the reflectance remains pretty low, and the convergence of iterative
# methods was more important.  So, I decided to specify the number of grid cells inside PML
# rather than the PML thickness.  This decision may change in the future.

export PMLParam
export create_stretched_∆l

mutable struct PMLParam
    m::Float  # degree of polynomial grading for σ and κ
    R::Float  # R: target reflection coefficient of PML
    κmax::Float  # maximum of real stretch factor
    amax::Float  # maximum of complex frequency shift
    ma::Float  # degree of polynomial grading for a
    PMLParam() = new(4.0, exp(-16), 1.0, 0.0, 4.0)
end

function create_stretched_∆l(ω::Number, ∆l::Tuple2{AbsVecReal}, l::Tuple2{AbsVecReal},
                             bounds::Tuple2{Real}, Npml::Tuple2{Integer}, pml::PMLParam=PMLParam())
    s∆l = create_stretched_∆l(ω, tuple.(∆l), tuple.(l), SVec{1}.(bounds), SVec{1}.(Npml), pml)

    return (s∆l[nPR][1], s∆l[nDL][1])
end

# Wrapper for converting AbstractVector's to SVec's
create_stretched_∆l(ω::Number, grid::Grid{K}, Npml::Tuple2{AbsVecInteger}, pml::PMLParam=PMLParam()) where {K} =
    create_stretched_∆l(ω, grid, SVec{K}.(Npml), pml)

create_stretched_∆l(ω::Number, ∆l::Tuple2{NTuple{K,AbsVecReal}}, l::Tuple2{NTuple{K,AbsVecReal}},
                    bounds::Tuple2{AbsVecReal}, Npml::Tuple2{AbsVecInteger}, pml::PMLParam=PMLParam()) where {K} =
    create_stretched_∆l(ω, ∆l, l, SVec{K}.(bounds), SVec{K}.(Npml), pml)

# Create vectors containig a stretched version of a given ∆l.
create_stretched_∆l(ω::Number, grid::Grid{K}, Npml::Tuple2{SInteger{K}}, pml::PMLParam=PMLParam()) where {K} =
    create_stretched_∆l(ω, grid.∆l, grid.l, grid.bounds, Npml, pml)

function create_stretched_∆l(ω::Number,  # angular frequency
                             ∆l::Tuple2{NTuple{K,AbsVecReal}},  # grid point spacings at primal and dual grid point locations
                             l::Tuple2{NTuple{K,AbsVecReal}},  # locations of primal and dual grid points w/o ghost points
                             bounds::Tuple2{SReal{K}},  # bounds[NEG][k] = boundary of domain at (-) end in k-direction
                             Npml::Tuple2{SInteger{K}},  # Npml[NEG][k] = number of cells inside PML at (-) end in k-direction
                             pml::PMLParam=PMLParam()  # PML parameters
                             ) where {K}
    lpml, Lpml = pml_loc(l[nPR], bounds, Npml)
    sfactor = create_sfactor(ω, l, lpml, Lpml, pml)

    s∆lprim = map((x,y)->x.*y, sfactor[nPR], ∆l[nPR])
    s∆ldual = map((x,y)->x.*y, sfactor[nDL], ∆l[nDL])

    s∆l = (s∆lprim, s∆ldual)

    return s∆l
end

# Wrapper for converting AbstractVector's to SVec's
create_sfactor(ω::Number, l::Tuple2{NTuple{K,AbsVecReal}}, lpml::Tuple2{AbsVecReal},
               Lpml::Tuple2{AbsVecReal}, pml::PMLParam=PMLParam()) where {K} =
    create_sfactor(ω, l, SVec{K}.(lpml), SVec{K}.(Lpml), pml)

# Create vectors of stretching factors.
function create_sfactor(ω::Number,  # angular frequency
                        l::Tuple2{NTuple{K,AbsVecReal}},  # locations of primal and dual grid points
                        lpml::Tuple2{SReal{K}},  # locations of PML interfaces
                        Lpml::Tuple2{SReal{K}},  # thicknesses of PML
                        pml::PMLParam=PMLParam()  # PML parameters
                        ) where {K}
    N = length.(l[nPR])  # (Nx, Ny, Nz)
    sfactor = (ones.(ComplexF,N), ones.(ComplexF,N))
    for k = 1:K
        lpmlₙ, lpmlₚ = lpml[nN][k], lpml[nP][k]
        Lpmlₙ, Lpmlₚ = Lpml[nN][k], Lpml[nP][k]
        for ngt = nPD
            r = l[ngt][k]
            s = sfactor[ngt][k]
            for i = 1:N[k]
                r[i] < lpmlₙ && (s[i] = calc_sfactor(ω, lpmlₙ-r[i], Lpmlₙ, pml))
                r[i] > lpmlₚ && (s[i] = calc_sfactor(ω, r[i]-lpmlₚ, Lpmlₚ, pml))
            end
        end
    end

    return sfactor
end

# Calculate the stretch factor at a given depth into PML.
function calc_sfactor(ω::Number,  # angular frequency
                      d::Real,  # depth into PML
                      Lpml::Real,  # thickness of PML
                      pml::PMLParam)  # PML parameters
    σmax = -(pml.m+1) * log(pml.R) /2Lpml  # -(m+1) ln(R) / (2 η Lpml), where η = 1 in units of η₀
    σ = σmax * (d/Lpml)^pml.m
    κ = 1 + (pml.κmax-1) * (d/Lpml)^pml.m
    a = pml.amax * (1 - d/Lpml)^pml.ma

    s_factor = κ + σ / (a + im*ω)  # s = κ + σ/(a + i ω ε), where ε = 1 in units of ε₀
end

# Wrapper for converting AbstractVector's to SVec's
pml_loc(lprim::NTuple{K,AbsVecReal}, bounds::Tuple2{AbsVecReal}, Npml::Tuple2{AbsVecInteger}) where {K} =
    pml_loc(lprim, SVec{K}.(bounds), SVec{K}.(Npml))

# Get the PML interface locations.
# The primal grid plane locations *without* the ghost plane locations are passed, so the
# positive-end boundaries of the computational domain cannot be retrieved from them.
# Therefore, the boundary information is passed separately.
function pml_loc(lprim::NTuple{K,AbsVecReal},  # lprim[k] = primal vertex locations in k-direction (excluding ghost location at (+) end)
                 bounds::Tuple2{SReal{K}},  # bounds[NEG][k] = boundary of domain at (-) end in k-direction
                 Npml::Tuple2{SInteger{K}}  # Npml[NEG][k] = number of cells inside PML at (-) end in k-direction
                 ) where {K}
    # Below, if Npml[nP] = 0, lprim[end+1-Npml[nP]] = lprim[end+1] is out of bound.  Because
    # lprim[end+1] points to the ghost point, we return the positive-end boundary in this
    # case.
    lpml = (map((v,n)->v[1+n], lprim, Npml[nN]), map((v,b,n)->(n==0 ? b : v[end+1-n]), lprim, bounds[nP], Npml[nP]))  # Tuple2{SFloat{K}}
    Lpml = (lpml[nN]-bounds[nN], bounds[nP]-lpml[nP])  # Tuple2{SFloat{K}}

    return lpml, Lpml
end


text:
using Documenter, Optimisers

DocMeta.setdocmeta!(Optimisers, :DocTestSetup, :(using Optimisers); recursive = true)

makedocs(modules = [Optimisers],
         doctest = VERSION == v"1.6",
         sitename = "Optimisers.jl",
         pages = ["Home" => "index.md",
                  "API" => "api.md"],
         format = Documenter.HTML(
             canonical = "https://fluxml.ai/Optimisers.jl/stable/",
             analytics = "UA-36890222-9",
             assets = ["assets/flux.css"],
             prettyurls = get(ENV, "CI", nothing) == "true"
         ),
)

deploydocs(
   repo = "github.com/FluxML/Optimisers.jl.git",
   target = "build",
   push_preview = true,
)


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ImageStorageProfile <: SwaggerModel
    osDisk::Any # spec type: Union{ Nothing, ImageOSDisk } # spec name: osDisk
    dataDisks::Any # spec type: Union{ Nothing, Vector{ImageDataDisk} } # spec name: dataDisks

    function ImageStorageProfile(;osDisk=nothing, dataDisks=nothing)
        o = new()
        validate_property(ImageStorageProfile, Symbol("osDisk"), osDisk)
        setfield!(o, Symbol("osDisk"), osDisk)
        validate_property(ImageStorageProfile, Symbol("dataDisks"), dataDisks)
        setfield!(o, Symbol("dataDisks"), dataDisks)
        o
    end
end # type ImageStorageProfile

const _property_map_ImageStorageProfile = Dict{Symbol,Symbol}(Symbol("osDisk")=>Symbol("osDisk"), Symbol("dataDisks")=>Symbol("dataDisks"))
const _property_types_ImageStorageProfile = Dict{Symbol,String}(Symbol("osDisk")=>"ImageOSDisk", Symbol("dataDisks")=>"Vector{ImageDataDisk}")
Base.propertynames(::Type{ ImageStorageProfile }) = collect(keys(_property_map_ImageStorageProfile))
Swagger.property_type(::Type{ ImageStorageProfile }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_ImageStorageProfile[name]))}
Swagger.field_name(::Type{ ImageStorageProfile }, property_name::Symbol) =  _property_map_ImageStorageProfile[property_name]

function check_required(o::ImageStorageProfile)
    (getproperty(o, Symbol("osDisk")) === nothing) && (return false)
    true
end

function validate_property(::Type{ ImageStorageProfile }, name::Symbol, val)
end


text:
### parsing.jl
#
# Copyright (C) 2017 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: uncertainty, error propagation, physics, plots
#
# This file is a part of Measurements.jl.
#
# License is MIT "Expat".
#
### Commentary:
#
# This file defines the recipes to plot Measurements vectors with Plots.jl package in 2D.
#
### Code:

using RecipesBase

@recipe function f(y::AbstractArray{<:Measurement})
    yerror := uncertainty.(y)
    value.(y)
end

@recipe function f(func::Function, x::AbstractArray{<:Measurement})
    y = func.(x)
    xerror := uncertainty.(x)
    yerror := uncertainty.(y)
    value.(x), value.(y)
end

@recipe function f(x::AbstractArray{<:Measurement}, y::AbstractArray{<:Measurement})
    xerror := uncertainty.(x)
    yerror := uncertainty.(y)
    value.(x), value.(y)
end

@recipe function f(x::AbstractArray{<:Measurement}, y::AbstractArray)
    xerror := uncertainty.(x)
    value.(x), y
end

@recipe function f(x::AbstractArray, y::AbstractArray{<:Measurement})
    yerror := uncertainty.(y)
    x, value.(y)
end


text:
module ElectrodermalActivity

using CImGui
using CImGui.CSyntax
using CImGui.CSyntax.CStatic
using CImGui.GLFWBackend
using CImGui.OpenGLBackend
using CImGui.GLFWBackend.GLFW
using CImGui.OpenGLBackend.ModernGL
using Printf
using DataFrames
using CSV

include("filedialog.jl")
include("gui.jl")

export launch

end # module


text:
module GigaFV
using WriteVTK
using Printf
using Logging
using LinearAlgebra
import YAML

include("configuration.jl")
include("basis.jl")
include("equations.jl")
include("grid.jl")
include("kernels/surface.jl")
include("kernels/time.jl")
include("plotters.jl")
include("error_writer.jl")
include("global_matrices.jl")

"""
    evaluate_rhs(eq, scenario, filter, globals, du, dofs, grid)

Evalutes the right-hand-side of the equation `eq` for
scenario `scenario`, with filter `filter`,
collection of global matrices `globals`, update
`du`, degrees of freedom `dofs` and grid `grid`.

Updates `du` in place.
"""
function evaluate_rhs(eq, scenario, globals, du, dofs, grid)
    buffers_face = Array{BuffersFaceIntegral,1}(undef, Threads.nthreads())
    for i=1:Threads.nthreads()
        buffers_face[i] = BuffersFaceIntegral(get_ndofs(eq))
    end

    nvar = get_ndofs(eq)
    du .= 0.0

    # compute max eigenvalue per thread, reduce later
    maxeigenvals = Array{Float64,1}(undef, Threads.nthreads())
    maxeigenvals .= -Inf

    # Evaluate all fluxes, compute eigenvalues
    Threads.@threads for i in eachindex(grid.cells)
        @views cell = grid.cells[i]
        @views data = dofs[:, cell.dataidx]
        @views flux = grid.flux[:, :,cell.dataidx]

        cureigenval = max(
            max_eigenval(eq, dofs, 1),
            max_eigenval(eq, dofs, 1)
        )
        maxeigenvals[Threads.threadid()] = max(maxeigenvals[Threads.threadid()], cureigenval)

        evaluate_flux(eq, data, flux)
    end
    Threads.@threads for i in eachindex(grid.cells)
        @views buffer_face = buffers_face[Threads.threadid()]

        @views cell = grid.cells[i]
        @views data = dofs[:, cell.dataidx]
        @views flux = grid.flux[:,:,cell.dataidx]

        # Here we also need to compute the maximum eigenvalue of each cell
        # and store it for each cell (needed for timestep restriction later!)
        faces = (left, top, right, bottom)
        # faces = instances(Face)
        for (i, neigh) in enumerate(cell.neighbors)
            @views dofsneigh = dofs[:,neigh.dataidx]
            @views fluxneigh = grid.flux[:,:,neigh.dataidx]

            isboundary = cell.facetypes[i] == boundary
            if isboundary
                normalidx = globals.normalidxs[faces[i]]

                # For boundary cells, we operate directly on the neighbour data
                evaluate_boundary(eq, scenario, faces[i], normalidx, data, buffer_face.boundary_state)
                evaluate_flux(eq, buffer_face.boundary_state, buffer_face.boundary_flux)
            end
            @views evaluate_face_integral(eq,
                globals,
                data,
                if isboundary buffer_face.boundary_state else dofsneigh end,
                flux,
                if isboundary buffer_face.boundary_flux else fluxneigh end,
                buffer_face,
                cell,
                faces[i],
                du[:,cell.dataidx])

        end

        @views du[:,cell.dataidx] .= du[:,cell.dataidx] / volume(cell)
    end

    grid.maxeigenval = maximum(maxeigenvals)
    @show grid.maxeigenval
end

"""
    main(configfile::String)

Runs a FV-simulation with configuration from `configfile`.
"""
function main(configfile::String)
    config = Configuration(configfile)
    eq = make_equation(config)
    scenario = make_scenario(config)
    grid = make_grid(config, eq, scenario)
    integrator = make_timeintegrator(config, grid)
    globals = GlobalMatrices(2)

    filename = "output/plot"

    # Init everything
    for cell in grid.cells
        @views interpolate_initial_dofs(eq, scenario, grid.dofs[:,cell.dataidx],cell)
    end

    plotter = VTKPlotter(eq, scenario, grid, filename)

    grid.time = 0
    timestep = 0
    next_plotted = config.plot_start

    while grid.time < config.end_time
        if timestep > 0
            time_start = time()
            dt = config.cellsize[1] * config.courant * 1/grid.maxeigenval
            # Only step up to either end or next plotting
            dt = min(dt, next_plotted-grid.time, config.end_time - grid.time)
            @assert dt > 0

            @info "Running timestep" timestep dt grid.time
            step(integrator, grid, dt) do du, dofs, time
                evaluate_rhs(eq, scenario, globals, du, dofs, grid)
            end

            grid.time += dt
            time_end = time()
            time_elapsed = time_end - time_start
            @info "Timestep took" time_elapsed
        else
            # Compute initial eigenvalue (needed for dt)
            grid.maxeigenval = -1
            for cell in grid.cells
                @views celldata = grid.dofs[:,cell.dataidx]
                for normalidx=1:2
                    cureigenval = max_eigenval(eq, celldata, normalidx)
                    grid.maxeigenval = max(grid.maxeigenval, cureigenval)
                end
            end
        end

        if abs(grid.time - next_plotted) < 1e-10
            @info "Writing output" grid.time
            plot(plotter)
            next_plotted = grid.time + config.plot_step
        end
        timestep += 1
    end
    save(plotter)
    if is_analytical_solution(eq, scenario)
        evaluate_error(eq, scenario, grid, grid.time)
    end
end

end


text:
#####
# This includes short-hands that are convenient but perhaps confusing
#####

export chebyshevt,chebyshevu,legendre,∫,⨜,⨍,∇,Δ,ChebyshevWeight,𝕀,ℝ,𝕌,𝒟

## Constructors

Fun()=Fun(identity)
Fun(d::Domain)=Fun(identity,d)
Fun(d::Space)=Fun(identity,d)

## Chebyshev & Legendre polynomials

chebyshevt{T<:Number}(n::Int,d::Interval{T}) = Fun([zeros(T,n);one(T)],Chebyshev(d))
chebyshevu{T<:Number}(n::Int,d::Interval{T}) = mod(n,2) == 1 ? Fun(interlace(zeros(T,div(n+2,2)),2ones(T,div(n+2,2))),Chebyshev(d)) : Fun(interlace(2ones(T,div(n+2,2)),zeros(T,div(n+2,2)))[1:n+1]-[one(T);zeros(T,n)],Chebyshev(d))
legendre{T<:Number}(n::Int,d::Interval{T}) = Fun([zeros(T,n);one(T)],Legendre(d))

for poly in (:chebyshevt,:chebyshevu,:legendre)
    @eval begin
        $poly{T<:Number}(n::Int,a::T,b::T) = $poly(n,Interval(a,b))
        $poly{T<:Number}(::Type{T},n::Int) = $poly(n,Interval{T}())
        $poly(n::Int) = $poly(Float64,n)
        $poly{T<:Number}(n::Range,d::Interval{T}) = map(i->$poly(i,d),n)
    end
end

ChebyshevWeight(d,k)=k==0?JacobiWeight(-0.5,-0.5,d):JacobiWeight(0.5,0.5,d)
ChebyshevWeight(d)=ChebyshevWeight(d,0)
ChebyshevWeight(k::Integer)=ChebyshevWeight(Interval(),k)
ChebyshevWeight()=ChebyshevWeight(0)

# shorthand for second order

ivp(d) = ivp(d,2)
bvp(d) = bvp(d,2)

## diff

# use conj(f.') for ArraySpace
Base.ctranspose(f::Fun)=differentiate(f)

∫(f::Fun)=integrate(f)
⨜(f::Fun)=cumsum(f)

for OP in (:Σ,:∮,:⨍,:⨎)
    @eval $OP(f::Fun)=sum(f)
end

∇(F::MultivariateFun) = grad(F)
Base.dot{M<:MultivariateFun}(∇::Function,F::Vector{M}) = div(F)
Base.cross{M<:MultivariateFun}(∇::Function,F::Vector{M}) = curl(F)

## Domains

const 𝕀 = Interval()
const ℝ = Line()
const 𝕌 = Circle()

𝒟 = Derivative()
Δ = Laplacian()


text:
# add one unit to the power limit
# SP+
export psp_incrementplimit

const INCREMENT_PLIMIT = [0x53, 0x49, 0x2B, 0x0D]

"""
    psp_incrementplimit(io)

Add one unit to the power limit.
"""
function psp_incrementplimit(io_psp::IO)
  write(io_psp, INCREMENT_PLIMIT)
  return nothing
end


text:
mutable struct ROCK2ConstantCache{T,T2,zType} <: OrdinaryDiffEqConstantCache
  ms::SVector{46, Int}
  fp1::SVector{46, T}
  fp2::SVector{46, T}
  recf::Vector{T2}
  zprev::zType
  mdegprev::Int
  mdeg::Int
  recind::Int
end
@cache struct ROCK2Cache{uType,rateType,uNoUnitsType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  gprev::uType
  gprev2::uType
  tmp::uType
  atmp::uNoUnitsType
  fsalfirst::rateType
  k::rateType
  k2::rateType
  constantcache::ROCK2ConstantCache
end

function alg_cache(alg::ROCK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  constantcache = ROCK2ConstantCache(uEltypeNoUnits, uEltypeNoUnits, u) # WIP: not sure about what type to use in here
  gprev = similar(u)
  gprev2 = similar(u)
  tmp = similar(u)
  atmp = similar(u,uEltypeNoUnits)
  fsalfirst = zero(rate_prototype)
  k = zero(rate_prototype)
  k2 = zero(rate_prototype)
  ROCK2Cache(u, uprev, gprev, gprev2, tmp, atmp, fsalfirst, k, k2, constantcache)
end

function alg_cache(alg::ROCK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  ROCK2ConstantCache(uEltypeNoUnits, uEltypeNoUnits, u) # WIP: not sure about what type to use in here
end

mutable struct ROCK4ConstantCache{T,T2,T3,T4,zType} <: OrdinaryDiffEqConstantCache
  ms::SVector{50, Int}
  fpa::Vector{T}
  fpb::Vector{T2}
  fpbe::Vector{T3}
  recf::Vector{T4}
  zprev::zType
  mdegprev::Int
  mdeg::Int
  recind::Int
end
@cache struct ROCK4Cache{uType,rateType,uNoUnitsType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  gprev::uType
  gprev2::uType
  gprev3::uType
  gprev4::uType
  gprev5::uType
  tmp::uType
  atmp::uNoUnitsType
  fsalfirst::rateType
  k::rateType
  k2::rateType
  k3::rateType
  k4::rateType
  k5::rateType
  constantcache::ROCK4ConstantCache
end

function alg_cache(alg::ROCK4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  constantcache = ROCK4ConstantCache(real(uEltypeNoUnits), real(uEltypeNoUnits), u)
  gprev = similar(u)
  gprev2 = similar(u)
  gprev3 = similar(u)
  gprev4 = similar(u)
  gprev5 = similar(u)
  tmp = similar(u)
  atmp = similar(u,uEltypeNoUnits)
  fsalfirst = zero(rate_prototype)
  k = zero(rate_prototype)
  k2 = zero(rate_prototype)
  k3 = zero(rate_prototype)
  k4 = zero(rate_prototype)
  k5 = zero(rate_prototype)
  ROCK4Cache(u, uprev, gprev, gprev2, gprev3, gprev4, gprev5, tmp, atmp, fsalfirst, k, k2, k3, k4, k5, constantcache)
end

function alg_cache(alg::ROCK4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  ROCK4ConstantCache(real(uEltypeNoUnits), real(uEltypeNoUnits),u)
end

mutable struct RKCConstantCache{zType} <: OrdinaryDiffEqConstantCache
  #to match the types to call maxeig!
  zprev::zType
end
@cache struct RKCCache{uType,rateType,uNoUnitsType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  gprev::uType
  gprev2::uType
  tmp::uType
  atmp::uNoUnitsType
  fsalfirst::rateType
  k::rateType
  constantcache::RKCConstantCache
end

function alg_cache(alg::RKC,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  constantcache = RKCConstantCache(u)
  gprev = similar(u)
  gprev2 = similar(u)
  tmp = similar(u)
  atmp = similar(u,uEltypeNoUnits)
  fsalfirst = zero(rate_prototype)
  k = zero(rate_prototype)
  RKCCache(u, uprev, gprev, gprev2, tmp, atmp, fsalfirst, k, constantcache)
end

function alg_cache(alg::RKC,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  RKCConstantCache(u)
end

@cache mutable struct IRKCConstantCache{uType,rateType,F,N} <: OrdinaryDiffEqConstantCache
  minm::Int64
  zprev::uType
  k2::rateType
  uf::F
  nlsolver::N
  du₁::rateType
  du₂::rateType
end

@cache mutable struct IRKCCache{uType,rateType,uNoUnitsType,JType,WType,UF,JC,N,F} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  gprev::uType
  gprev2::uType
  fsalfirst::rateType
  k::rateType
  du1::rateType
  f1ⱼ₋₁::rateType
  f1ⱼ₋₂::rateType
  f2ⱼ₋₁::rateType
  z::uType
  dz::uType
  b::uType
  tmp::uType
  atmp::uNoUnitsType
  J::JType
  W::WType
  uf::UF
  jac_config::JC
  linsolve::F
  nlsolver::N
  du₁::rateType
  du₂::rateType
  constantcache::IRKCConstantCache
end

function alg_cache(alg::IRKC,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  γ, c = 1//1, 1//1
  @oopnlsolve
  zprev = u
  k2  = rate_prototype
  du₁ = rate_prototype; du₂ = rate_prototype
  IRKCConstantCache(50,zprev,k2,uf,nlsolver,du₁,du₂)
end

function alg_cache(alg::IRKC,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  γ, c = 1//1, 1//1
  @iipnlsolve

  gprev = similar(u)
  gprev2 = similar(u)
  tmp = similar(u)
  atmp = similar(u,uEltypeNoUnits)
  fsalfirst = zero(rate_prototype)
  k  = zero(rate_prototype)
  k2 = zero(rate_prototype)
  zprev = similar(u)
  f1ⱼ₋₁ = zero(rate_prototype)
  f1ⱼ₋₂ = zero(rate_prototype)
  f2ⱼ₋₁ = zero(rate_prototype)
  du₁ = zero(rate_prototype)
  du₂ = zero(rate_prototype)
  constantcache = IRKCConstantCache(50,zprev,k2,uf,nlsolver,du₁,du₂)
  IRKCCache(u,uprev,gprev,gprev2,fsalfirst,k,du1,f1ⱼ₋₁,f1ⱼ₋₂,f2ⱼ₋₁,z,dz,b,tmp,atmp,J,W,uf,jac_config,linsolve,nlsolver,du₁,du₂,constantcache)
end

mutable struct ESERK5ConstantCache{T, zType} <: OrdinaryDiffEqConstantCache
  ms::SVector{49, Int}
  Cᵤ::SVector{5, Int}
  Cₑ::SVector{5, Int}
  zprev::zType
  Bᵢ::Vector{T}
  mdeg::Int
  start::Int
  internal_deg::Int
end

@cache struct ESERK5Cache{uType,rateType,uNoUnitsType} <: OrdinaryDiffEqMutableCache
  u::uType
  uprev::uType
  uᵢ::uType
  uᵢ₋₁::uType
  uᵢ₋₂::uType
  Sᵢ::uType
  tmp::uType
  atmp::uNoUnitsType
  fsalfirst::rateType
  k::rateType
  constantcache::ESERK5ConstantCache
end

function alg_cache(alg::ESERK5,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  constantcache = ESERK5ConstantCache(u)
  uᵢ = similar(u)
  uᵢ₋₁ = similar(u)
  uᵢ₋₂ = similar(u)
  Sᵢ   = similar(u)
  tmp = similar(u)
  atmp = similar(u,uEltypeNoUnits)
  fsalfirst = zero(rate_prototype)
  k = zero(rate_prototype)
  ESERK5Cache(u, uprev, uᵢ, uᵢ₋₁, uᵢ₋₂, Sᵢ, tmp, atmp, fsalfirst, k, constantcache)
end

function alg_cache(alg::ESERK5,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  ESERK5ConstantCache(u)
end


text:
# Regularized decomposition
# ------------------------------------------------------------
@with_kw mutable struct RDData{T <: AbstractFloat}
    Q̃::T = 1e10
    σ::T = 1.0
    incumbent::Int = 1
    major_iterations::Int = 0
    minor_iterations::Int = 0
end

@with_kw mutable struct RDParameters{T <: AbstractFloat}
    τ::T = 1e-6
    γ::T = 0.9
    σ::T = 1.0
    σ̅::T = 4.0
    σ̲::T = 0.5
end

"""
    RegularizedDecomposition

Functor object for using regularized decomposition regularization in an L-shaped algorithm. Create by supplying an [`RD`](@ref) object through `regularize ` in the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model.

...
# Parameters
- `σ::AbstractFloat = 1.0`: Initial value of regularization parameter. Controls the relative penalty of the deviation from the current major iterate.
- `σ̅::AbstractFloat = 4.0`: Maximum value of the regularization parameter.
- `σ̲::AbstractFloat = 0.5`: Minimum value of the regularization parameter.
- `log::Bool = true`: Specifices if L-shaped procedure should be logged on standard output or not.
- `penaltyterm::PenaltyTerm = Quadratic`: Specify penaltyterm variant ([`Quadratic`](@ref), [`Linearized`](@ref), [`InfNorm`](@ref), [`ManhattanNorm`][@ref])
...
"""
struct RegularizedDecomposition{T <: AbstractFloat, A <: AbstractVector, PT <: PenaltyTerm} <: AbstractRegularization
    data::RDData{T}
    parameters::RDParameters{T}

    ξ::A
    Q̃_history::A
    σ_history::A
    incumbents::Vector{Int}

    penaltyterm::PT

    function RegularizedDecomposition(ξ₀::AbstractVector, penaltyterm::PenaltyTerm; kw...)
        T = promote_type(eltype(ξ₀), Float32)
        ξ₀_ = convert(AbstractVector{T}, copy(ξ₀))
        A = typeof(ξ₀_)
        PT = typeof(penaltyterm)
        return new{T, A, PT}(RDData{T}(), RDParameters{T}(;kw...), ξ₀_, A(), A(), Vector{Int}(), penaltyterm)
    end
end

function initialize_regularization!(lshaped::AbstractLShapedSolver, rd::RegularizedDecomposition)
    rd.data.σ = rd.parameters.σ
    push!(rd.σ_history,rd.data.σ)
    # Initialize and update penalty
    initialize_penaltyterm!(rd.penaltyterm, lshaped.mastersolver, lshaped.x)
    c = vcat(get_obj(lshaped), active_model_objectives(lshaped))
    update_penaltyterm!(rd.penaltyterm, lshaped.mastersolver, c, 1/rd.data.σ, rd.ξ)
    return nothing
end

function log_regularization!(lshaped::AbstractLShapedSolver, rd::RegularizedDecomposition)
    @unpack Q̃,σ,incumbent = rd.data
    push!(rd.Q̃_history, Q̃)
    push!(rd.σ_history, σ)
    push!(rd.incumbents, incumbent)
    return nothing
end

function log_regularization!(lshaped::AbstractLShapedSolver, t::Integer, rd::RegularizedDecomposition)
    @unpack Q̃,σ,incumbent = rd.data
    rd.Q̃_history[t] = Q̃
    rd.σ_history[t] = σ
    rd.incumbents[t] =  incumbent
    return nothing
end

function take_step!(lshaped::AbstractLShapedSolver, rd::RegularizedDecomposition)
    @unpack Q,θ = lshaped.data
    @unpack τ = lshaped.parameters
    @unpack σ = rd.data
    @unpack γ,σ̅,σ̲ = rd.parameters
    t = timestamp(lshaped)
    σ̃ = incumbent_trustregion(lshaped, t, rd)
    Q̃ = incumbent_objective(lshaped, t, rd)
    need_update = false
    if abs(θ-Q) <= τ*(1+abs(θ)) || Q <= Q̃ + τ || rd.data.major_iterations == 0
        rd.ξ .= current_decision(lshaped)
        rd.data.Q̃ = copy(Q)
        need_update = true
        rd.data.incumbent = t
        rd.data.major_iterations += 1
    else
        rd.data.minor_iterations += 1
    end
    new_σ = if Q + τ <= (1-γ)*Q̃ + γ*θ
        max(σ, min(σ̅, 2*σ))
    elseif Q - τ >= γ*Q̃ + (1-γ)*θ
        min(σ, max(σ̲, 0.5*σ))
    else
        σ
    end
    if abs(new_σ - σ) > τ
        need_update = true
    end
    rd.data.σ = new_σ
    if need_update
        c = vcat(get_obj(lshaped), active_model_objectives(lshaped))
        update_penaltyterm!(rd.penaltyterm, lshaped.mastersolver, c, 1/rd.data.σ, rd.ξ)
    end
    return nothing
end

function solve_regularized_master!(lshaped::AbstractLShapedSolver, solver::LQSolver, rd::RegularizedDecomposition)
    solve_penalized!(rd.penaltyterm, lshaped.mastersolver, lshaped.mastervector, lshaped.x, rd.ξ)
    return nothing
end

# API
# ------------------------------------------------------------
"""
    RD

Factory object for [`RegularizedDecomposition`](@ref). Pass to `regularize ` in the `LShapedSolver` factory function. Equivalent factory calls: `RD`, `WithRD`, `RegularizedDecomposition`, `WithRegularizedDecomposition`. See ?RegularizedDecomposition for parameter descriptions.

"""
struct RD <: AbstractRegularizer
    penaltyterm::PenaltyTerm
    parameters::Dict{Symbol,Any}
end
RD(; penaltyterm = Quadratic(), kw...) = RD(penaltyterm, Dict{Symbol,Any}(kw))
WithRD(; penaltyterm = Quadratic(), kw...) = RD(penaltyterm, Dict{Symbol,Any}(kw))
RegularizedDecomposition(; penaltyterm = Quadratic(), kw...) = RD(penaltyterm, Dict{Symbol,Any}(kw))
WithRegularizedDecomposition(; penaltyterm = Quadratic(), kw...) = RD(penaltyterm, Dict{Symbol,Any}(kw))

function add_regularization_params!(regularizer::RD; kwargs...)
    push!(regularizer.parameters, kwargs...)
    for (k,v) in kwargs
        if k == :penaltyterm
            setfield!(regularizer, k, v)
            delete!(regularizer.parameters, k)
        end
    end
    return nothing
end

function (rd::RD)(x::AbstractVector)
    return RegularizedDecomposition(x, rd.penaltyterm; rd.parameters...)
end

function str(::RD)
    return "L-shaped using regularized decomposition"
end


text:
function module_dir()
    @__DIR__
end

using Pkg
Pkg.activate(module_dir())
using Dojo


text:
module ElectrodeCalculator

using Logging, LoggingExtras

function main()
  Base.eval(Main, :(const UserApp = ElectrodeCalculator))

  include(joinpath("..", "genie.jl"))

  Base.eval(Main, :(const Genie = ElectrodeCalculator.Genie))
  Base.eval(Main, :(using Genie))
end; main()

end


text:
"""
    bernoulli(n, x)

 Calculates Bernoulli polynomials ``B_n(x)``
 e.g., see

 + [https://en.wikipedia.org/wiki/Bernoulli_polynomials](https://en.wikipedia.org/wiki/Bernoulli_polynomials)
 + [http://dlmf.nist.gov/24](http://dlmf.nist.gov/24)

## Arguments
* ``n`` `::Integer`: the index into the series, ``n=0,1,2,3,...``
* ``x`` `::Real`: the point at which to calculate the polynomial

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> bernoulli(6, 1.2)
0.008833523809524735
```
"""
function bernoulli(n::Integer, x::Real)
    if n<0
        throw(DomainError(n))
    end
    if n == 0
        return 1 # zeta formula doesn't hold for n=0, so return explicit value
    elseif n == 1 # get some easy cases out of the way quickly
        return x-0.5
    elseif n == 2
        return x^2 - x + 1.0/6.0
    elseif n == 3
        return x^3 - 1.5*x^2 + 0.5*x
    elseif n == 4
        return x^4 - 2.0*x^3 +     x^2 - 1/30.0
    elseif n == 5
         return x^5 - 2.5*x^4 +(5.0/3.0)*x^3 - x/6.0
    end

 # from the Hurwitz-zeta function using

 # ``ζ(-n,x) = -B_{n+1}(x)/(n+1), for Re(x)>0,``

 # which is faster than direct calculation of the polynomial except for n<=4, but for
 # small n (n<=5) we use the exact polynomials. For negative x, we use the recursive
 # formula to push it into a postive range.
     # return -n*SpecialFunctions.zeta(1-n, x)
    # see https://carma.newcastle.edu.au/resources/jon/Preprints/Papers/Published-InPress/Oscillatory%20(Tapas%20II)/Papers/coffey-zeta.pdf, p.341
    # however, I am not really happy with the results I am getting out of this
    # so exploit symmetries in B_n(x) to compute recursively for x<=0
    # if x >= 0
    #      return -n*SpecialFunctions.zeta(1-n, x)
    # else
    #     return bernoulli(n, x+1) - n*x^(n-1)
    # end

    # total = 0.0
    # for k = 0:n-1
    #     total -= binomial(n+1, k) *  bernoulli(k,x)
    # end
    # total = total / (n+1)
    # total += x^n
    # return total

    # alternative, a little better, but relies on bernoulli numbers, which I only have to 35
    # and still has some non-negligable errors, so I still would like a better piece in here
    n = Int64(n)
    if n <= 35
        total = 0.0
        for k = 0:n
            k = Int64(k)
            total += binomial(n, k) *  bernoulli(n-k) * x^k
            # this will overflow for large n and k, but that migth be avoidable if we were more clever?
        end
        return total
    else
        # this is faster but less accurate, and runs into some weird problems, e.g.,
        # SpecialFunctions.zeta(-6, -1.0) = NaN (probably should be 2)
        return -n*SpecialFunctions.zeta(1-n, x)
    end
end


text:
abstract type AbstractTensor{S <: Tuple, T, N} <: AbstractArray{T, N} end

# aliases (too long name?)
const AbstractSecondOrderTensor{dim, T} = AbstractTensor{NTuple{2, dim}, T, 2}
const AbstractFourthOrderTensor{dim, T} = AbstractTensor{NTuple{4, dim}, T, 4}
const AbstractSymmetricSecondOrderTensor{dim, T} = AbstractTensor{Tuple{@Symmetry{dim, dim}}, T, 2}
const AbstractSymmetricFourthOrderTensor{dim, T} = AbstractTensor{NTuple{2, @Symmetry{dim, dim}}, T, 4}
const AbstractVec{dim, T} = AbstractTensor{Tuple{dim}, T, 1}
const AbstractMat{m, n, T} = AbstractTensor{Tuple{m, n}, T, 2}

const AbstractMatLike{T} = Union{
    AbstractMat{<: Any, <: Any, T},
    AbstractSymmetricSecondOrderTensor{<: Any, T},
    Transpose{T, <: AbstractVec{<: Any, T}},
}
const AbstractVecOrMatLike{T} = Union{AbstractVec{<: Any, T}, AbstractMatLike{T}}

# special name (not exported)
const AbstractSquareTensor{dim, T} = Union{AbstractTensor{Tuple{dim, dim}, T, 2},
                                           AbstractTensor{Tuple{@Symmetry{dim, dim}}, T, 2}}

# for AbstractArray interface
Base.IndexStyle(::Type{<: AbstractTensor}) = IndexLinear()

@pure Base.size(::Type{TT}) where {TT <: AbstractTensor} = tensorsize(Space(TT))
@inline function Base.size(TT::Type{<: AbstractTensor}, d::Int)
    S = size(TT)
    d > length(S) ? 1 : S[d]
end
@inline Base.size(x::AbstractTensor) = tensorsize(Space(x))
@pure Base.length(::Type{TT}) where {TT <: AbstractTensor} = prod(size(TT))

@pure Base.ndims(::Type{TT}) where {TT <: AbstractTensor} = length(size(TT))

@pure Base.axes(::Type{TT}) where {TT <: AbstractTensor} = map(Base.OneTo, size(TT))
function Base.axes(TT::Type{<: AbstractTensor}, d::Int)
    A = axes(TT)
    d > length(A) ? Base.OneTo(1) : A[d]
end

# indices
for Indices in (:CartesianIndices, :LinearIndices)
    @eval Base.$Indices(::Type{TT}) where {TT <: AbstractTensor} = $Indices(size(TT))
end
for func in (:independent_indices, :indices, :duplicates)
    @eval begin
        $func(::Type{TT}) where {S, TT <: AbstractTensor{S}} = $func(Space(S))
        $func(x::AbstractTensor) = $func(typeof(x))
    end
end

Space(::Type{TT}) where {S, TT <: AbstractTensor{S}} = Space(S)
Space(::AbstractTensor{S}) where {S} = Space(S)

# getindex_expr
function getindex_expr(x::Type{<: AbstractTensor}, ex::Union{Symbol, Expr}, i...)
    inds = independent_indices(x)
    :(Tuple($ex)[$(inds[i...])])
end

@generated function Base.getindex(x::AbstractTensor, I::Int...)
    ex = :()
    stride = 1
    for i in 1:length(I)
        if i == 1
            ex = :(I[1])
        else
            ex = :($ex + $stride * (I[$i] - 1))
        end
        stride *= size(x, i)
    end
    quote
        @_inline_meta
        @boundscheck checkbounds(x, I...)
        @inbounds x[$ex] # call getindex(x::AbstractTensor, i::Int)
    end
end

function Base.getindex(x::AbstractTensor, inds::Union{Int, StaticArray{<:Tuple, Int}, SOneTo, Colon}...)
    @_propagate_inbounds_meta
    Tensor(SArray(x)[inds...])
end

# to SArray
@generated function StaticArrays.SArray(x::AbstractTensor)
    exps = [getindex_expr(x, :x, i) for i in 1:length(x)]
    quote
        @_inline_meta
        @inbounds SArray{Tuple{$(size(x)...)}}($(exps...))
    end
end
function StaticArrays.SArray(x::Transpose{<: T, <: AbstractVec{dim, T}}) where {dim, T}
    SMatrix{1, dim}(Tuple(parent(x)))
end

# vcat
Base.vcat(a::AbstractVecOrMatLike) = a
Base.vcat(a::AbstractVecOrMatLike, b::AbstractVecOrMatLike) = Tensor(vcat(SArray(a), SArray(b)))
Base.vcat(a::AbstractVecOrMatLike, b::AbstractVecOrMatLike, c::AbstractVecOrMatLike...) = vcat(vcat(a, b), vcat(c...))
# hcat
Base.hcat(a::AbstractVec) = Tensor(hcat(SArray(a)))
Base.hcat(a::AbstractMatLike) = a
Base.hcat(a::AbstractVecOrMatLike, b::AbstractVecOrMatLike) = Tensor(hcat(SArray(a), SArray(b)))
Base.hcat(a::AbstractVecOrMatLike, b::AbstractVecOrMatLike, c::AbstractVecOrMatLike...) = hcat(hcat(a, b), hcat(c...))
for op in (:vcat, :hcat)
    _op = Symbol(:_, op)
    @eval @generated function $_op(xs::Union{AbstractVecOrMatLike, Real}...)
        exps = [xs[i] <: Real ? :(Vec(xs[$i])) : :(xs[$i]) for i in 1:length(xs)]
        quote
            $($op)($(exps...))
        end
    end
    for I in 0:10
        xs = map(i -> :($(Symbol(:x, i))::Real), 1:I)
        @eval Base.$op($(xs...), y::AbstractVecOrMatLike, zs::Union{AbstractVecOrMatLike, Real}...) = $_op($(xs...), y, zs...)
    end
end

if VERSION ≥ v"1.6"
    # reverse
    Base.reverse(x::AbstractTensor; dims = :) = Tensor(reverse(SArray(x); dims))
else
    Base.reverse(x::AbstractTensor) = Tensor(reverse(SArray(x)))
end


text:
packages = [
    "PyCall",
    "PyPlot",
    "JSON",
    "Docile",
    "ODE"
    ]

Pkg.init()
Pkg.update()

for package in packages
    Pkg.add(package)
end

max_stars_repo_name: novakgroupoxford/XPPjl

max_issues_repo_name: novakgroupoxford/XPPjl

max_forks_repo_name: novakgroupoxford/XPPjl


text:
"""
convals.jl
"""

"""
	ConVal{C,V,M,W}

Holds information about a constraint
"""
struct ConVal{S,C,Tc,Tic}
    con::C
    # constraint function value
    c::Tc
    # dual
    λ::Tc
    # penalty multiplier
    μ::Tc
    # active constraints
    a::Tc
    # index for this conval's `c` in
    # the global linearized constraint
    c_ginds::Tic
end

# constructors
function ConVal(con::AbstractConstraint{S}, n::Int, m::Int, c_ginds::Tic,
                M, V) where {S,Tic}
    p = length(con)
    c = V(zeros(p))
    λ = V(zeros(p))
    μ = V(fill(con.params.μ0, p))
    a = V(zeros(Bool, p))
    Tc = typeof(c)
    C = typeof(con)
    return ConVal{S,C,Tc,Tic}(con, c, λ, μ, a, c_ginds)
end

# methods
function update_active!(conval::ConVal{EQUALITY})
    conval.a .= true
    return nothing
end

function update_active!(conval::ConVal{INEQUALITY})
    a = conval.a
    c = conval.c
    λ = conval.λ
    a_tol = conval.con.params.a_tol
    con = conval.con
    for i = 1:length(a)
        a[i] = ((c[i] >= a_tol) | (abs(λ[i]) > a_tol))
    end
    return nothing
end

@inline violation(con::AbstractConstraint{EQUALITY}, c::AbstractVector) = (
    norm(c, Inf)
)

@inline violation(con::AbstractConstraint{INEQUALITY}, c::AbstractVector) = (
    max(0., maximum(c))
)

function update_dual_penalty!(convals::Vector{Vector{ConVal}})
    for (k, convals_) in enumerate(convals)
        for conval in convals_
            # update dual
            λ_max = conval.con.params.λ_max
            λ_min = sense(conval.con) == EQUALITY ? -λ_max : zero(λ_max)
            for i in eachindex(conval.λ)
                conval.λ[i] = clamp(conval.λ[i] + conval.μ[i] * conval.c[i], λ_min, λ_max)
            end
            # update penalty
            for i in eachindex(conval.μ)
                conval.μ[i] = clamp(conval.con.params.ϕ * conval.μ[i], 0, conval.con.params.μ_max)
            end
        end
    end
    return nothing
end

function max_violation_penalty(convals::Vector{Vector{ConVal}})
    max_violation = 0.
    max_penalty = 0.
    for (k, convals_) in enumerate(convals)
        for conval in convals_
            viol = violation(conval.con, conval.c)
            max_violation = max(max_violation, viol)
            max_penalty = max(max_penalty, maximum(conval.μ))
        end
    end
    return max_violation, max_penalty
end

function max_violation_info(convals::Vector{Vector{ConVal}})
    max_viol = -Inf
    info_str = ""
    for (k, convals_) in enumerate(convals)
        for conval in convals_
            max_viol_, info_str_ = max_violation_info(conval.con, conval.c, k)
            if max_viol_ > max_viol
                max_viol = max_viol_
                info_str = info_str_
            end
        end
    end
    return max_viol, info_str
end

# build list of convals from constraint list
function convals_from_constraint_list(cons::ConstraintList)
    # indices for tracking global concatenated constraint
    c_gind = 0
    convals = Vector{ConVal}[]
    # build a list of convals at each knot point
    for k = 1:cons.N
        convals_ = ConVal[]
        # iterate over each constraint
        for i in 1:length(cons)
            con = cons.constraints[i]
            knot_points = cons.inds[i]
            # if this constraint is active at the current knot point,
            # add a conval for it to the current list
            if k in knot_points
                p = length(con)
                c_ginds = cons.V((1:p) .+ c_gind)
                c_gind += p
                conval = ConVal(con, cons.n, cons.m, c_ginds,
                                cons.M, cons.V)
                push!(convals_, conval)
            end
        end
        push!(convals, convals_)
    end
    return convals
end