text:
# ALL TESTS THAT CONCERN MODFICIATION

# include
using LatticePhysics
using Base.Test

# begin the Modification testblock
modification_testset = @testset "Modifications of lattices and unitcells" begin

################################################################################
#
#   IMPLEMENTATIONS OF DIFFERENT LATTICE MODIFICATION FUNCTIONS
#   (MOST OF THEM ALSO WORK FOR UNITCELLS)
#
################################################################################

################################################################################
#
#   1) CONNECTIONS ADDING REMOVING
#       - Add a new connection
#       - remove connections (based on indices)
#       - remove connections (based on strength)
#
################################################################################

# begin the testset
@testset "Adding & removing connections" begin

    # First, Unitcell
    @testset "Adding and removing in Unitcells (honeycomb)" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)

        # First, try to add connections
        @testset "Adding connections" begin
            # try adding a random connection and testing the unitcell afterwards
            @test addConnection!(unitcell, 1, 1, "t1", (1,1)) == nothing
            @test testUnitcell(unitcell, 2,2)
            # test whether there are now 8 connections
            @test length(unitcell.connections) == 8
            # test whether there are now 4 connection strengths
            @test length(getConnectionStrengthList(unitcell)) == 4
            # test if one adds up the same connection again, the numbers stay the same
            @test addConnection!(unitcell, 1, 1, "t1", (1,1)) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(unitcell.connections) == 8
            @test length(getConnectionStrengthList(unitcell)) == 4
            # test if one adds up the another connection but with different wrap, the numbers change
            @test addConnection!(unitcell, 1, 1, "t1", (1,-1)) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(unitcell.connections) == 10
            @test length(getConnectionStrengthList(unitcell)) == 4
        end;

        # get a fresh test unitcell
        unitcell = getUnitcellHoneycomb(4)

        # Then, try to remove connections (based on index&strength)
        @testset "Removing connections based on index & strength" begin
            # try adding a random connection and testing the unitcell afterwards
            @test removeConnections!(unitcell, 1, 2, "tx") == nothing
            @test testUnitcell(unitcell, 2,2)
            # test whether there are now 4 connections
            @test length(unitcell.connections) == 4
            # test whether there are now 2 connection strengths
            @test length(getConnectionStrengthList(unitcell)) == 2
            # test if one adds up the same connection again, the numbers stay the same
            @test removeConnections!(unitcell, 1, 2, "tx") == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(unitcell.connections) == 4
            @test length(getConnectionStrengthList(unitcell)) == 2
        end;

        # get a fresh test unitcell
        unitcell = getUnitcellSquareOctagon(4)

        # Then, try to remove connections (based on index)
        @testset "Removing connections based on index" begin
            # try adding a random connection and testing the unitcell afterwards
            @test removeConnections!(unitcell, 1, 2) == nothing
            @test testUnitcell(unitcell, 2,2)
            # test whether there are now 10 connections
            @test length(unitcell.connections) == 10
            # test whether there are still 3 connection strengths
            @test length(getConnectionStrengthList(unitcell)) == 3
            # test if one adds up the same connection again, the numbers stay the same
            @test removeConnections!(unitcell, 1, 2) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(unitcell.connections) == 10
            @test length(getConnectionStrengthList(unitcell)) == 3
        end;

        # get a fresh test unitcell
        unitcell = getUnitcellSquareOctagon(4)

        # Then, try to remove connections (based on strength)
        @testset "Removing connections based on index" begin
            # try adding a random connection and testing the unitcell afterwards
            @test removeConnections!(unitcell, "tx") == nothing
            @test testUnitcell(unitcell, 2,2)
            # test whether there are now 8 connections
            @test length(unitcell.connections) == 8
            # test whether there are still 2 connection strengths
            @test length(getConnectionStrengthList(unitcell)) == 2
            # test if one adds up the same connection again, the numbers stay the same
            @test removeConnections!(unitcell, 1, 2) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(unitcell.connections) == 8
            @test length(getConnectionStrengthList(unitcell)) == 2
        end;

    # end the Unitcell testset here
    end;

    # Then, Lattice
    @testset "Adding and removing in Lattices (honeycomb 4x2 periodic)" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # get a test lattice
        lattice = getLattice(unitcell, [-4, -2])

        # First, try to add connections
        @testset "Adding connections" begin
            # try adding a random connection and testing the unitcell afterwards
            @test addConnection!(lattice, 2, 7, "t1", (1,1)) == nothing
            @test testLattice(lattice, 2,2)
            # test whether there are now 50 connections
            @test length(lattice.connections) == 50
            # test whether there are now 4 connection strengths
            @test length(getConnectionStrengthList(lattice)) == 4
            # test if one adds up the same connection again, the numbers stay the same
            @test addConnection!(lattice, 2, 7, "t1", (1,1)) == nothing
            @test testLattice(lattice, 2,2)
            @test length(lattice.connections) == 50
            @test length(getConnectionStrengthList(lattice)) == 4
            # test if one adds up the another connection but with different wrap, the numbers change
            @test addConnection!(lattice, 1, 1, "t1", (1,-1)) == nothing
            @test testLattice(lattice, 2,2)
            @test length(lattice.connections) == 52
            @test length(getConnectionStrengthList(lattice)) == 4
        end;

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # get a test lattice
        lattice = getLattice(unitcell, [-4, -2])

        # Then, try to remove connections (based on index&strength)
        @testset "Removing connections based on index & strength" begin
            # try adding a random connection and testing the unitcell afterwards
            @test removeConnections!(lattice, 1, 2, "tx") == nothing
            @test testLattice(lattice, 2,2)
            # test whether there are now 46 connections
            @test length(lattice.connections) == 46
            # test whether there are now 3 connection strengths
            @test length(getConnectionStrengthList(lattice)) == 3
            # test if one adds up the same connection again, the numbers stay the same
            @test removeConnections!(lattice, 1, 2, "tx") == nothing
            @test testLattice(lattice, 2,2)
            @test length(lattice.connections) == 46
            @test length(getConnectionStrengthList(lattice)) == 3
        end;

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # get a test lattice
        lattice = getLattice(unitcell, [-4, -2])

        # Then, try to remove connections (based on index)
        @testset "Removing connections based on index" begin
            # try adding a random connection and testing the unitcell afterwards
            @test removeConnections!(lattice, 6, 7) == nothing
            @test testLattice(lattice, 2,2)
            # test whether there are now 46 connections
            @test length(lattice.connections) == 46
            # test whether there are still 3 connection strengths
            @test length(getConnectionStrengthList(lattice)) == 3
            # test if one adds up the same connection again, the numbers stay the same
            @test removeConnections!(lattice, 6, 7) == nothing
            @test testLattice(lattice, 2,2)
            @test length(lattice.connections) == 46
            @test length(getConnectionStrengthList(lattice)) == 3
        end;

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # get a test lattice
        lattice = getLattice(unitcell, [-4, -2])

        # Then, try to remove connections (based on strength)
        @testset "Removing connections based on index" begin
            # try adding a random connection and testing the unitcell afterwards
            @test removeConnections!(lattice, "tx") == nothing
            @test testLattice(lattice, 2,2)
            # test whether there are now 32 connections
            @test length(lattice.connections) == 32
            # test whether there are still 2 connection strengths
            @test length(getConnectionStrengthList(lattice)) == 2
            # test if one adds up the same connection again, the numbers stay the same
            @test removeConnections!(lattice, 1, 2) == nothing
            @test testLattice(lattice, 2,2)
            @test length(lattice.connections) == 32
            @test length(getConnectionStrengthList(lattice)) == 2
        end;

    # end the Lattice testset here
    end;

# end the testset here
end;

################################################################################
#
#   2) SITES ADDING REMOVING
#       - add a new site
#       - remove a site (based on index)
#
################################################################################

# begin the testset
@testset "Adding & removing sites" begin

    # First, Unitcell
    @testset "Adding and removing in Unitcells (honeycomb)" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)

        # First, try to add site
        @testset "Adding sites" begin
            # try adding a random connection and testing the unitcell afterwards
            @test typeof(addSite!(unitcell, [1.0, 2.0])) == Int64
            @test testUnitcell(unitcell, 2,2)
        end;
        # Then, try to remove site
        @testset "Removing sites" begin
            # try removing the last (newly added) site and testing the unitcell afterwards
            @test removeSite!(unitcell, 3) == nothing
            @test testUnitcell(unitcell, 2,2)
            # try removing the first site and testing the unitcell afterwards
            @test removeSite!(unitcell, 1) == nothing
            @test testUnitcell(unitcell, 2,2)
        end;

    # end the Unitcell testset here
    end;

    # Then, Lattice
    @testset "Adding and removing in Lattices (honeycomb 4x2 periodic)" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # get a test lattice
        lattice = getLattice(unitcell, [-4, -2])

        # First, try to add site
        @testset "Adding sites" begin
            # try adding a random connection and testing the unitcell afterwards
            @test typeof(addSite!(lattice, [1.0, 2.0])) == Int64
            @test testLattice(lattice, 2,2)
        end;
        # Then, try to remove site
        @testset "Removing sites" begin
            # try removing the last (newly added) site and testing the unitcell afterwards
            @test removeSite!(lattice, 17) == nothing
            @test testLattice(lattice, 2,2)
            # try removing some intermediate site and testing the unitcell afterwards
            @test removeSite!(lattice, 5) == nothing
            @test testLattice(lattice, 2,2)
        end;

    # end the Lattice testset here
    end;

# end the testset here
end;

################################################################################
#
#   3) CONNECTION STRENGTH MODIFICATION
#       - all connections
#       - mapping of strengths
#       - evaluate strengths
#       - optimize connections
#
################################################################################

# begin the testset
@testset "Modifying connection strengths" begin

    # setting all connection strengths
    @testset "Setting all connection strengths" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # then, get a test lattice
        lattice  = getLattice(unitcell, [-4,-2])

        # First, Unitcell
        @testset "Setting in Unitcell" begin
            @test setAllInteractionStrengths!(unitcell, 1.0) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(getConnectionStrengthList(unitcell)) == 1
        end;
        # Then, Lattice
        @testset "Setting in Lattice" begin
            @test setAllInteractionStrengths!(lattice, 1.0) == nothing
            @test testLattice(lattice, 2,2)
            @test length(getConnectionStrengthList(lattice)) == 1
        end;

    # end the testset here
    end;

    # Only mapping desired connection strengths
    @testset "Mapping connection strengths" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # then, get a test lattice
        lattice  = getLattice(unitcell, [-4,-2])

        # First, Unitcell
        @testset "Mapping in Unitcell" begin
            # first, map only a part
            mapping = Dict("tx"=>"t", "ty"=>"t")
            @test mapInteractionStrengths!(unitcell, mapping) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(getConnectionStrengthList(unitcell)) == 2
            # then, map inside strings
            mapping = Dict("t"=>"J")
            @test mapInteractionStrengths!(unitcell, mapping, replace_in_strings=true) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(getConnectionStrengthList(unitcell)) == 2
            @test getConnectionStrengthList(unitcell)[1][1] == 'J'
            # then, map and evaluate
            mapping = Dict("J"=>1.0, "Jz"=>2.0)
            @test mapInteractionStrengths!(unitcell, mapping, replace_in_strings=false, evaluate=true) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(getConnectionStrengthList(unitcell)) == 2
            @test typeof(getConnectionStrengthList(unitcell)[1]) == Float64
        end;
        # Then, Lattice
        @testset "Mapping in Lattice" begin
            # first, map only a part
            mapping = Dict("tx"=>"t", "ty"=>"t")
            @test mapInteractionStrengths!(lattice, mapping) == nothing
            @test testLattice(lattice, 2,2)
            @test length(getConnectionStrengthList(lattice)) == 2
            # then, map inside strings
            mapping = Dict("t"=>"J")
            @test mapInteractionStrengths!(lattice, mapping, replace_in_strings=true) == nothing
            @test testLattice(lattice, 2,2)
            @test length(getConnectionStrengthList(lattice)) == 2
            @test getConnectionStrengthList(lattice)[1][1] == 'J'
            # then, map and evaluate
            mapping = Dict("J"=>1.0, "Jz"=>2.0)
            @test mapInteractionStrengths!(lattice, mapping, replace_in_strings=false, evaluate=true) == nothing
            @test testLattice(lattice, 2,2)
            @test length(getConnectionStrengthList(lattice)) == 2
            @test typeof(getConnectionStrengthList(lattice)[1]) == Float64
        end;

    # end the testset here
    end;

    # Only mapping desired connection strengths
    @testset "Evaluating connection strengths" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # then, get a test lattice
        lattice  = getLattice(unitcell, [-4,-2])

        # First, Unitcell
        @testset "Evaluating in Unitcell" begin
            # evaluate each kitaev parameter
            parameters = Dict("tx"=>1.0, "ty"=>1.0, "tz"=>2.0)
            @test evaluateInteractionStrengths!(unitcell, parameters) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(getConnectionStrengthList(unitcell)) == 2
            @test typeof(getConnectionStrengthList(unitcell)[1]) == Float64
        end;
        # Then, Lattice
        @testset "Evaluating in Lattice" begin
            # evaluate each kitaev parameter
            parameters = Dict("tx"=>1.0, "ty"=>1.0, "tz"=>2.0)
            @test evaluateInteractionStrengths!(lattice, parameters) == nothing
            @test testLattice(lattice, 2,2)
            @test length(getConnectionStrengthList(lattice)) == 2
            @test typeof(getConnectionStrengthList(lattice)[1]) == Float64
        end;

    # end the Unitcell testset here
    end;

    # Optimizing all connection, first Float64
    @testset "Optimizing connections (Float64)" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(1)
        # then, get a test lattice
        lattice  = getLattice(unitcell, [-4,-2])

        # Optimizing, Unitcell
        @testset "Optimizing in Unitcell" begin
            # add a redundant connection
            @test addConnection!(unitcell, 1,2, 2.0, (0,-1), overwrite=true) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(unitcell.connections) == 8
            # optimize
            @test optimizeConnections!(unitcell) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(getConnectionStrengthList(unitcell)) == 2
            @test length(unitcell.connections) == 6
            @test typeof(getConnectionStrengthList(unitcell)[1]) == Float64
        end;
        # Optimizing, Lattice
        @testset "Optimizing in Lattice" begin
            # add a redundant connection
            @test addConnection!(lattice, 1,2, 2.0, (0,0), overwrite=true) == nothing
            @test testLattice(lattice, 2,2)
            @test length(lattice.connections) == 50
            # optimize
            @test optimizeConnections!(lattice) == nothing
            @test testLattice(lattice, 2,2)
            @test length(getConnectionStrengthList(lattice)) == 2
            @test length(lattice.connections) == 48
            @test typeof(getConnectionStrengthList(lattice)[1]) == Float64
        end;

    # end the Unitcell testset here
    end;

    # Optimizing all connection, second String
    @testset "Optimizing connections (String)" begin

        # get a test unitcell
        unitcell = getUnitcellHoneycomb(4)
        # then, get a test lattice
        lattice  = getLattice(unitcell, [-4,-2])

        # Optimizing, Unitcell
        @testset "Optimizing in Unitcell" begin
            # add a redundant connection
            @test addConnection!(unitcell, 1,2, "t", (0,-1), overwrite=true) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(unitcell.connections) == 8
            # optimize
            @test optimizeConnections!(unitcell) == nothing
            @test testUnitcell(unitcell, 2,2)
            @test length(getConnectionStrengthList(unitcell)) == 3
            @test length(unitcell.connections) == 6
            @test typeof(getConnectionStrengthList(unitcell)[1]) == String
        end;
        # Optimizing, Lattice
        @testset "Optimizing in Lattice" begin
            # add a redundant connection
            @test addConnection!(lattice, 1,2, "t", (0,0), overwrite=true) == nothing
            @test testLattice(lattice, 2,2)
            @test length(lattice.connections) == 50
            # optimize
            @test optimizeConnections!(lattice) == nothing
            @test testLattice(lattice, 2,2)
            @test length(getConnectionStrengthList(lattice)) == 4
            @test length(lattice.connections) == 48
            @test typeof(getConnectionStrengthList(lattice)[1]) == String
        end;

    # end the Unitcell testset here
    end;

# end the testset here
end;

# end the Modification test block here
end;


text:
import  Base: ==
export  Shape,
        Quadtree,
        Point

immutable Shape
    ##Use UNION to make Float and Int together if necessary
    bounds::Array{Float64, 2}
    function Shape(bounds::Array{Int64, 2})
        new(convert(Array{Float64,2}, bounds))
    end
    function Shape(bounds::Array{Float64, 2})
        new(convert(Array{Float64,2}, bounds))
    end
    function Shape(x0::Real, y0::Real, x1::Real, y1::Real)
        new([x0 y0; x0 y1; x1 y1; x1 y0])
    end
end

## Implement == (isequal) function for Shapes
==(a::Shape, b::Shape) = a.bounds == b.bounds

type Quadtree
    x0::Real
    y0::Real
    x1::Real
    y1::Real
    objects::Vector{Shape}
    ne::Quadtree
    nw::Quadtree
    sw::Quadtree
    se::Quadtree

    local objects

    function Quadtree(bounds::Vector{Float64})
        objects = Shape[]
        new(bounds[1],
            bounds[2],
            bounds[3],
            bounds[4],
            objects)
    end

    function Quadtree(x0::Int64, y0::Int64, x1::Int64, y1::Int64)
        objects = Shape[]
        new(x0,
            y0,
            x1,
            y1,
            objects)
    end

end

type Point
    x
    y
    function Point(b::Array{Float64, 2})
        x = b[1]
        y = b[2]
        new(x,y)
    end
end

ext: jl


text:
""" Learning rate for scheme (τ+t)^-κ """
mutable struct InverseDecay <: GradDescent.Optimizer
    opt_type::String
    t::Int64
    κ::Float64
    τ::Int64
    ρ::Float64
end

""" Construct Learning Rate Scheme satisfying Robbins-Monro conditions"""
function InverseDecay(;τ::Int=100,κ::Real=0.51)
    InverseDecay("Inverse Decay",0,κ,τ,-1.0)
end

params(opt::InverseDecay) = "τ=$(opt.τ), κ=$(opt.κ)"

function GradDescent.update(opt::InverseDecay, g_t::AbstractArray{T,N}) where {T<:Real,N}
    opt.t += 1
    opt.ρ = (opt.t+opt.τ)^(-opt.κ)
    return opt.ρ * g_t
end


text:
using DifferentiableNAS
using Flux
using Flux: throttle, logitcrossentropy, onecold, onehotbatch
using Zygote: @nograd
using StatsBase: mean
using CUDA
include("../CIFAR10.jl")
@nograd onehotbatch
@nograd softmax

@testset "DARTS training test" begin
    using DifferentiableNAS
    using Flux
    using Flux: throttle, logitcrossentropy, onecold, onehotbatch, Optimiser
    using Zygote: @nograd
    using StatsBase: mean
    using Parameters
    using CUDA
    using Distributions
    using BSON
    using Dates

    include("../CIFAR10.jl")
    include("../training_utils.jl")

    argparams = trial_params(batchsize = 32)

    num_ops = length(PRIMITIVES)

    m = DARTSModel(num_cells = 4, channels = 4) |> gpu

    optimiser_α = Optimiser(WeightDecay(1e-3),ADAM(3e-4,(0.5,0.999)))
    optimiser_w = Optimiser(WeightDecay(3e-4),Momentum(0.025, 0.9))

    train, val = get_processed_data(argparams.val_split, argparams.batchsize, argparams.trainval_fraction)
    test = get_test_data(argparams.test_fraction)

    histepoch = historiessml()
    histbatch = historiessml()
    losses = [0.0, 0.0]

    datesnow = Dates.now()
    base_folder = string("test/models/darts_", datesnow)
    mkpath(base_folder)

    cbepoch = CbAll(CUDA.reclaim, histepoch, save_progress, CUDA.reclaim)
    cbbatch = CbAll(CUDA.reclaim, histbatch, CUDA.reclaim)

    DARTStrain1st!(loss, m, train, val, optimiser_α, optimiser_w, losses; cbepoch = cbepoch, cbbatch = cbbatch)

    loss1 = loss(m, train[1][1] |> gpu, train[1][2] |> gpu)
    acc1 = accuracy(m, test...)

    DARTStrain1st!(loss, m, train, val, optimiser; cb = evalcb)

    loss2 = loss(m, train[1][1] |> gpu, train[1][2] |> gpu)
    acc2 = accuracy(m, test...)

    @test loss2 < loss1
    @test acc2 > acc1
end

max_stars_repo_name: huangyxi/DifferentiableNAS

max_issues_repo_name: huangyxi/DifferentiableNAS

max_forks_repo_name: huangyxi/DifferentiableNAS


text:
using Logging
using LoggingExtras: TeeLogger
using TensorBoardLogger
using Dates

# I'm going to have one and only one MNIST and CIFAR model, so pretrained model
# only takes dataset name.
function maybe_train(model_file, model, train_fn)
    if isfile(model_file)
        @info "Already trained. Loading .." model_file
        @load model_file model
        return gpu(model)
    else
        @info "Pre-training .."
        train_fn(model)
        @info "saving .."
        @save model_file model=cpu(model)
        return model
    end
end

function maybe_load(model_file, model_fn)
    if isfile(model_file)
        @info "Loading from $model_file .."
        @load model_file model from_steps
        model = gpu(model)
        # NOTE: add 1 as starting step
        from_steps += 1
    else
        @info "Starting from scratch .."
        model = model_fn()
        from_steps = 1
    end
    return gpu(model), from_steps
end

function adv_exp_helper(expID, lr, total_steps, λ,
                        model_fn, ds_fn, pretrain_fn;
                        attack_fn,
                        print_steps, save_steps,
                        test_per_steps, test_run_steps,
                        test_attack_fn=attack_fn)
    model_file = "trained/$expID.bson"
    mkpath(dirname(model_file))

    # TODO record training @time? I need to use the average time instead of
    # accumulated time, because accumulated time introduces states that needs to
    # be saved. This is not very urgent.
    #
    # FIXME what should be my batch_size?
    ds, test_ds = ds_fn();
    x, y = next_batch!(ds) |> gpu;

    model, from_steps = maybe_load(model_file, model_fn)
    if from_steps == 1
        model = pretrain_fn(model)
    end

    @info "Progress" total_steps from_steps
    # stops here to avoid model warming up overhead
    if from_steps > total_steps return end

    logger = TBLogger("tensorboard_logs/$expID/train", tb_append, min_level=Logging.Info)
    test_logger = TBLogger("tensorboard_logs/$expID/test", tb_append, min_level=Logging.Info)

    # warm up the model
    model(x)

    # FIXME opt states
    opt = ADAM(lr);

    save_cb = create_save_cb(model_file, model, save_steps=save_steps)
    test_cb = create_adv_test_cb(model, test_ds,
                                 logger=test_logger,
                                 attack_fn=test_attack_fn,
                                 test_per_steps=test_per_steps, test_run_steps=test_run_steps)

    advtrain!(model, opt, attack_fn, ds,
              train_steps=total_steps,
              from_steps=from_steps,
              print_steps=print_steps,
              λ=λ,
              logger=logger,
              save_cb=save_cb,
              test_cb=test_cb)

end

function nat_exp_helper(expID, lr, total_steps,
                        model_fn, ds_fn)
    model_file = "trained/$expID.bson"
    mkpath(dirname(model_file))

    ds, test_ds = ds_fn();
    x, y = next_batch!(ds) |> gpu;

    model, from_steps = maybe_load(model_file, model_fn)

    @info "Progress" total_steps from_steps
    # stops here to avoid model warming up overhead
    if from_steps > total_steps return end

    logger = TBLogger("tensorboard_logs/$expID/train", tb_append, min_level=Logging.Info)
    test_logger = TBLogger("tensorboard_logs/$expID/test", tb_append, min_level=Logging.Info)

    # warm up the model
    model(x)

    # FIXME opt states
    opt = ADAM(lr);
    #
    # TODO WRN used SGD with Nesterov momentum 0.9 and weight decay 0.0005 start
    # lr=0.1, drop by 0.2 at 60,120,160 epoch
    #
    # FIXME weight decay?
    #
    # Nesterov(lr, 0.9)
    # Momentum(lr, 0.9)

    save_cb = create_save_cb(model_file, model, save_steps=40)
    test_cb = create_test_cb(model, test_ds, logger=test_logger,
                             test_per_steps=100, test_run_steps=20)

    train!(model, opt, ds,
           # DEBUG adding cifar augment. It does not seem to provide any improvement
           # augment=Augment(),
           train_steps=total_steps,
           from_steps=from_steps,
           print_steps=20,
           logger=logger,
           save_cb=save_cb,
           test_cb=test_cb)
end

function advae_exp_helper(expID, lr, total_steps,
                          ae_model_fn, ds_fn,
                          pretrain_fn, pretrained_cnn_fn;
                          λ, γ, β,
                          attack_fn,
                          print_steps, save_steps,
                          test_per_steps, test_run_steps)
    model_file = "trained/$expID.bson"
    @show model_file
    mkpath(dirname(model_file))

    ds, test_ds = ds_fn();
    x, y = next_batch!(ds) |> gpu;

    # here we have two models, AE and CNN
    ae, from_steps = maybe_load(model_file, ae_model_fn)
    if from_steps == 1
        ae = pretrain_fn(ae)
    end
    # load pretrained CNN
    cnn = pretrained_cnn_fn()
    # set test mode
    Flux.testmode!(cnn)
    # warm up the model
    cnn(ae(x))

    @info "Progress" total_steps from_steps
    if from_steps > total_steps return end

    logger = TBLogger("tensorboard_logs/$expID/train", tb_append, min_level=Logging.Info)
    test_logger = TBLogger("tensorboard_logs/$expID/test", tb_append, min_level=Logging.Info)

    opt = ADAM(lr);

    save_cb = create_save_cb(model_file, ae, save_steps=save_steps)
    test_cb = create_advae_test_cb(ae, cnn, test_ds,
                                   logger=test_logger,
                                   attack_fn=attack_fn,
                                   test_per_steps=test_per_steps, test_run_steps=test_run_steps)

    advae_train!(ae, cnn, opt, attack_fn, ds,
                 train_steps=total_steps,
                 from_steps=from_steps,
                 print_steps=print_steps,
                 λ=λ,
                 γ=γ,
                 logger=logger,
                 save_cb=save_cb,
                 test_cb=test_cb)

end

function free_exp_helper(expID, lr, total_steps, ε,
                         model_fn, ds_fn, pretrain_fn;
                         print_steps, save_steps,
                         test_per_steps, test_run_steps,
                         test_attack_fn=attack_fn)
    model_file = "trained/$expID.bson"
    mkpath(dirname(model_file))

    ds, test_ds = ds_fn();
    x, y = next_batch!(ds) |> gpu;

    model, from_steps = maybe_load(model_file, model_fn)
    if from_steps == 1
        model = pretrain_fn(model)
    end

    @info "Progress" total_steps from_steps
    # stops here to avoid model warming up overhead
    if from_steps > total_steps return end

    logger = TBLogger("tensorboard_logs/$expID/train", tb_append, min_level=Logging.Info)
    test_logger = TBLogger("tensorboard_logs/$expID/test", tb_append, min_level=Logging.Info)

    # warm up the model
    model(x)

    # FIXME opt states
    opt = ADAM(lr);

    save_cb = create_save_cb(model_file, model, save_steps=save_steps)
    test_cb = create_adv_test_cb(model, test_ds,
                                 logger=test_logger,
                                 attack_fn=test_attack_fn,
                                 test_per_steps=test_per_steps, test_run_steps=test_run_steps)

    free_train!(model, opt, ds, ε,
                train_steps=total_steps,
                from_steps=from_steps,
                print_steps=print_steps,
                logger=logger,
                save_cb=save_cb,
                test_cb=test_cb)

end

max_stars_repo_name: lihebi/AdvAE

max_issues_repo_name: lihebi/AdvAE

max_forks_repo_name: lihebi/AdvAE


text:
using Documenter, Test

include("test_synthetic_observations.jl")
include("test_forward_map.jl")
include("test_ensemble_column_models.jl")
include("test_turbulence_closure_parameters.jl")
include("test_eki.jl")

@testset "Doctests" begin
    doctest(OceanTurbulenceParameterEstimation)
end


text:
function getstructuregraphs()
    ontologies_api.OntologiesApi().get_structure_graphs()
end


text:
# Utils

# used in amcep in C
function phidf!(x::Cdouble, order, α, delay::StridedVector{Cdouble})
    if length(delay) != order+1
        throw(ArgumentError("inconsistent order or delay"))
    end

    ccall((:phidf, libSPTK), Void, (Cdouble, Cint, Cdouble, Ptr{Cdouble}),
          x, order, α, delay)
end

function lspcheck(lsp::StridedVector{Cdouble})
    ccall((:lspcheck, libSPTK),
          Cint, (Ptr{Cdouble}, Cint), lsp, length(lsp)-1)
end


text:
module V

const species = [
    # Birtwistle et al., MSB (2007)
    "E",
    "H",
    "E1",
    "E2",
    "E3",
    "E4",
    "E_E1",
    "H_E3",
    "H_E4",
    "E11",
    "E12",
    "E23",
    "E34",
    "E24",
    "E44",
    "E11P",
    "E12P",
    "E23P",
    "E34P",
    "E24P",
    "E44P",
    "G",
    "S",
    "I",
    "R",
    "O",
    "A",
    "E11G",
    "E11S",
    "E11R",
    "E12G",
    "E12S",
    "E12R",
    "E23G",
    "E23S",
    "E23I",
    "E23R",
    "E34G",
    "E34S",
    "E34I",
    "E34R",
    "E24G",
    "E24S",
    "E24I",
    "E24R",
    "E44G",
    "E44S",
    "E44I",
    "E44R",
    "sigmaG",
    "sigmaS",
    "sigmaI",
    "sigmaR",
    "sigmaA",
    "sigmaSP",
    "sigmaAP",
    "sigmaG_O",
    "sigmaG_A",
    "sigmaSP_G",
    "sigmaAP_S",
    "sigmaAP_I",
    "sigmaAP_R",
    "P3_A",
    "P2",
    "P3",
    "Akt",
    "Aktstar",
    "RsD",
    "RsT",
    "sigmaRP",
    "Raf",
    "Rafstar",
    "MEK",
    #"MEKstar",
    #"ERK",
    #"ERKstar",
    "OP",
    "AP",
    "A_sigmaG_O",
    "sigmaA_G",
    "sigmaA_G_O",
    "sigmaO",
    "E13",
    "E14",
    "E13P",
    "E14P",
    "E13G",
    "E13S",
    "E13I",
    "E13R",
    "E14G",
    "E14S",
    "E14I",
    "E14R",
    "fint",
    "T",
    "E11T",
    "E12T",
    "E23T",
    "E34T",
    "E24T",
    "E44T",
    "E13T",
    "E14T",
    "sigmaAP_T",
    "sigmaT",
    "E1_PT",
    "E2_PT",
    "E4_PT",
    "E_E1_PT",
    "H_E4_PT",
    # Nakakuki et al., Cell (2010)
    "ppMEKc",
    "CREBn",
    "pCREBn",
    "ERKc",
    "ERKn",
    "pERKc",
    "pERKn",
    "ppERKc",
    "ppERKn",
    "Elk1n",
    "pElk1n",
    "cFOSc",
    "cFOSn",
    "pcFOSc",
    "pcFOSn",
    "DUSPc",
    "DUSPn",
    "pDUSPc",
    "pDUSPn",
    "DUSPn_ERKn",
    "DUSPn_pERKn",
    "DUSPn_ppERKn",
    "pDUSPn_ERKn",
    "pDUSPn_pERKn",
    "pDUSPn_ppERKn",
    "RSKc",
    "pRSKc",
    "pRSKn",
    "PrecfosmRNAn",
    "PreduspmRNAn",
    "cfosmRNAc",
    "duspmRNAc",
    "Fc",
    "Fn",
    "FmRNAc",
    "PreFmRNAn",
]

for (idx,name) in enumerate(species)
    eval(Meta.parse("const $name = $idx"))
end

const n_species = length(species)

end  # module


text:
module UnicodeData

using Serialization

export udata

# Deserialized the dictionary built by the "build" step.
const serialized_path = joinpath(@__DIR__, "..", "deps", "unicode-dict.serial")
const CodepointDict = open(deserialize, serialized_path)

"""
udata(codepoint::T) where T <: Union{<:AbstractChar, <:Integer}) -> Union{String, Nothing}

Return a string containing information about `codepoint`. If `codepoint` has not been
assigned, return `nothing`.
"""
udata(x::UInt32) = get(CodepointDict, x, nothing)
udata(x::AbstractChar) = udata(codepoint(x))
udata(x::Integer) = udata(UInt32(x))

end # module


text:
toml_str(a; kwargs...) = sprint(io -> TOML.print(io, a; kwargs...))
toml_str(f, a; kwargs...) = sprint(io -> TOML.print(f, io, a; kwargs...))

@test toml_str(Dict("b" => 1, "c" => 2, "a" => 3); sorted=true) ==
    """
    a = 3
    b = 1
    c = 2
    """

@test toml_str(Dict("b" => 1, "ac" => 2, "abc" => 3); sorted=true, by=length) ==
    """
    b = 1
    ac = 2
    abc = 3
    """

struct MyStruct
    a::Int
end
@test_throws ErrorException toml_str(Dict("foo" => MyStruct(1)))
@test toml_str(Dict("foo" => MyStruct(1))) do x
        x isa MyStruct && return x.a
    end == """
        foo = 1
        """

@test toml_str(Dict("b" => SubString("foo"))) == "b = \"foo\"\n"


text:
#name: JuDecTreeApply
#meta.mlname: JuDecTree
#meta.mlrole: apply
#description: Custom Julia Apply func for Decision Tree
#language: julia
#input: blob model
#input: dataframe df
#input: string namesKeys [original features' names]
#input: string namesValues [new features' names]
#output: dataframe data_out

#>>> Load necessary packages
using Random
using BSON: @load
using DecisionTree

#>>> If original(namesKeys) and new(namesValues) passed, map original names to new
namesValues=split(namesValues, ",")
namesKeys=split(namesKeys, ",")
if length(namesKeys) > 0
    featuresNames = names(df)
    for i in 1:length(namesKeys)
        replace!(featuresNames, namesValues[i]  =>  namesKeys[i])
    end
end
testX=Matrix(select(df, Symbol.(featuresNames)))

#>>> Retrieve saved/trained model
@load file=model trained_model

#>>> Predict using trained model
predY = predict(trained_model, testX)
data_out = DataFrame([predY],[:out])

max_stars_repo_name: DotNetPart/public

max_issues_repo_name: DotNetPart/public

max_forks_repo_name: DotNetPart/public


text:
#=
function x0_volume(model::EoS,z; phase = :unknown)
    if phase == :unknown || is_liquid(phase)
        if typeof(model)<:SAFTgammaMie
            x0 = [log10(π/6*N_A*sum(z[i]*sum(model.group_multiplicities[i][k]*model.params.segment[k]*model.params.shapefactor[k]*model.params.sigma[k]^3 for k in @groups(i)) for i in @comps)/0.8)]
        elseif typeof(model)<:SAFT
            x0 = [log10(π/6*N_A*sum(z[i]*model.params.segment[i]*model.params.sigma[i]^3 for i in model.components)/0.8)]
        elseif typeof(model)<:Cubic
            x0 = [log10(sum(z[i]*z[j]*model.params.b[union(i,j)] for i in model.components for j in model.components)/0.8)]
        end
    elseif is_vapour(phase)
        if typeof(model)<:SAFTgammaMie
            x0 = [log10(π/6*N_A*sum(z[i]*sum(model.group_multiplicities[i][k]*model.params.segment[k]*model.params.shapefactor[k]*model.params.sigma[k]^3 for k in @groups(i)) for i in @comps)/1e-2)]
        elseif typeof(model)<:SAFT
            x0 = [log10(π/6*N_A*sum(z[i]*model.params.segment[i]*model.params.sigma[i]^3 for i in model.components)/1e-2)]
        elseif typeof(model)<:Cubic
            x0 = [log10(sum(z[i]*z[j]*model.params.b[union(i,j)] for i in model.components for j in model.components)/1e-2)]
        end
    elseif is_supercritical(phase)
        if typeof(model)<:SAFTgammaMie
            x0 = [log10(π/6*N_A*sum(z[i]*sum(model.group_multiplicities[i][k]*model.params.segment[k]*model.params.shapefactor[k]*model.params.sigma[k]^3 for k in @groups(i)) for i in @comps)/0.5)]
        elseif typeof(model)<:SAFT
            x0 = [log10(π/6*N_A*sum(z[i]*model.params.segment[i]*model.params.sigma[i]^3 for i in model.components)/0.5)]
        elseif typeof(model)<:Cubic
            x0 = [log10(sum(z[i]*z[j]*model.params.b[union(i,j)] for i in model.components for j in model.components)/0.5)]
        end
    end

    return x0
end
=#

#=x0_volume=#
# function x0_volume(model::SAFTgammaMie,z; phase = :unknown)
#     if phase == :unknown || is_liquid(phase)
#             x0 = [log10(π/6*N_A*sum(z[i]*sum(model.group_multiplicities[i][k]*model.params.segment[k]*model.params.shapefactor[k]*model.params.sigma[k]^3 for k in @groups(i)) for i in @comps)/0.8)]
#     elseif is_vapour(phase)
#             x0 = [log10(π/6*N_A*sum(z[i]*sum(model.group_multiplicities[i][k]*model.params.segment[k]*model.params.shapefactor[k]*model.params.sigma[k]^3 for k in @groups(i)) for i in @comps)/1e-2)]
#     elseif is_supercritical(phase)
#             x0 = [/0.5)]
#     end
#     return x0
# end

# function x0_volume(model::LJSAFT,z; phase = :unknown)
#     if phase == :unknown || is_liquid(phase)
#         x0 = [log10(π/6*sum(z[i]*model.params.segment[i]*model.params.b[i] for i in model.components)/0.8)]
#     elseif is_vapour(phase)
#         x0 = [log10(π/6*sum(z[i]*model.params.segment[i]*model.params.b[i] for i in model.components)/1e-2)]
#     elseif is_supercritical(phase)
#         x0 = [/0.5)]
#     end
#     return x0
# end

#=
lb_volume:

SAFTgammaMie:  log10(π/6*N_A*sum(z[i]*sum(model.group_multiplicities[i][k]*model.params.segment[k]*model.params.shapefactor[k]*model.params.sigma[k]^3 for k in @groups(i)) for i in @comps)
LJSAFT: log10(π/6*sum(z[i]*model.params.segment[i]*model.params.b[i] for i in @comps)
=#

"""
    x0_volume_liquid(model,T,z)

Returns an initial guess to the liquid volume, dependent on temperature and composition. by default is 1.25 times the [lower bound volume](@ref lb_volume).
"""
function x0_volume_liquid(model,T,z)
    v_lb = lb_volume(model,z)
    return v_lb*1.25
end

function x0_volume_liquid(model::SAFTVRMieModel,T,z)
    v_lb = lb_volume(model,z)
    return v_lb*1.5
end
"""
    x0_volume_gas(model,p,T,z)

Returns an initial guess to the gas volume, depending of pressure, temperature and composition. by default uses a [virial aproximation](@ref volume_virial)
"""
function x0_volume_gas(model,p,T,z)
    return volume_virial(model,p,T,z)
end

function x0_volume_sc(model,p,T,z)
    v_sc = lb_volume(model,z)
    return v_sc*2
end

"""
    x0_volume(model::EoSModel,p,T,z; phase = :unknown)

Returns an initial guess of the volume at a pressure, temperature, composition and suggested phase.

If the suggested phase is `:unkwown` or `:liquid`, calls [`x0_volume_liquid`](@ref).

If the suggested phase is `:gas`, calls [`x0_volume_gas`](@ref).

"""
function x0_volume(model::EoSModel,p,T,z; phase = :unknown)
    phase = Symbol(phase)
    if phase === :unknown || is_liquid(phase)
        return x0_volume_liquid(model,T,z)
    elseif is_vapour(phase)
        return x0_volume_gas(model,p,T,z)
    elseif is_supercritical(phase)
     else
        error("unreachable state on x0_volume")
    end
end

#=x0_sat_pure=#

# function x0_sat_pure(model::SAFTVRQMie)
#     x0    = [log10(π/6*N_A*model.params.segment[model.components[1]]*model.params.sigma[model.components[1]]^3/0.2),
#     log10(π/6*N_A*model.params.segment[model.components[1]]*model.params.sigma[model.components[1]]^3/1e-3)]
# end

# function x0_sat_pure(model::LJSAFT)
#     x0    = [log10(π/6*model.params.segment[model.components[1]]*model.params.b[model.components[1]]/0.5),
#     log10(π/6*model.params.segment[model.components[1]]*model.params.b[model.components[1]]/1e-3)]
# end

##=lb_volume=#
#
#lb_volume(model::LJSAFT,z; phase = :unknown) = [log10(π/6*sum(z[i]*model.params.segment[i]*model.params.b[i] for i in model.components)/1)]

"""
    lb_volume(model::EoSModel,z=SA[1.0])

Returns the lower bound volume.

It has different meanings depending on the Equation of State, but symbolizes the minimum allowable volume at a certain composition:

- SAFT EoS: the packing volume
- Cubic EoS, covolume (b) parameter

On empiric equations of state, the value is chosen to match the volume of the conditions at maximum pressure and minimum temperature
, but the equation itself normally can be evaluated at lower volumes.

On SAFT and Cubic EoS, volumes lower than `lb_volume` will likely error.

The lower bound volume is used for guesses of liquid volumes at a certain pressure, saturated liquid volumes and critical volumes.
"""
function lb_volume end

function lb_volume(model::SAFTModel, z = SA[1.0])
    seg = model.params.segment.values
    σᵢᵢ = model.params.sigma.diagvalues
    val = π/6*N_A*sum(z[i]*seg[i]*σᵢᵢ[i]^3 for i in 1:length(z))
    return val
end

function lb_volume(model::CubicModel,z = SA[1.0])
    n = sum(z)
    invn = one(n)/n
    b = model.params.b.values
    b̄ = dot(z,Symmetric(b),z)*invn*invn
    return b̄
end

function lb_volume(model::CPAModel,z = SA[1.0])
    n = sum(z)
    invn = one(n)/n
    b = model.params.b.values
    b̄ = dot(z,Symmetric(b),z)*invn*invn
    return b̄

end

function lb_volume(model::SAFTgammaMieModel, z = SA[1.0])
    vk  = model.igroups
    seg = model.params.segment.values
    S   = model.params.shapefactor.values
    σᵢᵢ = model.params.sigma.diagvalues
    val = π/6*N_A*sum(z[i]*sum(vk[i][k]*seg[k]*S[k]*σᵢᵢ[k]^3 for k in @groups(i)) for i in @comps)
    return val
end

#=scale_sat_pure=#

# function scale_sat_pure(model::SAFTgammaMie)
#     m̄  = sum(model.group_multiplicities[model.components[1]][k]*model.params.segment[k]*model.params.shapefactor[k] for k in @groups(model.components[1]))
#         σ̄3 = sum(model.group_multiplicities[model.components[1]][k]*model.group_multiplicities[model.components[1]][l]*
#                  model.params.segment[k]*model.params.segment[l]*
#                  model.params.shapefactor[k]*model.params.shapefactor[l]*
#                  model.params.sigma[union(k,l)]^3 for k in @groups(model.components[1]) for l in @groups(model.components[1]))/m̄^2
#         ϵ̄  = T_crit_pure(model)
#         p_scale    = σ̄3*N_A/R̄/ϵ̄
#         μ_scale    = 1/R̄/ϵ̄
#     return p_scale,μ_scale
# end

# function scale_sat_pure(model::LJSAFT)
#     p_scale    = model.params.b[model.components[1]]/N_A/R̄/model.params.T[model.components[1]]
#     μ_scale    = 1/R̄/model.params.T[model.components[1]]
#     return p_scale,μ_scale
# end

"""
    x0_sat_pure(model::EoSModel,T,z=SA[1.0])

Returns a 2-element vector corresponding to `[log10(Vₗ),log10(Vᵥ)]`, where Vₗ and Vᵥ are the liquid and vapor initial guesses.
Used in [`sat_pure`](@ref).
"""
function x0_sat_pure(model::EoSModel,T,z=SA[1.0])
    b = lb_volume(model,z)*one(T)
    B = second_virial_coefficient(model,T,z)
    x0v = -2*B + 2*b
    p = -0.25*R̄*T/B
    x0l = Ref(b)
    try
        x0l[] = volume_compress(model,p,T,z)
    catch
        x0l[] = b/0.25
    end

    if x0l[] >x0v || isnan(x0l[])
        x0l[] = b/0.25
    end
    #=here we solve the saturation with aproximate
    models of the EoS. on the gas side, we use
    a virial for logϕ, on the liquid side, we
    use an isothermal compressibility model based
    at the point P = P(-2B), vl(P)

    β = vt_isothermal_compressibility(model,x0l,T,z)
    fugv(P) = B*P/(R̄*T) #log(ϕv)
    P0 = p
    v0 = x0l
    fuglx(P) = -((v0*exp(β*P0)/R̄*T*β)*exp(-β*P) -log(P))
    fugl0 = fuglx(P0)
    fugl(P) = fuglx(P) - fugl0
    @show fugv(0.8P0)
    @show fugl(0.8P0)
    @show P0
    f0(P) = fugl(P) - fugv(P)
    @show Roots.find_zero(f0,0.9P0)
     =#
    x0  = [x0l[],x0v]

    return log10.(x0)
end
#=
appendix: logϕ for the isothermal compressibility aprox
from the volume_compress code:

    Δ(P) =  -(P-p0)*β = -βP + βP0
    v(P) = v0*exp(Δ(P)) #volume_compress uses a convergence modification

    Z-1 = v/RT - 1

    definition of fugacity coefficient

    logϕ = ∫(Z-1)/P dp from Psat to Pmax
    Z(P) = P*v(P)/RT
    logϕ = ∫v/RT -1/P dP
    logϕ = ∫v0*exp(βP0)*exp(-βP)/RT -1/P dP
    logϕ = (v0*exp(βP0)/RT)∫exp(-βP) dp -  ∫1/P dp
    logϕ = -(v0*exp(βP0)/RTβ)*exp(-βP) -log(P)

=#

function scale_sat_pure(model::EoSModel,z=SA[1.0])
    p    = 1/p_scale(model,z)
    μ    = 1/R̄/T_scale(model,z)
    return p,μ
end

#=x0_crit_pure=#

# x0_crit_pure(model::LJSAFT) = [1.5, log10(π/6*model.params.segment[model.components[1]]*model.params.b[model.components[1]]/0.3)]

"""
    x0_crit_pure(model::SAFTModel,z=SA[1.0])

Returns a 2-element vector corresponding to
    `[k,log10(Vc0)]`, where `k` is `Tc0/T_scale(model,z)`
"""
function x0_crit_pure end

function x0_crit_pure(model::SAFTModel,z=SA[1.0])
    lb_v = lb_volume(model,z)
    [2.0, log10(lb_v/0.3)]
end

function x0_crit_pure(model::CubicModel,z=SA[1.0])
    lb_v = lb_volume(model,z)
    [1.0, log10(lb_v/0.3)]
end

function x0_crit_pure(model::CPAModel,z=SA[1.0])
    lb_v = lb_volume(model,z)
    [2.0, log10(lb_v/0.3)]
end

function x0_crit_pure(model::EoSModel,z=SA[1.0])
    lb_v = lb_volume(model,z)
    [1.5, log10(lb_v/0.3)]
end

# x0_crit_pure(model::Cubic) = [1.0, log10(model.params.b[model.components[1]]/0.3)]
# x0_crit_pure(model::CPA) = [2, log10(model.params.b[model.components[1]]/0.3)]

#=T_crit_pure=#
# function T_crit_pure(model::SAFTgammaMie)
#     m̄ = sum(model.group_multiplicities[model.components[1]][k]*model.params.segment[k]*model.params.shapefactor[k] for k in @groups(model.components[1]))
#     return sum(model.group_multiplicities[model.components[1]][k]*model.group_multiplicities[model.components[1]][l]*
#                model.params.segment[k]*model.params.segment[l]*
#                model.params.shapefactor[k]*model.params.shapefactor[l]*
#                model.params.epsilon[union(k,l)] for k in @groups(model.components[1]) for l in @groups(model.components[1]))/m̄^2
# end
# T_crit_pure(model::LJSAFT) = model.params.T[model.components[1]]

# T_crit_pure(model::Cubic) = model.params.a[model.components[1]]/model.params.b[model.components[1]]/8.314*8/27

#=
 temperature scaling factor,
on critical based EoS, is a function of critical temperature
on SAFT EoS, is a function of ϵ
=#

"""
    T_scale(model::EoS,z=SA[1.0])

Represents a temperature scaling factor.

On any EoS based on Critical parameters (Cubic or Empiric EoS), the temperature scaling factor is chosen to be the critical temperature.

On SAFT or other molecular EoS, the temperature scaling factor is chosen to be a function of the potential depth ϵ.

Used as scaling factors in [`sat_pure`](@ref) and as input for solving [`crit_pure`](@ref)
"""
function T_scale end

function T_scale(model::SAFTModel,z=SA[1.0])
    ϵ = model.params.epsilon.diagvalues
    return prod(ϵ)^(1/length(ϵ))
end

#dont use αa, just a, to avoid temperature dependence
function T_scale(model::CubicModel,z=SA[1.0])
    n = sum(z)
    invn2 = one(n)/(n*n)
    Ωa,Ωb = ab_consts(model)
    _a = model.params.a.values
    _b = model.params.b.values
    a = dot(z, Symmetric(_a), z)*invn2/Ωa
    b = dot(z, Symmetric(_b), z)*invn2/Ωb
    return a/b/R̄
end

function T_scale(model::CPAModel,z=SA[1.0])
    n = sum(z)
    invn2 = one(n)/(n*n)
    Ωa,Ωb = ab_consts(model)
    _a = model.params.a.values
    _b = model.params.b.values
    a = dot(z, Symmetric(_a), z)*invn2/Ωa
    b = dot(z, Symmetric(_b), z)*invn2/Ωb
    return a/b/R̄

end

"""
    p_scale(model::SAFTModel,z=SA[1.0])

Represents a pressure scaling factor

On any EoS based on Critical parameters (Cubic or
Empiric EoS), the pressure scaling factor is
chosen to be a function of the critical pressure.

On SAFT or other molecular EoS, the temperature
scaling factor is chosen to a function of ∑(zᵢ*ϵᵢ*(σᵢᵢ)³)

Used as scaling factors in [`sat_pure`](@ref) and as input for solving [`crit_pure`](@ref)

"""
function p_scale(model::SAFTModel,z=SA[1.0])
    ϵ = model.params.epsilon.diagvalues
    σᵢᵢ = model.params.sigma.diagvalues
    val =  sum(z[i]*σᵢᵢ[i]^3/ϵ[i] for i in 1:length(z))*N_A/R̄
    return 1/val
end

function p_scale(model::SAFTgammaMieModel,z=SA[1.0])
    vk  = model.igroups
    seg = model.params.segment.values
    S   = model.params.shapefactor.values
    σ   = model.params.sigma.values
    m̄  = sum(vk[1][k]*seg[k]*S[k] for k in @groups(model.icomponents[1]))

    σ̄3 = sum(vk[1][k]*vk[1][l]*
                     seg[k]*seg[l]*
                     S[k]*S[l]*
                     σ[k,l]^3 for k in @groups(model.icomponents[1]) for l in @groups(model.icomponents[1]))/m̄^2
    ϵ̄ = T_scale(model,z)
    val    = σ̄3*N_A/R̄/ϵ̄
    return 1/val
end

function p_scale(model::CubicModel,z=SA[1.0])
    n = sum(z)
    invn2 = (1/n)^2
    Ωa,Ωb = ab_consts(model)
    _a = model.params.a.values
    _b = model.params.b.values
    a = invn2*dot(z, Symmetric(_a), z)/Ωa
    b = invn2*dot(z, Symmetric(_b), z)/Ωb
    return a/ (b^2) # Pc mean
end

function p_scale(model::CPAModel,z=SA[1.0])
    n = sum(z)
    invn2 = (1/n)^2
    Ωa,Ωb = ab_consts(model)
    _a = model.params.a.values
    _b = model.params.b.values
    a = invn2*dot(z, Symmetric(_a), z)/Ωa
    b = invn2*dot(z, Symmetric(_b), z)/Ωb
    return a/ (b^2) # Pc mean
end

#=
the following methods are fallbacks,
that require just the definition of T_scale,p_scale and lb_volume
respectively. if possible, each eos should define those
=#
function T_scales(model,z)
    n = length(z)
    x = zeros(n)
    res = zeros(n)
    for i = 1:n
        x[i] = 1.0
        res[i] = T_scale(model,x)
        x[i] = 0.0
    end
    return res
end

function lb_volumes(model,z)
    n = length(z)
    x = zeros(n)
    res = zeros(n)
    for i = 1:n
        x[i] = 1.0
        res[i] = lb_volume(model,x)
        x[i] = 0.0
    end
    return res
end

function p_scales(model,z)
    n = length(z)
    x = zeros(n)
    res = zeros(n)
    for i = 1:n
        x[i] = 1.0
        res[i] = p_scale(model,x)
        x[i] = 0.0
    end
end


text:
# Copied from
# https://github.com/tkf/Restacker.jl/blob/master/test/test_core.jl
module TestRestack

using Test
using Transducers: restack

struct ABC{A,B,C}
    a::A
    b::B
    c::C
end

struct CustomizedProperties{A,B,C}
    a::A
    b::B
    c::C
end

Base.getproperty(::CustomizedProperties, n::Symbol) = n

struct IncompleteInitialization{A,B}
    a::A
    b::B
    IncompleteInitialization{A,B}(a::A) where {A,B}= new{A,B}(a)
end

function testlabel(x)
    n = 50
    s = repr(x)
    length(s) < n && return s
    label = join(Iterators.take(s, n))
    return string(label, "…")
end

@testset "$(testlabel(x))" for x in Any[
    nothing,
    missing,
    1,
    1.0,
    NaN,
    1+2im,
    :symbol,
    "string",
    Int,
    Array,
    Vector,
    Vector{Int},
    ABC,
    ABC{Int},
    ABC{Int,Int,Int},
    (nothing, missing, 1, 1.0, NaN, 1 + 2im, :symbol, "string"),
    (Val(1), Val(2), nothing, missing, undef),
    (
        a = (nothing, missing, 1, 1.0, NaN, 1 + 2im, :symbol, "string"),
        b = (Val(1), Val(2), nothing, missing, undef),
    ),
    ABC(
        (nothing, missing, 1, 1.0, NaN, 1 + 2im, :symbol, "string"),
        (Val(1), Val(2), nothing, missing, undef),
        (d = 1, e = ABC(1, 2, 3)),
    ),
    CustomizedProperties(1, 2, 3),
    CustomizedProperties(1, 2, ABC(1, 2, CustomizedProperties(1, 2, 3))),
    IncompleteInitialization{Int,String}(2)
]
    @test restack(x) === x
end

@testset "CustomizedProperties" begin
    abc = CustomizedProperties(1, 2, 3)
    @test abc.a === :a
    @test getfield(abc, :a) === 1
end

end  # module


text:
using Oceananigans.TurbulenceClosures: VerstappenAnisotropicMinimumDissipation

interior(a, grid) = view(a, grid.Hx+1:grid.Nx+grid.Hx,
                            grid.Hy+1:grid.Ny+grid.Hy,
                            grid.Hz+1:grid.Nz+grid.Hz,
                        )

# Temporary until k index is reversed
const Nx, Ny, Nz = 16, 16, 16
const Hx, Hy, Hz = 1, 1, 1

function get_fields_from_checkpoint(filename)

    file = jldopen(filename)

    solution = (
        u = reverse(file["velocities/u"]; dims=3),
        v = reverse(file["velocities/v"]; dims=3),
        w = cat(zeros(Nx+2Hx, Ny+2Hy), reverse(file["velocities/w"][:, :, 2:Nz+2Hz]; dims=3); dims=3),
        T = reverse(file["tracers/T"]; dims=3),
        S = reverse(file["tracers/S"]; dims=3)
    )

    Gⁿ = (
        u = reverse(file["timestepper/Gⁿ/Gu"]; dims=3),
        v = reverse(file["timestepper/Gⁿ/Gv"]; dims=3),
        w = cat(zeros(Nx+2Hx, Ny+2Hy), reverse(file["timestepper/Gⁿ/Gw"][:, :, 2:Nz+2Hz]; dims=3); dims=3),
        T = reverse(file["timestepper/Gⁿ/GT"]; dims=3),
        S = reverse(file["timestepper/Gⁿ/GS"]; dims=3)
    )

    G⁻ = (
        u = reverse(file["timestepper/G⁻/Gu"]; dims=3),
        v = reverse(file["timestepper/G⁻/Gv"]; dims=3),
        w = cat(zeros(Nx+2Hx, Ny+2Hy), reverse(file["timestepper/G⁻/Gw"][:, :, 2:Nz+2Hz]; dims=3); dims=3),
        T = reverse(file["timestepper/G⁻/GT"]; dims=3),
        S = reverse(file["timestepper/G⁻/GS"]; dims=3)
    )

    close(file)

    return solution, Gⁿ, G⁻
end

function run_ocean_large_eddy_simulation_regression_test(arch, closure)
    name = "ocean_large_eddy_simulation_" * string(typeof(closure).name.wrapper)

    spinup_steps = 10000
      test_steps = 10
              Δt = 2.0

    # Parameters
      Qᵀ = 5e-5     # Temperature flux at surface
      Qᵘ = -2e-5    # Velocity flux at surface
    ∂T∂z = 0.005    # Initial vertical temperature gradient

    # Boundary conditions
    u_bcs = HorizontallyPeriodicBCs(    top = BoundaryCondition(Flux, Qᵘ)       )
    T_bcs = HorizontallyPeriodicBCs(    top = BoundaryCondition(Flux, Qᵀ),
                                     bottom = BoundaryCondition(Gradient, ∂T∂z) )
    S_bcs = HorizontallyPeriodicBCs(    top = BoundaryCondition(Flux, 5e-8)     )

    # Model instantiation
    model = Model(
             architecture = arch,
                 grid = RegularCartesianGrid(size=(16, 16, 16), length=(16, 16, 16)),
                 coriolis = FPlane(f=1e-4),
                 buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(α=2e-4, β=8e-4)),
                  closure = closure,
      boundary_conditions = BoundaryConditions(u=u_bcs, T=T_bcs, S=S_bcs)
    )

    ArrayType = typeof(model.velocities.u.data.parent)  # The type of the underlying data, not the offset array.

    #=
    # Initialize model: random noise damped at top and bottom
    Ξ(z) = randn() * z / model.grid.Lz * (1 + z / model.grid.Lz) # noise
    T₀(x, y, z) = 20 + ∂T∂z * z + ∂T∂z * model.grid.Lz * 1e-2 * Ξ(z)
    u₀(x, y, z) = sqrt(abs(Qᵘ)) * 1e-3 * Ξ(z)
    set!(model, u=u₀, w=u₀, T=T₀, S=35)

    time_step!(model, spinup_steps-test_steps, Δt)

    model.output_writers[:checkpointer] = Checkpointer(model;
                                                       frequency = test_steps,
                                                          prefix = name * "_",
                                                             dir = joinpath(dirname(@__FILE__), "data")
                                                      )

    time_step!(model, 2test_steps, Δt)
    pop!(model.output_writers, :checkpointer)
    =#

    initial_filename = joinpath(dirname(@__FILE__), "data", name * "_$spinup_steps.jld2")

    solution₀, Gⁿ₀, G⁻₀ = get_fields_from_checkpoint(initial_filename)

    model.velocities.u.data.parent .= ArrayType(solution₀.u)
    model.velocities.v.data.parent .= ArrayType(solution₀.v)
    model.velocities.w.data.parent .= ArrayType(solution₀.w)
    model.tracers.T.data.parent    .= ArrayType(solution₀.T)
    model.tracers.S.data.parent    .= ArrayType(solution₀.S)

    model.timestepper.Gⁿ.u.data.parent .= ArrayType(Gⁿ₀.u)
    model.timestepper.Gⁿ.v.data.parent .= ArrayType(Gⁿ₀.v)
    model.timestepper.Gⁿ.w.data.parent .= ArrayType(Gⁿ₀.w)
    model.timestepper.Gⁿ.T.data.parent .= ArrayType(Gⁿ₀.T)
    model.timestepper.Gⁿ.S.data.parent .= ArrayType(Gⁿ₀.S)

    model.timestepper.G⁻.u.data.parent .= ArrayType(G⁻₀.u)
    model.timestepper.G⁻.v.data.parent .= ArrayType(G⁻₀.v)
    model.timestepper.G⁻.w.data.parent .= ArrayType(G⁻₀.w)
    model.timestepper.G⁻.T.data.parent .= ArrayType(G⁻₀.T)
    model.timestepper.G⁻.S.data.parent .= ArrayType(G⁻₀.S)

    model.clock.time = spinup_steps * Δt
    model.clock.iteration = spinup_steps

    time_step!(model, test_steps, Δt; init_with_euler=false)

    final_filename = joinpath(dirname(@__FILE__), "data", name * "_$(spinup_steps+test_steps).jld2")

    solution₁, Gⁿ₁, G⁻₁ = get_fields_from_checkpoint(final_filename)

    for name in (:u, :v, :w, :T, :S)
        if name ∈ (:u, :v, :w)
            test_field = getproperty(model.velocities, name)
        else
            test_field = getproperty(model.tracers, name)
        end

        correct_field = getproperty(solution₁, name)

        Δ = Array(test_field.data.parent) .- correct_field

        Δ_min      = minimum(Δ)
        Δ_max      = maximum(Δ)
        Δ_mean     = mean(Δ)
        Δ_abs_mean = mean(abs, Δ)
        Δ_std      = std(Δ)

        @info(@sprintf("Δ%s: min=%.6g, max=%.6g, mean=%.6g, absmean=%.6g, std=%.6g\n",
                       name, Δ_min, Δ_max, Δ_mean, Δ_abs_mean, Δ_std))
    end

    @test all(Array(interior(solution₁.u, model.grid)) .≈ Array(interior(model.velocities.u)))
    @test all(Array(interior(solution₁.v, model.grid)) .≈ Array(interior(model.velocities.v)))
    @test all(Array(interior(solution₁.w, model.grid)) .≈ Array(interior(model.velocities.w)))
    @test all(Array(interior(solution₁.T, model.grid)) .≈ Array(interior(model.tracers.T)))
    @test all(Array(interior(solution₁.S, model.grid)) .≈ Array(interior(model.tracers.S)))

    return nothing
end


text:
module Staggered
using SparseArrays, LinearAlgebra
import sbp
using sbp.Operator: Operators2D
using sbp.Metrics: CovariantBasis
using sbp.Sparse: block_matrix, block_matrix_insert

function boundary_matrix_p(n, is_sparse=true)
        if is_sparse
                B = spzeros(n, n + 1)
        else
                B = zeros(n, n + 1)
        end

        B[1,1] = -1
        B[end,end] = 1

        return B
end

function boundary_matrix_m(n, is_sparse=true)
        if is_sparse
                B = spzeros(n, n - 1)
        else
                B = zeros(n, n - 1)
        end

        B[1,1] = -1
        B[end,end] = 1

        return B
end

function build_operators_2d(Dx::AbstractArray,
                            Dy::AbstractArray,
                            Px::AbstractArray,
                            Py::AbstractArray,
                            Hx::AbstractArray,
                            Hy::AbstractArray,
                            Bx::AbstractArray,
                            By::AbstractArray
                            )
        nx = size(Dx, 1)
        ny = size(Dy, 1)

        Ix = sparse(1.0I, nx, nx)
        Iy = sparse(1.0I, ny, ny)

        Dx2 = kron(Dx, Iy)
        Dy2 = kron(Ix, Dy)

        Px2 = kron(Px, Iy)
        Py2 = kron(Ix, Py)

        Hx2 = kron(Hx, Iy)
        Hy2 = kron(Ix, Hy)

        Bx2 = kron(Bx, Iy)
        By2 = kron(Ix, By)

        return Operators2D(nx, ny, Dx2, Dy2, Px2, Py2, Hx2, Hy2, Bx2, By2)

end

function build_all_operators_2d(builder::Function, nx::Int64, ny::Int64)
        xp, xm, h, Dxp, Dxm, Hxp, Hxm, Pxp, Pxm = builder(nx)
        yp, ym, h, Dyp, Dym, Hyp, Hym, Pyp, Pym = builder(ny)

        Bxp = boundary_matrix_p(nx)
        Bxm = boundary_matrix_m(nx+1)
        Byp = boundary_matrix_p(ny)
        Bym = boundary_matrix_m(ny+1)

        pp = build_operators_2d(Dxp, Dyp, Pxp, Pyp, Hxp, Hyp, Bxp, Byp)
        mm = build_operators_2d(Dxm, Dym, Pxm, Pym, Hxm, Hym, Bxm, Bym)
        pm = build_operators_2d(Dxp, Dym, Pxp, Pym, Hxp, Hym, Bxp, Bym)
        mp = build_operators_2d(Dxm, Dyp, Pxm, Pyp, Hxm, Hyp, Bxm, Byp)

        return pp, mm, pm, mp
end

"""
Construct the covariant basis at the cell centers (mm) using SBP operators.
The mapping function f(x,y) = (fx, fy) must be defined at the nodes (pp).
"""
function build_covariant_basis_mm(fx::AbstractArray, fy::AbstractArray,
                               mm::Operators2D, pm::Operators2D)
        x_r1 = mm.Dx * pm.Py * fx
        x_r2 = mm.Px * pm.Dy * fx
        y_r1 = mm.Dx * pm.Py * fy
        y_r2 = mm.Px * pm.Dy * fy

        return CovariantBasis(x_r1, y_r1, x_r2, y_r2)
end

"""
Construct the covariant basis at the v^1 components (pm) using SBP operators.
The mapping function f(x,y) = (fx, fy) must be defined at the nodes (pp).
"""
function build_covariant_basis_pm(fx::AbstractArray, fy::AbstractArray,
                               mm::Operators2D, pm::Operators2D,
                               mp::Operators2D)
        x_r1 = pm.Px * mm.Py * mp.Dx * fx
        x_r2 = pm.Dy * fx
        y_r1 = pm.Px * mm.Py * mp.Dx * fy
        y_r2 = pm.Dy * fy

        return CovariantBasis(x_r1, y_r1, x_r2, y_r2)
end

"""
Construct the covariant basis at the v^2 components (mp) using SBP operators.
The mapping function f(x,y) = (fx, fy) must be defined at the nodes (pp).
"""
function build_covariant_basis_mp(fx::AbstractArray, fy::AbstractArray,
                               mm::Operators2D, pm::Operators2D,
                               mp::Operators2D)
        x_r1 = mp.Dx * fx
        x_r2 = mp.Py * mm.Px * pm.Dy * fx
        y_r1 = mp.Dx * fy
        y_r2 = mp.Py * mm.Px * pm.Dy * fy

        return CovariantBasis(x_r1, y_r1, x_r2, y_r2)
end

end

max_stars_repo_name: ooreilly/sbpjl

max_issues_repo_name: ooreilly/sbpjl

max_forks_repo_name: ooreilly/sbpjl


text:
include("./imports.jl")

@testset ExtendedTestSet "Dummy learner convergence" begin
    learner = testlearner(coeff = 3)
    fit!(learner, repeat([TrainingPhase()], 5))
    @test  learner.model.coeff[1] ≈ 3 atol = 0.1
    @test learner.state.phase isa TrainingPhase
end

@testset ExtendedTestSet "Recorder" begin
    learner = testlearner(coeff = 3)
    h = learner.state.history
    @test h.epochs == 0
    @test h.nsteps == 0
    fit!(learner, 10)
    @test h.epochs == 10
    @test h.nstepsepoch == length(getdataloader(TrainingPhase(), learner))
    @test h.nsteps == 10 * length(getdataloader(TrainingPhase(), learner))

    idxs, losses = get(h.epochmetrics[ValidationPhase()], :loss)

    @test losses[1] > losses[end]
end

@testset ExtendedTestSet "Hyperparameter scheduling" begin
    nbatches = 16
    schedules = Schedules(Dict(LR => Schedule(2nbatches, 1e-4, 1e-6, anneal_linear)))
    learner = testlearner(coeff = 3, nbatches = nbatches, schedule = schedules)

    fit!(learner, TrainingPhase())
    @test getoptimparam(learner.opt, LR) ≈ ((1e-4 + 1e-6) / 2)

    fit!(learner, TrainingPhase())
    @test getoptimparam(learner.opt, LR) ≈ 1e-6
end


text:
module ProbitNodeTest

using Test
using ReactiveMP
using Random

@testset "ProbitNode" begin

    @testset "Creation" begin
        node = make_node(Probit)

        @test functionalform(node)          === Probit
        @test sdtype(node)                  === Stochastic()
        @test name.(interfaces(node))       === (:out, :in)
        @test factorisation(node)           === ((1, 2), )
        @test metadata(node)                === ProbitMeta(32)

        node = make_node(Probit, meta = 1)

        @test functionalform(node)          === Probit
        @test sdtype(node)                  === Stochastic()
        @test name.(interfaces(node))       === (:out, :in)
        @test factorisation(node)           === ((1, 2), )
        @test metadata(node)                === 1
    end

    @testset "Average energy" begin

        node = make_node(Probit)

        @test score(AverageEnergy(),
                    Probit,
                    Val{(:out, :in)},
                    (Marginal(Bernoulli(1), false, false), Marginal(NormalMeanVariance(0.0, 1.0), false, false)),
                    ProbitMeta()) ≈ 1.0

        @test score(AverageEnergy(),
                    Probit,
                    Val{(:out, :in)},
                    (Marginal(PointMass(1), false, false), Marginal(NormalMeanVariance(0.0, 1.0), false, false)),
                    ProbitMeta(100)) ≈ 1.0

        for k = 0:0.1:1
            @test score(AverageEnergy(),
                        Probit,
                        Val{(:out, :in)},
                        (Marginal(Bernoulli(k), false, false), Marginal(NormalMeanVariance(0.0, 1.0), false, false)),
                        ProbitMeta(100)) ≈ 1.0
        end

    end

end

end


text:
abstract type FixedLambdaGroupRidgeRegressor <: AbstractGroupRidgeRegressor end

"""
    SingleGroupRidgeRegressor(; λ,
                                decomposition = :default,
                                center = true,
                                scale = true)

Type representing vanilla Ridge regression with hyperparameter `λ`.
`center` and `scale` (default `true` for both) control whether the response and
features should be centered and scaled first (make sure that `center=true` if the
model is supposed to have an intercept!). `decomposition` can be one of `:default`,
`:cholesky` or `:woodbury` and determines how the linear system is solved.
"""
Base.@kwdef mutable struct SingleGroupRidgeRegressor{T,G} <: FixedLambdaGroupRidgeRegressor
    decomposition::Symbol = :default
    λ::T = 1.0
    groups::G = nothing
    center::Bool = true
    scale::Bool = true
end

_main_hyperparameter(::SingleGroupRidgeRegressor) = :λ
_main_hyperparameter_value(m) = getproperty(m, _main_hyperparameter(m))

function _default_hyperparameter_maximum(
    model::FixedLambdaGroupRidgeRegressor,
    fitted_machine,
)
    1000 * maximum(abs.(fitted_machine.cache.XtY))
end

function _default_param_min_ratio(ridge, fitted_machine)
    1e-6
end

function _default_scale(ridge, fitted_machine)
    :log10
end

function _groups(m::SingleGroupRidgeRegressor, p)
    isnothing(m.groups) ? GroupedFeatures([p]) : m.groups
end

function MMI.fit(m::FixedLambdaGroupRidgeRegressor, verb::Int, X, y)
    @unpack center, scale, groups = m
    Xmatrix = MMI.matrix(X)
    if center || scale
        x_transform = StatsBase.fit(
            ZScoreTransform,
            Xmatrix;
            dims = 1,
            center = center,
            scale = scale,
        )
        y_transform =
            StatsBase.fit(ZScoreTransform, y; dims = 1, center = center, scale = scale)
        Xmatrix = StatsBase.transform(x_transform, Xmatrix)
        y = StatsBase.transform(y_transform, y)
    else
        x_transform = nothing
        y_transform = nothing
    end
    p = size(Xmatrix, 2)
    groups = _groups(m, p)
    workspace = StatsBase.fit(m, Xmatrix, y, groups)  # see end_to_end.jl
    βs = StatsBase.coef(workspace)
    fitresult = (coef = βs, x_transform = x_transform, y_transform = y_transform)
    # return
    return fitresult, workspace, NamedTuple{}()
end

function MMI.update(
    model::AbstractGroupRidgeRegressor,
    verbosity::Int,
    old_fitresult,
    old_cache,
    X,
    y,
)
    new_λ = _main_hyperparameter_value(model)
    StatsBase.fit!(old_cache, new_λ)
    βs = StatsBase.coef(old_cache)
    fitresult = (
        coef = βs,
        x_transform = old_fitresult.x_transform,
        y_transform = old_fitresult.y_transform,
    )
    return fitresult, old_cache, NamedTuple{}()
end

function MMI.predict(model::AbstractGroupRegressor, fitresult, Xnew)
    Xnew = MMI.matrix(Xnew)
    @unpack coef, x_transform, y_transform = fitresult
    !isnothing(x_transform) && (Xnew = StatsBase.transform(x_transform, Xnew))
    ypred = Xnew * coef
    !isnothing(y_transform) && StatsBase.reconstruct!(y_transform, ypred)
    ypred
end

function range_and_grid(
    ridge::AbstractGroupRegressor,
    param_min,
    param_max,
    scale,
    resolution,
    n,
    rng
)
    param_symbol = _main_hyperparameter(ridge)
    param_range =
        range(ridge, param_symbol, lower = param_min, upper = param_max, scale = scale)
    model_grid =
        MLJTuning.grid(ridge, [param_symbol], [MLJ.iterator(param_range, resolution)])

    if length(model_grid) > n
        model_grid = sample(rng, model_grid, n; replace = false)
    end

    param_range, model_grid
end

"""
    MultiGroupRidgeRegressor(; decomposition, λ, groups)
"""
mutable struct MultiGroupRidgeRegressor{T,G<:GroupedFeatures} <:
               FixedLambdaGroupRidgeRegressor
    decomposition::Symbol
    λs::T   #Named tuple
    groups::G
    center::Bool
    scale::Bool
end

_main_hyperparameter(::MultiGroupRidgeRegressor) = :λs
_groups(m::MultiGroupRidgeRegressor, p) = m.groups

function MultiGroupRidgeRegressor(;
    groups::GroupedFeatures,
    λs::AbstractVector = ones(ngroups(groups)),
    decomposition = :default,
    center = true,
    scale = true,
)
    ngr = ngroups(groups)
    λ_expr = Tuple(Symbol.(:λ, Base.OneTo(ngr)))
    λ_tupl = MutableNamedTuple{λ_expr}(tuple(λs...))
    MultiGroupRidgeRegressor(decomposition, λ_tupl, groups, center, scale)
end

function range_and_grid(ridge::MultiGroupRidgeRegressor, λ_min, λ_max, scale, resolution, n, rng)
    λ_names = [Meta.parse("(λs.$λ)") for λ in keys(ridge.λs)]
    nparams = length(λ_names)
    λ_range =
        [range(ridge, λ, lower = λ_min, upper = λ_max, scale = scale) for λ in λ_names]
    λ_product_grid = MLJ.iterator.(λ_range, resolution)
    if nparams*log(resolution) > log(n)
        tmp_idx = zeros(Int, nparams)
        model_grid = [deepcopy(ridge) for i in Base.OneTo(n)]
        for i in Base.OneTo(n)
            sample!(rng, 1:resolution, tmp_idx)
            clone = model_grid[i]
            for k in eachindex(λ_names)
                MLJ.recursive_setproperty!(clone, λ_names[k], λ_product_grid[k][tmp_idx[k]])
            end
        end
    else
        model_grid = MLJTuning.grid(ridge, λ_names, λ_product_grid)
    end
    λ_range, model_grid
end

"""
    DefaultTuning(resolution, n, param_min_ratio, param_max, scale)

Determines the default set of hyperparameters to loop over when tuning a
`AbstractGroupRidgeRegressor` method.  Parameters are chosen on a grid
that is equidistant in `scale` (e.g. `:log10` or `:linear` or `:default`) with number
of points given by `resolution` (default `100`) that ranges from `param_min_ratio*param_max` to
`param_max`. Both `param_min_ratio` and `param_max` can be specified as `:default`,
in which case a method specific default choice will be used.

If there are multiple hyperparameters (say `d`),
then the above rules are used componentwise. `n` (default `1000`)
is the largest number of hyperparameters to explore (if `resolution^d > n`,
then the parameters are randomly subsampled to `n` of them).
"""
Base.@kwdef struct DefaultTuning{T,M}
    resolution::Int = 100
    n::Int = 5000
    param_min_ratio::M = :default
    param_max::T = :default
    scale = :default
end

function _tuning_grid(tuning::DefaultTuning, model, fitted_machine, rng)
    @unpack resolution, n, scale = tuning
    if tuning.param_max === :default
        param_max = _default_hyperparameter_maximum(model, fitted_machine)
    elseif isa(tuning.param_max, Number)
        param_max = tuning.param_max
    else
        error("param_max can be :default or a number only.")
    end

    if tuning.param_min_ratio === :default
        param_min_ratio = _default_param_min_ratio(model, fitted_machine)
    elseif isa(tuning.param_min_ratio, Number)
        param_min_ratio = tuning.param_min_ratio
    else
        error("param_min_ratio can be :default or a number only.")
    end

    if tuning.scale === :default
        _scale = _default_scale(model, fitted_machine)
    else
        _scale = tuning.scale
    end

    param_min = param_min_ratio * param_max
    param_range, model_grid = range_and_grid(model, param_min, param_max, _scale, resolution, n, rng)

    param_range, model_grid, param_max
end

"""
    LooRidgeRegressor(;ridge,
                       tuning = DefaultTuning(),
                       rng = Random.GLOBAL_RNG)

A MLJ model that wraps a `ridge` model such as `SigmaRidgeRegressor` and tunes
its parameters by leave-one-out-cross-validation with `tuning` settings defaulting to
[`DefaultTuning`](@ref). In case there is randomness in choosing the search space of
hyperparameters, then the `rng` may be specified (defaults to `Random.GLOBAL_RNG`).
"""
Base.@kwdef mutable struct LooRidgeRegressor{G,T} <: AbstractGroupRidgeRegressor
    ridge::G = SingleGroupRidgeRegressor()
    tuning::T = DefaultTuning()
    rng = Random.GLOBAL_RNG
end

_groups(loo::LooRidgeRegressor) = _groups(loo.ridge)

_workspace(wk::BasicGroupRidgeWorkspace) = wk
_workspace(wk) = wk.workspace

function MMI.fit(m::LooRidgeRegressor, verb::Int, X, y)
    ridge = m.ridge
    mach = MLJ.machine(ridge, X, y)
    fit!(mach)
    x_transform = mach.fitresult.x_transform
    y_transform = mach.fitresult.y_transform

    ridge_workspace = _workspace(mach.cache)
    param_range, model_grid, param_max = _tuning_grid(m.tuning, ridge, mach, m.rng)
    history = map(model_grid) do newm
        param = _main_hyperparameter_value(newm)
        mach.model = newm
        fit!(mach; verbosity = 0)
        λ = deepcopy(ridge_workspace.λs)
        loo = loo_error(ridge_workspace)
        (loo = loo, param = param, model = newm, λ = λ)
    end
    loos = [h.loo for h in history]
    params = [h.param for h in history]
    λs = [h.λ for h in history]
    best_model_idx = argmin(loos)
    best_model = model_grid[best_model_idx]
    best_param = params[best_model_idx]
    best_loo = loos[best_model_idx]
    best_λs = λs[best_model_idx]

    report = (
        best_model = best_model,
        best_param = best_param,
        best_λs = best_λs,
        best_loo = best_loo,
        loos = loos,
        λs = λs,
        params = params,
        param_max = param_max,
        param_range = param_range,
    )
    mach.model = best_model
    fit!(mach)
    βs = StatsBase.coef(ridge_workspace)
    fitresult = (coef = βs, x_transform = x_transform, y_transform = y_transform)
    # return
    return fitresult, ridge_workspace, report
end

Base.@kwdef mutable struct TunedRidgeRegressor{G,R,M,T} <: AbstractGroupRidgeRegressor
    ridge::G = SingleGroupRidgeRegressor(decomposition = :cholesky, λ = 1.0)
    tuning::T = DefaultTuning()
    resampling::R = MLJ.CV(nfolds = 5, shuffle=true)
    measure::M = MLJ.l2
end

function MMI.fit(m::TunedRidgeRegressor, verb::Int, X, y)
    ridge = m.ridge
    mach = MLJ.machine(ridge, X, y)
    fit!(mach)
    #x_transform = mach.fitresult.x_transform
    #y_transform = mach.fitresult.y_transform

    ridge_workspace = _workspace(mach.cache)
    param_range, model_grid, param_max = _tuning_grid(m.tuning, ridge, mach, m.resampling.rng)

    tuned_model = MLJ.TunedModel(
        model = ridge,
        ranges = model_grid,
        tuning = MLJ.Explicit(),
        resampling = m.resampling,
        measure = m.measure,
    )

    tuned_mach = MLJ.machine(tuned_model, X, y)
    fit!(tuned_mach)

    _fitresult = tuned_mach.fitresult.fitresult
    _cache = tuned_mach.fitresult.cache
    best_λs = deepcopy(_workspace(_cache).λs)

    tunedreport = tuned_mach.report
    best_param = _main_hyperparameter_value(tunedreport.best_model)

    tunedreport = (tunedreport..., best_param=best_param, best_λs = best_λs)
    # return
    return _fitresult, _cache, tunedreport
end


text:
"""
    predict(task, model, input[; device, context])

Predict a `target` from `input` using `model`. Optionally apply function `device`
to `x` before passing to `model` and use `context` instead of the default
context [`Inference`](#).
"""
function predict(task, model, input; device = cpu, undevice = cpu, context = Inference())
    if shouldbatch(task)
        return predictbatch(
            task,
            model,
            [input];
            device = device,
            undevice = undevice,
            context = context,
        ) |> only
    else
        return decodeypred(
            task,
            context,
            undevice(model(device(encodeinput(task, context, input)))),
        )
    end
end

"""
    predictbatch(task, model, inputs[; device, context])

Predict `targets` from a vector of `inputs` using `model` by batching them.
Optionally apply function `device` to batch before passing to `model` and
use `context` instead of the default [`Inference`](#).
"""
function predictbatch(
    task,
    model,
    inputs;
    device = cpu,
    undevice = cpu,
    context = Inference(),
)
    xs = device(
        MLUtils.batch([copy(encodeinput(task, context, input)) for input in inputs]),
    )
    ŷs = undevice(model(xs))
    targets = [decodeypred(task, context, ŷ) for ŷ in Datasets.unbatch(ŷs)]
    return targets
end


text:
"""
Plots the results of `mandel()` function
"""

function plotMandel( xrange::Tuple{Float64,Float64}, yrange::Tuple{Float64,Float64}; n=100 )
	M = rand(n,n);
	xval = range(xrange[1],xrange[2]; length=n );
	yval = range(yrange[1],yrange[2]; length=n );
	for i in 1:n, j in 1:n
		M[i,j] = mandel( xval[i] + im*yval[j]  )
	end
	heatmap( M, c=:thermal
		,aspect_ratio=1.0
		,xaxis=nothing
		,yaxis=nothing
	)
end

plotMandel( vrange::Tuple{Float64,Float64}; n=100 ) = plotMandel( vrange, vrange; n )


text:
## Package Logo

function printLogo()

    logo = string(raw"
                      ___________   __________   ___  ___   ___________   __________
                     /  _____   /  / ______  /  /   \/   \  \   _____  \  \  _______\
                    /  /____/  /  / /     / /  /          \  \  \____\  \  \ \________
                   /  _____   /  / /     / /  /            \  \   _____  \  \________ \
                  /  /    /  /  / /     / /  /   /\____/\   \  \  \    \  \          \ \
                 /  /____/  /  / /_____/ /  /   /        \   \  \  \____\  \   _______\ \
                /__________/  /_________/  /__ /          \___\  \__________\  \_________\

                ")
println(logo)

end

function versionBOMBS()

    logo = string(raw"
                                                  *
                                    **           * *           *                          **
                                   * *          *   *         * *                       **  *
                                  *  *         *    *        *  **                    **     *
                         *       *   *  *   ***     *     ***    **                 **        *      *
                        * *     *    * * * *         **  *         **      *      **           *    * *   ***
                       *   *   *     **   *            **            **   * * ****              ****   ***   ***
                   ** *     ***                                        * *   *                                  *
                  *  *                                                  *                                        *
                 *            ___________   __________   ___  ___   ___________   __________                      ***
                **           /  _____   /  / ______  /  /   \/   \  \   _____  \  \  _______\                        **
               *            /  /____/  /  / /     / /  /          \  \  \____\  \  \ \________                         *
             **            /  _____   /  / /     / /  /            \  \   _____  \  \________ \                       *
           **             /  /    /  /  / /     / /  /   /\____/\   \  \  \    \  \          \ \                    **
         **              /  /____/  /  / /_____/ /  /   /        \   \  \  \____\  \   _______\ \                 **
        *               /__________/  /_________/  /__ /          \___\  \__________\  \_________\ V 0.1.0       **
         ****                                                *                                                  *
             *                          ***                 * *                                                *
              *   ***                  *   *     *         *   *                  *                      **   *
               * *   **             ***     *   * ***     *     *   *          *** *           **       *  * *
                *      **      *   *         * *    *    *       * * **     ***     *         *  *     *    *
                         **   * * *           *     *   *         *    ** **         *       *    *   *
                           * *   *                   * *                 *            *     *      * *
                            *                         *                                *   *        *
                                                                                * *
                                                                                *
    ");

    println(logo)

end

versionBOMBS()

max_stars_repo_name: DavidGomezC/BOMBSbeta

max_issues_repo_name: DavidGomezC/BOMBSbeta

max_forks_repo_name: DavidGomezC/BOMBSbeta


text:
using FFTW # or using DSP

fft([1,1,1,1,0,0,0,0])

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# # Phillips model of Baroclinic Instability
#
#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/multilayerqg_2layer.ipynb).
#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/multilayerqg_2layer.ipynb).
#
# A simulation of the growth of barolinic instability in the Phillips 2-layer model
# when we impose a vertical mean flow shear as a difference ``\Delta U`` in the
# imposed, domain-averaged, zonal flow at each layer.

using FourierFlows, Plots, Printf

using FFTW: rfft, irfft
using Random: seed!
import GeophysicalFlows.MultilayerQG
import GeophysicalFlows.MultilayerQG: energies

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)
nothing # hide

# ## Numerical parameters and time-stepping parameters

n = 128                  # 2D resolution = n²
stepper = "FilteredRK4"  # timestepper
     dt = 6e-3           # timestep
 nsteps = 7000           # total number of time-steps
 nsubs  = 25             # number of time-steps for plotting (nsteps must be multiple of nsubs)
nothing # hide

# ## Physical parameters
 L = 2π                  # domain size
 μ = 5e-2                # bottom drag
 β = 5                   # the y-gradient of planetary PV

nlayers = 2              # number of layers
f₀, g = 1, 1             # Coriolis parameter and gravitational constant
 H = [0.2, 0.8]          # the rest depths of each layer
 ρ = [4.0, 5.0]          # the density of each layer

 U = zeros(nlayers) # the imposed mean zonal flow in each layer
 U[1] = 1.0
 U[2] = 0.0
nothing # hide

# ## Problem setup
# We initialize a `Problem` by providing a set of keyword arguments,
prob = MultilayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=U, dt=dt, stepper=stepper, μ=μ, β=β)
nothing # hide

# and define some shortcuts.
sol, clock, params, vars, grid = prob.sol, prob.clock, prob.params, prob.vars, prob.grid
x, y = grid.x, grid.y
nothing # hide

# ## Setting initial conditions

# Our initial condition is some small amplitude random noise. We smooth our initial
# condidtion using the `timestepper`'s high-wavenumber `filter`.

seed!(1234) # reset of the random number generator for reproducibility
q_i  = 4e-3randn((grid.nx, grid.ny, nlayers))
qh_i = prob.timestepper.filter .* rfft(q_i, (1, 2)) # only apply rfft in dims=1, 2
q_i  = irfft(qh_i, grid.nx, (1, 2)) # only apply irfft in dims=1, 2

MultilayerQG.set_q!(prob, q_i)
nothing # hide

# ## Diagnostics

# Create Diagnostics -- `energies` function is imported at the top.
E = Diagnostic(energies, prob; nsteps=nsteps)
diags = [E] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.
nothing # hide

# ## Output

# We choose folder for outputing `.jld2` files and snapshots (`.png` files).
filepath = "."
plotpath = "./plots_2layer"
plotname = "snapshots"
filename = joinpath(filepath, "2layer.jld2")
nothing # hide

# Do some basic file management
if isfile(filename); rm(filename); end
if !isdir(plotpath); mkdir(plotpath); end
nothing # hide

# And then create Output
get_sol(prob) = sol # extracts the Fourier-transformed solution
function get_u(prob)
  sol, params, vars, grid = prob.sol, prob.params, prob.vars, prob.grid

  @. vars.qh = sol
  streamfunctionfrompv!(vars.ψh, vars.qh, params, grid)
  @. vars.uh = -im * grid.l * vars.ψh
  invtransform!(vars.u, vars.uh, params)

  return vars.u
end

out = Output(prob, filename, (:sol, get_sol), (:u, get_u))
nothing # hide

# ## Visualizing the simulation

# We define a function that plots the potential vorticity field and the evolution
# of energy and enstrophy.

symlims(data) = maximum(abs.(extrema(data))) |> q -> (-q, q)

function plot_output(prob)
  Lx, Ly = prob.grid.Lx, prob.grid.Ly

  l = @layout Plots.grid(2, 3)
  p = plot(layout=l, size = (1000, 600))

  for m in 1:nlayers
    heatmap!(p[(m-1) * 3 + 1], x, y, vars.q[:, :, m]',
         aspectratio = 1,
              legend = false,
                   c = :balance,
               xlims = (-Lx/2, Lx/2),
               ylims = (-Ly/2, Ly/2),
               clims = symlims,
              xticks = -3:3,
              yticks = -3:3,
              xlabel = "x",
              ylabel = "y",
               title = "q_"*string(m),
          framestyle = :box)

    contourf!(p[(m-1) * 3 + 2], x, y, vars.ψ[:, :, m]',
              levels = 8,
         aspectratio = 1,
              legend = false,
                   c = :viridis,
               xlims = (-Lx/2, Lx/2),
               ylims = (-Ly/2, Ly/2),
               clims = symlims,
              xticks = -3:3,
              yticks = -3:3,
              xlabel = "x",
              ylabel = "y",
               title = "ψ_"*string(m),
          framestyle = :box)
  end

  plot!(p[3], 2,
             label = ["KE1" "KE2"],
            legend = :bottomright,
         linewidth = 2,
             alpha = 0.7,
             xlims = (-0.1, 2.35),
             ylims = (5e-10, 1e0),
            yscale = :log10,
            yticks = 10.0.^(-9:2:0),
            xlabel = "μt")

  plot!(p[6], 1,
             label = "PE",
            legend = :bottomright,
         linecolor = :red,
         linewidth = 2,
             alpha = 0.7,
             xlims = (-0.1, 2.35),
             ylims = (1e-10, 1e0),
            yscale = :log10,
            yticks = 10.0.^(-10:2:0),
            xlabel = "μt")

end
nothing # hide

# ## Time-stepping the `Problem` forward

# Finally, we time-step the `Problem` forward in time.

p = plot_output(prob)

startwalltime = time()

anim = @animate for j = 0:round(Int, nsteps / nsubs)
  if j % (1000 / nsubs) == 0
    cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy])

    log = @sprintf("step: %04d, t: %.1f, cfl: %.2f, KE1: %.3e, KE2: %.3e, PE: %.3e, walltime: %.2f min", clock.step, clock.t, cfl, E.data[E.i][1][1], E.data[E.i][1][2], E.data[E.i][2][1], (time()-startwalltime)/60)

    println(log)
  end

  for m in 1:nlayers
    p[(m-1) * 3 + 1][1][:z] = @. vars.q[:, :, m]
    p[(m-1) * 3 + 2][1][:z] = @. vars.ψ[:, :, m]
  end

  push!(p[3][1], μ * E.t[E.i], E.data[E.i][1][1])
  push!(p[3][2], μ * E.t[E.i], E.data[E.i][1][2])
  push!(p[6][1], μ * E.t[E.i], E.data[E.i][2][1])

  stepforward!(prob, diags, nsubs)
  MultilayerQG.updatevars!(prob)
end

gif(anim, "multilayerqg_2layer.gif", fps=18)

# ## Save

# Finally save the last snapshot.
savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step)
savefig(savename)