text:
module MultiFieldTests

using Test

@testset "MultiCellArrays" begin include("MultiCellArraysTests.jl") end

@testset "MultiCellMaps" begin include("MultiCellMapsTests.jl") end

@testset "MultiFESpaces" begin include("MultiFESpacesTests.jl") end

@testset "MultiFEFunctions" begin include("MultiFEFunctionsTests.jl") end

@testset "MultiFEBases" begin include("MultiFEBasesTests.jl") end

@testset "MultiAssemblers" begin include("MultiAssemblersTests.jl") end

@testset "MultiFEOperators" begin include("MultiFEOperatorsTests.jl") end

@testset "DGMultiFEOperators" begin include("DGMultiFEOperatorsTests.jl") end

@testset "VectorValuedMultiFEOperators" begin include("VectorValuedMultiFEOperatorsTests.jl") end

@testset "DarcyFEOperators" begin include("DarcyFEOperatorsTests.jl") end

@testset "MultiNonLinearFEOperators" begin include("MultiNonLinearFEOperatorsTests.jl") end

@testset "NSCavityTests" begin include("NSCavityTests.jl") end

end # module


text:
module TestRandomIncrements

using Test
using TSXPlayground:
    parallel_random_increments!,
    random_increments_problem_clusters,
    random_increments_problem_uniform,
    serial_random_increments!

ncounters = 2^20
nindices = 2^5

@testset "$label" for (label, (counters, indices)) in [
    ("uniform", random_increments_problem_uniform(ncounters, nindices)),
    ("cluster", random_increments_problem_clusters(ncounters, nindices, 5)),
]
    counters_seq = zero(counters)
    serial_random_increments!(counters_seq, indices)
    @testset for ntasks in 1:Threads.nthreads()
        counters_atomic = zero(counters)
        counters_tsx = zero(counters)
        parallel_random_increments!(counters_atomic, indices, ntasks)
        parallel_random_increments!(counters_tsx, indices, ntasks, Val(true), 10, 4)
        @test counters_atomic == counters_seq
        @test counters_tsx == counters_seq
    end
end

end  # module

ext: jl


text:
# This file is part of IntegerTriangles.jl.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module TrianglesPlot

"""
Plotting triangles.
"""
const ModuleTrianglesPlot = ""

using TrianglesBase, TrianglesExamples
using Plots; gr()

function PolyPlot(F, title)
    x = -2.0:0.05:1.0
    plot(x,  F[1], c="black",  linewidth=1)
    plot!(x, F[2], c="green",  line=:dash)
    plot!(x, F[3], c="blue",   linewidth=1)
    plot!(x, F[4], c="red",    line=:dash)
    plot!(x, F[5], c="orange", linewidth=1)
    plot!(x, F[6], c="blue",   line=:dash)
    plot!(x, F[7], c="black",  linewidth=1)
    title!(title)
    xlabel!("x")
    ylabel!("p(x)")

    # axis limit
    P = plot!(xlims=(-2.0, 1.0), ylims=(-5.0, 5.0))
    display(P)
end

# START-TEST-########################################################

function test()
end

function demo()
    # P = PolynomialFunction(LahTriangle(7))
    # PolyPlot(P, "Lah polynomials")
    # P = PolynomialFunction(EulerianTriangle(7))
    # PolyPlot(P, "Eulerian polynomials")
    # P = PolynomialFunction(LaguerreTriangle(7))
    # PolyPlot(P, "Laguerre polynomials")
    # P = PolynomialFunction(CatalanTriangle(7))
    # PolyPlot(P, "Catalan polynomials")
    # P = PolynomialFunction(MotzkinTriangle(7))
    # PolyPlot(P, "Motzkin polynomials")
    P = PolynomialFunction(NarayanaTriangle(7))
    PolyPlot(P, "Narayana polynomials")
    # P = PolynomialFunction(JacobsthalTriangle(7))
    # PolyPlot(P, "Jacobsthal polynomials")
    # P = PolynomialFunction(FibonacciTriangle(7))
    # PolyPlot(P, "Fibonacci polynomials")
end

function perf()
end

function main()
    test()
    # this might take a while
    demo()
    perf()
end

main()

end # module


text:
# This file was generated, do not modify it. # hide
b1 = select(boston, [:Crim, :Zn, :Indus])
first(b1, 2)


text:
# A GroupedRaft gathers "arrays" of rafts for grouped container-like access
type GroupedRaft <: AbstractRaft
    path::UTF8String
end


text:
include("../tests/test_lorenz63.jl")
include("../examples/lorenz63.jl")
using PyPlot
using JLD
function plot_sensitivity()
	vsh, dJds = test_lss()
	dJds = dJds[:]
	fig, ax = subplots(1,1)
	fig1, ax1 = subplots(1,1)
	vsh = vsh[:,100:end,10]'
	n = 0.005*(axes(vsh)[1] .- 1.)
	n_samples = size(dJds)[1]
	dJds_avg = sum(dJds)/n_samples
	var_dJds = sum(x -> x^2, dJds .- dJds_avg)/n_samples
	ax.plot(n, vsh[:,1], label=L"v_{\rm sh,x}")
	ax.plot(n, vsh[:,2], label=L"v_{\rm sh,y}")
	ax.plot(n, vsh[:,3], label=L"v_{\rm sh,z}")
	ax.set_xlabel("time",fontsize=28)
	ax.set_ylabel("tangent shadowing components",
					  fontsize=28)
	ax.xaxis.set_tick_params(labelsize=28)
	ax.yaxis.set_tick_params(labelsize=28)
	fig.legend(fontsize=28)
	ax1.set_ylabel("tangent sensitivities",fontsize=28)
	ax1.set_xlabel("sample number",fontsize=28)
	ax1.set_title("trajectory length = 3000", fontsize=28)
	ax1.xaxis.set_tick_params(labelsize=28)
	ax1.yaxis.set_tick_params(labelsize=28)
	ax1.grid(true)
	ax1.errorbar(x=1:n_samples,y=dJds,yerr=sqrt(var_dJds),
			 linestyle="none",ms=4)
	ax1.plot(sum(dJds)/n_samples*ones(n_samples),"b--")
	#mean value: 0.905328466, variance = 0.042, without
	# time dilation
end
function plot_condition_number()
	dJds, condnum = test_condition_number()
	fig, ax = subplots(1,1)
	n_samples = size(dJds)[1]
	n_arr = StepRange(500, 155, 5000)
	ax.plot(n_arr, condnum, ".", ms=10.0)
	ax.set_xlabel("trajectory length",fontsize=18)
	ax.set_ylabel("condition number",
					  fontsize=18)
	ax.xaxis.set_tick_params(labelsize=18)
	ax.yaxis.set_tick_params(labelsize=18)
end
function plot_data_assmln_err()
	data = load("../data/l63_asmln_ngd200_n2000.jld")
	z_prd = data["z_prd"]
	z_obs = data["z_obs"]
	msq_err = data["msq_err"]
	n_trj, n_gd, n_exps = size(z_prd)
	z_opt_prd = view(z_prd, :, n_gd, :)
	err = abs.((z_obs .- z_opt_prd)./z_obs)
	max_err = findmax(err, dims=1)[1][:]
	norm_max_err = max_err/maximum(max_err)
	clr_map = plt.get_cmap("Blues")
	clr_acc_err = clr_map(norm_max_err)
	times = dt*LinRange(0,n_trj,n_trj)
	times = times[100:end]
	err = err[100:end,:]
	fig, ax = subplots(1,1)
	for i = 1:n_exps
		ax.semilogy(times, err[:,i], ".",
					ms=0.5,color=clr_acc_err[i,:])
	end
	mean_err = (sum(err, dims=2)/n_exps)[:]
	std_err = sum((err .- mean_err).^2.0, dims=2)/n_exps
	std_err = sqrt.(std_err[:])
	ax.semilogy(times, mean_err, ".", color="k",ms=2.0)
	ax.xaxis.set_tick_params(labelsize=24)
	ax.yaxis.set_tick_params(labelsize=24)
	ax.set_xlabel("time", fontsize=24)
	ax.set_ylabel("Prediction error", fontsize=24)
	ax.grid(true)
end

max_stars_repo_name: nishaChandramoorthy/energies

max_issues_repo_name: nishaChandramoorthy/energies

max_forks_repo_name: nishaChandramoorthy/energies


text:
export Layout
module Layout

using CUDA
using GPUifyLoops: @unroll
using GemmKernels.Tiling
using StaticArrays

# ---------------------
# Customise computation
# ---------------------

@inline threadblock_condition(layout_a, layout_b, block_i, block_j, block_k, block_tile) = true

# ----------------------
# Explicit vectorisation
# ----------------------

struct Vec{N, T} end

@inline @generated function vloada(::Type{Vec{N, T}}, ptr::Core.LLVMPtr{T, AS}, i::Integer = 1) where {N, T, AS}
    alignment = sizeof(T) * N
    vec_len = (sizeof(T) * N) ÷ sizeof(Float32)

    return quote
        vec_ptr = Base.bitcast(Core.LLVMPtr{NTuple{$vec_len, VecElement{Float32}}, AS}, ptr)
        return unsafe_load(vec_ptr, (i-1) ÷ N + 1, Val($alignment))
    end
end

@inline @generated function vstorea!(::Type{Vec{N, T}}, ptr::Core.LLVMPtr{T, AS}, x, i::Integer = 1) where {N, T, AS}
    alignment = sizeof(T) * N
    vec_len = (sizeof(T) * N) ÷ sizeof(Float32)

    return quote
        vec_ptr = Base.bitcast(Core.LLVMPtr{NTuple{$vec_len, VecElement{Float32}}, AS}, ptr)
        return unsafe_store!(vec_ptr, x, (i-1) ÷ N + 1, Val($alignment))
    end
end

# -----------
# Layout base
# -----------

abstract type LayoutBase{T} end

@inline eltype(::Type{<:LayoutBase{T}}) where {T} = T
@inline physical_size(::Type{<:LayoutBase{T}}, logical_size::NamedTuple) where {T} = Tuple(logical_size)

# --------------
# Padded layouts
# --------------

struct Padded{L, P} end

@inline function pad_logical_coord(::Type{Padded{L, P}}, crd::NamedTuple) where {L, P}
    t = Tuple(crd)
    return typeof(crd)((Base.first(t) + P, Base.tail(t)...))
end

@inline eltype(::Type{Padded{L, P}}) where {L, P} = eltype(L)
@inline physical_size(::Type{Padded{L, P}}, logical_size::NamedTuple) where {L, P} = physical_size(L, pad_logical_coord(Padded{L, P}, logical_size))
@inline load(::Type{Padded{L, P}}, workspace, tile::Tile, logical_size::NamedTuple) where {L, P} = load(L, workspace, tile)
@inline store!(::Type{Padded{L, P}}, workspace, value, tile::Tile) where {L, P} = store!(L, workspace, value, tile::Tile)

# ---------------
# AlignedColMajor
# ---------------

struct AlignedColMajor{T} <: LayoutBase{T} end

@inline physical_size(::Type{Padded{AlignedColMajor{T}, P}}, logical_size::NamedTuple) where {T, P} = (logical_size[1] + P, logical_size[2])

@inline fragtype(::Type{AlignedColMajor{T}}, tile_size::NamedTuple) where {T} = NTuple{4, VecElement{Float32}}

@inline function load(::Type{AlignedColMajor{T}}, workspace, tile::Tile{size}) where {T, size}
    N = 16 ÷ sizeof(T)

    linear_base = linearise(tile.base, Base.size(workspace))
    linear_offset = linearise(tile.offset, Base.size(workspace))

    return vloada(Vec{N, T}, pointer(workspace), linear_base + linear_offset - 1)
end

@inline function store!(::Type{AlignedColMajor{T}}, workspace, value, tile::Tile{size}) where {T, size}
    N = 16 ÷ sizeof(T)

    linear_base = linearise(tile.base, Base.size(workspace))
    linear_offset = linearise(tile.offset, Base.size(workspace))

    vstorea!(Vec{N, T}, pointer(workspace), value, linear_base + linear_offset - 1)
end

# --------
# Diagonal
# --------

struct Diagonal{T} <: LayoutBase{T} end

@inline bitcast_helper(x::NTuple{8, VecElement{Float16}}) = Base.llvmcall(
    "
    %ret = bitcast <8 x i16> %0 to <4 x float>
    ret <4 x float> %ret
    ", NTuple{4, VecElement{Float32}}, Tuple{NTuple{8, VecElement{Float16}}}, x)

@inline function load(::Type{Diagonal{T}}, workspace, tile::Tile{size}) where {T, size}
    N = 16 ÷ sizeof(T)

    # The row index is given by t.index[1] + (k - 1), the column index is given by t.index[2] (0-based).
    # Only load on the diagonal, i.e. if row and column are equal.
    # Note that t.index[2] is 0-based, so we need to add 1 before loading from workspace.
    # TODO: Remove the <4 x float> everywhere, so we don't have to do this ugly casting all over the place.
    return bitcast_helper(ntuple(k -> VecElement{Float16}(tile.index[1] + k - 1 == tile.index[2] ? @inbounds(workspace[tile.index[2] + 1]) : 0), Val(8)))
end

@inline threadblock_condition(layout_a::Type{Diagonal{T}}, layout_b, block_i, block_j, block_k, block_tile) where {T} = abs(block_i - block_k) <= block_tile.size.K

# ---------------
# AlignedRowMajor
# ---------------

struct AlignedRowMajor{T} <: LayoutBase{T} end

@inline physical_size(::Type{Padded{AlignedRowMajor{T}, P}}, logical_size::NamedTuple) where {T, P} = (logical_size[2] + P, logical_size[1])

@inline fragtype(::Type{AlignedRowMajor{T}}, tile_size::NamedTuple) where {T} = NTuple{4, VecElement{Float32}}

@inline function load(::Type{AlignedRowMajor{T}}, workspace, tile::Tile{size}) where {T, size}
    N = 16 ÷ sizeof(T)

    linear_base = linearise(reverse(Tuple(tile.base)), Base.size(workspace))
    linear_offset = linearise(reverse(Tuple(tile.offset)), Base.size(workspace))

    return vloada(Vec{N, T}, pointer(workspace), linear_base + linear_offset - 1)
end

@inline function store!(::Type{AlignedRowMajor{T}}, workspace, value, tile::Tile{size}) where {T, size}
    N = 16 ÷ sizeof(T)

    linear_base = linearise(reverse(Tuple(tile.base)), Base.size(workspace))
    linear_offset = linearise(reverse(Tuple(tile.offset)), Base.size(workspace))

    vstorea!(Vec{N, T}, pointer(workspace), value, linear_base + linear_offset - 1)
end

# -------------------
# InterleavedColMajor
# -------------------

struct InterleavedColMajor{T} <: LayoutBase{T} end

@inline fragtype(::Type{InterleavedColMajor{T}}, tile_size::NamedTuple) where {T} = NTuple{tile_size[1] * tile_size[2], Complex{T}}

@inline function load(::Type{InterleavedColMajor{T}}, workspace, tile::Tile{size}) where {T, size}
    x = ntuple(i -> zero(Complex{T}), tile.size[1] * tile.size[2])

    @unroll for j = 1 : tile.size[2]
        @unroll for i = 1 : tile.size[1]
            t = translate_offset(tile, (i - 1, j - 1))
            @inbounds val = workspace[t.index[1] + 1, t.index[2] + 1]
            x = Base.setindex(x, val, (i - 1) * tile.size[2] + j)
        end
    end

    return x
end

@inline function store!(::Type{InterleavedColMajor{T}}, workspace, value, tile::Tile{size}) where {T, size}
    @unroll for j = 1 : tile.size[2]
        @unroll for i = 1 : tile.size[1]
            t = translate_offset(tile, (i - 1, j - 1))
            val = value[(i - 1) * tile.size[2] + j]
            @inbounds workspace[t.index[1] + 1, t.index[2] + 1] = val
        end
    end
end

# -------------------
# InterleavedRowMajor
# -------------------

struct InterleavedRowMajor{T} <: LayoutBase{T} end

@inline fragtype(::Type{InterleavedRowMajor{T}}, tile_size::NamedTuple) where {T} = NTuple{tile_size[1] * tile_size[2], Complex{T}}

@inline function load(::Type{InterleavedRowMajor{T}}, workspace, tile::Tile{size}) where {T, size}
    x = ntuple(i -> zero(Complex{T}), tile.size[1] * tile.size[2])

    @unroll for i = 1 : tile.size[1]
        @unroll for j = 1 : tile.size[2]
            t = translate_offset(tile, (i - 1, j - 1))
            @inbounds val = workspace[t.index[2] + 1, t.index[1] + 1]
            x = Base.setindex(x, val, (i - 1) * tile.size[2] + j)
        end
    end

    return x
end

@inline function store!(::Type{InterleavedRowMajor{T}}, workspace, value, tile::Tile{size}) where {T, size}
    @unroll for i = 1 : tile.size[1]
        @unroll for j = 1 : tile.size[2]
            t = translate_offset(tile, (i - 1, j - 1))
            val = value[(i - 1) * tile.size[2] + j]
            @inbounds workspace[t.index[2] + 1, t.index[1] + 1] = val
        end
    end
end

# -------------
# SplitColMajor
# -------------

struct SplitColMajor{T} <: LayoutBase{T} end

@inline function physical_size(::Type{SplitColMajor{T}}, logical_size::NamedTuple) where {T}
    t = Tuple(logical_size)
    return (t..., 2)
end

@inline fragtype(::Type{SplitColMajor{T}}, tile_size::NamedTuple) where {T} = NTuple{tile_size[1] * tile_size[2], Complex{T}}

@inline function load(::Type{SplitColMajor{T}}, workspace, tile::Tile{size}) where {T, size}
    x = ntuple(i -> zero(Complex{T}), tile.size[1] * tile.size[2])

    @unroll for j = 1 : tile.size[2]
        @unroll for i = 1 : tile.size[1]
            t = translate_offset(tile, (i - 1, j - 1))
            @inbounds val = workspace[t.index[1] + 1, t.index[2] + 1, 1] + im *
                            workspace[t.index[1] + 1, t.index[2] + 1, 2]
            x = Base.setindex(x, val, (i - 1) * tile.size[2] + j)
        end
    end

    return x
end

@inline function store!(::Type{SplitColMajor{T}}, workspace, value, tile::Tile{size}) where {T, size}
    @unroll for j = 1 : tile.size[2]
        @unroll for i = 1 : tile.size[1]
            t = translate_offset(tile, (i - 1, j - 1))
            val = value[(i - 1) * tile.size[2] + j]
            @inbounds workspace[t.index[1] + 1, t.index[2] + 1, 1] = val.re
            @inbounds workspace[t.index[1] + 1, t.index[2] + 1, 2] = val.im
        end
    end
end

# -------------
# SplitRowMajor
# -------------

struct SplitRowMajor{T} <: LayoutBase{T} end

@inline function physical_size(::Type{Padded{SplitRowMajor{T}, P}}, logical_size::NamedTuple) where {T, P}
    return (logical_size[2] + P, logical_size[1], 2)
end

@inline fragtype(::Type{SplitRowMajor{T}}, tile_size::NamedTuple) where {T} = NTuple{tile_size[1] * tile_size[2], Complex{T}}

@inline function load(::Type{SplitRowMajor{T}}, workspace, tile::Tile{size}) where {T, size}
    x = ntuple(i -> zero(Complex{T}), tile.size[1] * tile.size[2])

    @unroll for i = 1 : tile.size[1]
        @unroll for j = 1 : tile.size[2]
            t = translate_offset(tile, (i - 1, j - 1))
            @inbounds val = workspace[t.index[2] + 1, t.index[1] + 1, 1] + im *
                            workspace[t.index[2] + 1, t.index[1] + 1, 2]
            x = Base.setindex(x, val, (i - 1) * tile.size[2] + j)
        end
    end

    return x
end

@inline function store!(::Type{SplitRowMajor{T}}, workspace, value, tile::Tile{size}) where {T, size}
    @unroll for i = 1 : tile.size[1]
        @unroll for j = 1 : tile.size[2]
            t = translate_offset(tile, (i - 1, j - 1))
            val = value[(i - 1) * tile.size[2] + j]
            @inbounds workspace[t.index[2] + 1, t.index[1] + 1, 1] = val.re
            @inbounds workspace[t.index[2] + 1, t.index[1] + 1, 2] = val.im
        end
    end
end

end


text:
struct SemObsCommon{
        A <: Union{AbstractArray, Nothing},
        B <: AbstractArray,
        C <: Union{AbstractArray, Nothing},
        D <: AbstractFloat,
        O <: Union{AbstractFloat, Nothing},
        R <: Union{AbstractArray, Nothing}} <: SemObs
    data::A
    obs_cov::B
    obs_mean::C
    n_man::D
    n_obs::O
    data_rowwise::R
end

function SemObsCommon(;
        data = nothing,
        obs_cov = nothing,
        meanstructure = false,
        rowwise = false)
    # if no cov. matrix was given, compute one
    if isnothing(obs_cov) obs_cov = Statistics.cov(data) end
    isnothing(data) ? n_obs = nothing : n_obs = convert(Float64, size(data, 1))
    n_man = Float64(size(obs_cov, 1))
    # if a meanstructure is needed, compute observed means
    meanstructure ?
        obs_mean = vcat(Statistics.mean(data, dims = 1)...) :
        obs_mean = nothing
    rowwise ?
        data_rowwise = [data[i, :] for i = 1:convert(Int64, n_obs)] :
        data_rowwise = nothing
    return SemObsCommon(data, obs_cov, obs_mean, n_man, n_obs, data_rowwise)
end

struct SemObsMissing{
        A <: AbstractArray,
        D <: AbstractFloat,
        O <: AbstractFloat,
        P <: Vector,
        P2 <: Vector,
        R <: Vector,
        PD <: AbstractArray,
        PO <: AbstractArray,
        PVO <: AbstractArray,
        A2 <: AbstractArray,
        A3 <: AbstractArray
        } <: SemObs
    data::A
    n_man::D
    n_obs::O
    patterns::P # missing patterns
    patterns_not::P2
    rows::R # coresponding rows in data_rowwise
    data_rowwise::PD # list of data
    pattern_n_obs::PO # observed rows per pattern
    pattern_nvar_obs::PVO # number of non-missing variables per pattern
    obs_mean::A2
    obs_cov::A3
end

function SemObsMissing(data)

    # remove persons with only missings
    keep = Vector{Int64}()
    for i = 1:size(data, 1)
        if any(.!ismissing.(data[i, :]))
            push!(keep, i)
        end
    end
    data = data[keep, :]

    n_obs = size(data, 1)
    n_man = size(data, 2)

    # compute and store the different missing patterns with their rowindices
    missings = ismissing.(data)
    patterns = [missings[i, :] for i = 1:size(missings, 1)]

    patterns_cart = findall.(!, patterns)
    data_rowwise = [data[i, patterns_cart[i]] for i = 1:n_obs]
    data_rowwise = convert.(Array{Float64}, data_rowwise)

    remember = Vector{BitArray{1}}()
    rows = [Vector{Int64}(undef, 0) for i = 1:size(patterns, 1)]
    for i = 1:size(patterns, 1)
        unknown = true
        for j = 1:size(remember, 1)
            if patterns[i] == remember[j]
                push!(rows[j], i)
                unknown = false
            end
        end
        if unknown
            push!(remember, patterns[i])
            push!(rows[size(remember, 1)], i)
        end
    end
    rows = rows[1:length(remember)]
    n_patterns = size(rows, 1)

    # sort by number of missings
    sort_n_miss = sortperm(sum.(remember))
    remember = remember[sort_n_miss]
    remember_cart = findall.(!, remember)
    remember_cart_not = findall.(remember)
    rows = rows[sort_n_miss]

    pattern_n_obs = size.(rows, 1)
    pattern_nvar_obs = length.(remember_cart)

    cov_mean = [cov_and_mean(data_rowwise[rows]) for rows in rows]
    obs_cov = [cov_mean[1] for cov_mean in cov_mean]
    obs_mean = [cov_mean[2] for cov_mean in cov_mean]

    return SemObsMissing(data, Float64(n_man), Float64(n_obs), remember_cart,
    remember_cart_not,
    rows, data_rowwise, Float64.(pattern_n_obs), Float64.(pattern_nvar_obs),
    obs_mean, obs_cov)
end


text:
using SymPyCall
using Test

include("tests.jl")
include("test-math.jl")
include("test-matrix.jl")
include("test-ode.jl")
include("test-logical.jl")
include("test-specialfuncs.jl")
#include("test-permutations.jl")
#include("test-physics.jl")
#include("test-external-module.jl")
include("test-latexify.jl")

max_stars_repo_name: jverzani/SymPyCall

max_issues_repo_name: jverzani/SymPyCall

max_forks_repo_name: jverzani/SymPyCall


text:
using  Distributions
using Test

d = Kolmogorov()
# compare to Smirnov's tabulated values
@test round(cdf(d,0.28), digits=6) == .000_001
@test round(cdf(d,0.50), digits=6) == .036_055
@test round(cdf(d,0.99), digits=6) == .719_126
@test round(cdf(d,1.00), digits=6) == .730_000
@test round(cdf(d,1.01), digits=6) == .740_566
# @test round(cdf(d,1.04),digits=6) == .770_434 # table value appears to be wrong
@test round(cdf(d,1.50), digits=6) == .977_782
@test round(cdf(d,2.00), digits=6) == .999_329
@test round(cdf(d,2.50), digits=7) == .999_992_5
@test round(cdf(d,3.00), digits=8) == .999_999_97


text:
#==
This example shows multiobjective Paretofronts with 3 objectives.
The front is convex. The NBI method does not catch the full Pareto front,
e.g. as seen by the missing arcs between f1 and f2, whilst f3 is 0.
To capture this area, we need some of the elements in β to be negative.

Taken from
Il Yong Kim and OL De Weck.
Adaptive weighted sum method for multiobjective optimization:
a new method for pareto front generation.
Structural and Multidisciplinary Optimization, 31(2):105–116, 2006.
==#

using MultiJuMP, JuMP
using Ipopt

m = MultiModel(solver = IpoptSolver())

@variable(m, x[i=1:3] >= 0)
@NLexpression(m, f1, x[1])
@NLexpression(m, f2, x[2])
@NLexpression(m, f3, x[3])
@NLconstraint(m, x[1]^4+2x[2]^3+5x[3]^2<=1)

obj1 = SingleObjective(f1, sense = :Max)#,
                       #iv = Dict{Symbol, Any}(:x => [0, 0., 1.2]))

obj2 = SingleObjective(f2, sense = :Max)#,
                       #iv = Dict{Symbol, Any}(:x => [0, 0., 1.2]))

obj3 = SingleObjective(f3, sense = :Max)#,
                       #iv = Dict{Symbol, Any}(:x => [0, 0., 2.]))

md = getMultiData(m)
md.objectives = [obj1, obj2, obj3]
#md.objectives = [obj1, obj2]
md.pointsperdim = 10
solve(m, method = :NBI)

using Plots
pyplot()
plot(md)


text:
"""
    Base.:-(a::UncertainDataset, b::UncertainDataset;
        n = 10000)

Subtraction operator for uncertain value datasets.

To obtain the difference for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th subtracted uncertain value is a kernel density estimate to those differences, i.e.
`subt_val = UncertainValue(UnivariateKDE, sample_xᵢ .- sample_yᵢ)`.
"""
function Base.:-(a::UncertainDataset, b::UncertainDataset; n = 10000)
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Dataset lengths do not match ($N, $n_vals_b)"))
    end

    sum_dists = [resample(a[i], n) .- resample(b[i], n) for i = 1:N]
    UncertainDataset([UncertainValue(sum_dists[i]) for i = 1:N])
end

"""
    Base.:-(a::Real, b::UncertainDataset;
        n = 10000)

Subtraction operator for scalars and uncertain value datasets.

To obtain the difference for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th subtracted uncertain value is a kernel density estimate to those differences, i.e.
`subt_val = UncertainValue(UnivariateKDE, sample_xᵢ .- sample_yᵢ)`.
"""
function Base.:-(a::T, b::UncertainDataset; n = 10000) where {T <: Number}
    N = length(a)

    UncertainDataset([UncertainValue(a .- resample(b[i], n)) for i = 1:N])
end

"""
    Base.:-(a::Vector{Real}, b::UncertainDataset;
        n = 10000)

Subtraction operator for scalars and uncertain value datasets.

To obtain the difference for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th subtracted uncertain value is a kernel density estimate to those differences, i.e.
`subt_val = UncertainValue(UnivariateKDE, sample_xᵢ .- sample_yᵢ)`.
"""
function Base.:-(a::Vector{T}, b::UncertainDataset; n = 10000) where {T <: Number}
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Lengths do not match ($N, $n_vals_b)"))
    end

    UncertainDataset([UncertainValue(a[i] .- resample(b[i], n)) for i = 1:N])
end

"""
    Base.:-(a::UncertainDataset, b::Real;
        n = 10000)

Subtraction operator for scalars and uncertain value datasets.

To obtain the difference for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th subtracted uncertain value is a kernel density estimate to those differences, i.e.
`subt_val = UncertainValue(UnivariateKDE, sample_xᵢ .- sample_yᵢ)`.
"""
function Base.:-(a::UncertainDataset, b::T; n = 10000) where {T <: Number}
    N = length(a)

    UncertainDataset([UncertainValue(resample(a[i], n) .- b) for i = 1:N])
end

"""
    Base.:-(a::UncertainDataset, b::Real;
        n = 10000)

Subtraction operator for scalars and uncertain value datasets.

To obtain the difference for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th subtracted uncertain value is a kernel density estimate to those differences, i.e.
`subt_val = UncertainValue(UnivariateKDE, sample_xᵢ .- sample_yᵢ)`.
"""
function Base.:-(a::UncertainDataset, b::Vector{T}; n = 10000) where {T <: Number}
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Lengths do not match ($N, $n_vals_b)"))
    end

    UncertainDataset([UncertainValue(resample(a[i], n) .- b[i]) for i = 1:N])
end


text:
export cycle_detection_floyd, cycle_detection_brent

# https://en.wikipedia.org/wiki/Cycle_detection#Algorithms

cycle_detection_floyd(f, x0) = begin
    # Main phase of algorithm: finding a repetition x_i = x_2i.
    # The hare moves twice as quickly as the tortoise and
    # the distance between them increases by 1 at each step.
    # Eventually they will both be inside the cycle and then,
    # at some point, the distance between them will be
    # divisible by the period λ.
    tortoise = f(x0) # f(x0) is the element/node next to x0.
    hare = (f ∘ f)(x0)
    while tortoise ≠ hare
        tortoise = f(tortoise)
        hare = (f ∘ f)(hare)
    end

    # At this point the tortoise position, ν, which is also equal
    # to the distance between hare and tortoise, is divisible by
    # the period λ. So hare moving in circle one step at a time,
    # and tortoise (reset to x0) moving towards the circle, will
    # intersect at the beginning of the circle. Because the
    # distance between them is constant at 2ν, a multiple of λ,
    # they will agree as soon as the tortoise reaches index μ.

    # Find the position μ of first repetition.
    μ = 0
    tortoise = x0
    while tortoise ≠ hare
        tortoise = f(tortoise)
        hare = f(hare)   # Hare and tortoise move at same speed
        μ += 1
    end

    # Find the length of the shortest cycle starting from x_μ
    # The hare moves one step at a time while tortoise is still.
    # λ is incremented until λ is found.
    λ = 1
    hare = f(tortoise)
    while tortoise ≠ hare
        hare = f(hare)
        λ += 1
    end

    λ, μ
end

cycle_detection_brent(f, x0) = begin
    # main phase: search successive powers of two
    power = λ = 1
    tortoise = x0
    hare = f(x0)  # f(x0) is the element/node next to x0.
    while tortoise ≠ hare
        if power == λ  # time to start a new power of two?
            tortoise = hare
            power *= 2
            λ = 0
        end
        hare = f(hare)
        λ += 1
    end

    # Find the position of the first repetition of length λ
    μ = 0
    tortoise = hare = x0
    for i in 0:λ-1
        hare = f(hare)
    end
    # The distance between the hare and tortoise is now λ.

    # Next, the hare and tortoise move at same speed until they agree
    while tortoise ≠ hare
        tortoise = f(tortoise)
        hare = f(hare)
        μ += 1
    end

    λ, μ
end


text:
# ---
# title: 355. Design Twitter
# id: problem355
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Hash Table, Heap, Design
# link: <https://leetcode.com/problems/design-twitter/description/>
# hidden: true
# ---
#
# Design a simplified version of Twitter where users can post tweets,
# follow/unfollow another user and is able to see the 10 most recent tweets in
# the user's news feed. Your design should support the following methods:
#
#   1. **postTweet(userId, tweetId)** : Compose a new tweet.
#   2. **getNewsFeed(userId)** : Retrieve the 10 most recent tweet ids in the user's news feed. Each item in the news feed must be posted by users who the user followed or by the user herself. Tweets must be ordered from most recent to least recent.
#   3. **follow(followerId, followeeId)** : Follower follows a followee.
#   4. **unfollow(followerId, followeeId)** : Follower unfollows a followee.
#
# **Example:**
#
#
#
#     Twitter twitter = new Twitter();
#
#     // User 1 posts a new tweet (id = 5).
#     twitter.postTweet(1, 5);
#
#     // User 1's news feed should return a list with 1 tweet id -> [5].
#     twitter.getNewsFeed(1);
#
#     // User 1 follows user 2.
#     twitter.follow(1, 2);
#
#     // User 2 posts a new tweet (id = 6).
#     twitter.postTweet(2, 6);
#
#     // User 1's news feed should return a list with 2 tweet ids -> [6, 5].
#     // Tweet id 6 should precede tweet id 5 because it is posted after tweet id 5.
#     twitter.getNewsFeed(1);
#
#     // User 1 unfollows user 2.
#     twitter.unfollow(1, 2);
#
#     // User 1's news feed should return a list with 1 tweet id -> [5],
#     // since user 1 is no longer following user 2.
#     twitter.getNewsFeed(1);
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
@testset "formatter" begin
    @static if VERSION < v"1.4"
        @test_broken !isempty(Atom.FORMAT_TEXT_KWARGS)
    else
        @test !isempty(Atom.FORMAT_TEXT_KWARGS)
    end
end


text:
#############################################################
# steven varga 2015 march, Toronto
# licence: MIT Licence
#############################################################

###########################################################
# pretty print result layer by layer
#
import Base.show
export test_gradients

function show(model::Net, e )
    #show( model )
    println("\n*************************************************")
    println("gradient check with finite difference method")
    println("passed: '.'   failed: 'x' ")
    println("*************************************************")
    pos = 1
    for l = 1:length(model.layers)
        if Mocha.has_param(model.layers[l])
            println()
            println( model.layers[l] )
            println( "---------------------------------------------------" )
            w = model.states[l].W.data
            b = model.states[l].b.data
            W = reshape( e[pos:(pos+length(w)-1)], size(w) );
            pos += length(w);  println("weights:");show( W ); println("\n bias:")
            B = reshape( e[pos:(pos+length(b)-1)], size(b) );
            pos += length(b); show( B' );  println()
        end
    end
end

################################################################
# model paramter unrolling into a vector
# suboptimal;
# optimal implemantation constructs model params in a cont.
# space, where unrolling becomes trivial
# it left as an exercise for the reader
function unroll_parameters( model::Net )
    theta = Array{Float64}() # initial state is one length ??? weird
    for l = 1:length(model.layers)
        if Mocha.has_param(model.layers[l])

            for m in model.states[l].parameters
                θ = m.blob.data
                size =  length(θ)
                theta = [ theta; reshape(θ, size, 1)]
            end
        end
    end
    # note the begining, work around the first element
    # is fluke
    return theta[2:end]
end

function unroll_gradients( model::Net )
    theta = Array{Float64}() # initial state is one length ??? weird
    for l = 1:length(model.layers)
        if Mocha.has_param(model.layers[l])
            for m in model.states[l].parameters
                θ = m.gradient.data
                size =  length(θ)
                theta = [ theta; reshape(θ, size, 1)]
            end
        end
    end
    # note the begining, work around as the first element
    # is fluke
    return theta[2:end]
end

#################################################
# updates model paramaters by copying
#
function update_θ!(model,θᵢ)
    pos = 1
    for l = 1:length(model.layers)
        if Mocha.has_param(model.layers[l])

            for m in model.states[l].parameters
                θ = m.blob.data
                len =  length(θ)
                Base.copy!(θ, reshape( θᵢ[pos:(pos+len-1)], size(θ) ))
                pos += len
            end
        end
    end

end

# (J,θ,▽) =
function  hypothesis_and_gradient( model::Net )
  θ = unroll_parameters(model)
  # set_model_params!(model,θ)
  # use closure to embedd params
  function J( θᵢ )
    update_θ!(model,θᵢ)           # update model params
    forward( model )              # compute cost
    return model.states[end].loss # which is stored in the last node or model state
  end

  function grad!(θᵢ,∇ᵢ)
    update_θ!(model,θᵢ)           # update model parameters
    backward(model)               # compute gradients
    ∇ = unroll_gradients(model)   # retrieve them from model state
    Base.copy!(∇ᵢ,∇)              # and update them
  end

  return (J,θ,grad!)
end

#########################################################
#
# J cost function
# g! computes and updates the gradient of the model given θ model params
function compute_finite_difference( J::Function, g!::Function, θ::Vector{Float64}; ε = 1e-6, digit=8 )
  ∇ᵋ = similar(θ);  # gradient/slope with 2ε wiggle room
  ∇  = similar(θ);  # gradient/slope at midpoint
  # compute C cost and ∇ gradient at mid point
  C = J(θ);  	 g!(θ,∇);

  ## define the two sided derivative
  θ⁺ = similar(θ); θ⁻ = similar(θ);

  # iterate through cost function and calculate slope
  for i=1:length(θ)
    Base.copy!(θ⁺,θ);  θ⁺[i] = θ⁺[i] + ε;
    Base.copy!(θ⁻,θ);  θ⁻[i] = θ⁻[i] - ε;
    ∇ᵋ[i] = ( J(θ⁺) - J(θ⁻) ) / 2ε
  end
  return (∇ᵋ,∇)
end

# ############################################################
# gradient check with two sided finite difference method
# based on prof. Andrew Ng,  Machine learning, Coursera
#
#  -- this part could use more work, as an example if gradient fails
#  indicate the positon in the layer itself; to aid debugging.  Since author
#  ported this code from his implentation, further fitting is needed

function gradient_check(model::Net, epsilon::Float64, digit::Int, visual::Bool)
    # create objective that computes grad( θ ), and cost( θ )
    (J,θ, grad) = hypothesis_and_gradient( model::Net )
    ∇ᵋ,∇ = compute_finite_difference( J, grad, θ )

    # do actual comparison with `digit` numerical percision
    # ∇⁺ = round(∇⁺, 4); 	∇ = round(∇, 4)
    idx = round.( abs.(∇ᵋ - ∇), digit ) .!= 0
    if visual
        δ = Array{Char}(length(idx));  fill!(δ,'.')
        δ[idx] = 'x'
        show(model, δ)
        #show(model,round(∇ᵋ,digit) ); show(model,round(∇,digit))
    end
    # return false if fail at any point
    # TODO: check if correct
    sum( round.( abs.(∇ᵋ - ∇), digit) ) < epsilon
end

###############################################################
# exported method
#
# make sure to set datalayer so the total size is same as  batch_size
# one line of data+label is sufficient; but works with minibatches as well
#
# your milage may vary; function of epsilon ~ 1e-6 - 1e-10
# signifacant digit 6 - 8 higher is better gradient
# if your fanout is big theninitial weights may be set small
# when gradient fails reduce fan-out and test again; your
# gradients may be ok
###############################################################
function test_gradients(net::Net; epsilon=1e-6, digit=6, visual=true )
    return  typeof(net.backend) == Mocha.CPUBackend  ?
         gradient_check( net, epsilon, digit, visual ) :false
end

#= test driver; used code from existing example
detach and remove leading closing comment

use_cuda = false

using Mocha
srand(12345678)

############################################################
# Prepare Random Data
############################################################
N = 5    # works with arbitrary minibatch size as long as
         # N == batch_size in MemoryDataLayer so it cycles through
         # and gets the same data during forward()
M = 10
P = 4

X = rand(M, N)
W = rand(M, P)
B = rand(P, 1)

Y = (W'*X .+ B)
Y = Y + 0.01*randn(size(Y))

############################################################
# Define network
############################################################
 backend = CPUBackend()
init(backend)

data_layer = MemoryDataLayer(batch_size=N, data=Array[X,Y])

w1 = InnerProductLayer(neuron=Neurons.Sigmoid(), name="ip1",output_dim=20, tops=[:a], bottoms=[:data])
w2 = InnerProductLayer(neuron=Neurons.Identity(), name="ip2",output_dim=4, tops=[:b], bottoms=[:a])
loss_layer = SquareLossLayer(name="loss", bottoms=[:b, :label] )

net = Net("TEST", backend, [w1,w2, loss_layer, data_layer])
println(net)

# epsilon:     milage may vary 1e-4 - 1e-8
# digit:       compare this many digits to check for 'identity'
# visualize:   prints out correct and failed positions of gradients
test_gradients(net, epsilon=1e-8, digit=6, visual=true )

shutdown(backend)
=#


text:
spg_niggli_reduce!(lattice::Array{Float64, 2},
                symprec::Float64=1e-5) = ccall((:spg_niggli_reduce, libsymspg), Int32,
                                              (Ptr{Float64}, Float64),
                                              lattice, symprec)

spg_delaunay_reduce!(lattice::Array{Float64, 2},
                    symprec::Float64=1e-5) = ccall((:spg_delaunay_reduce, libsymspg), Int32,
                                                  (Ptr{Float64}, Float64),
                                                  lattice, symprec)


text:
using ITensors
using Test
import Combinatorics

digits(::Type{T},i,j,k) where {T} = T(i*10^2+j*10+k)

@testset "ITensor $T Contractions" for T ∈ (Float64,ComplexF64)
  mi,mj,mk,ml,mα = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  α = Index(mα,"alpha")
  @testset "Test contract ITensors" begin
    A = randomITensor(T)
    B = randomITensor(T)
    Ai = randomITensor(T,i)
    Bi = randomITensor(T,i)
    Aj = randomITensor(T,j)
    Aij = randomITensor(T,i,j)
    Bij = randomITensor(T,i,j)
    Aik = randomITensor(T,i,k)
    Ajk = randomITensor(T,j,k)
    Ajl = randomITensor(T,j,l)
    Akl = randomITensor(T,k,l)
    Aijk = randomITensor(T,i,j,k)
    Ajkl = randomITensor(T,j,k,l)
    Aikl = randomITensor(T,i,k,l)
    Aklα = randomITensor(T,k,l,α)
    Aijkl = randomITensor(T,i,j,k,l)
    @testset "Test contract ITensor (Scalar*Scalar -> Scalar)" begin
      C = A*B
      @test scalar(C)≈scalar(A)*scalar(B)
    end
    @testset "Test contract ITensor (Scalar*Vector -> Vector)" begin
      C = A*Ai
      @test array(C)≈scalar(A)*array(Ai)
    end
    @testset "Test contract ITensor (Vector*Scalar -> Vector)" begin
      C = Aj*A
      @test array(C)≈scalar(A)*array(Aj)
    end
    @testset "Test contract ITensors (Vectorᵀ*Vector -> Scalar)" begin
      C = Ai*Bi
      CArray = transpose(array(Ai))*array(Bi)
      @test CArray≈scalar(C)
    end
    @testset "Test contract ITensors (Vector*Vectorᵀ -> Matrix)" begin
      C = Ai*Aj
      for ii ∈ 1:dim(i), jj ∈ 1:dim(j)
        @test C[i(ii),j(jj)] ≈ Ai[i(ii)]*Aj[j(jj)]
      end
    end
    @testset "Test contract ITensors (Matrix*Scalar -> Matrix)" begin
      Aij = permute(Aij,i,j)
      C = Aij*A
      @test array(permute(C,i,j))≈scalar(A)*array(Aij)
    end
    @testset "Test contract ITensors (Matrix*Vector -> Vector)" begin
      Aij = permute(Aij,i,j)
      C = Aij*Aj
      CArray = array(permute(Aij,i,j))*array(Aj)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrixᵀ*Vector -> Vector)" begin
      Aij = permute(Aij,j,i)
      C = Aij*Aj
      CArray = transpose(array(Aij))*array(Aj)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Vector*Matrix -> Vector)" begin
      Aij = permute(Aij,i,j)
      C = Ai*Aij
      CArray = transpose(transpose(array(Ai))*array(Aij))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Vector*Matrixᵀ -> Vector)" begin
      Aij = permute(Aij,j,i)
      C = Ai*Aij
      CArray = transpose(transpose(array(Ai))*transpose(array(Aij)))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrix*Matrix -> Scalar)" begin
      Aij = permute(Aij,i,j)
      Bij = permute(Bij,i,j)
      C = Aij*Bij
      CArray = LinearAlgebra.tr(array(Aij)*transpose(array(Bij)))
      @test CArray≈scalar(C)
    end
    @testset "Test contract ITensors (Matrix*Matrix -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      CArray = array(Aij)*array(Ajk)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrixᵀ*Matrix -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      CArray = transpose(array(Aij))*array(Ajk)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrix*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      CArray = array(Aij)*transpose(array(Ajk))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrixᵀ*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      CArray = transpose(array(Aij))*transpose(array(Ajk))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrix⊗Matrix -> 4-tensor)" begin
      C = Aij*Akl
      for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k), ll ∈ 1:dim(l)
        @test C[i(ii),j(jj),k(kk),l(ll)] ≈ Aij[i(ii),j(jj)]*Akl[k(kk),l(ll)]
      end
    end
    @testset "Test contract ITensors (3-Tensor*Scalar -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aijk*A
      @test array(permute(C,i,j,k)) ≈ scalar(A)*array(Aijk)
    end
    @testset "Test contract ITensors (3-Tensor*Vector -> Matrix)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aijk*Ai
      CArray = reshape(reshape(array(permute(Aijk,j,k,i)),dim(j)*dim(k),dim(i))*array(Ai),dim(j),dim(k))
      @test CArray≈array(permute(C,j,k))
    end
    @testset "Test contract ITensors (Vector*3-Tensor -> Matrix)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aj*Aijk
      CArray = reshape(transpose(array(Aj))*reshape(array(permute(Aijk,j,i,k)),dim(j),dim(i)*dim(k)),dim(i),dim(k))
      @test CArray≈array(permute(C,i,k))
    end
    @testset "Test contract ITensors (3-Tensor*Matrix -> Vector)" begin
      Aijk = permute(Aijk,i,j,k)
      Aik = permute(Aik,i,k)
      C = Aijk*Aik
      CArray = reshape(array(permute(Aijk,j,i,k)),dim(j),dim(i)*dim(k))*vec(array(Aik))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (3-Tensor*Matrix -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajl = permute(Ajl,j,l)
      C = Aijk*Ajl
      CArray = reshape(reshape(array(permute(Aijk,i,k,j)),dim(i)*dim(k),dim(j))*array(Ajl),dim(i),dim(k),dim(l))
      @test CArray≈array(permute(C,i,k,l))
    end
    @testset "Test contract ITensors (Matrix*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Akl = permute(Akl,k,l)
      C = Akl*Aijk
      CArray = reshape(array(permute(Akl,l,k))*reshape(array(permute(Aijk,k,i,j)),dim(k),dim(i)*dim(j)),dim(l),dim(i),dim(j))
      @test CArray≈array(permute(C,l,i,j))
    end
    @testset "Test contract ITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajkl = permute(Ajkl,j,k,l)
      C = Aijk*Ajkl
      CArray = reshape(array(permute(Aijk,i,j,k)),dim(i),dim(j)*dim(k))*reshape(array(permute(Ajkl,j,k,l)),dim(j)*dim(k),dim(l))
      @test CArray≈array(permute(C,i,l))
    end
    @testset "Test contract ITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijk ∈ Combinatorics.permutations([i,j,k]), inds_jkl ∈ Combinatorics.permutations([j,k,l])
        Aijk = permute(Aijk,inds_ijk...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijk
        CArray = reshape(array(permute(Ajkl,l,j,k)),dim(l),dim(j)*dim(k))*reshape(array(permute(Aijk,j,k,i)),dim(j)*dim(k),dim(i))
        @test CArray≈array(permute(C,l,i))
      end
    end
    @testset "Test contract ITensors (4-Tensor*3-Tensor -> 1-Tensor)" begin
      for inds_ijkl ∈ Combinatorics.permutations([i,j,k,l]), inds_jkl ∈ Combinatorics.permutations([j,k,l])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijkl
        CArray = reshape(array(permute(Ajkl,j,k,l)),1,dim(j)*dim(k)*dim(l))*reshape(array(permute(Aijkl,j,k,l,i)),dim(j)*dim(k)*dim(l),dim(i))
        @test vec(CArray)≈array(permute(C,i))
      end
    end
    @testset "Test contract ITensors (4-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijkl ∈ Combinatorics.permutations([i,j,k,l]), inds_klα ∈ Combinatorics.permutations([k,l,α])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Aklα = permute(Aklα,inds_klα...)
        C = Aklα*Aijkl
        CArray = reshape(reshape(array(permute(Aklα,α,k,l)),dim(α),dim(k)*dim(l))*reshape(array(permute(Aijkl,k,l,i,j)),dim(k)*dim(l),dim(i)*dim(j)),dim(α),dim(i),dim(j))
        @test CArray≈array(permute(C,α,i,j))
      end
    end
  end # End contraction testset
end

@testset "Contraction conversions" begin

  @testset "Real scalar * Complex ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    x = rand(Float64)
    A = randomITensor(ComplexF64,i,j)
    B = x*A
    for ii ∈ dim(i), jj ∈ dim(j)
      @test B[i(ii),j(jj)] == x*A[i(ii),j(jj)]
    end
  end
  @testset "Complex scalar * Real ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    x = rand(ComplexF64)
    A = randomITensor(Float64,i,j)
    B = x*A
    for ii ∈ dim(i), jj ∈ dim(j)
      @test B[i(ii),j(jj)] == x*A[i(ii),j(jj)]
    end
  end
  @testset "Real ITensor * Complex ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    k = Index(2,"k")
    A = randomITensor(Float64,i,j)
    B = randomITensor(ComplexF64,j,k)
    C = A*B
    @test array(permute(C,i,k)) ≈ array(A)*array(B)
  end
  @testset "Complex ITensor * Real ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    k = Index(2,"k")
    A = randomITensor(ComplexF64,i,j)
    B = randomITensor(Float64,j,k)
    C = A*B
    @test array(permute(C,i,k)) ≈ array(A)*array(B)
  end

  @testset "Outer Product Real ITensor * Complex ITensor" begin
    i = Index(2, "i")
    j = Index(2, "j")
    A = randomITensor(Float64, i)
    B = randomITensor(ComplexF64, j)
    C = A*B
    @test array(permute(C, i, j)) ≈ kron(array(A), transpose(array(B)))
  end

  @testset "Outer Product: Complex ITensor * Real ITensor" begin
    i = Index(2, "i")
    j = Index(2, "j")
    A = randomITensor(ComplexF64, i)
    B = randomITensor(Float64, j)
    C = A*B
    @test array(permute(C, i, j)) ≈ kron(array(A), transpose(array(B)))
  end

end

nothing


text:
"""
    LinearizedModel{M,L,T} <: AbstractLinearModel

A container for the linearized model that holds the full nonlinear model, the linearized model, and the
trajectory of linearization points. The same dynamics and jacobian functions can still be called on the
`LinearizedModel` type.

# Constructors
    LinearizedModel(nonlinear_model::AbstractModel, Z::AbstractTrajectory; kwargs...)
    LinearizedModel(nonlinear_model::AbstractModel, [z::AbstractKnotPoint]; kwargs...)

Linearizes `nonlinear_model` about the trajectory `Z` or a single point `z`. If not specified,
`z` is defined as the state and control defined by `zeros(nonlinear_model)`.

Linearization is, by default, on the continuous system.

# Keyword Arguments
* `is_affine` - Linearize the system with an affine term, such that the new state is the same as the original state. See below for more details.
* `dt` - Time step. If not provided, defaults to the value in `Z` or `z`. Must be specified and non-zero for the system to be discretized. If `dt = NaN`, then the dt will be inferred from the trajectory (useful for variable step sizes).
* `integration` - An explicit integration method. Must also specify a non-zero dt.

If `is_affine = false`, the dynamics are defined as:
``f(x,u) \\approx f(x_0, u_0) + A \\delta x + B \\delta u``

which defines the error state ``\\delta x = x - x_0``
    as the state of the linearized system. Here ``A`` and ``B`` are the partial derivative
    of the dynamics with respect to the state and control, respectively.

If `is_affine = true`, the form is an affine function of the form
``f(x,u) \\approx A x + B u + d``

where ``d = f(x_0,u_0) - A x_0 - B u_0``, which maintains the same definition of the state.
"""
struct LinearizedModel{M, Q, L, T, D} <: AbstractLinearModel
    model::M
    linmodel::L
    Z::T
    F::D
    function LinearizedModel(::Type{Q}, model::AbstractModel, linmodel::LinearModel,
            Z::AbstractTrajectory) where Q <: QuadratureRule
        if !is_discrete(linmodel)
            @assert Q == Continuous
        end
        n,m = size(linmodel)
        F = DynamicsJacobian(model)
        new{typeof(model), Q, typeof(linmodel), typeof(Z), typeof(F)}(
            model, linmodel, Z, F
        )
    end
end

for method in (:state_dim, :control_dim, :is_discrete, :is_affine, :is_timevarying)
    @eval $(method)(model::LinearizedModel) = $(method)(model.linmodel)
end
integration(::LinearizedModel{<:Any,Q}) where Q = Q

get_A(model::LinearizedModel, k=1) = get_A(model.linmodel, k)
get_B(model::LinearizedModel, k=1) = get_B(model.linmodel, k)
get_k(model::LinearizedModel, t) = get_k(model.linmodel, t)

"""
"""
function LinearizedModel(nonlinear_model::AbstractModel, Z::AbstractTrajectory;
        integration::Type{Q}=DEFAULT_Q,
        dt=0.0,
        is_affine=false
    ) where {Q<:Explicit}
    n,m = state_dim(nonlinear_model), control_dim(nonlinear_model)

    is_discrete = dt != zero(dt)
    if is_discrete
        if !isnan(dt)
            set_dt!(Z, dt)
        end
        Q0 = Q
    else
        Q0 = Continuous
        if Q != DEFAULT_Q
            @warn "Discarding input for integration, since the system in continuous. Must specify a non-zero dt if you want discretize the system."
        end
    end

    if !isnan(dt) && is_discrete
        set_dt!(Z, dt)
    elseif !isnan(dt) && !is_discrete
    end

    times = length(Z) > 1 ? get_times(Z) : 1:0

    linmodel = LinearModel(n, m; is_affine=is_affine, times=times, dt=dt)

    model = LinearizedModel(Q0, nonlinear_model, linmodel, Z)

    linearize!(Q0, model)

    model
end

# single point linearization
function LinearizedModel(nonlinear_model::AbstractModel,
        z::KnotPoint=KnotPoint(zeros(nonlinear_model)...,0.,0.);
        kwargs...
    ) where {Q<:Explicit}
    LinearizedModel(nonlinear_model, Traj([z]); kwargs...)
end

# Pass dynamics function through to linear model
@inline dynamics(model::LinearizedModel, x, u, t=0.0) = dynamics(model.linmodel, x, u, t)
@inline discrete_dynamics(::Type{PassThrough}, model::LinearizedModel, x, u, t, dt) =
    discrete_dynamics(PassThrough, model.linmodel, x, u, t, dt)

@inline jacobian!(∇f::AbstractMatrix, model::LinearizedModel, z::AbstractKnotPoint) =
    jacobian!(∇f, model.linmodel, z)

@inline discrete_jacobian!(::Type{PassThrough}, ∇f, model::LinearizedModel, z::AbstractKnotPoint) =
    discrete_jacobian!(PassThrough, ∇f, model.linmodel, z)

"""
    update_trajectory!(model::LinearizedModel, Z::AbstractTrajectory, integration::=DEFAULT_Q)

Updates the trajectory inside of the `model` and relinearizes (and discretizes for discrete
    models) the model about the new trajectory.
"""
function update_trajectory!(model::LinearizedModel{<:Any,Q}, Z::AbstractTrajectory) where {Q<:QuadratureRule}
    model.Z .= Z
    linearize!(Q, model)
end

# linearize a single knot point
"""
    linearize!(model::LinearizedModel, integration=DEFAULT_Q)

Linearize the nonlinear model `model.model` about `model.Z`, storing the result in
`model.linmodel`.

Linearization defaults to the form:
``f(x,u) \\approx f(x_0, u_0) + A \\delta x + B \\delta u``

if `is_affine(model.linmodel == false`, which defines the error state ``\\delta x = x - x_0``
    as the state of the linearized system. Here ``A`` and ``B`` are the partial derivative
    of the dynamics with respect to the state and control, respectively.

Otherwise, the form is an affine function of the form
``f(x,u) \\approx A x + B u + d``

where ``d = f(x_0,u_0) - A x_0 - B u_0``, which maintains the same definition of the state.

"""
function linearize!(::Type{Q}, model::LinearizedModel) where Q <: QuadratureRule
    nlmodel = model.model
    linmodel = model.linmodel

    for k in eachindex(model.Z)
        if is_discrete(linmodel)
            discretize!(Q, model, k)
        else
            @assert Q == Continuous
            z = model.Z[k]
            F = model.F
            jacobian!(F, nlmodel, z)
            if is_affine(model)
                linmodel.d[k] .= dynamics(model, z) - F * z.z
            end
            linmodel.A[k] .= get_A(F)
            linmodel.B[k] .= get_B(F)
        end
    end
end

"""
    discretize!(::Type{Q}, model::LinearizedModel, k)

Discretize the linearized model at time step k, using integration `Q`.
"""
function discretize!(::Type{Q}, model::LinearizedModel, k::Int) where Q <: Explicit
    z = model.Z[k]
    discrete_jacobian!(Q, model.F, model.model, z)
    if is_affine(model)
        d = model.linmodel.d[k]
        d .= discrete_dynamics(Q, model.model, z) - model.F * z.z
    end

    # Copy to LinearModel
    model.linmodel.A[k] .= model.F.A
    model.linmodel.B[k] .= model.F.B
end

function discretize!(::Type{Exponential}, model::LinearizedModel, k::Int)
    n,m = size(model.linmodel)
    z = model.Z[k]
    E = model.linmodel.E
    F = model.F

    # Calculate continuous Jacobian of nonlinear system
    jacobian!(F, model.model, z)

    # Calculate Matrix exponential
    matrix_exponential!(E, get_A(F), get_B(F), z.dt)

    # Copy discrete system
    ix,iu = 1:n, n .+ (1:m)
    A = model.linmodel.A[k]
    B = model.linmodel.B[k]
    A .= E[ix,ix]
    B .= E[ix,iu]

    if is_affine(model)
        D = E[ix, (n+m) .+ (1:n)]
        d0 = dynamics(model.model, z)  # nonlinear dynamics
        d0 = d0 - F * z.z              # continuous affine term

        d = model.linmodel.d[k]
        mul!(d, D, d0)
    end
end


text:
import Printf.@sprintf
export dm_params

function dm_params()
    rodent = keys(DM_LOCOMOTION_RODENT)
    humanoid = keys(DM_LOCOMOTION_HUMANOID)
    dm_control = keys(DM_CONTROL_SUITE_SIZE)

    game = Dict("rodent" => rodent, "humanoid" => humanoid, "dm_control" => dm_control)
    shards = 0:4

    @info game shards
end

const DM_LOCOMOTION_RODENT = Dict{String, String}(
    "rodent_gaps" => "dm_locomotion/rodent_gaps/seq2",
    "rodent_escape" => "dm_locomotion/rodent_bowl_escape/seq2",
    "rodent_two_touch" => "dm_locomotion/rodent_two_touch/seq40",
    "rodent_mazes" => "dm_locomotion/rodent_mazes/seq40"
)

const DM_LOCOMOTION_RODENT_SIZE = Dict{String, Tuple}(
    "observation/walker/actuator_activation" => (38,),
    # "observation/walker/sensors_torque" => (),
    # "observation/walker/sensors_force" => (),
    "observation/walker/body_height" => (1,),
    "observation/walker/end_effectors_pos" => (12,),
    "observation/walker/joints_pos"=> (30,),
    "observation/walker/joints_vel"=> (30,),
    "observation/walker/tendons_pos"=> (8,),
    "observation/walker/tendons_vel"=> (8,),
    "observation/walker/appendages_pos"=> (15,),
    "observation/walker/world_zaxis"=> (3,),
    "observation/walker/sensors_accelerometer"=> (3,),
    "observation/walker/sensors_velocimeter"=> (3,),
    "observation/walker/sensors_gyro"  => (3,),
    "observation/walker/sensors_touch"=> (4,),
    "observation/walker/egocentric_camera"=> (64, 64, 3),
    "action"=> (38,),
    "discount"=> (),
    "reward"=> (),
    "step_type"=> ()
)

const DM_LOCOMOTION_HUMANOID = Dict{String, String}(
    "humanoid_corridor" => "dm_locomotion/humanoid_corridor/seq2",
    "humanoid_gaps" => "dm_locomotion/humanoid_gaps/seq2",
    "humanoid_walls" => "dm_locomotion/humanoid_walls/seq40"
)

const DM_LOCOMOTION_HUMANOID_SIZE = Dict{String, Tuple}(
    # "observation/walker/actuator_activation" => (0,),
    "observation/walker/sensors_torque" => (6,),
    # "observation/walker/sensors_force" => (),
    "observation/walker/joints_vel"=> (56,),
    "observation/walker/sensors_velocimeter"=> (3,),
    "observation/walker/sensors_gyro"=> (3,),
    "observation/walker/joints_pos"=> (56,),
    "observation/walker/appendages_pos" => (15,),
    "observation/walker/world_zaxis"=> (3,),
    "observation/walker/body_height"=> (1,),
    "observation/walker/sensors_accelerometer"=> (3,),
    "observation/walker/end_effectors_pos"=> (12,),
    "observation/walker/egocentric_camera"=> (
        64,
        64,
        3,
    ),
    "action"=> (56,),
    "discount"=> (),
    "reward"=> (),
    # "episodic_reward"=> (),
    "step_type"=> ()
)

const DM_CONTROL_SUITE_SIZE = Dict{String, Dict{String, Tuple}}(
    "cartpole_swingup" => Dict{String, Tuple}(
        "observation/position"=> (3,),
        "observation/velocity"=> (2,),
        "action"=> (1,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    ),
    "cheetah_run" => Dict{String, Tuple}(
        "observation/position"=> (8,),
        "observation/velocity"=> (9,),
        "action"=> (6,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    ),
    "finger_turn_hard" => Dict{String, Tuple}(
        "observation/position"=> (4,),
        "observation/velocity"=> (3,),
        "observation/touch"=> (2,),
        "observation/target_position"=> (2,),
        "observation/dist_to_target"=> (1,),
        "action"=> (2,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    ),
    "fish_swim" => Dict{String, Tuple}(
        "observation/target"=> (3,),
        "observation/velocity"=> (13,),
        "observation/upright"=> (1,),
        "observation/joint_angles"=> (7,),
        "action"=> (5,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    ),
    "humanoid_run" => Dict{String, Tuple}(
        "observation/velocity"=> (27,),
        "observation/com_velocity"=> (3,),
        "observation/torso_vertical"=> (3,),
        "observation/extremities"=> (12,),
        "observation/head_height"=> (1,),
        "observation/joint_angles"=> (21,),
        "action"=> (21,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    ),
    "manipulator_insert_ball" => Dict{String, Tuple}(
        "observation/arm_pos"=> (16,),
        "observation/arm_vel"=> (8,),
        "observation/touch"=> (5,),
        "observation/hand_pos"=> (4,),
        "observation/object_pos"=> (4,),
        "observation/object_vel"=> (3,),
        "observation/target_pos"=> (4,),
        "action"=> (5,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    ),
    "manipulator_insert_peg" => Dict{String, Tuple}(
        "observation/arm_pos"=> (16,),
        "observation/arm_vel"=> (8,),
        "observation/touch"=> (5,),
        "observation/hand_pos"=> (4,),
        "observation/object_pos"=> (4,),
        "observation/object_vel"=> (3,),
        "observation/target_pos"=> (4,),
        "episodic_reward"=> (),
        "action"=> (5,),
        "discount"=> (),
        "reward"=> (),
        "step_type"=> ()
    ),
    "walker_stand" => Dict{String, Tuple}(
        "observation/orientations"=> (14,),
        "observation/velocity"=> (9,),
        "observation/height"=> (1,),
        "action"=> (6,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    ),
    "walker_walk" => Dict{String, Tuple}(
        "observation/orientations"=> (14,),
        "observation/velocity"=> (9,),
        "observation/height"=> (1,),
        "action"=> (6,),
        "discount"=> (),
        "reward"=> (),
        "episodic_reward"=> (),
        "step_type"=> ()
    )
)

const DM_LOCOMOTION = merge(DM_LOCOMOTION_HUMANOID, DM_LOCOMOTION_RODENT)

num_shards = 100

function dm_init()
    repo = "rl-unplugged-dm"
    for task in keys(DM_CONTROL_SUITE_SIZE)
        for index in 0:num_shards-1
            register(
                DataDep(
                    "$repo-$task-$index",
                    """
                    Dataset: RL Unplugged dm-control
                    Credits: https://arxiv.org/abs/2006.13888
                    Url: https://github.com/deepmind/deepmind-research/tree/master/rl_unplugged
                    Authors: Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Tom Le Paine,
                    Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal,
                    Josh Merel, Daniel Mankowitz, Cosmin Paduraru, Gabriel
                    Dulac-Arnold, Jerry Li, Mohammad Norouzi, Matt Hoffman,
                    Ofir Nachum, George Tucker, Nicolas Heess, Nando deFreitas
                    Year: 2020

                    DeepMind Control Suite Tassa et al., 2018 is a set of control tasks implemented in MuJoCo Todorov et al., 2012.
                    We consider a subset of the tasks provided in the suite that cover a wide range of difficulties.
                    Most of the datasets in this domain are generated using D4PG. For the environments Manipulator insert
                    ball and Manipulator insert peg we use V-MPO Song et al., 2020 to generate the data as D4PG is unable to
                    solve these tasks. Datasets are available for 9 dm-control-suite tasks. For details on how the dataset was generated,
                    please refer to the paper. DeepMind Control Suite is a traditional continuous action RL benchmark. In particular, it is recommended
                    that you test your approach in DeepMind Control Suite if you are interested in comparing against other state of the art offline RL methods.
                    """,
                    "gs://rl_unplugged/dm_control_suite/$task/"*@sprintf("train-%05i-of-%05i", index, num_shards);
                    fetch_method = fetch_gc_file,
                )
            )
        end
    end
    for task in keys(DM_LOCOMOTION)
        for index in 0:num_shards-1
            register(
                DataDep(
                    "$repo-$task-$index",
                    """
                    Dataset: RL Unplugged dm-locomotion
                    Credits: https://arxiv.org/abs/2006.13888
                    Url: https://github.com/deepmind/deepmind-research/tree/master/rl_unplugged
                    Authors: Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Tom Le Paine,
                    Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal,
                    Josh Merel, Daniel Mankowitz, Cosmin Paduraru, Gabriel
                    Dulac-Arnold, Jerry Li, Mohammad Norouzi, Matt Hoffman,
                    Ofir Nachum, George Tucker, Nicolas Heess, Nando deFreitas
                    Year: 2020

                    These tasks are made up of the corridor locomotion tasks involving the CMU Humanoid, for which prior efforts
                    have either used motion capture data Merel et al., 2019a, Merel et al., 2019b or training from scratch Song et
                    al., 2020. In addition, the DM Locomotion repository contains a set of tasks adapted to be suited to a virtual
                    rodent Merel et al., 2020. We emphasize that the DM Locomotion tasks feature the combination of challenging high-DoF
                    continuous control along with perception from rich egocentric observations. For details on how the dataset was
                    generated, please refer to the paper.

                    It is recommended that you to try offline RL methods on DeepMind Locomotion dataset, if you are interested in very challenging
                    offline RL dataset with continuous action space.
                    """,
                    "gs://rl_unplugged/$(DM_LOCOMOTION[task])/"*@sprintf("train-%05i-of-%05i", index, num_shards);
                    fetch_method = fetch_gc_file
                )
            )
        end
    end
end


text:
#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.

# hygiene.jl
# Make sure that our macros have good hygiene
using Base.Test

module M
import JuMP

mymod = JuMP.Model()
mysense = :Min
JuMP.@variable(mymod, x >= 0)
r = 3:5
JuMP.@variable(mymod, y[i=r] <= i)
JuMP.@constraint(mymod, x + sum{ j*y[j], j=r } <= 1)
JuMP.@constraint(mymod, sum{ y[j], j=r ; j == 4} <= 1)
JuMP.@constraint(mymod, -1 <= x + y[3] <= 1)
JuMP.@objective(mymod, mysense, y[4])
JuMP.@NLconstraint(mymod, y[3] == 1)

end


text:
include("../setupTests.jl")
# build domain and true image
domain = [0.0 1.0 0.0 1.0]
n      = [16,16]
Minv   = getRegularMesh(domain,n)
xc     = getCellCenteredGrid(Minv)
xtrue = sin(2*pi*xc[:,1]).*sin(pi*xc[:,2])

# get noisy data
A     = sparse(1.0I,Minv.nc,Minv.nc)
ids   = sort(rand(1:Minv.nc,round(Int64,Minv.nc*.8)))
A     = A[ids,:]
btrue = A*xtrue
bdata = btrue + .1*randn(size(btrue))/norm(btrue,Inf)
Wd    = ones(length(btrue))

# generate misfit param
i1     = (1:round(Int64,size(A,1)/2))
i2     = (round(Int64,size(A,1)/2)+1:size(A,1))
pFor1  = LSparam(A[i1,:],[])
pFor2  = LSparam(A[i2,:],[])
sigmaBack = zeros(length(xtrue))
Iact         = sparse(1.0I,Minv.nc,Minv.nc)
gl1          = getGlobalToLocal(1.0,sigmaBack)
gl2          = getGlobalToLocal(Iact,sigmaBack)
bd1          = bdata[i1]
bd2          = bdata[i2]
Wd1          = Wd[i1]
Wd2          = Wd[i2]

pMisRefs    = Array(RemoteChannel,2)
pMisRefs[1] = initRemoteChannel(getMisfitParam,workers()[1],pFor1,Wd1,bd1,SSDFun,fMod,gl1)
pMisRefs[2] = initRemoteChannel(getMisfitParam,workers()[min(2,nworkers())],pFor2,Wd2,bd2,SSDFun,fMod,gl2)

pMis    = Array{MisfitParam}(2)
pMis[1] = getMisfitParam(pFor1,Wd1,bd1,SSDFun,fMod,gl1)
pMis[2] = getMisfitParam(pFor2,Wd2,bd2,SSDFun,fMod,gl2)

# setup pInv
alpha        = 1.
x0           = mean(xtrue)*ones(Minv.nc)
boundsLow    = -.5*ones(Minv.nc)
boundsHigh   = 0.5*ones(Minv.nc)
sigmaBack    = zeros(Minv.nc)

print("\t== compare with projGN...")
pInv         = getInverseParam(Minv,fMod,diffusionReg,alpha,x0,boundsLow,boundsHigh)
pInv.maxIter = 5
x1t, = projGN(copy(x0),pInv,pMis)
pInv.maxIter = 5
x2t, = projGN(copy(x0),pInv,pMisRefs,out=1)
@test norm(x1t-x2t)/norm(x1t) < 1e-12
@test all(x1t.>=boundsLow)
@test all(x2t.>=boundsLow)
@test all(x1t.<=boundsHigh)
@test all(x2t.<=boundsHigh)
print("passed! ===")

print("\t== test projGN with projCG solver...")
pInv.maxIter = 5
x1, = projGN(copy(x0),pInv,pMis)
pInv.maxIter = 5
x2, = projGN(copy(x0),pInv,pMisRefs,out=1)
@test norm(x1-x2)/norm(x1) < 1e-12
@test norm(x1t-x1)/norm(x1t) < 1e-12
@test norm(x2t-x2)/norm(x2t) < 1e-12
@test all(x1.>=boundsLow)
@test all(x2.>=boundsLow)
@test all(x1.<=boundsHigh)
@test all(x2.<=boundsHigh)
print("passed! ===")

print("\t== test projGN with direct solver...")
pInv.maxIter = 5
x1, = projGN(copy(x0),pInv,pMis,solveGN=projGNexplicit)
pInv.maxIter = 5
x2, = projGN(copy(x0),pInv,pMisRefs,solveGN=projGNexplicit,out=1)
@test norm(x1-x2)/norm(x1) < 1e-12
@test norm(x1t-x1)/norm(x1t) < 1e-2
@test norm(x2t-x2)/norm(x2t) < 1e-2
@test all(x1.>=boundsLow)
@test all(x2.>=boundsLow)
@test all(x1.<=boundsHigh)
@test all(x2.<=boundsHigh)

print("\t== test projGN with singular system...")
pInv.alpha = 0.
x2, = projGN(copy(x0),pInv,pMis[1],solveGN=projGNexplicit,out=1)
@test all(x2.>=boundsLow)
@test all(x2.<=boundsHigh)

print("passed! ===")


text:
using MarketData

facts("TimeVector construction") do

    context("TimeVectors are constructed") do
        @fact true --> true
    end
end


text:
# Set tests
@test Set(["a","c","d","e"]) - Set(["b","a","e","x"]) == Set(["c","d"])
@test Set(["a","c","d","e"]) - Set(["b","x"]) == Set(["a","c","d","e"])
@test Set(["a","c","d","e"]) - Set{String}([]) == Set(["a","c","d","e"])
@test Set{String}([]) - Set(["a","c","d","e"]) == Set([])
@test Set{String}([]) - Set{String}([]) == Set([])
@test Set(["a","c"]) × Set(["b"]) == Set([("a","b"),("c","b")])
@test Set(["a","c"]) ⊗ Set(["b"]) == Set([("a","b"),("b","c")])
@test Set(["a","c"]) × Set(["b","d"]) == Set([("a","b"),("a","d"),("c","b"),("c","d")])
@test Set(["a","c"]) ⊗ Set(["b","d"]) == Set([("a","b"),("a","d"),("b","c"),("c","d")])
@test Set(["a","c","e"]) × Set(["b","d"]) == Set([("a","b"),("a","d"),("c","b"),("c","d"),("e","b"),("e","d")])
@test Set(["a","c","e"]) ⊗ Set(["b","d"]) == Set([("a","b"),("a","d"),("b","c"),("c","d"),("b","e"),("d","e")])
@test Set(["a","b"]) × Set(["b"]) == Set([("a","b"),("b","b")])
@test Set(["a","b"]) ⊗ Set(["b"]) == Set([("a","b")])
@test Set(["a","b"]) × Set{String}([]) == Set([])
@test Set(["a","b"]) ⊗ Set{String}([]) == Set([])
@test Set{String}([]) × Set(["a","b"]) == Set([])
@test Set{String}([]) ⊗ Set(["a","b"]) == Set([])
@test Set{String}([]) × Set{String}([]) == Set([])
@test Set{String}([]) ⊗ Set{String}([]) == Set([])


text:
print("String? ")
y = readline()
println("Your input was \"", y, "\".\n")
print("Integer? ")
y = readline()
try
    y = parse(Int, y)
    println("Your input was \"", y, "\".\n")
catch
    println("Sorry, but \"", y, "\" does not compute as an integer.")
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
function solve(model::AbstractModel,
               probtype::AbstractDynamicExperiment,
               alg::AbstractODEAlgorithm=DP5();
               kwargs...)
    prob = build_problem(model, probtype; kwargs...)
    sol = solve(prob, alg; kwargs...)
    return sol
end

function solve(model::AbstractModel,
               probtype::AbstractStochasticExperiment,
               alg::AbstractSDEAlgorithm=SRIW1();
               kwargs...)
    prob = build_problem(model, probtype; kwargs...)
    sol = solve(prob, alg; kwargs...)
    return sol
end

function solve(prob::AbstractNLOProblem, alg; kwargs...)
    sol = DiffEqBase.solve(prob.prob, alg; kwargs...)
    return DynamicNLOSolution(sol, prob)
end
function solve(prob::DynamicIkedaProblem, alg; kwargs...)
    sol = DiffEqBase.solve(prob.prob, alg; callback=prob.ikeda_callback, kwargs...)
    DynamicNLOSolution(sol, prob)
end


text:
# LSAC JUNE 2007 Form 8LSN75
# https://www.lsac.org/sites/default/files/legacy/docs/default-source/jd-docs/sampleptjune.pdf

# Question 1

#A company employee generates a series of five-digit product
#codes in accordance with the following rules:
#The codes use the digits 0, 1, 2, 3, and 4, and no others.
Ω,℧ = ABparse(["a, b, c, d, e  ∈  0:4"]); Ω[℧]

#Each digit occurs exactly once in any code.
Ω,℧ = ABparse("{{i}} != {{!i}}", Ω,℧); Ω[℧]

#The second digit has a value exactly twice that of the
#first digit.
Ω,℧ = ABparse("b == a*2", Ω,℧); Ω[℧]

#The value of the third digit is less than the value of the
#fifth digit
Ω,℧ = ABparse("c < e", Ω,℧); Ω[℧]

#1. If the last digit of an acceptable product code is 1, it
#must be true that the
Ω1,℧1 = ABparse("e == 1 ==> a==2", Ω,℧); Ω1[℧1]
#(A) first digit is 2
Ω1,℧1 = ABparse("e == 1 ==> b==0", Ω,℧); Ω1[℧1]
#(B) second digit is 0
Ω1,℧1 = ABparse("e == 1 ==> c==3", Ω,℧); Ω1[℧1]
#(C) third digit is 3
Ω1,℧1 = ABparse("d == 1 ==> d==4", Ω,℧); Ω1[℧1]
#(D) fourth digit is 4
Ω1,℧1 = ABparse("d == 1 ==> d==0", Ω,℧); Ω1[℧1]
#(E) fourth digit is 0

#2. Which one of the following must be true about any
#acceptable product code?
range(Ω, ℧)
Ω[℧]

#(A) The digit 1 appears in some position before the
#digit 2.
Ω2,℧2 = ABparse("{{i}} == 1 &&& {{>i}} == 2 {{1}}", Ω,℧); Ω2[℧2]

#(B) The digit 1 appears in some position before the
#digit 3.
Ω2,℧2 = ABparse("{{i}} == 1 &&& {{>i}} == 3 {{1}}", Ω,℧); Ω2[℧2]

#(C) The digit 2 appears in some position before the
#digit 3.
Ω2,℧2 = ABparse("{{i}} == 2 &&& {{>i}} == 3 {{1}}", Ω,℧); Ω2[℧2]

#(D) The digit 3 appears in some position before the
#digit 0.
Ω2,℧2 = ABparse("{{i}} == 0 &&& {{>i}} == 3 {{1}}", Ω,℧); Ω2[℧2]

#(E) The digit 4 appears in some position before the
#digit 3.
Ω2,℧2 = ABparse("{{i}} == 4 &&& {{>i}} == 3 {{1}}", Ω,℧); Ω2[℧2]

#3. If the third digit of an acceptable product code is not 0,
#which one of the following must be true?

Ω3,℧3 = ABparse("c != 0", Ω,℧); Ω3[℧3]
range(Ω3,℧3)

#(A) The second digit of the product code is 2.
#(B) The third digit of the product code is 3.
#(C) The fourth digit of the product code is 0.
#(D) The fifth digit of the product code is 3. <true>
#(E) The fifth digit of the product code is 1.

Ω[℧]

#4. Any of the following pairs could be the third and
#fourth digits, respectively, of an acceptable product
#code, EXCEPT:
# (A) 0, 1
ABparse(["c == 0", "d == 1"], Ω=Ω,℧=℧)
# (B) 0, 3
ABparse(["c == 0", "d == 3"], Ω=Ω,℧=℧)
#(C) 1, 0
ABparse(["c == 1", "d == 0"], Ω=Ω,℧=℧)
#(D) 3, 0
ABparse(["c == 3", "d == 0"], Ω=Ω,℧=℧)
#(E) 3, 4
ABparse(["c == 3", "d == 4"], Ω=Ω,℧=℧) # zero feasible outcomes

Ω[℧]

#5. Which one of the following must be true about any
#acceptable product code?
#(A) There is exactly one digit between the digit 0
#and the digit 1.
ABparse(["{{i}} == 1|0 <=> {{i+2}} == 1|0"], Ω=Ω,℧=℧) # zero feasible outcomes

#(B) There is exactly one digit between the digit 1
#and the digit 2.
ABparse(["{{i}} == 1|2 <=> {{i+2}} == 1|2"], Ω=Ω,℧=℧) # zero feasible outcomes

#(C) There are at most two digits between the digit 1 ??? Not working
#and the digit 3.
ABparse(["{{i}} == 1|3 <=> {{i+3}} == 1|3 |||| {{i}} == 1|3 <=> {{i+4}} == 1|3"], Ω=Ω,℧=℧) # zero feasible outcomes

ABparse("{{i}} != {{i+1}}", Ω,℧) # zero feasible outcomes

(D) There are at most two digits between the digit 2
and the digit 3.
ABparse(["{{i}} == 2|3 <==> {{i+3}} == 2|3"], Ω=Ω,℧=℧) # zero feasible outcomes

(E) There are at most two digits between the digit 2
and the digit 4.
ABparse(["""{{i}} == 2|4 <==> {{i+3}} == 2|4 ||||
            {{i}} == 2|4 <==> {{i+2}} == 2|4 ||||
            {{i}} == 2|4 <==> {{i+1}} == 2|4 """], Ω=Ω,℧=℧) # zero feasible outcomes

Ω[℧]


text:
module ImageClipboard

using ImageIO, FileIO, ColorTypes

export clipboard_img

# linux-X11
include("_xclip.jl")
# linux-wayland
include("_wlclipboard.jl")
# mac
include("_osascript.jl")
# windows
include("_powershell.jl")

"""
    clipboard_img() -> AbstractMatrix{<:Colorant}

Paste an image from clipboard
"""
function clipboard_img()
    if Sys.islinux() && get(ENV, "XDG_SESSION_TYPE", "") == "wayland"
        if _isavailable_wlclipboard()
            img = _wlclipboard()
        elseif _isavailable_xclip()
            img = _xclip()
        else
            error("Please install wlclipboard or xclip to your system")
        end
    elseif Sys.islinux()
        if _isavailable_xclip()
            img = _xclip()
        else
            error("Please install xclip to your system")
        end
    elseif Sys.iswindows()
        img = _powershell()
    elseif Sys.isapple()
        img = _osascript()
    else
        error("Only Linux, Windows, and Mac are supported")
    end
    return img
end

"""
    clipboard_img(img::AbstractMatrix{<:Colorant})

Copy an image to clipboard
"""
function clipboard_img(img::AbstractMatrix{<:Colorant})
    if Sys.islinux() && get(ENV, "XDG_SESSION_TYPE", "") == "wayland"
        if _isavailable_wlclipboard()
            _wlclipboard(img)
        elseif _isavailable_xclip()
            _xclip(img)
        else
            error("Please install wlclipboard or xclip to your system")
        end
    elseif Sys.islinux()
        if _isavailable_xclip()
            _xclip(img)
        else
            error("Please install xclip to your system")
        end
    elseif Sys.iswindows()
        _powershell(img)
    elseif Sys.isapple()
        _osascript(img)
    else
        error("Only Linux, Windows, and Mac are supported")
    end
    return nothing
end

end # module


text:
immutable ZeroOneLoss <: Loss
end

@inline function loss_value(loss::ZeroOneLoss, p::Real, y::Real)
    return Float64(round(Int, p) != y)
end

Base.show(io::IO, loss::ZeroOneLoss) = print(io, "0-1 Loss")


text:
using PyPlot

srand(1)
N_simul = 150
cp = ConsumptionProblem()

c_bar, b1, b2 = consumption_complete(cp)
debt_complete = [b1, b2]

c_path, debt_path, y_path, s_path = consumption_incomplete(cp, N_simul=N_simul)

fig, ax = subplots(1, 2, figsize=(15, 5))

ax[1][:set_title]("Consumption paths")
ax[1][:plot](1:N_simul, c_path, label="incomplete market")
ax[1][:plot](1:N_simul, c_bar * ones(N_simul), label="complete market")
ax[1][:plot](1:N_simul, y_path, label="income", lw=2, alpha=.6, ls="--")
ax[1][:legend]()
ax[1][:set_xlabel]("Periods")
ax[1][:set_ylim]([1.4, 2.1])

ax[2][:set_title]("Debt paths")
ax[2][:plot](1:N_simul, debt_path, label="incomplete market")

ax[2][:plot](1:N_simul, debt_complete[s_path], label="complete market")
ax[2][:plot](1:N_simul, y_path, label="income", alpha=.6, ls="--")
ax[2][:legend]()
ax[2][:axhline](0, color="k", ls="--")
ax[2][:set_xlabel]("Periods")


text:
function sparse_reorder(A)
    (m,n)=size(A);
    U=Array(Array{Int64},(m,));
    for i=1:m
        U[i]=A.rowval[A.colptr[i]:A.colptr[i+1]-1];
    end
    p=sortperm(U,lt=lexless);
    return p;
end

max_stars_repo_name: ReidAtcheson/incns

max_issues_repo_name: ReidAtcheson/incns

max_forks_repo_name: ReidAtcheson/incns


text:
using Test, Pkg
using CairoMakie

path = normpath(joinpath(dirname(pathof(AbstractPlotting)), "..", "test", "ReferenceTests"))
Pkg.develop(PackageSpec(path=path))
using ReferenceTests
using ReferenceTests: unique_name
CairoMakie.activate!(type = "png")

excludes = Set((
    "examples3d:65-Colored Mesh",
    "examples3d:466-Line GIF",
    "short_tests:115",
    "examples2d:205-Streamplot animation",
    "examples2d:220-Line changing colour",
    "examples3d:272-Axis + Surface",
    "examples3d:518-Streamplot 3D",
    "examples3d:134-Meshscatter Function",
    "examples2d:330-Hollow pie chart",
    "examples3d:150-Record Video",
    "examples3d:12-Image on Geometry (Earth)",
    "layouting:16-Comparing contours, image, surfaces and heatmaps",
    "examples3d:56-Textured Mesh",
    "examples2d:321-Simple pie chart",
    "examples3d:393-Animated surface and wireframe",
    "examples2d:339-Open pie chart",
    "examples3d:376-image scatter",
    "examples3d:244-surface + contour3d",
    "short_tests:111",
    "examples3d:18-Orthographic Camera",
    "documentation:23-Legend",
    "attributes:43-rotation",
    "examples3d:168-3D Contour with 2D contour slices",
    "examples3d:113-Surface with image",
    "examples2d:4-Test heatmap + image overlap",
    "examples2d:164-Text Annotation",
    "step-2",
    "examples2d:40-FEM polygon 2D.png",
    "examples2d:174-Text rotation",
    "examples3d:211-Image on Surface Sphere",
    "examples2d:66-FEM mesh 2D",
    "examples2d:145-Hbox",
    "examples3d:427-Stars",
    "examples2d:108-Subscenes",
    "examples3d:196-Arrows 3D",
    "layouting:1-Layouting",
    # sigh this is actually super close,
    # but doesn't interpolate the values inside the
    # triangles, so looks pretty different
    "examples2d:40-FEM polygon 2D"
))

database = ReferenceTests.load_database()

filter!(database) do (name, entry)
    !(unique_name(entry) in excludes) &&
    !(:volume in entry.used_functions) &&
    !(:volume! in entry.used_functions)
end

files, recorded = ReferenceTests.record_tests(database)

ReferenceTests.reference_tests(recorded)


text:
# Biryani -- A conversion and validation toolbox
# By: Emmanuel Raviart <emmanuel@raviart.com>
#
# Copyright (C) 2015 Emmanuel Raviart
# https://github.com/eraviart/Biryani.jl
#
# This file is part of Biryani.
#
# 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.

function input_to_json(convertible::Convertible)
  """Convert a JSON string to Julia data.

  .. warning:: Like most converters, a ``nothing`` value is not converted.
  """
  if convertible.error !== nothing || convertible.value === nothing
    return convertible
  end
  try
    value = JSON.parse(convertible.value)
    return Convertible(value, convertible.context)
  catch exception
    if isa(exception, BoundsError) || isa(exception, ErrorException)
      return Convertible(convertible.value, convertible.context, N_("Invalid JSON"))
    else
      rethrow()
    end
  end
end


text:
export JacobianConvexparameters,
    CharpolyCoeff,
    JacobianDeterminantConvexparameters

function JacobianConvexparameters(S::AbstractAlgebra.Ring,
                                  N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix,
                                  lpos::UnitRange{T}=1:size(E, 2),
                                  hpos::UnitRange{T}=.+(length(lpos), 1:size(N, 1))) where {T <: Integer}
    ## Migrating to AbstractAlgebra
    Nnemo = Nemo.matrix(S, N)
    Ynemo = Nemo.matrix(S, transpose(Y))
    ## Square matrix with h's diagonal
    H = Diagonal(S, Nemo.gens(S)[hpos])
    ## Diagonal with coordinates of E
    D = Diagonal(S, E * Nemo.gens(S)[lpos])
    ## Return the jacobian under convex parameters
    return Nnemo * D * Ynemo * H
end

function CharpolyCoeff(R::AbstractAlgebra.Ring, M::MatElem, ncoeff::Integer)
    P, x = Nemo.PolynomialRing(R, "x")
    # q = coeff(divexact(charpoly(P,Jhl),x^mindeg),0);
    return Nemo.coeff(charpoly(P, M), ncoeff)
end

function JacobianDeterminantConvexparameters(R::AbstractAlgebra.Ring,
                                             N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix,
                                             lpos::UnitRange{T}=1:size(E, 2),
                                             hpos::UnitRange{T}=.+(length(lpos), 1:size(N, 1)),
                                             mindeg::Integer=size(N, 1) - rank(Nemo.matrix(R, N))) where {T <: Integer}
    return CharpolyCoeff(R, JacobianConvexparameters(R, N, Y, E, lpos, hpos), mindeg)
end

function JacobianConvexparameters(H::AbstractAlgebra.Ring,
                                  L::AbstractAlgebra.Ring,
                                  N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix)
    ## Migrating to AbstractAlgebra
    Nnemo = Nemo.matrix(H, N)
    Ynemo = Nemo.matrix(H, transpose(Y))
    ## Square matrix with h's diagonal
    digH = Diagonal(H, Nemo.gens(H))
    ## Diagonal with coordinates of E
    digL = Diagonal(H, E * Nemo.gens(L))
    ## Return the jacobian under convex parameters
    return Nnemo * digL * Ynemo * digH
end

function JacobianDeterminantConvexparameters(H::AbstractAlgebra.Ring,
                                             L::AbstractAlgebra.Ring,
                                             N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix,
                                             mindeg::Integer=size(N, 1) - rank(Nemo.matrix(H, N)))
    return CharpolyCoeff(H, JacobianConvexparameters(H, L, N, Y, E), mindeg)
end


text:
###############################################################################
#
# Abstract ERA5 dataset type
#
###############################################################################
"""
    abstract type AbstractERA5Data{FT}

Hierachy of AbstractUngriddedData
- [`ERA5LandHourly`](@ref)
"""
abstract type AbstractERA5Data end

###############################################################################
#
# ERA5 data types
#
###############################################################################
"""
    struct ERA5LandHourly <: AbstractERA5Data

ERA5 Land hourly data from 1981 to present:
    cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-land
"""
struct ERA5LandHourly <: AbstractERA5Data end

"""
    struct ERA5SingleLevelsHourly <: AbstractERA5Data

ERA5 Single Level hourly data from 1981 to present:
    cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-single-levels
"""
struct ERA5SingleLevelsHourly <: AbstractERA5Data end


text:
include(joinpath(@__DIR__, "imports.jl"))  # For precompile.
markers = [:CD3z, :EOMES, :Perforin, :Granzyme_A, :Siglec7]

using Distributed
rmprocs(workers())
addprocs(length(markers))
@everywhere include(joinpath(@__DIR__, "imports.jl"))

postprocess_only = false

# Read command line args.
if length(ARGS) > 1
  resultsdir = ARGS[1]
  awsbucket = ARGS[2]
  println("ARGS: ", ARGS); flush(stdout)
  nsamps = 10000
  nburn = 2000
else
  resultsdir = "results/test/"
  awsbucket = nothing
  nsamps = 10
  nburn = 20

  if postprocess_only
    SCRATCH_DIR = ENV["SCRATCH_DIR"]
    SIMNAME = "run0"
    resultsdir = "$(SCRATCH_DIR)/cde/datastudy/$(SIMNAME)"
    awsbucket = "s3://cytof-density-estimation/datastudy/$(SIMNAME)"
  end
end

# Read data.
donor = 1
path_to_data = "../../../data/TGFBR2/cytof-data/donor$(donor).csv"
subsample_size = 20000
data = CDE.Util.subsample(DataFrame(CSV.File(path_to_data)), subsample_size,
                          seed=0)
configs = [let
             _awsbucket = (awsbucket == nothing) ? awsbucket : "$(awsbucket)/donor$(donor)/$(marker)"
             yC, yT = CDE.Util.partition(data, marker, log_response=true)
             (yC=yC, yT=yT, K=5, nsamps=nsamps, nburn=nburn, marker=marker,
              awsbucket=_awsbucket,
              resultsdir=joinpath(resultsdir, "donor$(donor)/$(marker)"))
           end for marker in markers]

# Parallel run.
println("Starting runs ...")
if !postprocess_only  # Do run as well.
  res = pmap(run, configs, on_error=identity)
else
  # Post processing only.
  res = pmap(resdir -> let
               out = BSON.load("$(resdir)/results.bson")
               imgdir = joinpath(resdir, "img"); mkpath(imgdir)
               postprocess(out[:chain], out[:laststate], out[:summarystats],
                           out[:data].yC, out[:data].yT, imgdir; bw_postpred=0.2,
                           ygrid=collect(range(-8, 8, length=1000)),
                           density_legend_pos=:topleft)
             end, getfield.(configs, :resultsdir), on_error=identity)
end

println("Status of runs:")
foreach(z -> println(z[2], " => ", z[1]),
        zip(res, getfield.(configs, :resultsdir)))

println("DONE WITH ALL!"); flush(stdout)

max_stars_repo_name: luiarthur/CytofDensityEstimation

max_issues_repo_name: luiarthur/CytofDensityEstimation

max_forks_repo_name: luiarthur/CytofDensityEstimation