text:
using Test
using CLIMA.MoistThermodynamics
using NCDatasets
using Random
MT = MoistThermodynamics
using CLIMA.PlanetParameters
using LinearAlgebra

using CLIMAParameters

float_types = [Float32, Float64]

include("data_tests.jl")

@testset "moist thermodynamics - isentropic processes" begin
    for FT in [Float64]
        # for FT in float_types
        e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice =
            MT.tested_convergence_range(50, FT)
        Φ = FT(1)
        Random.seed!(15)
        perturbation = FT(0.1) * rand(length(T))

        # TODO: Use reasonable values for ambient temperature/pressure
        T∞, p∞ = T .* perturbation, p .* perturbation
        @test air_temperature.(p, θ_liq_ice, Ref(DryAdiabaticProcess())) ≈
              (p ./ FT(MSLP)) .^ (FT(R_d) / FT(cp_d)) .* θ_liq_ice
        @test air_pressure_given_θ.(θ_liq_ice, Φ, Ref(DryAdiabaticProcess())) ≈
              FT(MSLP) .*
              (1 .- Φ ./ (θ_liq_ice .* FT(cp_d))) .^ (FT(cp_d) / FT(R_d))
        @test air_pressure.(T, T∞, p∞, Ref(DryAdiabaticProcess())) ≈
              p∞ .* (T ./ T∞) .^ (FT(1) / FT(kappa_d))
    end
end

@testset "moist thermodynamics - correctness" begin
    FT = Float64
    # ideal gas law
    @test air_pressure(FT(1), FT(1), PhasePartition(FT(1))) === FT(R_v)
    @test air_pressure(FT(1), FT(2), PhasePartition(FT(1), FT(0.5), FT(0))) ===
          FT(R_v)
    @test air_pressure(FT(1), FT(1)) === FT(R_d)
    @test air_pressure(FT(1), FT(2)) === 2 * FT(R_d)
    @test air_density(FT(1), FT(1)) === 1 / FT(R_d)
    @test air_density(FT(1), FT(2)) === 2 / FT(R_d)

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

    @test cp_m(PhasePartition(FT(0))) === FT(cp_d)
    @test cp_m(PhasePartition(FT(1))) === FT(cp_v)
    @test cp_m(PhasePartition(FT(1), FT(1))) === FT(cp_l)
    @test cp_m(PhasePartition(FT(1), FT(0), FT(1))) === FT(cp_i)
    @test cp_m(FT) == FT(cp_d)

    @test cv_m(PhasePartition(FT(0))) === FT(cp_d - R_d)
    @test cv_m(PhasePartition(FT(1))) === FT(cp_v - R_v)
    @test cv_m(PhasePartition(FT(1), FT(1))) === FT(cv_l)
    @test cv_m(PhasePartition(FT(1), FT(0), FT(1))) === FT(cv_i)
    @test cv_m(FT) == FT(cv_d)

    # speed of sound
    @test soundspeed_air(T_0 + 20, PhasePartition(FT(0))) ==
          sqrt(cp_d / cv_d * R_d * (T_0 + 20))
    @test soundspeed_air(T_0 + 100, PhasePartition(FT(1))) ==
          sqrt(cp_v / cv_v * R_v * (T_0 + 100))

    # specific latent heats
    @test latent_heat_vapor(FT(T_0)) ≈ LH_v0
    @test latent_heat_fusion(FT(T_0)) ≈ LH_f0
    @test latent_heat_sublim(FT(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(FT(T_triple), Liquid()) ≈ press_triple
    @test saturation_vapor_pressure(FT(T_triple), Ice()) ≈ press_triple

    @test q_vap_saturation(FT(T_triple), ρ, PhasePartition(FT(0))) ==
          ρ_v_triple / ρ
    @test q_vap_saturation(FT(T_triple), ρ, PhasePartition(q_tot, q_tot)) ==
          ρ_v_triple / ρ

    @test q_vap_saturation_generic(FT(T_triple), ρ; phase = Liquid()) ==
          ρ_v_triple / ρ
    @test q_vap_saturation_generic(FT(T_triple), ρ; phase = Ice()) ==
          ρ_v_triple / ρ
    @test q_vap_saturation_generic.(FT(T_triple - 20), ρ; phase = Liquid()) >=
          q_vap_saturation_generic.(FT(T_triple - 20), ρ; phase = Ice())

    @test saturation_excess(FT(T_triple), ρ, PhasePartition(q_tot)) ==
          q_tot - ρ_v_triple / ρ
    @test saturation_excess(FT(T_triple), ρ, PhasePartition(q_tot / 1000)) ==
          0.0

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

    @test air_temperature(
        cv_m(PhasePartition(FT(0))) * (T - T_0),
        PhasePartition(FT(0)),
    ) === FT(T)
    @test air_temperature(
        cv_m(PhasePartition(FT(q_tot))) * (T - T_0) + q_tot * e_int_v0,
        PhasePartition(q_tot),
    ) ≈ FT(T)

    @test total_energy(FT(e_kin), FT(e_pot), FT(T_0)) === FT(e_kin) + FT(e_pot)
    @test total_energy(FT(e_kin), FT(e_pot), FT(T)) ≈
          FT(e_kin) + FT(e_pot) + cv_d * (T - T_0)
    @test total_energy(FT(0), FT(0), FT(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)
    @test liquid_fraction(T) === FT(0)
    @test liquid_fraction(T, PhasePartition(q_tot, q_liq, q_liq)) === FT(0.5)
    q = PhasePartition_equil(T, ρ, q_tot)
    @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(T) === FT(1)
    @test liquid_fraction(T, PhasePartition(q_tot, q_liq, q_liq / 2)) ===
          FT(2 / 3)
    q = PhasePartition_equil(T, ρ, q_tot)
    @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
    tol_T = 1e-1
    q_tot = FT(0)
    ρ = FT(1)
    @test MT.saturation_adjustment_SecantMethod(
        internal_energy_sat(300.0, ρ, q_tot),
        ρ,
        q_tot,
        10,
        1e-2,
    ) ≈ 300.0
    @test abs(
        MT.saturation_adjustment(
            internal_energy_sat(300.0, ρ, q_tot),
            ρ,
            q_tot,
            10,
            1e-2,
        ) - 300.0,
    ) < tol_T

    q_tot = FT(0.21)
    ρ = FT(0.1)
    @test MT.saturation_adjustment_SecantMethod(
        internal_energy_sat(200.0, ρ, q_tot),
        ρ,
        q_tot,
        10,
        1e-2,
    ) ≈ 200.0
    @test abs(
        MT.saturation_adjustment(
            internal_energy_sat(200.0, ρ, q_tot),
            ρ,
            q_tot,
            10,
            1e-2,
        ) - 200.0,
    ) < tol_T
    q = PhasePartition_equil(T, ρ, q_tot)
    @test q.tot - q.liq - q.ice ≈
          vapor_specific_humidity(q) ≈
          q_vap_saturation(T, ρ)

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

    # potential temperatures
    T = FT(300)
    @test liquid_ice_pottemp_given_pressure(T, FT(MSLP)) === T
    @test liquid_ice_pottemp_given_pressure(T, FT(MSLP) / 10) ≈
          T * 10^(R_d / cp_d)
    @test liquid_ice_pottemp_given_pressure(
        T,
        FT(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 dry_pottemp_given_pressure(T, p, PhasePartition(q_tot)) isa typeof(p)
    @test air_temperature_from_liquid_ice_pottemp_given_pressure(
        dry_pottemp_given_pressure(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 exner_given_pressure(p, PhasePartition(q_tot)) isa typeof(p)
end

@testset "moist 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

    for FT in float_types
        rtol = FT(1e-2)
        e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice =
            MT.tested_convergence_range(50, FT)

        # PhaseEquil
        ts_exact = PhaseEquil.(e_int, ρ, q_tot, 100, FT(1e-4))
        ts = PhaseEquil.(e_int, ρ, q_tot)
        # Should be machine accurate (because ts contains `e_int`,`ρ`,`q_tot`):
        @test all(
            getproperty.(PhasePartition.(ts), :tot) .≈
            getproperty.(PhasePartition.(ts_exact), :tot),
        )
        @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,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol,
        ))

        # PhaseEquil
        ts_exact =
            PhaseEquil.(
                e_int,
                ρ,
                q_tot,
                100,
                FT(1e-4),
                MT.saturation_adjustment_SecantMethod,
            )
        ts =
            PhaseEquil.(
                e_int,
                ρ,
                q_tot,
                30,
                FT(1e-1),
                MT.saturation_adjustment_SecantMethod,
            ) # Needs to be in sync with default
        # Should be machine accurate (because ts contains `e_int`,`ρ`,`q_tot`):
        @test all(
            getproperty.(PhasePartition.(ts), :tot) .≈
            getproperty.(PhasePartition.(ts_exact), :tot),
        )
        @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,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol,
        ))

        # LiquidIcePotTempSHumEquil
        ts_exact = LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot, 40, FT(1e-3))
        ts = LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot)
        # Should be machine accurate:
        @test all(air_density.(ts) .≈ air_density.(ts_exact))
        @test all(
            getproperty.(PhasePartition.(ts), :tot) .≈
            getproperty.(PhasePartition.(ts_exact), :tot),
        )
        # Approximate (temperature must be computed via saturation adjustment):
        @test all(isapprox.(
            internal_energy.(ts),
            internal_energy.(ts_exact),
            rtol = rtol,
        ))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            liquid_ice_pottemp.(ts_exact),
            rtol = rtol,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol,
        ))

        # LiquidIcePotTempSHumEquil_given_pressure
        ts_exact =
            LiquidIcePotTempSHumEquil_given_pressure.(
                θ_liq_ice,
                p,
                q_tot,
                40,
                FT(1e-3),
            )
        ts = LiquidIcePotTempSHumEquil_given_pressure.(θ_liq_ice, p, q_tot)
        # Should be machine accurate:
        @test all(
            getproperty.(PhasePartition.(ts), :tot) .≈
            getproperty.(PhasePartition.(ts_exact), :tot),
        )
        # Approximate (temperature must be computed via saturation adjustment):
        @test all(isapprox.(
            air_density.(ts),
            air_density.(ts_exact),
            rtol = rtol,
        ))
        @test all(isapprox.(
            internal_energy.(ts),
            internal_energy.(ts_exact),
            rtol = rtol,
        ))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            liquid_ice_pottemp.(ts_exact),
            rtol = rtol,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol,
        ))

        # LiquidIcePotTempSHumNonEquil
        ts_exact =
            LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt, 40, FT(1e-3))
        ts = LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt)
        # Should be machine accurate:
        @test all(
            getproperty.(PhasePartition.(ts), :tot) .≈
            getproperty.(PhasePartition.(ts_exact), :tot),
        )
        @test all(
            getproperty.(PhasePartition.(ts), :liq) .≈
            getproperty.(PhasePartition.(ts_exact), :liq),
        )
        @test all(
            getproperty.(PhasePartition.(ts), :ice) .≈
            getproperty.(PhasePartition.(ts_exact), :ice),
        )
        @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),
            rtol = rtol,
        ))
        @test all(isapprox.(
            liquid_ice_pottemp.(ts),
            liquid_ice_pottemp.(ts_exact),
            rtol = rtol,
        ))
        @test all(isapprox.(
            air_temperature.(ts),
            air_temperature.(ts_exact),
            rtol = rtol,
        ))

    end

end

@testset "moist thermodynamics - constructor consistency" begin

    # Make sure `ThermodynamicState` arguments are returned unchanged

    for FT in float_types
        rtol = FT(1e-2)
        e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice =
            MT.tested_convergence_range(50, FT)

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

        ts = PhaseDry_given_pT.(p, T)
        @test all(internal_energy.(ts) .≈ internal_energy.(T))
        @test all(air_density.(ts) .≈ ρ)

        # PhaseEquil
        ts =
            PhaseEquil.(
                e_int,
                ρ,
                q_tot,
                30,
                FT(1e-1),
                Ref(MT.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.(e_int, ρ, q_tot)
        @test all(internal_energy.(ts) .≈ e_int)
        @test all(getproperty.(PhasePartition.(ts), :tot) .≈ q_tot)
        @test all(air_density.(ts) .≈ ρ)

        # PhaseNonEquil
        ts = PhaseNonEquil.(e_int, ρ, q_pt)
        @test all(internal_energy.(ts) .≈ e_int)
        @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),
        )
        @test all(air_density.(ts) .≈ ρ)

        # air_temperature_from_liquid_ice_pottemp_given_pressure-liquid_ice_pottemp inverse
        θ_liq_ice_ = liquid_ice_pottemp_given_pressure.(T, p, q_pt)
        @test all(
            air_temperature_from_liquid_ice_pottemp_given_pressure.(
                θ_liq_ice_,
                p,
                q_pt,
            ) .≈ T,
        )

        # liquid_ice_pottemp-air_temperature_from_liquid_ice_pottemp_given_pressure inverse
        T =
            air_temperature_from_liquid_ice_pottemp_given_pressure.(
                θ_liq_ice,
                p,
                q_pt,
            )
        @test all(liquid_ice_pottemp_given_pressure.(T, p, q_pt) .≈ θ_liq_ice)

        # Accurate but expensive `LiquidIcePotTempSHumNonEquil` constructor (Non-linear temperature from θ_liq_ice)
        T_non_linear =
            air_temperature_from_liquid_ice_pottemp_non_linear.(
                θ_liq_ice,
                ρ,
                10,
                FT(1e-3),
                q_pt,
            )
        T_expansion =
            air_temperature_from_liquid_ice_pottemp.(θ_liq_ice, ρ, q_pt)
        @test all(isapprox.(T_non_linear, T_expansion, rtol = rtol))
        e_int_ = internal_energy.(T_non_linear, q_pt)
        ts = PhaseNonEquil.(e_int_, ρ, q_pt)
        @test all(T_non_linear .≈ air_temperature.(ts))
        @test all(θ_liq_ice .≈ liquid_ice_pottemp.(ts))

        # LiquidIcePotTempSHumEquil
        ts = LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot, 40, FT(1e-3))
        @test all(isapprox.(liquid_ice_pottemp.(ts), θ_liq_ice, atol = 1e-1))
        @test all(isapprox.(air_density.(ts), ρ, rtol = rtol))
        @test all(getproperty.(PhasePartition.(ts), :tot) .≈ q_tot)

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

        # LiquidIcePotTempSHumEquil_given_pressure
        ts =
            LiquidIcePotTempSHumEquil_given_pressure.(
                θ_liq_ice,
                p,
                q_tot,
                40,
                FT(1e-3),
            )
        @test all(isapprox.(liquid_ice_pottemp.(ts), θ_liq_ice, atol = 1e-1))
        @test all(
            getproperty.(PhasePartition.(ts), :tot) .≈ getproperty.(q_pt, :tot),
        )
        @test all(isapprox.(air_pressure.(ts), p, atol = FT(MSLP) * 2e-2))

        # LiquidIcePotTempSHumNonEquil_given_pressure
        ts = LiquidIcePotTempSHumNonEquil_given_pressure.(θ_liq_ice, p, q_pt)
        @test all(liquid_ice_pottemp.(ts) .≈ θ_liq_ice)
        @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),
        )

        # LiquidIcePotTempSHumNonEquil
        ts = LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt, 5, FT(1e-3))
        @test all(θ_liq_ice .≈ liquid_ice_pottemp.(ts))
        @test all(air_density.(ts) .≈ ρ)
        @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),
        )
    end

end

@testset "moist 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.
    FT = Float32
    e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(50, FT)

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

    ts_dry = PhaseDry.(e_int, ρ)
    ts_dry_pT = PhaseDry_given_pT.(p, T)
    ts_eq = PhaseEquil.(e_int, ρ, q_tot, 15, FT(1e-1))
    ts_T =
        TemperatureSHumEquil.(
            air_temperature.(ts_dry),
            air_pressure.(ts_dry),
            q_tot,
        )
    ts_neq = PhaseNonEquil.(e_int, ρ, q_pt)
    ts_θ_liq_ice_eq =
        LiquidIcePotTempSHumEquil.(θ_liq_ice, ρ, q_tot, 40, FT(1e-3))
    ts_θ_liq_ice_eq_p =
        LiquidIcePotTempSHumEquil_given_pressure.(
            θ_liq_ice,
            p,
            q_tot,
            40,
            FT(1e-3),
        )
    ts_θ_liq_ice_neq = LiquidIcePotTempSHumNonEquil.(θ_liq_ice, ρ, q_pt)
    ts_θ_liq_ice_neq_p =
        LiquidIcePotTempSHumNonEquil_given_pressure.(θ_liq_ice, p, q_pt)

    for ts in (
        ts_dry,
        ts_dry_pT,
        ts_eq,
        ts_T,
        ts_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.(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.(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.(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.(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.(virtual_pottemp.(ts)) == typeof.(e_int)
        @test eltype.(gas_constants.(ts)) == typeof.(e_int)
        @test typeof.(getproperty.(PhasePartition.(ts), :tot)) == typeof.(e_int)
    end

end

@testset "moist thermodynamics - dry limit" begin

    FT = Float64
    e_int, ρ, q_tot, q_pt, T, p, θ_liq_ice = MT.tested_convergence_range(50, FT)

    # PhasePartition test is noisy, so do this only once:
    ts_dry = PhaseDry(first(e_int), first(ρ))
    ts_eq = PhaseEquil(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.(e_int, ρ)
    ts_eq = PhaseEquil.(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(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(soundspeed_air.(ts_eq) .≈ soundspeed_air.(ts_dry))
    @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(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, Ref(Ice())) .≈
        saturation_vapor_pressure.(ts_dry, Ref(Ice())),
    )
    @test all(
        saturation_vapor_pressure.(ts_eq, Ref(Liquid())) .≈
        saturation_vapor_pressure.(ts_dry, Ref(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

max_stars_repo_name: yairchn/CLIMA

max_issues_repo_name: yairchn/CLIMA

max_forks_repo_name: yairchn/CLIMA


text:
export DataSource, LabeledDataSource, InMemoryLabeledDataSource
export EncodedInMemoryLabeledDataSource, DataFrameLabeledDataSource
export nobs, nvar, features, targets, groundtruth, bias, splitTrainTest!
export dataSource, encodeDataSource

using ArrayViews
using DataFrames

import StatsBase.nobs
import Base.shuffle!

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

abstract DataSource

nobs(source::DataSource) = @_not_implemented
nvar(source::DataSource) = @_not_implemented
features(source::DataSource) = @_not_implemented
features(source::DataSource, offset::Int, length::Int) = @_not_implemented

# labeled sources

abstract LabeledDataSource <: DataSource
abstract InMemoryLabeledDataSource <: LabeledDataSource

groundtruth(source::LabeledDataSource) = @_not_implemented
targets(source::LabeledDataSource) = @_not_implemented
targets(source::LabeledDataSource, offset::Int, length::Int) = @_not_implemented
bias(source::LabeledDataSource) = @_not_implemented
nclasses(source::LabeledDataSource) = @_not_implemented
labels(source::LabeledDataSource) = @_not_implemented
classDistribution(source::LabeledDataSource) = @_not_implemented

# ==========================================================================
# In-memory labeled sources

immutable EncodedInMemoryLabeledDataSource{E<:ClassEncoding,G<:Any,N} <: InMemoryLabeledDataSource
  features::AbstractArray{Float64,2}
  targets::AbstractArray{Float64,N}
  groundtruth::AbstractArray{G,1}
  encoding::E
  bias::Float64

  function EncodedInMemoryLabeledDataSource(features::AbstractArray{Float64,2},
                                            targets::AbstractArray{Float64,1},
                                            groundtruth::AbstractArray{G,1},
                                            encoding::E,
                                            bias::Float64)
    typeof(encoding) <: OneOfKClassEncoding && throw(ArgumentError("Can't have OneOutOfK-Encoding with a vector as target"))
    size(features, 2) == length(targets) || throw(DimensionMismatch("Features and targets have to have the same number of observations"))
    new(features, targets, groundtruth, encoding, bias)
  end

  function EncodedInMemoryLabeledDataSource(features::AbstractArray{Float64,2},
                                            targets::AbstractArray{Float64,2},
                                            groundtruth::AbstractArray{G,1},
                                            encoding::OneOfKClassEncoding,
                                            bias::Float64)
    nclasses(encoding) == size(targets, 1) || throw(DimensionMismatch("Targets have to have the same number of rows as the encoding has labels"))
    size(features, 2) == size(targets, 2) || throw(DimensionMismatch("Features and targets have to have the same number of observations"))
    new(features, targets, groundtruth, encoding, bias)
  end
end

function EncodedInMemoryLabeledDataSource{E<:ClassEncoding,G<:Any}(
    features::AbstractArray{Float64,2},
    targets::AbstractArray{Float64,1},
    groundtruth::AbstractArray{G,1},
    encoding::E,
    bias::Float64 = 1.)
  EncodedInMemoryLabeledDataSource{E,G,1}(features, targets, groundtruth, encoding, bias)
end

function EncodedInMemoryLabeledDataSource{E<:ClassEncoding}(
    features::AbstractArray{Float64,2},
    targets::AbstractArray{Float64,1},
    encoding::E,
    bias::Float64 = 1.)
  typeof(encoding) <: OneOfKClassEncoding && throw(ArgumentError("Can't have OneOutOfK-Encoding with a vector as target"))
  size(features, 2) == length(targets) || throw(DimensionMismatch("Features and targets have to have the same number of observations"))
  EncodedInMemoryLabeledDataSource(features, targets, labeldecode(encoding, targets), encoding, bias)
end

function EncodedInMemoryLabeledDataSource{G<:Any}(
    features::AbstractArray{Float64,2},
    targets::AbstractArray{Float64,2},
    groundtruth::AbstractArray{G,1},
    encoding::OneOfKClassEncoding,
    bias::Float64 = 1.)
  EncodedInMemoryLabeledDataSource{OneOfKClassEncoding,G,2}(features, targets, groundtruth, encoding, bias)
end

function EncodedInMemoryLabeledDataSource(
    features::AbstractArray{Float64,2},
    targets::AbstractArray{Float64,2},
    encoding::OneOfKClassEncoding,
    bias::Float64 = 1.)
  nclasses(encoding) == size(targets, 1) || throw(DimensionMismatch("Targets have to have the same number of rows as the encoding has labels"))
  size(features, 2) == size(targets, 2) || throw(DimensionMismatch("Features and targets have to have the same number of observations"))
  EncodedInMemoryLabeledDataSource(features, targets, labeldecode(encoding, targets), encoding, bias)
end

groundtruth{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  source.groundtruth

nobs{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  size(source.features, 2)

nvar{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  size(source.features, 1)

features{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  source.features

features{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}, offset::Int, length::Int) =
  view(source.features, :, offset:(offset+length-1))

targets{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  source.targets

targets{E<:ClassEncoding,G}(source::EncodedInMemoryLabeledDataSource{E,G,1}, offset::Int, length::Int) =
  view(source.targets, offset:(offset + length - 1))

targets{E<:ClassEncoding,G}(source::EncodedInMemoryLabeledDataSource{E,G,2}, offset::Int, length::Int) =
  view(source.targets, :, offset:(offset + length - 1)
       )
bias{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  source.bias

nclasses{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  nclasses(source.encoding)

labels{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  labels(source.encoding)

classDistribution{E<:ClassEncoding,G,N}(source::EncodedInMemoryLabeledDataSource{E,G,N}) =
  classDistribution(source.encoding, labeldecode(source.encoding, source.targets))

function shuffle!{G}(X::Array{Float64,2}, t::Array{Float64,1}, g::Array{G,1})
  rows = size(X, 1)
  cols = size(X, 2)
  for c = 1:cols
    i = rand(c:cols)
    for r = 1:rows
      @inbounds X[r,c], X[r,i] = X[r,i], X[r,c]
    end
    @inbounds t[c], t[i] = t[i], t[c]
    @inbounds g[c], g[i] = g[i], g[c]
  end
  nothing
end

function shuffle!{G}(X::Array{Float64,2}, t::Array{Float64,2}, g::Array{G,1})
  rows = size(X, 1)
  cols = size(X, 2)
  rowsT = size(t, 1)
  for c = 1:cols
    i = rand(c:cols)
    for r = 1:rows
      @inbounds X[r,c], X[r,i] = X[r,i], X[r,c]
    end
    for r = 1:rowsT
      @inbounds t[r,c], t[r,i] = t[r,i], t[r,c]
    end
    @inbounds g[c], g[i] = g[i], g[c]
  end
  nothing
end

function splitTrainTest!{E<:ClassEncoding,G,N}(
    source::EncodedInMemoryLabeledDataSource{E,G,N};
    p_train::FloatingPoint = .7)
  @assert 0. < p_train < 1.
  X = source.features
  t = source.targets
  g = source.groundtruth
  shuffle!(X, t, g)
  ce = source.encoding
  bias = source.bias
  n = nobs(source)
  sn = safeFloor(n * p_train)
  if N == 1
    trainData = dataSource(view(X, :, 1:sn),
                           view(t, 1:sn),
                           view(g, 1:sn),
                           ce, bias = bias)
    testData = dataSource(view(X, :, (sn+1):n),
                          view(t, (sn+1):n),
                          view(g, (sn+1):n),
                          ce, bias = bias)
    return trainData, testData
  else
    trainData = dataSource(view(X, :, 1:sn),
                           view(t, :, 1:sn),
                           view(g, 1:sn),
                           ce, bias = bias)
    testData = dataSource(view(X, :, (sn+1):n),
                          view(t, :, (sn+1):n),
                          view(g, (sn+1):n),
                          ce, bias = bias)
    return trainData, testData
  end
end

# ==========================================================================
# DataFrame labeled sources

immutable DataFrameLabeledDataSource <: LabeledDataSource
  dataFrame::DataFrame
  formula::Formula
end

# ==========================================================================
# Encode a DataFrameLabeledDataSource

function encodeDataSource{E<:ClassEncoding}(source::DataFrameLabeledDataSource, ::Type{E})
  mf = ModelFrame(source.formula, source.dataFrame)
  mm = ModelMatrix(mf)
  dfBias = in(0, mm.assign)
  X = dfBias ? mm.m[:,2:end]: mm.m
  extBias = dfBias * 1.0
  t = convert(Vector{Float64}, model_response(mf))
  ce = E(t)
  t_enc = labelencode(ce, t)
  dataSource(X', t_enc, ce, bias=extBias)
end

# ==========================================================================
# Choose best DataSource for the parameters

function dataSource{E<:ClassEncoding, G, N}(
    features::AbstractArray{Float64,2},
    targets::AbstractArray{Float64,N},
    groundtruth::AbstractArray{G,1},
    encoding::E;
    bias::Float64 = 1.)
  EncodedInMemoryLabeledDataSource(features, targets, groundtruth, encoding, bias)
end

function dataSource{E<:ClassEncoding, N}(
    features::AbstractArray{Float64,2},
    targets::AbstractArray{Float64,N},
    encoding::E;
    bias::Float64 = 1.)
  EncodedInMemoryLabeledDataSource(features, targets, encoding, bias)
end

function dataSource(formula::Formula, data::DataFrame)
  DataFrameLabeledDataSource(data, formula)
end

function dataSource{E<:ClassEncoding}(formula::Formula, data::DataFrame, ::Type{E})
  encodeDataSource(DataFrameLabeledDataSource(data, formula), E)
end

ext: jl


text:
module Phi0
function f(q::Vector{T}, q_dot::Vector{T}, xi::T) where T
    l1 = q[1]
    l2 = q[2]
    l3 = q[3]
    l1_dot = q_dot[1]
    l2_dot = q_dot[2]
    l3_dot = q_dot[3]

    [-160351958.315909*xi.^10.*(2*l1 - l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^4 + 5073636.18108931*xi.^8.*(2*l1 - l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^3 - 99887.2123151958*xi.^6.*(2*l1 - l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^2 + 1053.49794238683*xi.^4.*(2*l1 - l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2) - 4.44444444444444*xi.^2.*(2*l1 - l2 - l3).*(l1 + l2 + l3 + 0.45), -160351958.315909*sqrt(3)*xi.^10.*(l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^4 + 5073636.18108931*sqrt(3)*xi.^8.*(l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^3 - 99887.2123151958*sqrt(3)*xi.^6.*(l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^2 + 1053.49794238683*sqrt(3)*xi.^4.*(2*l1 - l2 - l3).*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2) - 4.44444444444444*sqrt(3)*xi.^2.*(l2 - l3).*(l1 + l2 + l3 + 0.45), 60131984.3684659*xi.^9.*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^4 - 1522090.85432679*xi.^7.*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^3 + 22474.6227709191*xi.^5.*(l1 + l2 + l3 + 0.45).*(l1.^2 - l1.*l2 - l1.*l3 + l2.^2 - l2.*l3 + l3.^2).^2 - 79.0123456790123*xi.^3.*(l1 + l2 + l3 + 0.45).*(2*l1.^2 - 2*l1.*l2 - 2*l1.*l3 + 2*l2.^2 - 2*l2.*l3 + 2*l3.^2) + xi.*(l1 + l2 + l3 + 0.45)/3]
end
end


text:
@testset "C functions {$ST}" begin

 v_ptr = GridapDistributedPETScWrappers.C.VecCreate(ST)
 chk(GridapDistributedPETScWrappers.C.VecSetType(v_ptr, GridapDistributedPETScWrappers.C.VECMPI))
 b = Vec{ST}(v_ptr)
 resize!(b, mlocal=sys_size)
 global_indices = first(localpart(b))-1:last(localpart(b))-1  # zero based
  for i=1:sys_size
    idxm = [global_indices[i] ]
    val = [ rhs[i] ]
    GridapDistributedPETScWrappers.C.SetValues(b.p, idxm, val)
  end

  GridapDistributedPETScWrappers.C.AssemblyBegin(b.p)
  GridapDistributedPETScWrappers.C.AssemblyEnd(b.p)

   for i=1:sys_size
    idxm = [global_indices[i] ]
    val = ST[ 0.0 ]
    GridapDistributedPETScWrappers.C.GetValues(b.p, idxm, val)
    @test val[1] == rhs[i]
  end

  # matrix
  A = Mat(ST, mlocal=sys_size, nlocal=sys_size)
  row_range, col_range = localranges(A)

  for i=1:sys_size
    for j=1:sys_size
      row_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[i] - 1 ]
      col_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[j] - 1 ]
      val = ST[ A_julia[i, j] ]
      GridapDistributedPETScWrappers.C.SetValues(A.p, row_idx, col_idx, val)
    end
  end

  GridapDistributedPETScWrappers.C.AssemblyBegin(A.p)
  GridapDistributedPETScWrappers.C.AssemblyEnd(A.p)

  for i=1:sys_size
    for j=1:sys_size
      row_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[i] - 1 ]
      col_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[j] - 1 ]
      val = ST[ 0.0 ]
      GridapDistributedPETScWrappers.C.GetValues(A.p, row_idx, col_idx, val)
      @test val[1] == A_julia[i, j]
    end
  end

  idx = collect(first(row_range)-1:last(row_range)-1)
  vals = collect(1:(sys_size*sys_size))
  vals2 = convert(Array{ST, 1}, vals)
  GridapDistributedPETScWrappers.C.SetValues(A.p, idx, idx, vals2)
  GridapDistributedPETScWrappers.C.AssemblyBegin(A.p)
  GridapDistributedPETScWrappers.C.AssemblyEnd(A.p)

  for i=1:sys_size
    for j=1:sys_size
      row_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[i] - 1 ]
      col_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[j] - 1 ]
      val = ST[ 0.0 ]
      GridapDistributedPETScWrappers.C.GetValues(A.p, row_idx, col_idx, val)
      @test val[1] == (i + (j-1)*sys_size)
    end
  end

  # block matrix
  B = Mat(ST, mlocal=sys_size, nlocal=sys_size, bs=sys_size, mtype=GridapDistributedPETScWrappers.C.MATMPIBAIJ)
  idx = GridapDistributedPETScWrappers.C.PetscInt[comm_rank]
  GridapDistributedPETScWrappers.C.SetValues(B.p, idx, idx, vals2)
  GridapDistributedPETScWrappers.C.AssemblyBegin(B.p)
  GridapDistributedPETScWrappers.C.AssemblyEnd(B.p)
  for i=1:sys_size
    for j=1:sys_size
      row_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[i] - 1 ]
      col_idx = GridapDistributedPETScWrappers.C.PetscInt[ row_range[j] - 1 ]
      val = ST[ 0.0 ]
      GridapDistributedPETScWrappers.C.GetValues(A.p, row_idx, col_idx, val)
      @test val[1] == (i + (j-1)*sys_size)
    end
  end

  # shell matrix
  # if ST == Float64
  #   ctx = (1, 2, 3)
  #   ctx_ptr = pointer_from_objref(ctx)
  #   c_ptr = GridapDistributedPETScWrappers.C.MatCreateShell(sys_size, sys_size, GridapDistributedPETScWrappers.C.PETSC_DETERMINE, GridapDistributedPETScWrappers.C.PETSC_DETERMINE, ctx_ptr)
  #   C = Mat{ST}(c_ptr)
  #
  #
  #   f_ptr = cfunction(mymult, GridapDistributedPETScWrappers.C.PetscErrorCode, (GridapDistributedPETScWrappers.C.Mat{ST}, GridapDistributedPETScWrappers.C.Vec{ST}, GridapDistributedPETScWrappers.C.Vec{ST}))
  #   GridapDistributedPETScWrappers.C.MatShellSetOperation(C.p, GridapDistributedPETScWrappers.C.MATOP_MULT, f_ptr)
  #
  #   ctx_ret = GridapDistributedPETScWrappers.C.MatShellGetContext(C.p)
  #   @test unsafe_pointer_to_objref(ctx_ret) == ctx
  #
  #   x = Vec(ST[1:sys_size;])
  #   xlocal = LocalVector(x)
  #   b = Vec(zeros(ST, sys_size))
  #   *(C, x, b)
  #   gc()  # avoid a finalizer problem
  #   blocal = LocalVector(b)
  #   for i=1:length(blocal)
  #     @test blocal[i] == ST(i*i)
  #   end
  #
  #   restore(blocal)
  # end
end


text:
using Documenter, BartIO

makedocs(
         sitename = "BartIO.jl",
         modules  = [BartIO],
         pages=[
                "Home" => "index.md"
               ])

deploydocs(;
               repo="github.com/aTrotier/BartIO.jl",
               versions = nothing
           )


text:
# joRadon - Radon transform (experimental)

## helper module
module joRadon_etc
    using FFTW
    using JOLI: jo_convert
    # linear and parabolic Radon transform and its adjoint.
    #
    # Tristan van Leeuwen, 2012 (tleeuwen@eos.ubc.ca)
    # Modified by Henryk Modzelewski (Nov 2018)
    #
    # use:
    #   out = fwd_lpradon(input,t,h,q,power)
    #   out = adj_lpradon(input,t,h,q,power)
    #
    # input:
    #   input - vectorized matrix of size (length(t) x length(h)) for forward, (length(t) x length(q)) for adjoint)
    #   t     - time vector in seconds
    #   h     - offset vecror in meters
    #   q     - radon parameter
    #   power - 1: linear radon, 2: parabolic radon
    #   mode  - 1: forward, -1: adjoint
    #
    # output:            println(size(L))
    #   vectorized matrix of size (length(t) x length(q)) for forward, (length(t) x length(h)) for adjoint)
    #
    #
    function fwd_lpradon(in::Vector{idt}, t::Vector, h::Vector, q::Vector, power::Integer, rdt::DataType) where idt<:Real

        (idt<:Real && rdt<:Real) || throw(joLinearFunctionException("joRadon: input and output must be real"))

        dt   = t[2] - t[1]
        tN   = length(t)
        hN   = length(h)
        qN   = length(q)
        fftN = 2*nextpow(2,tN)
        two::idt = 2

        input = reshape(in,tN,hN)
        h=jo_convert(idt,h)
        q=jo_convert(idt,q)

        input_padded = Matrix{idt}(undef,fftN, hN)
        input_padded[1:tN,:] = input[:,:]
        input_padded[tN+1:fftN,:] .= zero(idt)

        input_padded  = fft(input_padded, 1)
        out = zeros(Complex{rdt}, fftN, qN)

        for k = 2:Int(floor(fftN/2))
            f = two*pi*(k-1)/fftN/dt
            L = exp.(im*f*(h.^power)*q')
            tmp = transpose(input_padded[k,:])*L
            out[k,:] = tmp
            out[fftN+2-k,:] = conj(tmp)
        end
        out = real(ifft(out, 1))
        out = jo_convert(rdt,vec(out[1:tN,:]))

        return out
    end
    function adj_lpradon(in::Vector{idt}, t::Vector, h::Vector, q::Vector, power::Integer, rdt::DataType) where idt<:Real

        (idt<:Real && rdt<:Real) || throw(joLinearFunctionException("joRadon: input and output must be real"))

        dt   = t[2] - t[1]
        tN   = length(t)
        hN   = length(h)
        qN   = length(q)
        fftN = 2*nextpow(2,tN)
        two::idt = 2

        input = reshape(in,tN,qN)
        h=jo_convert(idt,h)
        q=jo_convert(idt,q)

        input_padded = Matrix{idt}(undef, fftN, qN)
        input_padded[1:tN,:] = input[:,:]
        input_padded[tN+1:fftN,:] .= zero(idt)

        input_padded  = fft(input_padded,1)
        out = zeros(Complex{rdt}, fftN, hN)

        for k = 2:Int(floor(fftN/2))
            f = two*pi*(k-1)/fftN/dt
            L = exp.(im*f*(h.^power)*q')
            tmp = transpose(input_padded[k,:])*L'
            out[k,:] = tmp
            out[fftN+2-k,:] = conj(tmp)
        end

        out = real(ifft(out,1))
        out = jo_convert(rdt,vec(out[1:tN,:]))

        return out
    end
end
using .joRadon_etc

export joRadon
"""
    julia> joRadon(t,h,q,power)

2d linear and parabolic Radon transform

# Signature

    joRadon(t::Vector,h::Vector,q::Vector,power::Integer;DDT::DataType=joComplex,RDT::DataType=DDT)

# Arguments

- `t`: time vector [s]
- `h`: offset vector [m]
- `q`: Radon parameter
- `power`: 1: linear radon, 2: parabolic radon

# Notes

- joRadon is dimensionalized; it is not standard units-independent operator
- 2D array must be in vectorized form of array(t,h) of size (length(t),length(q))

"""
function joRadon(t::Vector,h::Vector,q::Vector,power::Integer;DDT::DataType=joComplex,RDT::DataType=DDT)
    (DDT<:Real && RDT<:Real) || throw(joLinearFunctionException("joRadon: domain and range must be real"))
    (power>0 && power<3) || throw(joLinearFunctionException("joRadon: power can only be either 1 or 2"))
    nt=length(t)
    nh=length(h)
    nq=length(q)
    return joLinearFunctionFwd_A(nt*nq,nt*nh,
        v1->joRadon_etc.fwd_lpradon(vec(v1),t,h,q,power,RDT),
        v2->joRadon_etc.adj_lpradon(vec(v2),t,h,q,power,DDT),
        DDT,RDT;
        name="joRadon")
end


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

"""
Test a given sequence of actions (for example, NMAC rule sequence).
Returns fitness score as well as the expr.
Deprecated: use ExprSearch.RefExprSearch and GrammarExpts.ACASX_Ref
"""
module ACASX_Test_Seq

export setup, playsequence, nmacrule

using ACASXProblem
using DerivationTrees

const NMAC_SEQUENCE = Int64[2, 2, 7, 10, 7, 7, 5, 3, 10]
#results in :(F((D[:,76] .< 100) & (D[:,77] .< 500))) which is 38 chars

function setup(; runtype::Symbol=:nmacs_vs_nonnmacs,
               data::AbstractString="dasc",
               manuals::AbstractString="dasc_manual",
               clusterdataname::AbstractString="",

               maxsteps::Int64=20)

  problem = ACASXClustering(runtype, data, manuals, clusterdataname)

  grammar = get_grammar(problem)
  tree_params = DerivTreeParams(grammar, maxsteps)
  tree = DerivationTree(tree_params)

  return problem, tree
end

function playsequence{T}(problem::ACASXClustering{T}, tree::DerivationTree,
                      sequence::Vector{Int64}=NMAC_SEQUENCE)
  play!(tree, sequence)
  expr = get_expr(tree)
  fitness = get_fitness(problem, expr)

  return tree, fitness, expr
end

function nmacrule()
  problem, tree = setup()
  tree, fitness, expr = playsequence(problem, tree)
  @show fitness
  @show expr
  return tree, fitness, expr
end

end #module


text:
using Cassette
import Base:~
import Distributions.logpdf
export sample, unif, pointwise, <|, rt, Ω, logpdf
const ID = NTuple{N, Int} where N

# Random variables are functions from a sample space to some realization space.
# They are either primitive, or pointwise functions of primitives, i.e., derived
# The sample space object is a mapping from identifiers to outputs of primitives

"Sample space object"
struct Ω{T}
  data::T
end

# A projection of ω to some index id is simply ω[id]
# ΩProj is a symbolic representation of this projection
# It's useful to maintain a reference to the parent ω

"id'th element of ω, with refeence to parent"
struct ΩProj{OM}
  ω::OM
  id::ID
end

"Project `ω` onto id"
proj(ω, id) = ΩProj(ω, id)

# # Sampling
# "Sample a random ω ∈ Ω"
# sample(::Type{Ω{T}}) where T = Ω(T())
# sample(f) = f(sample(Ω))

# Primitives

# FIXME: get! should be omega type dependent
function (T::Type{<:Distributions.Distribution})(ωπ::ΩProj, args...)
  get!(ωπ.ω.data, ωπ.id, rand(T(args...)))
end

(T::Type{<:Distributions.Distribution})(ω::Ω, args...) =
  T(proj(ω, (1,)), args...)

# Conditioning
struct Observation{X, V}
  x::X
  v::V
end
observe(x, v) = Observation(x, v)
Base.:|(x, y::Observation) = Observation()

# log pdf
"Wraps val, and is mutable"
mutable struct Wrapper{T}
  val::T
  seen::Set{ID}
end
Cassette.@context LogPdfCtx

function logpdf(rv, ω)
  # FIXME: Don't double count
  ld = Wrapper(0.0, Set{ID}()) # log 1
  ctx = LogPdfCtx(metadata = ld)
  Cassette.overdub(ctx, rv, ω)
  ld.val
end

function Cassette.posthook(ctx::LogPdfCtx, ret, T::Type{<:Distributions.Distribution}, ωπ::ΩProj, args...)
  # @show ωπ.id
  # println("")
  if ωπ.id ∉ ctx.metadata.seen
    # @show ret
    # @show args
    @show ctx.metadata.val += Distributions.logpdf(T(args...), ret)
    push!(ctx.metadata.seen, ωπ.id)
  end
  println("")
  ret
end

# This is necessary to compose contexts, in particular logpdf and hadctx
function Cassette.posthook(ctx::LogPdfCtx, ret, ::typeof(Cassette.overdub), ictx, T::Type{<:Distributions.Distribution}, ωπ::ΩProj, args...)
  println("")?
  if ωπ.id ∉ ctx.metadata.seen
    # @show ret
    # @show args
    @show ctx.metadata.val += Distributions.logpdf(T(args...), ret)
    push!(ctx.metadata.seen, ωπ.id)
  end
  println("")
  ret
end

# (Conditional) independence
Cassette.@context CIIDCtx # Use cassette to augment enivonrment with extra state

"""Conditionally independent and identically distributed given `shared` parents

Returns the `id`th element of an (exchangeable) sequence of random variables
that are identically distributed with `f` but conditionally independent given
random variables in `shared`.

This is essentially constructs a plate where all variables in shared are shared,
and all other parents are not.
"""
function ciid(f, id::ID, shared)
  let ctx = CIIDCtx(metadata = (shared = shared, id = id))
    ω -> withctx(ctx, f, ω)
  end
end

withctx(ctx, f, ω) = Cassette.overdub(ctx, f, ω)
function Cassette.overdub(ctx::CIIDCtx, ::typeof(withctx), ctxinner, f, ω)
  # Merge the scope
  id = (ctxinner.metadata.id..., ctx.metadata.id...)
  shared = (ctxinner.metadata.shared..., ctx.metadata.shared...)
  Cassette.overdub(CIIDCtx(metadata = (shared = shared, id = id)), f, ω)
end

"I.I.D.: `ciid` with nothing shar ed"
ciid(f, id::ID) = ciid(f, id, ())
ciid(f, id::Integer) = ciid(f, (id,))
ciid(f, id::Integer, shared) = ciid(f, (id,), shared)

function Cassette.overdub(ctx::CIIDCtx, T::Type{<:Distributions.Distribution}, ω::Ω, args...)
  @show proj(ω, ctx.metadata.id)
  T(proj(ω, ctx.metadata.id), args...)
end

function Cassette.overdub(ctx::CIIDCtx, x, ω::Ω)
  # @show x
  # @show ctx.metadata.shared
  # TODO: This is using julia function equivalence which is a bit shaky eh
  if x in ctx.metadata.shared
    x(ω)
  else
    Cassette.recurse(ctx, x, ω)
  end
end

# Syntactic Sugar (to make model-building nicer)

"Random tuple"
rt(fs...) = ω -> map(f -> f(ω), fs)

"""Supports notation `i ~ x <| (y,z,)`
which is the ith element of an (exchangeable) sequence of random variables that are
identically distributed with x but conditionally independent given y and z.
"""
struct Plate{F, S}
  f::F
  shared::S
end

@inline <|(f, shared::Tuple) = Plate(f, shared)
~(id, f) = ciid(f, id)
~(id, plate::Plate) = ciid(plate.f, id, plate.shared)

# Pointwise
Cassette.@context PWCtx
Lifted = Union{map(typeof, (+, -, /, *))...}
Cassette.overdub(::PWCtx, op::Lifted, x::Function) = ω -> op(x(ω))
Cassette.overdub(::PWCtx, op::Lifted, x::Function, y::Function) = ω -> op(x(ω), y(ω))
Cassette.overdub(::PWCtx, op::Lifted, x::Function, y) = ω -> op(x(ω), y)
Cassette.overdub(::PWCtx, op::Lifted, x, y::Function) = ω -> op(x, y(ω))
pointwise(f) = Cassette.overdub(PWCtx(), f)


text:
__precompile__(false)
module PlugAndPlot

using GroupedErrors
using TextParse, IndexedTables, IterableTables
using QML, StatPlots
import JuliaDB
gr()

export choose_data, build_window

include("types.jl")
include("select_functions.jl")
include("build_window.jl")
include("build_plot.jl")
end # module


text:
using Distributions, Random

function simulateLR(;Nd = 4, Nc = 1,β0 = 1., β = fill(.5, Nc), σ = 1, kwargs...)
    x = rand(Normal(10,5),Nd,Nc)
    y = β0 .+ x*β .+ rand(Normal(0,σ),Nd)
    return (x=x, y=y, Nd=Nd, Nc=Nc)
 end


text:
# Error type and decoding functionality

export CuError

"""
    CuError(code)
    CuError(code, meta)

Create a CUDA error object with error code `code`. The optional `meta` parameter indicates
whether extra information, such as error logs, is known.
"""
struct CuError <: Exception
    code::CUresult
    meta::Any

    CuError(code, meta=nothing) = new(CUresult(code), meta)
end

Base.convert(::Type{CUresult}, err::CuError) = err.code

Base.:(==)(x::CuError,y::CuError) = x.code == y.code

"""
    name(err::CuError)

Gets the string representation of an error code.

This name can often be used as a symbol in source code to get an instance of this error.
For example:

```jldoctest
julia> using CUDAdrv

julia> err = CuError(1)
CuError(1, ERROR_INVALID_VALUE)

julia> name(err)
"ERROR_INVALID_VALUE"

julia> CUDAdrv.ERROR_INVALID_VALUE
CuError(1, ERROR_INVALID_VALUE)
```
"""
function name(err::CuError)
    if err.code == -1%UInt32
        "ERROR_USING_STUBS"
    else
        str_ref = Ref{Cstring}()
        cuGetErrorName(err, str_ref)
        unsafe_string(str_ref[])[6:end]
    end
end

"""
    description(err::CuError)

Gets the string description of an error code.
"""
function description(err::CuError)
    if err.code == -1%UInt32
        "Cannot use the CUDA stub libraries."
    else
        str_ref = Ref{Cstring}()
        cuGetErrorString(err, str_ref)
        unsafe_string(str_ref[])
    end
end

function Base.showerror(io::IO, err::CuError)
    print(io, "CUDA error: $(description(err)) (code $(reinterpret(Int32, err.code)), $(name(err)))")

    if err.meta != nothing
        print(io, "\n")
        print(io, err.meta)
    end
end

Base.show(io::IO, err::CuError) = print(io, "CuError($(err.code))")

# define shorthands that give CuErryr objects
for code in instances(cudaError_enum)
    name = String(Symbol(code))
    shorthand = Symbol(name[6:end]) # strip the CUDA_ prefix
    @eval const $shorthand = CuError($code)
end

## API call wrapper

const apicall_hook = Ref{Union{Nothing,Function}}(nothing)

# TODO: for the next breaking version
# const call_hooks = Set{Function}()
# @noinline do_call_hooks(fun) = filter!(hook->hook(fun), call_hooks)

macro check(ex)
    # check is used in front of `ccall` or `@runtime_ccall`s that work on a tuple (fun, lib)
    if Meta.isexpr(ex, :call)
        @assert ex.args[1] == :ccall
        @assert Meta.isexpr(ex.args[2], :tuple)
        fun = String(ex.args[2].args[1].value)
    elseif Meta.isexpr(ex, :macrocall)
        @assert ex.args[1] == Symbol("@runtime_ccall")
        @assert Meta.isexpr(ex.args[3], :tuple)
        fun = String(ex.args[3].args[1].value)
    else
        error("@check should prefix ccall or @runtime_ccall")
    end

    # strip any version tag (e.g. cuEventDestroy_v2 -> cuEventDestroy)
    m = match(r"_v\d+$", fun)
    if m !== nothing
        fun = fun[1:end-length(m.match)]
    end

    quote
        # NOTE: this hook is used by CUDAnative.jl to initialize upon the first API call
        apicall_hook[] !== nothing && apicall_hook[]($(QuoteNode(Symbol(fun))))

        res::CUresult = $(esc(ex))
        if res != CUDA_SUCCESS
            throw(CuError(res))
        end

        return
    end
end


text:
using DataFrames
using Distributions
using NearestNeighbors
using StatsBase

using Base.Threads
using Dates: now, DateTime

import LightGraphs

function drop_unused_components!(data::BmmData; min_n_samples::Int=2)
    non_noise_ids = findall(data.assignment .> 0)
    n_mols_per_cell = num_of_molecules_per_cell(data) # Can't use comp.n_samples because it's updated only on the maximization step
    existed_ids = findall(n_mols_per_cell .>= min_n_samples)
    id_map = zeros(Int, length(data.components))
    id_map[existed_ids] .= 1:length(existed_ids)

    data.assignment[non_noise_ids] .= id_map[data.assignment[non_noise_ids]]
    data.components = data.components[existed_ids]
end

adjacent_component_ids(assignment::Array{Int, 1}, adjacent_points::Array{Int, 1})::Array{Int, 1} =
    [x for x in Set(assignment[adjacent_points]) if x > 0]

"""
...
# Arguments
- `assignment::Array{Int, 1}`:
- `adjacent_points::Array{Int, 1}`:
- `adjacent_weights::Array{Float64, 1}`: weights here mean `1 / distance` between two adjacent points
"""
@inline function adjacent_component_weights!(comp_weights::Vector{Float64}, comp_ids::Vector{Int}, component_weights::Dict{Int, Float64},
        assignment::Vector{Int}, adjacent_points::Vector{Int}, adjacent_weights::Vector{Float64}, confidences::Vector{Float64})
    empty!(component_weights)
    empty!(comp_weights)
    empty!(comp_ids)
    zero_comp_weight = 0.0

    @inbounds @simd for i in 1:length(adjacent_weights)
        c_point = adjacent_points[i]
        c_id = assignment[c_point]
        c_conf = confidences[c_point]
        cw = adjacent_weights[i]
        if c_id == 0
            zero_comp_weight += cw * (1 - c_conf)
        else
            component_weights[c_id] = get(component_weights, c_id, 0.0) + cw * c_conf
        end
    end

    for (k, v) in component_weights
        push!(comp_ids, k)
        push!(comp_weights, v)
    end

    return zero_comp_weight
end

function adjust_densities_by_prior_segmentation!(denses::Vector{Float64}, segment_id::Int, largest_cell_id::Int,
        data::BmmData, adj_classes::Vector{Int}, adj_weights::Vector{Float64})
    seg_prior_pow = data.prior_seg_confidence * exp(3*data.prior_seg_confidence)
    seg_size = data.n_molecules_per_segment[segment_id]
    largest_cell_size = min(get(data.components[largest_cell_id].n_molecules_per_segment, segment_id, 0) + 1, seg_size)

    for j in eachindex(adj_weights)
        c_adj = adj_classes[j]
        if c_adj == largest_cell_id
            continue
        end

        main_seg = data.main_segment_per_cell[c_adj]
        n_cell_mols_per_seg = min(get(data.components[c_adj].n_molecules_per_segment, segment_id, 0) + 1, seg_size)

        if (segment_id == main_seg) || (main_seg == 0) # Part against over-segmentation
            denses[j] *= (1. - data.prior_seg_confidence)^0.5 * (n_cell_mols_per_seg / largest_cell_size)^data.prior_seg_confidence
        else # Part against under-segmentation and overlap
            denses[j] *= (1. - data.prior_seg_confidence)^0.5 * (1. - n_cell_mols_per_seg / seg_size)^seg_prior_pow
        end
    end
end

function expect_dirichlet_spatial!(data::BmmData; stochastic::Bool=true, noise_density_threshold::Float64=1e-100)
    component_weights = Dict{Int, Float64}();
    adj_classes = Int[]
    adj_weights = Float64[]
    denses = Float64[]

    has_seg_prior = !isempty(data.segment_per_molecule)
    adj_classes_global = get_global_adjacent_classes(data)

    for i in 1:size(data.x, 1)
        x::Float64 = position_data(data)[1,i]
        y::Float64 = position_data(data)[2,i]
        gene::Int = composition_data(data)[i]
        confidence::Float64 = data.confidence[i]
        mol_cluster::Int = isempty(data.cluster_per_molecule) ? 0 : data.cluster_per_molecule[i]
        segment_id = has_seg_prior ? data.segment_per_molecule[i] : 0

        # Looks like it's impossible to optimize further, even with vectorization. It means that creating vectorized version of expect_dirichlet_spatial makes few sense
        zero_comp_weight = adjacent_component_weights!(adj_weights, adj_classes, component_weights, data.assignment,
            data.adjacent_points[i], data.adjacent_weights[i], data.confidence)

        if i in keys(adj_classes_global)
            n1 = length(adj_classes)
            append!(adj_classes, adj_classes_global[i])
            append!(adj_weights, ones(length(adj_classes) - n1) .* data.real_edge_weight)
        end

        empty!(denses)
        largest_cell_id = 0
        largest_cell_size = 0
        for j in eachindex(adj_weights)
            c_adj = adj_classes[j]
            cc = data.components[c_adj]
            c_dens = confidence * adj_weights[j] * cc.prior_probability * pdf(cc, x, y, gene, use_smoothing=data.use_gene_smoothing)
            if (c_adj > 0) && (c_adj < length(data.cluster_per_cell)) && (data.cluster_per_cell[c_adj] != mol_cluster)
                c_dens *= data.cluster_penalty_mult
            end

            if (segment_id > 0) && (data.main_segment_per_cell[c_adj] in (segment_id, 0))
                # Find an adjacent cell with the largest number of molecules per cell after assigning the current point to it
                cur_mols_per_seg = min(get(cc.n_molecules_per_segment, segment_id, 0) + 1, data.n_molecules_per_segment[segment_id])
                if (cur_mols_per_seg > largest_cell_size) || ((cur_mols_per_seg == largest_cell_size) && (cc.n_samples > data.components[largest_cell_id].n_samples))
                    largest_cell_size = cur_mols_per_seg
                    largest_cell_id = c_adj
                end
            end

            push!(denses, c_dens)
        end

        if (segment_id > 0) & (largest_cell_id > 0)
            adjust_densities_by_prior_segmentation!(denses, segment_id, largest_cell_id, data, adj_classes, adj_weights)
        end

        if sum(denses) < noise_density_threshold
            assign!(data, i, 0) # Noise class
            continue
        end

        if confidence < 1.0
            push!(denses, (1 - confidence) * fmax(data.real_edge_weight, zero_comp_weight) * data.noise_density)
            push!(adj_classes, 0)
        end

        if !stochastic
            assign!(data, i, adj_classes[findmax(denses)[2]])
            continue
        end

        assign!(data, i, fsample(adj_classes, denses))
    end
end

function update_prior_probabilities!(components::Array{Component, 1}, new_component_weight::Float64)
    c_weights = [max(c.n_samples, new_component_weight) for c in components]
    prior_probs = rand(Distributions.Dirichlet(c_weights))
    # t_mmc = 150.0;
    # t_dist = Normal(t_mmc, t_mmc * 3)
    # prior_probs = pdf.(t_dist, c_weights)
    # prior_probs[c_weights .< t_mmc] .= c_weights[c_weights .< t_mmc] ./ t_mmc .* pdf(t_dist, t_mmc)

    for (c, p) in zip(components, prior_probs)
        c.prior_probability = p
    end
end

function maximize!(data::BmmData)
    ids_by_assignment = split_ids(data.assignment .+ 1, max_factor=length(data.components)+1)[2:end]

    @inbounds @views for i in 1:length(data.components)
        p_ids = ids_by_assignment[i]
        maximize!(data.components[i], position_data(data)[:, p_ids], composition_data(data)[p_ids], confidence(data)[p_ids])
    end

    data.noise_density = estimate_noise_density_level(data)
    if isinf(data.noise_density)
        error("Infinity noise density")
    end

    if length(data.cluster_per_molecule) > 0
        data.cluster_per_cell = [isempty(ids) ? 0 : mode(data.cluster_per_molecule[ids]) for ids in ids_by_assignment]
    end
end

function noise_composition_density(data::BmmData)::Float64
#     mean([mean(c.composition_params.counts[c.composition_params.counts .> 0] ./ c.composition_params.sum_counts) for c in data.components]);
    acc = 0.0 # Equivalent to the above commented expression if sum_counts == sum(counts)
    n_comps = 0.0
    for (ci, c) in enumerate(data.components)
        if c.composition_params.sum_counts < 1e-3
            continue
        end

        inn_acc = 0
        for v in c.composition_params.counts
            inn_acc += (v > 0)
        end

        acc += 1.0 / inn_acc
        n_comps += 1.0
    end

    return acc / fmax(n_comps, 1.0)
end

function estimate_noise_density_level(data::BmmData)::Float64
    composition_density = noise_composition_density(data)

    std_vals = data.distribution_sampler.shape_prior.std_values;
    position_density = pdf(MultivariateNormal([0.0, 0.0], diagm(0 => std_vals.^2)), 3 .* std_vals)

    return position_density * composition_density
end

append_empty_component!(data::BmmData) =
    push!(data.components, sample_distribution!(data))[end]

function append_empty_components!(data::BmmData, new_component_frac::Float64)
    for i in 1:round(Int, new_component_frac * length(data.components))
        append_empty_component!(data)
    end
end

function get_global_adjacent_classes(data::BmmData)::Dict{Int, Array{Int, 1}}
    adj_classes_global = Dict{Int, Array{Int, 1}}()
    for (cur_id, comp) in enumerate(data.components)
        if comp.n_samples > 0
            continue
        end

        nearest_id = knn(data.position_knn_tree, comp.position_params.μ, 1)[1][1]
        for t_id in data.adjacent_points[nearest_id]
            if t_id in keys(adj_classes_global)
                push!(adj_classes_global[t_id], cur_id)
            else
                adj_classes_global[t_id] = [cur_id]
            end
        end

        if nearest_id in keys(adj_classes_global)
            push!(adj_classes_global[nearest_id], cur_id)
        else
            adj_classes_global[nearest_id] = [cur_id]
        end
    end

    return adj_classes_global
end

function build_cell_graph(assignment::Vector{Int}, adjacent_points::Array{Vector{Int}, 1}, mol_ids::Vector{Int}, cell_id::Int; confidence::Union{Vector{Float64}, Nothing}=nothing, confidence_threshold::Float64=0.5)
    cur_adj_point = Vector{Int}[]
    if confidence === nothing
        cur_adj_point = filter.(p -> assignment[p] == cell_id, adjacent_points[mol_ids])
    else
        mol_ids = mol_ids[confidence[mol_ids] .> confidence_threshold]
        cur_adj_point = filter.(p -> (assignment[p] == cell_id) & (confidence[p] .> confidence_threshold), view(adjacent_points, mol_ids))
    end

    vert_id_per_mol_id = Dict(mi => vi for (vi, mi) in enumerate(mol_ids))

    for points in cur_adj_point
        for j in 1:length(points)
            points[j] = vert_id_per_mol_id[points[j]]
        end
    end

    neg = LightGraphs.SimpleGraphs.cleanupedges!(cur_adj_point)
    return LightGraphs.SimpleGraph(neg, cur_adj_point), mol_ids
end

function get_connected_components_per_label(assignment::Vector{Int}, adjacent_points::Array{Vector{Int}, 1}, min_molecules_per_cell::Int; kwargs...)
    mol_ids_per_cell = split(1:length(assignment), assignment .+ 1)[2:end]
    real_cell_ids = findall(length.(mol_ids_per_cell) .>= min_molecules_per_cell)
    graph_per_cell = [build_cell_graph(assignment, adjacent_points, mol_ids_per_cell[ci], ci; kwargs...)[1] for ci in real_cell_ids];
    return LightGraphs.connected_components.(graph_per_cell), real_cell_ids, mol_ids_per_cell
end

function split_cells_by_connected_components!(data::BmmData; add_new_components::Bool, min_molecules_per_cell::Int)
    conn_comps_per_cell, real_cell_ids, mol_ids_per_cell = get_connected_components_per_label(data.assignment, data.adjacent_points,
        min_molecules_per_cell; confidence=data.confidence)

    for (cell_id, conn_comps) in zip(real_cell_ids, conn_comps_per_cell)
        if length(conn_comps) < 2
            continue
        end

        largest_cc_id = findmax(length.(conn_comps))[2]
        for (ci, c_ids) in enumerate(conn_comps)
            if ci == largest_cc_id
                continue
            end

            mol_ids = mol_ids_per_cell[cell_id][c_ids]

            if add_new_components
                append_empty_component!(data)
                data.assignment[mol_ids] .= length(data.components)
            else
                data.assignment[mol_ids] .= 0
            end
        end
    end
end

log_em_state(data::BmmData, iter_num::Int, time_start::DateTime) =
    @info "BMM part done for $(now() - time_start) in $iter_num iterations. #Components: $(sum(num_of_molecules_per_cell(data) .> 0)). " *
        "Noise level: $(round(mean(data.assignment .== 0) * 100, digits=3))%"

track_progress!(progress::Nothing) = nothing
track_progress!(progress::Progress) = next!(progress)

function bmm!(data::BmmData; min_molecules_per_cell::Int, n_iters::Int=1000, log_step::Int=4, verbose=true,
              new_component_frac::Float64=0.05, new_component_weight::Float64=0.3,
              assignment_history_depth::Int=0, trace_components::Bool=false, progress::Union{Progress, Nothing}=nothing,
              component_split_step::Int=3)
    time_start = now()

    if verbose
        println("BMM fit started")
    end

    if (assignment_history_depth > 0) && !(:assignment_history in keys(data.tracer))
        data.tracer[:assignment_history] = Vector{Int}[]
    end

    if trace_components && !(:component_history in keys(data.tracer))
        data.tracer[:component_history] = Vector{Component}[]
    end

    it_num = 0
    trace_n_components!(data, min_molecules_per_cell);

    maximize!(data)

    for i in 1:n_iters
        append_empty_components!(data, new_component_frac)
        update_prior_probabilities!(data.components, new_component_weight)
        update_n_mols_per_segment!(data)

        expect_dirichlet_spatial!(data)

        if (i % component_split_step == 0) || (i == n_iters)
            split_cells_by_connected_components!(data; add_new_components=(new_component_frac > 1e-10), min_molecules_per_cell=(i == n_iters ? 0 : min_molecules_per_cell))
        end

        drop_unused_components!(data)
        maximize!(data)

        it_num = i
        if verbose && i % log_step == 0
            log_em_state(data, it_num, time_start)
        end

        trace_n_components!(data, min_molecules_per_cell);
        trace_assignment_history!(data, assignment_history_depth; use_guids=!trace_components)
        if trace_components
            trace_component_history!(data)
        end

        track_progress!(progress)
    end

    if verbose
        log_em_state(data, it_num, time_start)
    end

    return data
end

run_bmm_parallel(bm_data_arr, args...; kwargs...) =
    run_bmm_parallel!(deepcopy(bm_data_arr), args...; kwargs...)

function run_bmm_parallel!(bm_data_arr::Array{BmmData, 1}, n_iters::Int; min_molecules_per_cell::Int, kwargs...)::BmmData
    @info "Algorithm start"
    p = Progress(n_iters * length(bm_data_arr))
    @threads for i in 1:length(bm_data_arr)
        bmm!(bm_data_arr[i]; n_iters=n_iters, min_molecules_per_cell=min_molecules_per_cell, verbose=false, progress=p, kwargs...)
    end

    bm_data_merged = merge_bm_data(bm_data_arr)

    if :assignment_history in keys(bm_data_merged.tracer)
        bm_data_merged.assignment = estimate_assignment_by_history(bm_data_merged)[1];
        maximize!(bm_data_merged)
    end

    drop_unused_components!(bm_data_merged; min_n_samples=1)
    maximize!(bm_data_merged)

    @info "Done!"

    return bm_data_merged
end

max_stars_repo_name: YinuoJin/Baysor

max_issues_repo_name: YinuoJin/Baysor

max_forks_repo_name: YinuoJin/Baysor


text:
## test Point
using Iris
p = Point(1)
p = Point(1,2)
p = Point((1,2,3))

p.x
p.y
p.ϕ
p.r

@code_warntype Point(1)
@code_warntype Point(1,2)
@code_warntype Point((1,2,3))
f(p) = p.x
@code_warntype f(p)

## test BoundaryLayers
using Iris
PML{1}(.3)
cPML{2}(.3)
noBL{1}(.4)
conj(PML{3}(.3))

pml = PML{1}(.3,0,.3)
@code_warntype pml(.1)

nbl = noBL{1}(.3,0,.3)
@code_warntype nbl(.1)

## test Shapes1D
using Iris
interval = Interval(0,.3)
@code_warntype interval(.2)

## test DielectricFunction
using Iris
DielectricFunction()
de = DielectricFunction(1,.1)
desin(p::Point,params::Dict) = complex(params[:n1],params[:n2]*sin(params[:period]*p.x))
DielectricFunction(desin,2,.3)
de2 = DielectricFunction(desin,Dict(:period=>π,:n1=>2,:n2=>.1))
de2(Point(2.1))
F = PumpFunction()
F(Point(1.4))
fsin(p::Point,params::Dict) = params[:f0] + params[:f1]*sin(params[:period]*p.x)
F2 = PumpFunction(fsin,Dict(:period=>π,:f0=>1,:f1=>.1))
F2(Point(1.1))

@code_warntype DielectricFunction()
@code_warntype DielectricFunction(1,.1)
@code_warntype DielectricFunction(desin,2,.3)
@code_warntype DielectricFunction(desin,Dict{Symbol,Float64}(:period=>π,:n1=>2,:n2=>.1))
@code_warntype de2(Point(2.1))
@code_warntype PumpFunction()
@code_warntype F(Point(1.4))
@code_warntype PumpFunction(fsin,Dict(:period=>π,:f0=>1,:f1=>.1))
@code_warntype F2(1.1)

## lattices
using Iris
lat = Lattice(.01)
# Lattice(1.0,1.0;type=:Polar)
# lat = Lattice(1.0,1.0,1.0;angles=(.1,.1,0.))

@code_warntype Lattice(1.0)
@code_warntype Lattice(1.0,1.0)
@code_warntype Lattice(1.0,1.0,1.0;angles=(.1,.1,0.))
f(lat) = lat.r0
@code_warntype f(lat)

## Boundary
using Iris
interval = Interval(0,1)
bcs = (DirichletBC{1}(),DirichletBC{2}())
bls = (PML{1}(.3),PML{2}(.3))
bnd = Boundary(interval,bcs,bls)

@code_warntype Boundary(interval,bcs,bls)

## Domain
using Iris
lat = Lattice(.013)
interval = Interval(0,1)
bcs = (DirichletBC{1}(),DirichletBC{2}())
bls = (PML{1}(.3),PML{2}(.3))
bnd = Boundary(interval,bcs,bls)

dom = Domain(bnd,lat,DielectricFunction(1),PumpFunction(0))
@code_warntype Domain(bnd,lat,DielectricFunction(1),PumpFunction(0))

## Simulation
using Iris
lat = Lattice(.013)
interval1 = Interval(0,2)
interval2 = Interval(.5,1)
bcs = (DirichletBC{1}(),DirichletBC{2}())
bls = (PML{1}(.2),PML{2}(.2))
bnd1 = Boundary(interval1,bcs,bls)
bnd2 = Boundary(interval2,bcs,bls)

dom1 = Domain(bnd1,lat,DielectricFunction(1),PumpFunction(0))
dom2 = Domain(bnd2,lat,DielectricFunction(2),PumpFunction(0))
sim = Simulation(10,dom1,dom2)

##
using Iris
using NLsolve
lat = Lattice(2π/30/15/3)
interval1 = Interval(0,3)
interval2 = Interval(1,2)
bcs1 = (MatchedBC{1}(in=[1]),MatchedBC{2}(in=[1]))
bls1 = (noBL{1}(4*2π/30),noBL{2}(4*2π/30))
bcs2 = (noBC{1}(),noBC{2}())
bls2 = (noBL{1}(),noBL{2}())
bnd1 = Boundary(interval1,bcs1,bls1)
bnd2 = Boundary(interval2,bcs2,bls2)

dom1 = Domain(bnd1,lat,DielectricFunction(1),PumpFunction(0))
dom2 = Domain(bnd2,lat,DielectricFunction(3),PumpFunction(1),TwoLevelSystem(30,-.15,2))
sim = Simulation(30,dom2,dom1)
@time ms = maxwell_salt(sim,1)
ωl,ψl = eig_kl(sim,30;nev=1)
@time nlsolve(ms,real(ωl),.1ψl;show_trace=true);
@time SALT(sim,real(ωl),ψl;show_trace=true)

@btime eig_kl(sim,30);
@code_warntype eig_kl(sim,30)
using NonlinearEigenproblems
nep = maxwell_nep(sim)
ω,ψ = contour_beyn(nep;σ=8,radius=.3,N=75,neigs=2,k=3); println(ω)
@btime eig_knl(sim,30; maxit=40);
@code_warntype eig_knl(sim,30; maxit=40)

η,u = eig_cf(sim,30)
@btime eig_cf(sim,30);
@code_warntype eig_cf(sim,30)

perm = sortperm(map(a->a.vec[1],sim.x))
plot(map(a->a.vec[1],sim.x)[perm],100abs2.(ψl[end÷3+1:2end÷3,1])[perm])
plot(map(a->a.vec[1],sim.x)[perm],100abs2.(ψnl[end÷3+1:2end÷3,1])[perm])
plot(map(a->a.vec[1],sim.x)[perm],100abs2.(u[end÷3+1:2end÷3,1])[perm])

## dispersions
using Iris
using NonlinearEigenproblems
lat = Lattice(2π/30/30/3)
interval1 = Interval(0,3)
interval2 = Interval(1,2)
bcs1 = (MatchedBC{1}(out=[1]),MatchedBC{2}(out=[1]))
bls1 = (noBL{1}(4*2π/30),noBL{2}(4*2π/30))
bcs2 = (noBC{1}(),noBC{2}())
bls2 = (noBL{1}(),noBL{2}())
bnd1 = Boundary(interval1,bcs1,bls1)
bnd2 = Boundary(interval2,bcs2,bls2)

dom1 = Domain(bnd1,lat,DielectricFunction(1),PumpFunction(0))
dom2 = Domain(bnd2,lat,DielectricFunction(3),PumpFunction(1),TwoLevelSystem(31,.01,2))
sim = Simulation(30,dom2,dom1)

nep1 = maxwell_nep(sim)
nep2 = SPMF_NEP(nep1.A[1:7],nep1.fi[1:7])
ω,ψ = contour_beyn(nep1;σ=32,radius=1,neigs=3,k=5,logger=2,N=75)
ω2,ψ2 = contour_beyn(nep2;σ=32,radius=1,neigs=3,k=5,logger=2,N=75)
# newton(nep1;λ=32,maxit=50,logger=2)
# newton(nep2;λ=32,maxit=50,logger=2)
# augnewton(nep1;λ=32,maxit=50,logger=2)
# augnewton(nep2;λ=32,maxit=50,logger=2)
resinv(nep1;λ=32,logger=2)
resinv(nep2;λ=32,logger=2)
quasinewton(nep1;λ=32,logger=2)
quasinewton(nep2;λ=32,logger=2)
mslp(nep1;λ=32,logger=2)
mslp(nep2;λ=32,logger=2)
# implicitdet(nep1;λ=32,maxit=150,logger=2)
# implicitdet(nep2;λ=32,maxit=150,logger=2)
# nlar(nep1;λ=32,logger=2,neigs=2)
# nlar(nep2;λ=32,logger=2,neigs=2)
# jd_betcke(nep1;λ=32,logger=2)
# jd_betcke(nep2;λ=32,logger=2)
# jd_effenberger(nep2;λ=32,logger=2)
# jd_effenberger(nep2;λ=32,logger=2)
# contour_block_SS(nep1;σ=32,neigs=1,radius=1,k=6,logger=2,N=150)
# contour_block_SS(nep2;σ=32,neigs=3,k=5,logger=2,N=75)
# iar(nep1;σ=32,logger=1,neigs=2)
iar(nep2;σ=32,logger=1,neigs=2)
# iar_chebyshev(nep1;σ=32,logger=1,neigs=2)
iar_chebyshev(nep2;σ=32,logger=1,neigs=2)
# tiar(nep1;σ=32,logger=1,neigs=2)
tiar(nep2;σ=32,logger=1,neigs=2)
# ilan(nep1;σ=32,logger=1,neigs=1)
# ilan(nep2;σ=32,logger=1,neigs=1)

## test jacobian

using LinearAlgebra
using SparseArrays

tls = TwoLevelSystem(10,1,1)
ωs = [9. 11.]

dω = .00001
v = [1 5 .1]
Dω = dω*v

γ1 = Iris.Common.Dispersions.TwoLevelSystems.γ(tls,ωs,ψl)
γ2 = Iris.Common.Dispersions.TwoLevelSystems.γ(tls,ωs+Dω,ψl)
dγ = (γ2-γ1)/dω
∂γ = Iris.Common.Dispersions.TwoLevelSystems.∂γ∂ω(tls,ωs,ψl).*v
norm(dγ-∂γ)

Γ1 = Iris.Common.Dispersions.TwoLevelSystems.Γ(tls,ωs,ψl)
Γ1-abs2.(γ1)
Γ2 = Iris.Common.Dispersions.TwoLevelSystems.Γ(tls,ωs+Dω,ψl)
dΓ = (Γ2-Γ1)/dω
∂Γ = Iris.Common.Dispersions.TwoLevelSystems.∂Γ∂ω(tls,ωs,ψl).*v
norm(dΓ-∂Γ)

h1 = Iris.Common.Dispersions.TwoLevelSystems.H(tls,ωs,ψl)
h2 = Iris.Common.Dispersions.TwoLevelSystems.H(tls,ωs+Dω,ψl)
dh = (h2 - h1)/dω
∂h = Iris.Common.Dispersions.TwoLevelSystems.∂H∂ω(tls,ωs,ψl)
norm(∂h[:,1]*v[1]+∂h[:,2]*v[2]+∂h[:,3]*v[3] - dh)

χ1 = Iris.Common.Dispersions.TwoLevelSystems.χ(tls,ωs,ψl)
χ2 = Iris.Common.Dispersions.TwoLevelSystems.χ(tls,ωs+Dω,ψl)
dχψ = ((χ2-χ1)/dω)*ψl[:]
∂χψ = Iris.Common.Dispersions.TwoLevelSystems.∂χ∂ωψ(tls,ωs,ψl)
norm(∂χψ[1][:]*v[1] + ∂χψ[2][:]*v[2] + ∂χψ[3][:]*v[3] - dχψ)

####

using Iris

lat = Lattice(2π/30/10/3)
    interval1 = Interval(0,3)
    interval2 = Interval(1,2)
    bcs1 = (MatchedBC{1}(out=[1]),MatchedBC{2}(out=[1]))
    bls1 = (noBL{1}(4*2π/30),noBL{2}(4*2π/30))
    bcs2 = (noBC{1}(),noBC{2}())
    bls2 = (noBL{1}(),noBL{2}())
    bnd1 = Boundary(interval1,bcs1,bls1)
    bnd2 = Boundary(interval2,bcs2,bls2)

    dom1 = Domain(bnd1,lat,DielectricFunction(1),PumpFunction(0))
    dom2 = Domain(bnd2,lat,DielectricFunction(3),PumpFunction(1))
    sim = Simulation(30,dom2,dom1)

using Random
ωs = [8. 9. 11.]
ψ = rand(MersenneTwister(1),ComplexF64,3length(sim.x),length(ωs))
d = .0001
dψ = d*rand(ComplexF64,3length(sim.x),length(ωs))
tls = TwoLevelSystem(10,.2,1)
χ1 = susceptability(tls,ωs,ψ)
χ2 = susceptability(tls,ωs,ψ+dψ)
dχ = (χ2-χ1)
J = jacobian_lasing(tls,ωs,ψ)
J*vcat(real(dψ[:]),imag(dψ[:]))


text:
struct LinearPolicyWithBasis{TP, TB} <: AbstractPolicy where {TP, TB}
    π::TP
    ϕ::TB
end

function params(π::LinearPolicyWithBasis)
    return params(π.π)
end

function (π::LinearPolicyWithBasis)(s)
    feats = π.ϕ(s)
    return π.π(feats)
end

function logpdf(π::LinearPolicyWithBasis, s, a)
    feats = π.ϕ(s)
    return logpdf(π.π, feats, a)
end

function grad_logpdf!(ψ, π::LinearPolicyWithBasis, s, a)
    feats = π.ϕ(s)
    return grad_logpdf!(ψ, π.π, feats, a)
end

function grad_logpdf(π::LinearPolicyWithBasis, s, a)
    feats = π.ϕ(s)
    logp, ψ = grad_logpdf(π.π, feats, a)
    return logp, ψ
end

function sample_with_trace!(ψ, action, π::LinearPolicyWithBasis, s)
    feats = π.ϕ(s)
    return sample_with_trace!(ψ, action, π.π, feats)
end

function sample_with_trace(π::LinearPolicyWithBasis, s)
    feats = π.ϕ(s)
    return sample_with_trace(π.π, feats)
end


text:
using Microbiome
using MicrobiomePlots
using Colors
using Random
using Test

@testset "Abundances" begin
    Random.seed!(1)
    M = rand(100, 10)

    abund = abundancetable(
        M, ["sample_$x" for x in 1:10],
        ["feature_$x" for x in 1:100])

    @test typeof(abundanceplot(abund, topabund=5)) <: Plots.Plot
    @test typeof(abundanceplot(abund, srt=collect(10:-1:1))) <: Plots.Plot

    @test_skip typeof(abundanceplot(abund, srt=:feature_1)) <: Plots.Plot # Needs method feature sorting

    ann = annotationbar(parse.(Color, ["red", "white", "blue"]))

    @test typeof(plot(ann)) <: Plots.Plot
end

max_issues_repo_name: BioJulia/MicrobiomePlots

max_forks_repo_name: BioJulia/MicrobiomePlots


text:
Dict(#
    :replication => 1000,
    :output_directory => "results",
    :nprocs => 4,
    :overwrite => false,
    :seed_offset => 0,
)

max_stars_repo_name: ndgnuh/RiceBPH

max_issues_repo_name: ndgnuh/RiceBPH

max_forks_repo_name: ndgnuh/RiceBPH


text:
function ptdf_networkflow(ps_m::CanonicalModel,
                          branches::Array{Br,1},
                          buses::Array{PSY.Bus,1},
                          expression::Symbol,
                          PTDF::AxisArrays.AxisArray,
                          time_range::UnitRange{Int64}) where {Br <: PSY.Branch}

    ps_m.constraints[:network_flow] = JuMP.Containers.DenseAxisArray{JuMP.ConstraintRef}(undef, [b.name for b in branches], time_range)
    ps_m.constraints[:nodal_balance] = JuMP.Containers.DenseAxisArray{JuMP.ConstraintRef}(undef, [bn.name for bn in buses], time_range)

     _remove_undef!(ps_m.expressions[expression])

    for t in time_range
        for b in branches
            ps_m.constraints[:network_flow][b.name,t] = JuMP.@constraint(ps_m.JuMPmodel, ps_m.variables[:Fbr][b.name,t] == PTDF[b.name,:].data'*ps_m.expressions[expression].data[:,t])
        end

        for b in branches
            _add_to_expression!(ps_m.expressions[expression], b.connectionpoints.from.number, t, ps_m.variables[:Fbr][b.name,t], -1)
            _add_to_expression!(ps_m.expressions[expression], b.connectionpoints.to.number, t, ps_m.variables[:Fbr][b.name,t], 1)
        end

        for b in buses
            ps_m.constraints[:nodal_balance][b.name, t] = JuMP.@constraint(ps_m.JuMPmodel, ps_m.expressions[expression][b.number,t] == 0)
        end

    end

    return

end

#=
The previous implementation of the PTDF constraints showed to be faster in the 5 bus system. I have added
another implementatio for further testing with larger systems

 ─────────────────────────────────────────────────────────────────────────────

The first implementation, above was tested this way with the respective results

to = TimerOutput()

@timeit to "canonical_model"   ps_m = PSI.CanonicalModel(Model(ipopt_optimizer),
                              Dict{String, JuMP.Containers.DenseAxisArray{JuMP.VariableRef}}(),
                              Dict{String, JuMP.Containers.DenseAxisArray}(),
                              nothing,
                              Dict{Symbol, PSI.JuMPAffineExpressionArray}(:var_active => PSI.JuMPAffineExpressionArray(undef, 5, 24),
                                                                         :var_reactive => PSI.JuMPAffineExpressionArray(undef, 5, 24)),
                              nothing);
@timeit to "build_thermal"    PSI.construct_device!(ps_m, PSY.ThermalGen, PSI.ThermalDispatch, PM.StandardACPForm, sys5b);
@timeit to "build_load"    PSI.construct_device!(ps_m, PSY.PowerLoad, PSI.StaticPowerLoad, PM.StandardACPForm, sys5b);
@timeit to "add_flow"      PSI.flow_variables(ps_m, PM.DCPlosslessForm, branches5, 1:24)
@timeit to "PTDF cons" begin
    @timeit to "allocate_space" ps_m.constraints["Flow_con1"] = JuMP.Containers.DenseAxisArray{JuMP.ConstraintRef}(undef, [b.name for b in branches5], 1:24)
    @timeit to "make constraints" for t in 1:24
                                    for b in branches5
                                        ps_m.constraints["Flow_con1"][b.name,t] = JuMP.@constraint(ps_m.JuMPmodel, ps_m.variables[:Fbr][b.name,t] == PTDF[b.name,:].data'*ps_m.expressions[:var_active][:,t])
                                    end
                                end
                            end

 ─────────────────────────────────────────────────────────────────────────────
                                      Time                   Allocations
                              ──────────────────────   ───────────────────────
       Tot / % measured:           305ms / 12.3%           4.78MiB / 91.0%

 Section              ncalls     time   %tot     avg     alloc   %tot      avg
 ─────────────────────────────────────────────────────────────────────────────
 PTDF cons                 1   35.7ms  95.2%  35.7ms   3.30MiB  75.8%  3.30MiB
   allocate_space          1   33.3ms  88.6%  33.3ms   2.82MiB  64.8%  2.82MiB
   make constraints        1   2.43ms  6.47%  2.43ms    488KiB  11.0%   488KiB
 build_thermal             1   1.39ms  3.71%  1.39ms    904KiB  20.3%   904KiB
 add_flow                  1    221μs  0.59%   221μs   64.8KiB  1.45%  64.8KiB
 canonical_model           1    133μs  0.36%   133μs   46.3KiB  1.04%  46.3KiB
 build_load                1   62.3μs  0.17%  62.3μs   63.8KiB  1.43%  63.8KiB
 ─────────────────────────────────────────────────────────────────────────────

 ─────────────────────────────────────────────────────────────────────────────

 The second implementation and the results are as follows:

 to = TimerOutput()

@timeit to "canonical_model"   ps_m = PSI.CanonicalModel(Model(ipopt_optimizer),
                              Dict{String, JuMP.Containers.DenseAxisArray{JuMP.VariableRef}}(),
                              Dict{String, JuMP.Containers.DenseAxisArray}(),
                              nothing,
                              Dict{Symbol, PSI.JuMPAffineExpressionArray}(:var_active => PSI.JuMPAffineExpressionArray(undef, 5, 24),
                                                                         :var_reactive => PSI.JuMPAffineExpressionArray(undef, 5, 24)),
                              nothing);
@timeit to "build_thermal"    PSI.construct_device!(ps_m, PSY.ThermalGen, PSI.ThermalDispatch, PM.StandardACPForm, sys5b);
@timeit to "build_load"    PSI.construct_device!(ps_m, PSY.PowerLoad, PSI.StaticPowerLoad, PM.StandardACPForm, sys5b);
@timeit to "add_flow"      PSI.flow_variables(ps_m, PM.DCPlosslessForm, branches5, 1:24)
@timeit to "PTDF cons" begin
    @timeit to "allocate_space" ps_m.constraints["Flow_con2"] = JuMP.Containers.DenseAxisArray{JuMP.ConstraintRef}(undef, [b.name for b in branches5], 1:24)
    @timeit to "make constraints" begin for t in 1:24
            for b in branches5
                expr = JuMP.AffExpr(0.0, ps_m.variables[:Fbr][b.name,t] => 1.0)
                    for n in nodes5
                        JuMP.add_to_expression!(expr, -1*PTDF[b.name,:].data[n.number]*ps_m.expressions[:var_active][n.number,t])
                    end
                ps_m.constraints["Flow_con2"][b.name,t] = JuMP.@constraint(ps_m.JuMPmodel, expr == 0.0)
        end
    end
    end
end

 ─────────────────────────────────────────────────────────────────────────────
                                      Time                   Allocations
                              ──────────────────────   ───────────────────────
       Tot / % measured:           1.25s / 2.48%           5.10MiB / 90.9%

 Section              ncalls     time   %tot     avg     alloc   %tot      avg
 ─────────────────────────────────────────────────────────────────────────────
 PTDF cons                 1   29.0ms  93.9%  29.0ms   3.58MiB  77.3%  3.58MiB
   allocate_space          1   24.8ms  80.2%  24.8ms   2.82MiB  60.8%  2.82MiB
   make constraints        1   4.19ms  13.5%  4.19ms    781KiB  16.4%   781KiB
 build_thermal             1   1.40ms  4.54%  1.40ms    904KiB  19.0%   904KiB
 add_flow                  1    267μs  0.86%   267μs   64.8KiB  1.36%  64.8KiB
 canonical_model           1    136μs  0.44%   136μs   46.3KiB  0.98%  46.3KiB
 build_load                1   94.2μs  0.30%  94.2μs   63.8KiB  1.34%  63.8KiB
 ─────────────────────────────────────────────────────────────────────────────

 =#


text:
import MathProgBase
const MPB = MathProgBase.SolverInterface
import SemidefiniteModels
const SDM = SemidefiniteModels

export PARAMETER_DEFAULT, PARAMETER_UNSTABLE_BUT_FAST, PARAMETER_STABLE_BUT_SLOW

export SDPAMathProgModel, SDPASolver

struct SDPASolver <: MPB.AbstractMathProgSolver
    options::Dict{Symbol,Any}
end
SDPASolver(;kwargs...) = SDPASolver(Dict{Symbol,Any}(kwargs))

mutable struct SDPAMathProgModel <: SDM.AbstractSDModel
    problem::SDPAProblem
    options::Dict{Symbol,Any}
    function SDPAMathProgModel(; kwargs...)
        new(SDPAProblem(), Dict{Symbol, Any}(kwargs))
    end
end
SDM.SDModel(s::SDPASolver) = SDPAMathProgModel(; s.options...)
MPB.ConicModel(s::SDPASolver) = SDM.SDtoConicBridge(SDM.SDModel(s))
MPB.LinearQuadraticModel(s::SDPASolver) = MPB.ConicToLPQPBridge(MPB.ConicModel(s))

MPB.supportedcones(s::SDPASolver) = [:Free,:Zero,:NonNeg,:NonPos,:SOC,:RSOC,:SDP]
function MPB.setvartype!(m::SDPAMathProgModel, vtype, blk, i, j)
    if vtype != :Cont
        error("Unsupported variable type $vtype by SDPA")
    end
end

function MPB.loadproblem!(m::SDPAMathProgModel, filename::AbstractString)
    error("not supported yet")
end
#writeproblem(m, filename::String)
function MPB.loadproblem!(m::SDPAMathProgModel, blkdims::Vector{Int}, constr::Int)
    m.problem = SDPAProblem()
    setParameterType(m.problem, PARAMETER_DEFAULT)
    setparameters!(m.problem, m.options)
    inputConstraintNumber(m.problem, constr)
    inputBlockNumber(m.problem, length(blkdims))
    for (i, blkdim) in enumerate(blkdims)
        inputBlockSize(m.problem, i, blkdim)
        inputBlockType(m.problem, i, blkdim < 0 ? LP : SDP)
    end
    initializeUpperTriangleSpace(m.problem)
end

function SDM.setconstrB!(m::SDPAMathProgModel, val, constr::Integer)
    @assert constr > 0
    inputCVec(m.problem, constr, val)
end
function SDM.setconstrentry!(m::SDPAMathProgModel, coef, constr::Integer, blk::Integer, i::Integer, j::Integer)
    @assert constr > 0
    inputElement(m.problem, constr, blk, i, j, float(coef), false)
end
function SDM.setobjentry!(m::SDPAMathProgModel, coef, blk::Integer, i::Integer, j::Integer)
    inputElement(m.problem, 0, blk, i, j, float(coef), false)
end

function MPB.optimize!(m::SDPAMathProgModel)
    SDPA.initializeUpperTriangle(m.problem, false)
    SDPA.initializeSolve(m.problem)
    SDPA.solve(m.problem)
end

function MPB.status(m::SDPAMathProgModel)
    status = getPhaseValue(m.problem)
    if status == noINFO
        return :Uninitialized
    elseif status == pFEAS
        return :Suboptimal
    elseif status == dFEAS
        return :Suboptimal
    elseif status == pdFEAS
        return :Optimal
    elseif status == pdINF
        return :Infeasible
    elseif status == pFEAS_dINF
        return :Unbounded
    elseif status == pINF_dFEAS
        return :Infeasible
    elseif status == pdOPT
        return :Optimal
    elseif status == pUNBD
        return :Unbounded
    elseif status == dUNBD
        return :Infeasible
    end
end

function MPB.getobjval(m::SDPAMathProgModel)
    getPrimalObj(m.problem)
end
function MPB.getsolution(m::SDPAMathProgModel)
    PrimalSolution(m.problem)
end
function MPB.getdual(m::SDPAMathProgModel)
    unsafe_wrap(Array, getResultXVec(m.problem), getConstraintNumber(m.problem))
end
function MPB.getvardual(m::SDPAMathProgModel)
    VarDualSolution(m.problem)
end


text:
#
# Denoise an SPD Example with Cyclic Proximal Point applied to the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data as for the corresponding CP algorithm
#
using Manopt
using Images, CSV, DataFrames, LinearAlgebra, JLD2
#
# Settings
ExportResult = true
ExportOrig = true
ExportResultVideo = false
ExportTable = true
#
# Manifold and Data
f = artificialSPDImage2(32)
pixelM = SymmetricPositiveDefinite(3)
resultsFolder = "src/examples/Total_Variation/SPD_TV/"
experimentName = "ImageCPPA"
if !isdir(resultsFolder)
    mkdir(resultsFolder)
end
if ExportOrig
    renderAsymptote(resultsFolder*experimentName*"-orig.asy"; data=f, scaleAxes=(7.5,7.5,7.5))
end
#
# Parameters
α = 6.
maxIterations = 4000
#
# Build Problem for L2-TV
M = Power(pixelM,size(f))
d = length(size(f))
rep(d) = (d>1) ? [ones(Int,d)...,d] : d
fidelity(x) = 1/2*distance(M,x,f)^2
Λ(x) = forwardLogs(M,x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM),getValue(repeat(x,rep(d)...)),getValue(Λ(x))),1)
#
# Setup and Optimize
cost(x) =  fidelity(x) + α*prior(x)
proximalMaps = [(λ,x) -> proxDistance(M,λ,f,x,2), (λ,x) -> proxTV(M,α*λ,x,1)]
x0 = f
@time y, yRec = cyclicProximalPoint(M,cost,proximalMaps,x0;
debug = [:Iteration," | ", DebugProximalParameter()," | ", :Change, " | ", :Cost, "\n",100,:Stop],
record = [:Iteration, :Iterate, :Cost],
stoppingCriterion = stopAfterIteration(maxIterations)
)
#
# Results
if ExportResult
    renderAsymptote(resultsFolder*experimentName*"-result-$(maxIterations)-α$(replace(string(α), "." => "-")).asy"; data=y, render=4, scaleAxes=(7.5,7.5,7.5) )
end
if ExportTable
    A = cat( [ y[1] for y in yRec], [y[3] for y in yRec]; dims=2 )
    CSV.write(string(resultsFolder*experimentName*"ResultCost.csv"),
    DataFrame(A), writeheader=false);
    save(resultsFolder*experimentName*"-CostValue.jld2",
    Dict("compareCostFunctionValue" => last(yRec)[3])
    )
end