text:
using SerializedElementArrays
using Test

using SerializedElementArrays: disk

@testset "SerializedElementArrays.jl" begin
  @testset "SerializedElementArrays.SerializedElementVector" begin
    d = SerializedElementArrays.SerializedElementVector()
    @test size(d) == (0,)
    @test length(d) == 0
    @test ndims(d) == 1
    @test eltype(d) == Any
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    @test !isfile(SerializedElementArrays.filename(d, 1))

    d = SerializedElementArrays.SerializedElementVector(undef, 4)
    @test size(d) == (4,)
    @test length(d) == 4
    @test ndims(d) == 1
    @test eltype(d) == Any
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test !isfile(SerializedElementArrays.filename(d, n))
      @test !isassigned(d, n)
    end
    d[1] = 1
    @test isfile(SerializedElementArrays.filename(d, 1))
    @test d[1] == 1
    d[2] = 1.2
    @test isfile(SerializedElementArrays.filename(d, 2))
    @test d[2] == 1.2
    d[3] = "XXX"
    @test d[3] == "XXX"
    @test isfile(SerializedElementArrays.filename(d, 3))
    @test !isassigned(d, 4)
    @test_throws UndefRefError d[4]

    d = SerializedElementArrays.SerializedElementVector{Int}(undef, 4)
    @test size(d) == (4,)
    @test length(d) == 4
    @test ndims(d) == 1
    @test eltype(d) == Int
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test !isfile(SerializedElementArrays.filename(d, n))
      @test !isassigned(d, n)
    end
    d[1] = 1
    @test isfile(SerializedElementArrays.filename(d, 1))
    @test isassigned(d, 1)
    @test d[1] == 1
    @test_throws InexactError d[2] = 1.2

    d = SerializedElementArrays.SerializedElementVector(1:6)
    @test size(d) == (6,)
    @test length(d) == 6
    @test ndims(d) == 1
    @test eltype(d) == Int
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test isfile(SerializedElementArrays.filename(d, n))
      @test isassigned(d, n)
      @test d[n] == n
    end

    d = SerializedElementArrays.SerializedElementVector{Float64}(1:6)
    @test size(d) == (6,)
    @test length(d) == 6
    @test ndims(d) == 1
    @test eltype(d) == Float64
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test isfile(SerializedElementArrays.filename(d, n))
      @test isassigned(d, n)
      @test d[n] == n
    end

    a = Array(d)
    @test a isa Vector{Float64}
    @test length(a) == 6
    for n in 1:length(a)
      @test a[n] == n
    end

    a = Vector(d)
    @test a isa Vector{Float64}
    @test length(a) == 6
    for n in 1:length(a)
      @test a[n] == n
    end

    a = Vector{ComplexF64}(d)
    @test a isa Vector{ComplexF64}
    @test length(a) == 6
    for n in 1:length(a)
      @test a[n] == complex(n)
    end

    @test_throws MethodError Matrix(d)

    d2 = disk(a)
    @test d2 isa SerializedElementArrays.SerializedElementVector{ComplexF64}
    @test size(d2) == (6,)
    @test SerializedElementArrays.pathname(d) ≠ SerializedElementArrays.pathname(d2)
  end

  @testset "SerializedElementArrays.SerializedElementMatrix" begin
    d = SerializedElementArrays.SerializedElementMatrix()
    @test size(d) == (0, 0)
    @test length(d) == 0
    @test ndims(d) == 2
    @test eltype(d) == Any

    d = SerializedElementArrays.SerializedElementMatrix(undef, 4, 5)
    @test size(d) == (4, 5)
    @test length(d) == 20
    @test ndims(d) == 2
    @test eltype(d) == Any
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test !isfile(SerializedElementArrays.filename(d, n))
      @test !isassigned(d, n)
    end
    d[1] = 1
    @test isfile(SerializedElementArrays.filename(d, 1))
    @test d[1] == 1
    d[2] = 1.2
    @test isfile(SerializedElementArrays.filename(d, 2))
    @test d[2] == 1.2
    d[3] = "XXX"
    @test d[3] == "XXX"
    @test isfile(SerializedElementArrays.filename(d, 3))
    d[1, 2] = "YYY"
    @test d[1, 2] == "YYY"
    @test d[5] == "YYY"
    @test isfile(SerializedElementArrays.filename(d, 1, 2))
    @test isfile(SerializedElementArrays.filename(d, 5))
    @test !isassigned(d, 4)
    @test_throws UndefRefError d[4]

    d = SerializedElementArrays.SerializedElementMatrix{Int}(undef, 4, 5)
    @test size(d) == (4, 5)
    @test length(d) == 20
    @test ndims(d) == 2
    @test eltype(d) == Int
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test !isfile(SerializedElementArrays.filename(d, n))
      @test !isassigned(d, n)
    end
    d[1, 2] = 1
    @test isfile(SerializedElementArrays.filename(d, 5))
    @test isfile(SerializedElementArrays.filename(d, 1, 2))
    @test isassigned(d, 5)
    @test isassigned(d, 1, 2)
    @test d[1, 2] == 1
    @test d[5] == 1
    @test_throws InexactError d[2, 2] = 1.2

    d = SerializedElementArrays.SerializedElementMatrix(reshape(1:6, 2, 3))
    @test size(d) == (2, 3)
    @test length(d) == 6
    @test ndims(d) == 2
    @test eltype(d) == Int
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test isfile(SerializedElementArrays.filename(d, n))
      @test isassigned(d, n)
      @test d[n] == n
      @test d[n] isa Int64
    end

    d = SerializedElementArrays.SerializedElementMatrix{Float64}(reshape(1:6, 2, 3))
    @test size(d) == (2, 3)
    @test length(d) == 6
    @test ndims(d) == 2
    @test eltype(d) == Float64
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test isfile(SerializedElementArrays.filename(d, n))
      @test isassigned(d, n)
      @test d[n] == n
      @test d[n] isa Float64
    end

    a = Array(d)
    @test a isa Matrix{Float64}
    @test length(a) == 6
    for n in 1:length(a)
      @test a[n] == n
    end

    a = Matrix(d)
    @test a isa Matrix{Float64}
    @test length(a) == 6
    for n in 1:length(a)
      @test a[n] == n
    end

    d2 = disk(a)
    @test d2 isa SerializedElementArrays.SerializedElementMatrix{Float64}
    @test size(d2) == (2, 3)
    @test SerializedElementArrays.pathname(d) ≠ SerializedElementArrays.pathname(d2)
  end

  @testset "SerializedElementArrays.SerializedElementArray" begin
    @test_throws MethodError SerializedElementArrays.SerializedElementArray()
    @test_throws MethodError SerializedElementArrays.SerializedElementArray{Float64}()

    # TODO: this also fails for Array, should we support it?
    @test_throws MethodError SerializedElementArrays.SerializedElementArray(undef, 2, 3, 4)
    d = SerializedElementArrays.SerializedElementArray{Any}(undef, 2, 3, 4)
    @test size(d) == (2, 3, 4)
    @test length(d) == 24
    @test ndims(d) == 3
    @test eltype(d) == Any
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test !isfile(SerializedElementArrays.filename(d, n))
      @test !isassigned(d, n)
    end
    d[1] = 1
    @test isfile(SerializedElementArrays.filename(d, 1))
    @test d[1] == 1
    d[2] = 1.2
    @test isfile(SerializedElementArrays.filename(d, 2))
    @test d[2] == 1.2
    d[3] = "XXX"
    @test d[3] == "XXX"
    @test isfile(SerializedElementArrays.filename(d, 3))
    d[1, 2, 3] = "YYY"
    @test d[1, 2, 3] == "YYY"
    @test isfile(SerializedElementArrays.filename(d, 1, 2, 3))
    @test !isassigned(d, 4)
    @test_throws UndefRefError d[4]

    d = SerializedElementArrays.SerializedElementArray{Int}(undef, 2, 3, 4)
    @test size(d) == (2, 3, 4)
    @test length(d) == 24
    @test ndims(d) == 3
    @test eltype(d) == Int
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test !isfile(SerializedElementArrays.filename(d, n))
      @test !isassigned(d, n)
    end
    d[1, 2, 3] = 1
    @test isfile(SerializedElementArrays.filename(d, 1, 2, 3))
    @test isassigned(d, 1, 2, 3)
    @test d[1, 2, 3] == 1
    @test_throws InexactError d[2, 2, 4] = 1.2

    d = SerializedElementArrays.SerializedElementArray(reshape(1:24, 2, 3, 4))
    @test size(d) == (2, 3, 4)
    @test length(d) == 24
    @test ndims(d) == 3
    @test eltype(d) == Int
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test isfile(SerializedElementArrays.filename(d, n))
      @test isassigned(d, n)
      @test d[n] == n
      @test d[n] isa Int64
    end

    d = SerializedElementArrays.SerializedElementArray{Float64}(reshape(1:24, 2, 3, 4))
    @test size(d) == (2, 3, 4)
    @test length(d) == 24
    @test ndims(d) == 3
    @test eltype(d) == Float64
    @test startswith(SerializedElementArrays.pathname(d), tempdir())
    @test ispath(SerializedElementArrays.pathname(d))
    for n in 1:length(d)
      @test isfile(SerializedElementArrays.filename(d, n))
      @test isassigned(d, n)
      @test d[n] == n
      @test d[n] isa Float64
    end

    a = Array(d)
    @test a isa Array{Float64,3}
    @test length(a) == 24
    for n in 1:length(a)
      @test a[n] == n
    end

    a = Array{<:Any,3}(d)
    @test a isa Array{Float64,3}
    @test length(a) == 24
    for n in 1:length(a)
      @test a[n] == n
    end

    d2 = disk(a)
    @test d2 isa SerializedElementArrays.SerializedElementArray{Float64,3}
    @test size(d2) == (2, 3, 4)
    @test SerializedElementArrays.pathname(d) ≠ SerializedElementArrays.pathname(d2)
  end

  @testset "Composite element type" begin
    v = [randn(3, 4) for n in 1:2, m in 1:3]
    d = disk(v)
    @test size(d) == (2, 3)
    @test length(d) == (6)
    @test eltype(d) == Matrix{Float64}
    d11 = d[1, 1]
    @test size(d11) == (3, 4)
    @test d11 isa Matrix{Float64}
    d23 = d[end, end]
    @test size(d23) == (3, 4)
    @test d23 isa Matrix{Float64}
  end

  @testset "Undefined references" begin
    a = Array{Matrix{Float64}}(undef, 2, 3)
    d = disk(a)
    @test size(d) == (2, 3)
    @test !isassigned(a, 1, 2)
    @test !isassigned(d, 1, 2)
    @test isempty(readdir(SerializedElementArrays.pathname(d)))
    x = randn(5, 5)
    d[1, 2] = x
    @test x == d[1, 2]
    @test length(readdir(SerializedElementArrays.pathname(d))) == 1
    y = randn(3, 4)
    d[2, 3] = y
    @test y == d[2, 3]
    @test length(readdir(SerializedElementArrays.pathname(d))) == 2
  end

  if VERSION >= v"1.4"
    @testset "Custom paths" begin
      a = randn(3, 4)
      d = disk(a; path="my_tmp")
      @test size(d) == (3, 4)
      @test length(readdir(SerializedElementArrays.pathname(d))) == 12
      @test dirname(SerializedElementArrays.pathname(d)) == "my_tmp"
    end
  end

  @testset "disk(::SerializedElementArray)" begin
    a = randn(2, 3)
    d1 = disk(a)
    d2 = disk(a)
    @test SerializedElementArrays.pathname(d1) ≠ SerializedElementArrays.pathname(d2)
    d3 = disk(d2)
    @test SerializedElementArrays.pathname(d2) == SerializedElementArrays.pathname(d3)
  end

  @testset "disk(::Function)" begin
    n, m = 4, 5
    d1 = disk(n, m) do n, m
      [i + j for i in 1:n, j in 1:m]
    end
    d2 = disk(n, m; full=true) do n, m
      [i + j for i in 1:n, j in 1:m]
    end
    d3 = disk(n, m; force_gc=false) do n, m
      [i + j for i in 1:n, j in 1:m]
    end
    @test d1[1, 2] == 1 + 2
    @test d2[1, 2] == 1 + 2
    @test d3[1, 2] == 1 + 2
  end
end


text:
@testset "Environment" begin
	# 2D
	surf = x -> 0.05 * sin(π * x[1])
	surf_grad = x -> [0.05 * π * cos(π * x[1])]
	x0 = ones(10)

	env = ContactImplicitMPC.environment_2D(surf)
	@test norm(env.surf(x0) - surf(x0)) < 1e-10
	@test norm(env.surf_grad(x0) - surf_grad(x0)) < 1e-10
	@test size(env.surf(x0)) == ()
	@test size(env.surf_grad(x0)) == (1,)

	# 3D
	surf = x -> 0.05 * sin(π * x[1])
	surf_grad = x -> [0.05 * π * cos(π * x[1]); 0.0]
	x0 = ones(10)

	env = ContactImplicitMPC.environment_3D(surf)
	@test norm(env.surf(x0) - surf(x0)) < 1e-10
	@test norm(env.surf_grad(x0) - surf_grad(x0)) < 1e-10
	@test size(env.surf(x0)) == ()
	@test size(env.surf_grad(x0)) == (2,)
end

ext: jl


text:
module BringRadicals
export bringrad

bringrad(a::Integer) = bringrad(float(a))
function bringrad(a::T) where T<Real
    x = abs(a)
    if x < .22
        at4 = a^4
        return a*evalpoly(at4, T(1, -1, 5, -35, 285))
    end
    if x<2
        guess = x*evalpoly(x, T(1.0098200485727955, -0.1885556510194153, 1.2609058087320157,
                              -3.7000253271311796, 3.4819520619583963, -1.109350971521824))
    else
        guess = @fastmath exp2(log2(x)*inv(T(5))))-inv(x+10)
    end
    tol = sqrt(eps(guess))
    guess = copysign(guess, -a)
    while true
        gt4 = (guess^2)^2 # avoid expense of floating point pow
        newguess = guess - muladd(guess, gt4, guess+a)/muladd(5, gt4, 1)
        abs(newguess - guess) <= tol && (guess = newguess; break)
        guess = newguess
    end
    gt4 = guess^4
    return guess - fma(guess, gt4, guess+a)/fma(5, gt4, 1)
end
end


text:
# Autogenerated wrapper script for XLA_Tools_jll for x86_64-linux-gnu-cxx11
export convert_computation, dumped_computation_to_operation_list, dumped_computation_to_text, show_literal, show_signature, show_text_literal

JLLWrappers.@generate_wrapper_header("XLA_Tools")
JLLWrappers.@declare_executable_product(convert_computation)
JLLWrappers.@declare_executable_product(dumped_computation_to_operation_list)
JLLWrappers.@declare_executable_product(dumped_computation_to_text)
JLLWrappers.@declare_executable_product(show_literal)
JLLWrappers.@declare_executable_product(show_signature)
JLLWrappers.@declare_executable_product(show_text_literal)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_executable_product(
        convert_computation,
        "bin/convert_computation",
    )

    JLLWrappers.@init_executable_product(
        dumped_computation_to_operation_list,
        "bin/dumped_computation_to_operation_list",
    )

    JLLWrappers.@init_executable_product(
        dumped_computation_to_text,
        "bin/dumped_computation_to_text",
    )

    JLLWrappers.@init_executable_product(
        show_literal,
        "bin/show_literal",
    )

    JLLWrappers.@init_executable_product(
        show_signature,
        "bin/show_signature",
    )

    JLLWrappers.@init_executable_product(
        show_text_literal,
        "bin/show_text_literal",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using LinearAlgebra: diagm, eigen

"""
Inertial instability version of build_matrix
"""

function build_matrix(k, m, grid, physics, jet)

    N2     =  physics.N^2
    Z      = zeros(grid.Ny+1, grid.Ny+1)
    I      =  diagm(0 => ones(grid.Ny+1))
    Q      =  diagm(0 =>   physics.fz .- jet.dU)
    By     =  diagm(0 => - physics.fy .* jet.dU)
    U      = k * diagm(0=>jet.U) - 1im * physics.ν * m^2 * I
    fyomDy = physics.fy / m * grid.D

    # Set up eigenvalue problem: [u, v, b]
    A = [                             U                (1im*Q + fyomDy)[:, 2:end-1]   -1im*k/m*I;
        -(1im*physics.fz*I + fyomDy)[2:end-1,:]                  U[2:end-1,2:end-1]  -1/m*grid.D[2:end-1,:] ;
        1im*N2*k/m*I                             (N2/m*grid.D - 1im* By)[:,2:end-1]            U];

    return A
end

function compute_spectrum(A, Neigs)

    λ_in, λvec_in = eigen(A);
            index = sortperm(imag(λ_in), rev=true);
             λvec = λvec_in[:,index];
                λ = λ_in[index];

               ωs = λ[1:Neigs];
            modes = λvec[:,1:Neigs];

        return ωs, modes
    end

    # Must modify this for this particular problem


text:
@testset "Power modulo" begin
    R = ResidueRing(ZZ,7)
    Rx,x = R["x"]
    f = Base.rand(Rx,10)
    g = x^4+x+1
    e = fmpz(10)^20
    @test powmod(f,e,g) == powmod(f,e%(7^degree(g)-1),g)
    @test powmod(f,fmpz(0),g) == 1
end


text:
using Documenter, Gasdynamics1D

#ENV["GKSwstype"] = "nul" # removes GKS warnings during plotting
ENV["GKSwstype"] = "100" # removes GKS warnings during plotting

makedocs(
    sitename = "Gasdynamics1D.jl",
    doctest = true,
    clean = true,
    pages = [
        "Home" => "index.md",
        "Manual" => ["manual/0-BasicGasDynamics.md",
                     "manual/1-IsentropicGasDynamics.md",
                     "manual/2-ConvergingDivergingNozzle.md",
                     "manual/3-NormalShocks.md",
                     "manual/4-FannoFlow.md",
                     "manual/5-RayleighFlow.md"
                     ]
        #"Internals" => [ "internals/properties.md"]
    ],
    #format = Documenter.HTML(assets = ["assets/custom.css"])
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
        mathengine = MathJax(Dict(
            :TeX => Dict(
                :equationNumbers => Dict(:autoNumber => "AMS"),
                :Macros => Dict()
            )
        ))
    ),
    #assets = ["assets/custom.css"],
    #strict = true
)

#if "DOCUMENTER_KEY" in keys(ENV)
deploydocs(
     repo = "github.com/UCLAMAEThreads/Gasdynamics1D.jl.git",
     target = "build",
     devbranch = "main",
     deps = nothing,
     make = nothing
     #versions = "v^"
)
#end


text:
"""
    detrend(x, y, poly_order)

Subtract polynomial from `y` data.
"""
detrend(x, y, poly_order) = poly_order == 0 ? y .- mean(y) : y - fit(x, y, poly_order).(x)
detrend(c::Curve, poly_order) = Curve(c.x, detrend(c.x, c.y, poly_order))

"""
    NoiseSample <: AbstractCurve

Curve holding a noise sample. At minimum, a constant offset
is removed from the noise sample upon construction.
If provided, a polynomial of order `n` is subtracted.

# Fields
- `x::Vector{Float64}` : x-grid
- `y::Vector{Float64}` : y-value

# Constructors

    NoiseSample(x::Vector{T}, y::Vector{T})
    NoiseSample(x::Vector{T}, y::Vector{T}, n::Int)
    NoiseSample(s::Curve)
    NoiseSample(s::Curve, n::Int)

# Notes

Noise samples should not have a slow trend on top of the noise,
otherwise noise parameters cannot be reliably infered.
Construct `NoiseSample` objects with a high enough order
`n` so as to remove any slow variations.
"""
struct NoiseSample{T} <: AbstractCurve
    x::Vector{T}
    y::Vector{T}
    function NoiseSample{T}(x, y, n::Int=0) where T
        verify_same_length(x, y)
        return new{T}(x, detrend(x, y, n))
    end
end

NoiseSample(x::Vector{T}, y::Vector{T}, n::Int=0) where T = NoiseSample{T}(x, y, n)
NoiseSample(y::Vector{T}, n::Int=0) where T = NoiseSample(collect(T, 1:length(y)), y, n)
NoiseSample(s::Curve, n::Int=0) = NoiseSample(s.x, s.y, n)

"""
    std(ns::NoiseSample; kw...)

Retrieve the standard deviation of the `NoiseSample` `ns`.
"""
Statistics.std(ns::NoiseSample; kw...) = Statistics.std(ns.y; kw...)

# -------------------------------------
# Noise models
# -------------------------------------

"""
    AbstractNoiseModel

Supertype of noise models.
"""
abstract type AbstractNoiseModel end

"""
    GaussianNoiseModel{T <: Real} <: AbstractNoiseModel

Model to describe noise following a univariate Gaussian distribution
(uncorrelated noise).

# Fields

- `σ::T` : standard deviation
"""
struct GaussianNoiseModel{T <: Real} <: AbstractNoiseModel
    σ::T
end

Base.eltype(::GaussianNoiseModel{T}) where T = T

function Base.show(io::IO, nm::GaussianNoiseModel)
    println(io, "GaussianNoiseModel(σ = $(nm.σ))")
end

"""
    MvGaussianNoiseModel{T <: Real} <: AbstractNoiseModel

Model to describe noise following a multivariate Gaussian distribution
(correlated noise).

# Fields
- `δx::T`: noise grid spacing
- `α::T` : autocovariance amplitude
- `λ::T` : autocovaraiance lag
"""
struct MvGaussianNoiseModel{T <: Real} <: AbstractNoiseModel
    δx::T
    α::T
    λ::T
end

Base.eltype(::MvGaussianNoiseModel{T}) where T = T

function Base.show(io::IO, nm::MvGaussianNoiseModel)
    println(io, "MvGaussianNoiseModel(α = $(nm.α), λ = $(nm.λ))")
end

"""
    gauss_kernel(Δx, β)

Fit function for noise autocorrelation.

    β = [α, λ]
    f(δx) = α^2 * exp(-0.5 * (Δx/λ)^2)
"""
function gauss_kernel(Δx, β)
    α, λ = β
    return @. α^2 * exp(-0.5 * (Δx/λ)^2)
end

"""
    get_cov(δx::T, n::Int, α::T, λ::T) where {T <: Real}

Return covariance matrix `Σ` with size `n` × `n` for the exponentiated quadratic kernel
with amplitude `α` and lag `λ` calculated on a grid of spacing `δx`.
"""
function get_cov(δx::T, n::Int, α::T, λ::T) where {T <: Real}
    β = [α, λ]
    Σ = Array{T}(undef, n, n)
    for i = (1:n), j = (1:n)
        Δx = δx * (i - j)
        Σ[i,j] = gauss_kernel(Δx, β)
        if i == j
            Σ[i,j] *= 1.0000001 # for numerical stability (?)
        end
    end
    return Σ
end
get_cov(nm::MvGaussianNoiseModel{T}, n::Int) where {T <: Real} = get_cov(nm.δx, n, nm.α, nm.λ)

function estimate_autocov(n::NoiseSample)
    # check that spectral grid is equally spaced
    if !allapproxequal(diff(n.x))
        throw(ArgumentError("Noise analysis only works on an evenly spaced spectral grid."))
    end
    δx = abs(n.x[1]-n.x[2]) # x-values must be evenly spaced
    lag_units = begin
        N = convert(Integer, round(length(n)/2))
        collect(1:N)
    end
    lags = δx .* lag_units
    acov = autocov(n.y, lag_units)
    return lags, acov
end

function _fit_noise(lags::Vector{T}, autocov::Vector{T}; α_guess=1.0, λ_guess=1.0) where {T <: Real}
    ma = maximum(autocov) # rescale autocov to make fit more robust
    fit = curve_fit(gauss_kernel, lags, autocov./ma, [α_guess, λ_guess])
    α, λ = fit.param
    return MvGaussianNoiseModel(lags[1], α * √ma, λ)
end

"""
    fit_noise(ns::NoiseSample; α_guess=1.0, λ_guess=1.0)

Build a `MvGaussianNoiseModel` for the `NoiseSample` `ns`.

Estimates parameters `α` and `λ` by fitting the exponentiated quadratic covariance function

    k(Δx) = α² exp(-0.5 (Δx/λ)²)

to the autocovariance (k) vs. the lag (Δx) of the noise sample.

Change initial guess if fit does not converge to sensible result.
"""
function fit_noise(n::NoiseSample; α_guess=1.0, λ_guess=1.0)
    lags, acov = estimate_autocov(n)
    return _fit_noise(lags, acov; α_guess=α_guess, λ_guess=λ_guess)
end

"""
Generate correlated noise from a noise sample or noise model.
"""
function generate_noise(nm::MvGaussianNoiseModel{T}, len::Int, samples::Int) where {T <: Real}
    Σ = get_cov(nm, len)
    return Particles(samples, MvNormal(Σ))
end
function generate_noise(ns::NoiseSample, samples::Int; kw...)
    nm = fit_noise(ns; kw...)
    δx = ns.x[2] - ns.x[1]
    return generate_noise(nm, length(ns), samples)
end

"""
Generate uncorrelated noise from a GaussianNoiseModel.
"""
function generate_noise(nm::GaussianNoiseModel{T}, len::Int, samples::Int) where {T <: Real}
    return Particles(samples, MvNormal(zeros(T, len), nm.σ))
end

"""
    add_noise(c::Curve, nm::AbstractNoiseModel, N::Int=10_000)

Add noise from noise model to curve, retrieve `UncertainCurve` with `N` samples.
"""
function add_noise(c::Curve, nm::GaussianNoiseModel, N::Int=10_000)
    return UncertainCurve(c.x, Particles(N, MvNormal(c.y, nm.σ)))
end
function add_noise(c::Curve, nm::MvGaussianNoiseModel, N::Int=10_000)
    δx = c.x[2] - c.x[1]
    Σ = get_cov(nm, length(c))
    return UncertainCurve(c.x, Particles(N, MvNormal(c.y, Σ)))
end


text:
# 01-blstm.jl
#
# Julia implementation Copyright (c) 2018 Matthew C. Kelley
#
# An implementation of the bidirectional LSTM neural net for
# speech recognition defind by Graves & Schmidhuber ([Graves, A., &
# Schmidhuber, J. (2005). Framewise phoneme classification with
# bidirectional LSTM and other neural network architectures. Neural
# Networks, 18(5-6), 602-610.]). The best results the authors report
# is after training this network using a weighted error function, then
# retraining it for 17 epochs with the categorical cross entropy
# error function. The present implementation is of their network that
# does not use retraining, for which the best results were found for
# training for around 20 epochs.

using Flux
using Flux: crossentropy, softmax, flip, sigmoid, LSTM
using JLD

# Paths to the training and test data directories
traindir = "train"
testdir = "test"

# Component layers of the bidirectional LSTM layer
forward = LSTM(26, 93)
backward = LSTM(26, 93)
output = Dense(186, 61)

"""
    BLSTM(x)

BLSTM layer using above LSTM layers

# Parameters
* **x** A 2-tuple containing the forward and backward time samples;
the first is from processing the sequence forward, and the second
is from processing it backward

# Returns
* The concatenation of the forward and backward LSTM predictions
"""
BLSTM(x) = vcat.(forward.(x), flip(backward, x))

"""
    model(x)

The chain of functions representing the trained model.

# Parameters
* **x** The utterance that the model should process

# Returns
* The model's predictions for each time step in `x`
"""
model(x) = softmax.(output.(BLSTM(x)))

"""
   loss(x, y)

Calculates the categorical cross-entropy loss for an utterance

# Parameters
* **x** Iterable containing the frames to classify
* **y** Iterable containing the labels corresponding to the frames
in `x`

# Returns
* The calculated loss value

# Side-effects
* Resets the state in the BLSTM layer
"""
function loss(x, y)
    l = sum(crossentropy.(model(x), y))
    Flux.reset!((forward, backward))
    return l
end

"""
    read_data(data_dir)

Reads in the data contained in a specified directory

# Parameters
* **data_dir** String of the path to the directory containing the data

# Return
* **Xs** Vector where each element is a vector of the frames for
one utterance
* **Ys** A vector where each element is a vector of the labels for
the frames for one utterance
"""
function read_data(data_dir)
    fnames = readdir(data_dir)

    Xs = Vector()
    Ys = Vector()

    for (i, fname) in enumerate(fnames)
        print(string(i) * "/" * string(length(fnames)) * "\r")
        x, y = load(joinpath(data_dir, fname), "x", "y")
        x = [x[i,:] for i in 1:size(x,1)]
        y = [y[i,:] for i in 1:size(y,1)]
        push!(Xs, x)
        push!(Ys, y)
    end

    return (Xs, Ys)
end

"""
    predict(x)

Make predictions on the data using the model defined above

# Parameters
* **x** An iterable containing the frames for a single utterance

# Returns
* The predicted scores for each phoneme class for each frame in `x`

# Side effects
* Resets the state in the BLSTM layer after making predictions
"""
function predict(x)
    ŷ = model(x)
    Flux.reset!((forward, backward))
    return ŷ
end

"""
    evaluate_accuracy(data)

Evaluates the accuracy of the model on a set of data; can be used
either for validation or test accuracy

# Parameters
* **data** An iterable of paired values where the first element is
all the frames for a single utterance, and the second is the
associated frame labels to compare the model's predictions against

# Returns
* The predicted accuracy value as a proportion of the number of
correct predictions over the total number of predictions made
"""
function evaluate_accuracy(data)
    correct = Vector()
    for (x, y) in data
        y = indmax.(y)
        ŷ = indmax.(predict(x))
        correct = vcat(correct,
                        [ŷ_n == y_n for (ŷ_n, y_n) in zip(ŷ, y)])
    end
    sum(correct) / length(correct)
end

println("Loading files")
Xs, Ys = read_data(traindir)
data = collect(zip(Xs, Ys))

# Move 5% (184 files) of the TIMIT data into a validation set
val_data = data[1:184]
data = data[185:length(data)]

# Begin training
println("Beginning training")

opt = Momentum(params((forward, backward, output)), 10.0^-5; ρ=0.9)
epochs = 20

for i in 1:epochs
    println("Epoch " * string(i) * "/" * string(epochs))
    data = data[shuffle(1:length(data))]
    val_data = val_data[shuffle(1:length(val_data))]

    Flux.train!(loss, data, opt)

    print("Validating\r")
    val_acc = evaluate_accuracy(val_data)
    println("Val acc. " * string(val_acc))
    println()
end

# Clearn up some memory
val_data = 0
data = 0
Xs = 0
Ys = 0
gc()

# Test model
print("Testing\r")
Xs_test, Ys_test = read_data(testdir)
test_data = collect(zip(Xs_test, Ys_test))
test_acc = evaluate_accuracy(test_data)
println("Test acc. " * string(test_acc))
println()

max_stars_repo_name: maetshju/flux-blstm-implementation

max_issues_repo_name: maetshju/flux-blstm-implementation

max_forks_repo_name: maetshju/flux-blstm-implementation


text:
using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.DGBalanceLawDiscretizations
using CLIMA.DGBalanceLawDiscretizations.NumericalFluxes
using CLIMA.MPIStateArrays
using CLIMA.LowStorageRungeKuttaMethod
using CLIMA.ODESolvers
using CLIMA.GenericCallbacks
using CLIMA.ParametersType
using LinearAlgebra
using StaticArrays
using Logging, Printf, Dates
using CLIMA.VTK
using DelimitedFiles
using Dierckx
using Random

using TimerOutputs

const to = TimerOutput()

if haspkg("CuArrays")
    using CUDAdrv
    using CUDAnative
    using CuArrays
    CuArrays.allowscalar(false)
    const ArrayType = CuArray
else
    const ArrayType = Array
end

# Prognostic equations: ρ, (ρu), (ρv), (ρw), (ρe_tot), (ρq_tot)
# For the dry example shown here, we load the moist thermodynamics module
# and consider the dry equation set to be the same as the moist equations but
# with total specific humidity = 0.
using CLIMA.MoistThermodynamics
using CLIMA.PlanetParameters: R_d, cp_d, grav, cv_d, MSLP, T_0, Omega

# State labels
const _nstate = 6
const _ρ, _U, _V, _W, _E, _QT = 1:_nstate
const stateid = (ρid = _ρ, Uid = _U, Vid = _V, Wid = _W, Eid = _E, QTid = _QT)
const statenames = ("RHO", "U", "V", "W", "E", "QT")

# Viscous state labels
const _nviscstates = 16
const _τ11, _τ22, _τ33, _τ12, _τ13, _τ23, _qx, _qy, _qz, _Tx, _Ty, _Tz, _θx, _θy, _θz, _SijSij = 1:_nviscstates

# Gradient state labels
# Gradient state labels
const _states_for_gradient_transform = (_ρ, _U, _V, _W, _E, _QT)

const _nauxstate = 15
const _a_x, _a_y, _a_z, _a_sponge, _a_02z, _a_z2inf, _a_rad, _a_ν_e, _a_LWP_02z, _a_LWP_z2inf,_a_q_liq,_a_θ, _a_P,_a_T, _a_soundspeed_air = 1:_nauxstate

if !@isdefined integration_testing
  const integration_testing =
    parse(Bool, lowercase(get(ENV,"JULIA_CLIMA_INTEGRATION_TESTING","false")))
end

# Problem constants (TODO: parameters module (?))
@parameter Prandtl 71 // 100 "Prandtl (molecular)"
@parameter  μ_sgs 100.0 "Constant dynamic viscosity"
@parameter Prandtl_t 1//3 "Prandtl_t"
@parameter cp_over_prandtl cp_d / Prandtl_t "cp_over_prandtl"

# Random number seed
const seed = MersenneTwister(0)

function global_max(A::MPIStateArray, states=1:size(A, 2))
  host_array = Array ∈ typeof(A).parameters
  h_A = host_array ? A : Array(A)
  locmax = maximum(view(h_A, :, states, A.realelems))
  MPI.Allreduce([locmax], MPI.MAX, A.mpicomm)[1]
end

function global_mean(A::MPIStateArray, states=1:size(A,2))
  host_array = Array ∈ typeof(A).parameters
  h_A = host_array ? A : Array(A)
  (Np, nstate, nelem) = size(A)
  numpts = (nelem * Np) + 1
  localsum = sum(view(h_A, :, states, A.realelems))
  MPI.Allreduce([localsum], MPI.SUM, A.mpicomm)[1] / numpts
end

# User Input
const numdims = 2
const Npoly = 4

# Define grid size
Δx    = 250
Δy    = 250
Δz    = 250

#
# OR:
#
# Set Δx < 0 and define  Nex, Ney, Nez:
#
(Nex, Ney, Nez) = (5, 5, 5)

# Physical domain extents
const (xmin, xmax) = (0, 10000)
const (ymin, ymax) = (0,  6400)
const (zmin, zmax) = (0,  6400)

#Get Nex, Ney from resolution
const Lx = xmax - xmin
const Ly = ymax - ymin
const Lz = zmax - ymin

if ( Δx > 0)
    #
    # User defines the grid size:
    #
    Nex = ceil(Int64, Lx / (Δx * Npoly))
    Ney = ceil(Int64, Ly / (Δy * Npoly))
    Nez = ceil(Int64, Lz / (Δz * Npoly))
else
    #
    # User defines the number of elements:
    #
    Δx = Lx / (Nex * Npoly)
    Δy = Ly / (Ney * Npoly)
    Δz = Lz / (Nez * Npoly)
end

DoF = (Nex*Ney*Nez)*(Npoly+1)^numdims*(_nstate)
DoFstorage = (Nex*Ney*Nez)*(Npoly+1)^numdims*(_nstate + _nviscstates + _nauxstate + CLIMA.Mesh.Grids._nvgeo) +
             (Nex*Ney*Nez)*(Npoly+1)^(numdims-1)*2^numdims*(CLIMA.Mesh.Grids._nsgeo)

# Smagorinsky model requirements : TODO move to SubgridScaleTurbulence module
@parameter C_smag 0.15 "C_smag"
# Equivalent grid-scale
#Δ = (Δx * Δy * Δz)^(1/3)
Δ = min(Δx, Δy)
const Δsqr = Δ * Δ

# -------------------------------------------------------------------------
# Preflux calculation: This function computes parameters required for the
# DG RHS (but not explicitly solved for as a prognostic variable)
# In the case of the rising_thermal_bubble example: the saturation
# adjusted temperature and pressure are such examples. Since we define
# the equation and its arguments here the user is afforded a lot of freedom
# around its behaviour.
# The preflux function interacts with the following
# Modules: NumericalFluxes.jl
# functions: wavespeed, cns_flux!, bcstate!
# -------------------------------------------------------------------------
@inline function preflux(Q, VF, aux, _...)
  @inbounds begin
    ρ, U, V, W = Q[_ρ], Q[_U], Q[_V], Q[_W]
    ρinv = 1 / ρ
    u, v, w = ρinv * U, ρinv * V, ρinv * W
  end
end

#-------------------------------------------------------------------------
#md # Soundspeed computed using the thermodynamic state TS
# max eigenvalue
@inline function wavespeed(n, Q, aux, t, u, v, w)
  @inbounds begin
    (n[1] * u + n[2] * v + n[3] * w) + aux[_a_soundspeed_air]
  end
end

# -------------------------------------------------------------------------
# ### read sounding
#md #
#md # The sounding file contains the following quantities along a 1D column.
#md # It needs to have the following structure:
#md #
#md # z[m]   theta[K]  q[g/kg]   u[m/s]   v[m/s]   p[Pa]
#md # ...      ...       ...      ...      ...      ...
#md #
#md #
# -------------------------------------------------------------------------
function read_sounding()
  #read in the original squal sounding
  fsounding  = open(joinpath(@__DIR__, "../soundings/SOUNDING_PYCLES_Z_T_P.dat"))
  sounding = readdlm(fsounding)
  close(fsounding)
  (nzmax, ncols) = size(sounding)
  if nzmax == 0
    error("SOUNDING ERROR: The Sounding file is empty!")
  end
  return (sounding, nzmax, ncols)
end

# -------------------------------------------------------------------------
# ### Physical Flux (Required)
#md # Here, we define the physical flux function, i.e. the conservative form
#md # of the equations of motion for the prognostic variables ρ, U, V, W, E, QT
#md # $\frac{\partial Q}{\partial t} + \nabla \cdot \boldsymbol{F} = \boldsymbol {S}$
#md # $\boldsymbol{F}$ contains both the viscous and inviscid flux components
#md # and $\boldsymbol{S}$ contains source terms.
#md # Note that the preflux calculation is splatted at the end of the function call
#md # to cns_flux!
# -------------------------------------------------------------------------
cns_flux!(F, Q, VF, aux, t) = cns_flux!(F, Q, VF, aux, t, preflux(Q,VF, aux)...)
@inline function cns_flux!(F, Q, VF, aux, t, u, v, w)
    @inbounds begin
        FT = eltype(F)
        D_subsidence = 3.75e-6
        ρ, U, V, W, E, QT = Q[_ρ], Q[_U], Q[_V], Q[_W], Q[_E], Q[_QT]
        P = aux[_a_P]
        xvert = aux[_a_y]

        # Inviscid contributions
        F[1, _ρ],  F[2, _ρ],  F[3, _ρ]  = U          , V          , W
        F[1, _U],  F[2, _U],  F[3, _U]  = u * U  + P , v * U      , w * U
        F[1, _V],  F[2, _V],  F[3, _V]  = u * V      , v * V + P  , w * V
        F[1, _W],  F[2, _W],  F[3, _W]  = u * W      , v * W      , w * W + P
        F[1, _E],  F[2, _E],  F[3, _E]  = u * (E + P), v * (E + P), w * (E + P)
        F[1, _QT], F[2, _QT], F[3, _QT] = u * QT     , v * QT     , w * QT

        #Derivative of T and Q:
        vqx, vqy, vqz = VF[_qx], VF[_qy], VF[_qz]
        vTx, vTy, vTz = VF[_Tx], VF[_Ty], VF[_Tz]

        SijSij = VF[_SijSij]

        #Dynamic eddy viscosity from Smagorinsky:
        ν_e = 75.0  #sqrt(2SijSij) * C_smag^2 * Δsqr
        D_e = 75.0  #ν_e / Prandtl_t

        # Multiply stress tensor by viscosity coefficient:
        τ11, τ22, τ33 = VF[_τ11] * ν_e, VF[_τ22]* ν_e, VF[_τ33] * ν_e
        τ12 = τ21 = VF[_τ12] * ν_e
        τ13 = τ31 = VF[_τ13] * ν_e
        τ23 = τ32 = VF[_τ23] * ν_e

        # Viscous velocity flux (i.e. F^visc_u in Giraldo Restelli 2008)
        F[1, _U] -= τ11; F[2, _U] -= τ12; F[3, _U] -= τ13
        F[1, _V] -= τ21; F[2, _V] -= τ22; F[3, _V] -= τ23
        F[1, _W] -= τ31; F[2, _W] -= τ32; F[3, _W] -= τ33

        # Viscous Energy flux (i.e. F^visc_e in Giraldo Restelli 2008)
        F[1, _E] -= u * τ11 + v * τ12 + w * τ13 + cp_over_prandtl * vTx * ν_e
        F[2, _E] -= u * τ21 + v * τ22 + w * τ23 + cp_over_prandtl * vTy * ν_e
        F[3, _E] -= u * τ31 + v * τ32 + w * τ33 + cp_over_prandtl * vTz * ν_e

        # Viscous contributions to mass flux terms
        F[1, _QT] -=  0.0 #vqx * D_e
        F[2, _QT] -=  0.0 #vqy * D_e
        F[3, _QT] -=  0.0 #vqz * D_e
    end
end

# -------------------------------------------------------------------------
#md # Here we define a function to extract the velocity components from the
#md # prognostic equations (i.e. the momentum and density variables). This
#md # function is not required in general, but provides useful functionality
#md # in some cases.
# -------------------------------------------------------------------------
# Compute the velocity from the state
const _ngradstates = 6
gradient_vars!(gradient_list, Q, aux, t, _...) = gradient_vars!(gradient_list, Q, aux, t, preflux(Q,~,aux)...)
@inline function gradient_vars!(gradient_list, Q, aux, t, u, v, w)
  @inbounds begin
    T = aux[_a_T]
    θ = aux[_a_θ]
    ρ, QT =Q[_ρ], Q[_QT]
    # ordering should match states_for_gradient_transform
    gradient_list[1], gradient_list[2], gradient_list[3] = u, v, w
    gradient_list[4], gradient_list[5], gradient_list[6] = θ, QT, T
  end
end

# -------------------------------------------------------------------------
#md ### Viscous fluxes.
#md # The viscous flux function compute_stresses computes the components of
#md # the velocity gradient tensor, and the corresponding strain rates to
#md # populate the viscous flux array VF. SijSij is calculated in addition
#md # to facilitate implementation of the constant coefficient Smagorinsky model
#md # (pending)
@inline function compute_stresses!(VF, grad_mat, _...)
  @inbounds begin
    dudx, dudy, dudz = grad_mat[1, 1], grad_mat[2, 1], grad_mat[3, 1]
    dvdx, dvdy, dvdz = grad_mat[1, 2], grad_mat[2, 2], grad_mat[3, 2]
    dwdx, dwdy, dwdz = grad_mat[1, 3], grad_mat[2, 3], grad_mat[3, 3]
    # compute gradients of moist vars and temperature
    dθdx, dθdy, dθdz = grad_mat[1, 4], grad_mat[2, 4], grad_mat[3, 4]
    dqdx, dqdy, dqdz = grad_mat[1, 5], grad_mat[2, 5], grad_mat[3, 5]
    dTdx, dTdy, dTdz = grad_mat[1, 6], grad_mat[2, 6], grad_mat[3, 6]
    # virtual potential temperature gradient: for richardson calculation
    # strains
    # --------------------------------------------
    # SMAGORINSKY COEFFICIENT COMPONENTS
    # --------------------------------------------
    S11 = dudx
    S22 = dvdy
    S33 = dwdz
    S12 = (dudy + dvdx) / 2
    S13 = (dudz + dwdx) / 2
    S23 = (dvdz + dwdy) / 2
    # --------------------------------------------
    # SMAGORINSKY COEFFICIENT COMPONENTS
    # --------------------------------------------
    # FIXME: Grab functions from module SubgridScaleTurbulence
    SijSij = S11^2 + S22^2 + S33^2 + 2S12^2 + 2S13^2 + 2S23^2

    #--------------------------------------------
    # deviatoric stresses
    # Fix up index magic numbers
    VF[_τ11] = 2 * (S11 - (S11 + S22 + S33) / 3)
    VF[_τ22] = 2 * (S22 - (S11 + S22 + S33) / 3)
    VF[_τ33] = 2 * (S33 - (S11 + S22 + S33) / 3)
    VF[_τ12] = 2 * S12
    VF[_τ13] = 2 * S13
    VF[_τ23] = 2 * S23

    # TODO: Viscous stresse come from SubgridScaleTurbulence module
    VF[_qx], VF[_qy], VF[_qz] = dqdx, dqdy, dqdz
    VF[_Tx], VF[_Ty], VF[_Tz] = dTdx, dTdy, dTdz
    VF[_θx], VF[_θy], VF[_θz] = dθdx, dθdy, dθdz
    VF[_SijSij] = SijSij
  end
end

# -------------------------------------------------------------------------
#md ### Auxiliary Function (Not required)
#md # In this example the auxiliary function is used to store the spatial
#md # coordinates. This may also be used to store variables for which gradients
#md # are needed, but are not available through teh prognostic variable
#md # calculations. (An example of this will follow - in the Smagorinsky model,
#md # where a local Richardson number via potential temperature gradient is required)
# -------------------------------------------------------------------------
@inline function auxiliary_state_initialization!(aux, x, y, z)
  @inbounds begin
      FT = eltype(aux)
      xvert = y
      aux[_a_y] = xvert
  end
end

# -------------------------------------------------------------------------
# generic bc for 2d , 3d

@inline function bcstate!(QP, VFP, auxP, nM, QM, VFM, auxM, bctype, t, uM, vM, wM)
    @inbounds begin

        x, y, z = auxM[_a_x], auxM[_a_y], auxM[_a_z]
        xvert = y

        ρM, UM, VM, WM, EM, QTM = QM[_ρ], QM[_U], QM[_V], QM[_W], QM[_E], QM[_QT]

        # No flux boundary conditions
        # No shear on walls (free-slip condition)
        UnM = nM[1] * UM + nM[2] * VM + nM[3] * WM
        QP[_U] = UM - 2 * nM[1] * UnM
        QP[_V] = VM - 2 * nM[2] * UnM
        QP[_W] = WM - 2 * nM[3] * UnM
        #QP[_ρ] = ρM   #this is:  dρ/dn = 0  i.e. ρ+ = ρ-
        #QP[_E] = EM   #this is:  dE/dn = 0  i.e. E+ = E-
        #VFP   .= VFM
        #VFP   .= 0.0    #This means that stress tau at the boundary is zero (notice
        #  that we are solving a viscous problem (nu=75) with a slip boundary; clearly this is physically incosistent but it will do for the sake of this benchmark (Straka 1993).
        Pr = 0.7
        ν = 75
        VFP[_Ty] = -grav*Pr/(ν*cv_d*cp_d)
        VFP[_Tz] = -grav*Pr/(ν*cv_d*cp_d)

        #=if xvert < 0.0001
        #if bctype  CODE_BOTTOM_BOUNDARY  FIXME: THIS NEEDS TO BE CHANGED TO CODE-BASED B.C. FOR TOPOGRAPHY
            #Dirichelt on T:
            SST    = 292.5
            q_tot  = QP[_QT]/QP[_ρ]
            q_liq  = auxM[_a_q_liq]
            e_int  = internal_energy(SST, PhasePartition(q_tot, q_liq, 0.0))
            e_kin  = 0.5*(QP[_U]^2/ρM^2 + QP[_V]^2/ρM^2 + QP[_W]^2/ρM^2)
            e_pot  = grav*xvert
            E      = ρM * total_energy(e_kin, e_pot, SST, PhasePartition(q_tot, q_liq, 0.0))
            QP[_E] = E
        end
        =#
        nothing
    end
end
# -------------------------------------------------------------------------
"""
 Neumann boundary conditions on
 all states, and on T for viscous problems:

 dQ/dn = 0
 dT/dn = -g/cv_d

"""
@inline function stresses_boundary_penalty!(VF, nM, gradient_listM, QM, aM, gradient_listP, QP, aP, bctype, t)
    VF .= 0

    stresses_penalty!(VF, nM, gradient_listM, QM, aM, gradient_listP, QP, aP, t)
end

@inline function stresses_penalty!(VF, nM, gradient_listM, QM, aM, gradient_listP, QP, aP, t)
  @inbounds begin
    n_Δgradient_list = similar(VF, Size(3, _ngradstates))
    for j = 1:_ngradstates, i = 1:3
      n_Δgradient_list[i, j] = nM[i] * (gradient_listP[j] - gradient_listM[j]) / 2
    end
    compute_stresses!(VF, n_Δgradient_list)
  end
end
# -------------------------------------------------------------------------

@inline function source!(S,Q,aux,t)
  # Initialise the final block source term
  S .= 0

  # Typically these sources are imported from modules
  @inbounds begin
    source_geopot!(S, Q, aux, t)
  end
end

@inline function source_geopot!(S,Q,aux,t)
  @inbounds S[_V] += - Q[_ρ] * grav
end

function preodefun!(disc, Q, t)
  DGBalanceLawDiscretizations.dof_iteration!(disc.auxstate, disc, Q) do R, Q, QV, aux
    @inbounds let
      ρ, U, V, W, E, QT = Q[_ρ], Q[_U], Q[_V], Q[_W], Q[_E], Q[_QT]
      xvert = aux[_a_y]
      e_int = (E - (U^2 + V^2+ W^2)/(2*ρ) - ρ * grav * xvert) / ρ
      q_tot = QT / ρ
      TS = PhaseEquil(e_int, q_tot, ρ)
      T = air_temperature(TS)
      P = air_pressure(TS) # Test with dry atmosphere
      q_liq = PhasePartition(TS).liq

      R[_a_T] = T
      R[_a_P] = P
      R[_a_q_liq] = q_liq
      R[_a_soundspeed_air] = soundspeed_air(TS)
      R[_a_θ] = virtual_pottemp(TS)
    end
  end
end

# initial condition
"""
    User-specified. Required.
    This function specifies the initial conditions
    for the dycoms driver.
"""
# NEW FUNCTION
# initial condition
function dc!(dim, Q, t, x, y, z, _...)
    FT                = eltype(Q)
    R_gas::FT         = R_d
    c_p::FT           = cp_d
    c_v::FT           = cv_d
    p0::FT            = MSLP
    gravity::FT       = grav
    # initialise with dry domain
    q_tot::FT         = 0
    q_liq::FT         = 0
    q_ice::FT         = 0
    # perturbation parameters for rising bubble
    rx                    = 4000
    ry                    = 2000
    xc                    = 0
    yc                    = 3000
    r                     = sqrt( (x - xc)^2/rx^2 + (y - yc)^2/ry^2)
    θ_ref::FT         = 300
    θ_c::FT           = -15.0
    Δθ::FT            = 0.0
    if r <= 1
        Δθ = θ_c * (1 + cospi(r))/2
    end
    qvar                  = PhasePartition(q_tot)
    θ                     = θ_ref + Δθ # potential temperature
    π_exner               = 1.0 - gravity / (c_p * θ) * y # exner pressure
    ρ                     = p0 / (R_gas * θ) * (π_exner)^ (c_v / R_gas) # density

    P                     = p0 * (R_gas * (ρ * θ) / p0) ^(c_p/c_v) # pressure (absolute)
    T                     = P / (ρ * R_gas) # temperature
    U, V, W               = 0.0 , 0.0 , 0.0  # momentum components
    # energy definitions
    e_kin                 = (U^2 + V^2 + W^2) / (2*ρ)/ ρ
    e_pot                 = gravity * y
    e_int                 = internal_energy(T, qvar)
    E                     = ρ * (e_int + e_kin + e_pot)  #* total_energy(e_kin, e_pot, T, q_tot, q_liq, q_ice)
    @inbounds Q[_ρ], Q[_U], Q[_V], Q[_W], Q[_E], Q[_QT]= ρ, U, V, W, E, ρ * q_tot
end

function run(mpicomm, dim, Ne, N, timeend, FT, dt)

    ##
      #-----------------------------------------------------------------
    # build physical range to be stratched
    #-----------------------------------------------------------------
    x_range = grid1d(xmin, xmax; nelem=Ne[1])
    y_range = grid1d(ymin, ymax; nelem=Ne[2])

    #-----------------------------------------------------------------
    # Build grid stretching along whichever direction
    # (ONLY Z for now. We need to decide what function we want to use for x and y)
    #-----------------------------------------------------------------
    y_range = grid1d(ymin, ymax, SingleExponentialStretching(2.5), nelem=Ne[2])

    #-----------------------------------------------------------------
    # END grid stretching
    #-----------------------------------------------------------------

    brickrange = (x_range, y_range)
    #-----------------------------------------------------------------
    #Build grid:
    #-----------------------------------------------------------------

  # User defined periodicity in the topl assignment
  # brickrange defines the domain extents
  @timeit to "Topo init" topl = StackedBrickTopology(mpicomm, brickrange, periodicity=(false,false))

  @timeit to "Grid init" grid = DiscontinuousSpectralElementGrid(topl,
                                                                 FloatType = FT,
                                                                 DeviceArray = ArrayType,
                                                                 polynomialorder = N)

  numflux!(x...) = NumericalFluxes.rusanov!(x..., cns_flux!, wavespeed, preflux)
  numbcflux!(x...) = NumericalFluxes.rusanov_boundary_flux!(x..., cns_flux!, bcstate!, wavespeed, preflux)

  # spacedisc = data needed for evaluating the right-hand side function
  @timeit to "Space Disc init" spacedisc = DGBalanceLaw(grid = grid,
                                                        length_state_vector = _nstate,
                                                        flux! = cns_flux!,
                                                        numerical_flux! = numflux!,
                                                        numerical_boundary_flux! = numbcflux!,
                                                        number_gradient_states = _ngradstates,
                                                        states_for_gradient_transform =
                                                        _states_for_gradient_transform,
                                                        number_viscous_states = _nviscstates,
                                                        gradient_transform! = gradient_vars!,
                                                        viscous_transform! = compute_stresses!,
                                                        viscous_penalty! = stresses_penalty!,
                                                        viscous_boundary_penalty! = stresses_boundary_penalty!,
                                                        auxiliary_state_length = _nauxstate,
                                                        auxiliary_state_initialization! = (x...) ->
                                                        auxiliary_state_initialization!(x...),
                                                        source! = source!,
                                                        preodefun! = preodefun!)

  # This is a actual state/function that lives on the grid
  @timeit to "IC init" begin
    initialcondition(Q, x...) = dc!(Val(dim), Q, FT(0), x...)
    Q = MPIStateArray(spacedisc, initialcondition)
  end

  @timeit to "Time stepping init" begin
    lsrk = LSRK54CarpenterKennedy(spacedisc, Q; dt = dt, t0 = 0)

    #=eng0 = norm(Q)
    @info @sprintf """Starting
    norm(Q₀) = %.16e""" eng0
    =#
    # Set up the information callback
    starttime = Ref(now())
    cbinfo = GenericCallbacks.EveryXWallTimeSeconds(10, mpicomm) do (s=false)
      if s
        starttime[] = now()
      else
        #energy = norm(Q)
        #globmean = global_mean(Q, _ρ)
        @info @sprintf("""Update
                       simtime = %.16e
                       runtime = %s""",
                       ODESolvers.gettime(lsrk),
                       Dates.format(convert(Dates.DateTime,
                                            Dates.now()-starttime[]),
                                    Dates.dateformat"HH:MM:SS")) #, energy )#, globmean)
      end
    end

    npoststates = 3
    _o_u, _o_v, _o_w = 1:npoststates
    postnames = ("u", "v", "w")
    postprocessarray = MPIStateArray(spacedisc; nstate=npoststates)

    step = [0]
    cbvtk = GenericCallbacks.EveryXSimulationSteps(1000) do (init=false)
      DGBalanceLawDiscretizations.dof_iteration!(postprocessarray, spacedisc, Q) do R, Q, QV, aux
        @inbounds let
          u, v, w = preflux(Q, QV, aux)
          R[_o_u] = u
          R[_o_v] = v
          R[_o_w] = w
        end
      end

      outprefix = @sprintf("./CLIMA-output-scratch/stretching/dy_%dD_mpirank%04d_step%04d", dim,
                           MPI.Comm_rank(mpicomm), step[1])
      @debug "doing VTK output" outprefix
      writevtk(outprefix, Q, spacedisc, statenames, postprocessarray, postnames)

      step[1] += 1
      nothing
    end
  end

  @info @sprintf """ ---- COMPLETE: Grid built successfully ----"""

  # Initialise the integration computation. Kernels calculate this at every timestep??
  @timeit to "solve" solve!(Q, lsrk; timeend=timeend, callbacks=(cbinfo, cbvtk))
end

using Test
let
  MPI.Initialized() || MPI.Init()
  Sys.iswindows() || (isinteractive() && MPI.finalize_atexit())
  mpicomm = MPI.COMM_WORLD
  ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
  loglevel = ll == "DEBUG" ? Logging.Debug :
      ll == "WARN"  ? Logging.Warn  :
      ll == "ERROR" ? Logging.Error : Logging.Info
  logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull
  global_logger(ConsoleLogger(logger_stream, loglevel))
  @static if haspkg("CUDAnative")
      device!(MPI.Comm_rank(mpicomm) % length(devices()))
  end
  # User defined number of elements
  # User defined timestep estimate
  # User defined simulation end time
  # User defined polynomial order
  numelem = (Nex, Ney)
  dt = 0.025
  timeend = dt
  polynomialorder = Npoly
  FT = Float64
  dim = numdims

  if MPI.Comm_rank(mpicomm) == 0
    @info @sprintf """ ------------------------------------------------------"""
    @info @sprintf """   ______ _      _____ __  ________                    """
    @info @sprintf """  |  ____| |    |_   _|  ...  |  __  |                 """
    @info @sprintf """  | |    | |      | | |   .   | |  | |                 """
    @info @sprintf """  | |    | |      | | | |   | | |__| |                 """
    @info @sprintf """  | |____| |____ _| |_| |   | | |  | |                 """
    @info @sprintf """  | _____|______|_____|_|   |_|_|  |_|                 """
    @info @sprintf """                                                       """
    @info @sprintf """ ------------------------------------------------------"""
    @info @sprintf """ MWE_grid_stretching                                   """
    @info @sprintf """   Resolution:                                         """
    @info @sprintf """     (Δx, Δy) = (%.2e, %.2e)                           """ Δx Δy
    @info @sprintf """     (Nex, Ney) = (%d, %d)                             """ Nex Ney
    @info @sprintf """     DoF = %d                                          """ DoF
    @info @sprintf """     Minimum necessary memory to run this test: %g GBs """ (DoFstorage * sizeof(FT))/1000^3
    @info @sprintf """     Time step dt: %.2e                                """ dt
    @info @sprintf """     End time  t : %.2e                                """ timeend
    @info @sprintf """ ------------------------------------------------------"""
  end

  engf_eng0 = run(mpicomm, dim, numelem[1:dim], polynomialorder, timeend,
                  FT, dt)

  show(to)
end

isinteractive() || MPI.Finalize()

nothing

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
#=
    maybe_floattype(T)

To keep consistant with Base `diff` that Int array outputs Int array. It is sometimes
useful to only promote types for Bool and FixedPoint. In most of the time, `floattype`
should be the most reliable way.
=#
maybe_floattype(::Type{T}) where {T} = T
maybe_floattype(::Type{Bool}) = floattype(Bool)
maybe_floattype(::Type{T}) where {T<:FixedPoint} = floattype(T)
maybe_floattype(::Type{CT}) where {CT<:Color} = base_color_type(CT){maybe_floattype(eltype(CT))}

#=
A helper to eagerly check if particular function makes sense for `extreme_filter` semantics.
For instance, `max(::RGB, ::RGB)` is not well-defined and we should early throw errors so that
our users don't get encrypted error messages.
=#
function require_select_function(f, ::Type{T}) where {T}
    if !_is_select_function(f, T)
        hint = "does `f(x::T, y::T)` work as expected?"
        throw(ArgumentError("function `$f` is not a well-defined select function on type `$T`: $hint"))
    end
end
_is_select_function(f, ::Type{T}) where {T} = _is_select_function_trial(f, T)
function _is_select_function(f, ::Type{T}) where {T<:Real}
    f in (min, max) && return true
    return _is_select_function_trial(f, T)
end
function _is_select_function(f, ::Type{CT}) where {CT<:AbstractGray}
    f in (min, max) && return true
    return _is_select_function_trial(f, CT)
end
function _is_select_function(f, ::Type{CT}) where {CT<:Colorant}
    # min/max is not well-defined on generic color space
    f in (min, max) && return false
    return _is_select_function_trial(f, CT)
end
function _is_select_function_trial(f, ::Type{T}) where {T}
    # for generic case, just run a trial and see if it doesn't error
    v = zero(T)
    try
        f(v, v)
        return true
    catch
        return false
    end
end


text:
module PartialTorus

import SimpleMolecule: Molecule
import MolecularBuild: PartialFibonacciTorus
import MolecularBuild: BuildPartialTorus

function assemble(systems::Array{Molecule,1}, params::Dict)
    return BuildPartialTorus.build(
        PartialFibonacciTorus,
        systems;
        number=convert(Int, params["number"]),
        radius=float(params["radius"]),
        a=float(params["a"]),
        b=float(params["b"]),
        zmin=float(params["zmin"]),
        zmax=float(params["zmax"]),
        center=[float(x) for x in params["center"]],
        old_orientation=Float64[float(x) for x in params["orientation"]],
        aligned=true,
        inverted=params["inverted"]
    )
end

export assemble
end

max_stars_repo_name: yuhangwang/SimShape

max_issues_repo_name: yuhangwang/SimShape

max_forks_repo_name: yuhangwang/SimShape


text:
using Test
using MultivariateMoments

@testset "VectorizedHermitianMatrix" begin
    Q = MultivariateMoments.vectorized_hermitian_matrix(Int, (i, j) -> i == j ? i * 2 - 1 : 2 - im, 2, 1:2)
    @test Q[1:2, 1:2] isa Matrix{Complex{Int}}
    @test Q[1:2, 1:2] == [1 2 - im; 2 + im 3]
    @test square_getindex(Q, 1:2) isa VectorizedHermitianMatrix{Int, Bool, Complex{Int}}
    @test square_getindex(Q, 1:2).Q == [1, 2, 3, -1]
    @test square_getindex(Q, 2:-1:1) isa VectorizedHermitianMatrix{Int, Bool, Complex{Int}}
    @test square_getindex(Q, 2:-1:1).Q == [3, 2, 1, -1]
    @test square_getindex(Q, 1:1) isa VectorizedHermitianMatrix{Int, Bool, Complex{Int}}
    @test square_getindex(Q, 1:1).Q == [1]
    @test square_getindex(Q, 2:2) isa VectorizedHermitianMatrix{Int, Bool, Complex{Int}}
    @test square_getindex(Q, 2:2).Q == [3]
    @test eltype(Q) == Complex{Int}
    @test 1 == @inferred Q[1, 1]
    @test 2 - im == @inferred Q[1, 2]
    @test 2 + im == @inferred Q[2, 1]
    @test 3 == @inferred Q[2, 2]
    symmetric_setindex!(Q, 4, 1, 1)
    @test Q.Q == [4, 2, 3, -1]
    symmetric_setindex!(Q, 5 - 2im, 1, 2)
    @test Q.Q == [4, 5, 3, -2]
    P = copy(Q)
    @test P.n == 2
    symmetric_setindex!(P, 6, 2, 2)
    @test Q.Q == [4, 5, 3, -2]
    @test P.n == 2
    @test P.Q == [4, 5, 6, -2]
    R = map(i -> i - 1, P)
    @test R.n == 2
    @test R.Q == [3, 4, 5, -2]
    @test P.Q == [4, 5, 6, -2]
    S = map(conj, R)
    @test S.n == 2
    @test S.Q == [3, 4, 5, 2]
    @test_throws ErrorException symmetric_setindex!(S, im, 1, 1)
    M = [1       2 + 3im 4 + 5im
         2 - 3im 6       7 + 8im
         4 - 5im 7 - 8im 9]
    N = MultivariateMoments.vectorized_hermitian_matrix(Int, (i, j) -> M[i, j], 3, 3:-1:1)
    @test Matrix(N) == M[3:-1:1, 3:-1:1]
    symmetric_setindex!(N, 4 - 5im, 3, 1)
    @test Matrix(N) != M[3:-1:1, 3:-1:1]
    M[3, 1] = 4 + 5im
    M[1, 3] = 4 - 5im
    @test Matrix(N) == M[3:-1:1, 3:-1:1]
    for QQ in [similar(N, (2, 2)), similar(N, 2, 2), similar(typeof(N), (2, 2)), similar(typeof(N), 2, 2)]
        @test typeof(QQ) == typeof(N)
        symmetric_setindex!(QQ, 2, 1, 1)
        @test QQ[1, 1] == 2
        symmetric_setindex!(QQ, 3 + 4im, 1, 2)
        @test QQ[1, 2] == 3 + 4im
        @test QQ[2, 1] == 3 - 4im
        symmetric_setindex!(QQ, 5, 2, 2)
        @test QQ[2, 2] == 5
    end
end


text:
using Test
using ExproniconLite

@testset "x function" begin
    @test_expr xtuple(1, :x) == :((1, x))
    @test_expr xnamedtuple(;x=2, y=3) == :((x = 2, y = 3))
    @test_expr xcall(Base, :sin, 1; x=2) == :($Base.sin(1; x = 2))
    @test_expr xpush(:coll, :x) == :($Base.push!(coll, x))
    @test_expr xfirst(:coll) == :($Base.first(coll))
    @test_expr xlast(:coll) == :($Base.last(coll))
    @test_expr xprint(:coll) == :($Base.print(coll))
    @test_expr xprintln(:coll) == :($Base.println(coll))
    @test_expr xmap(:f, :coll) == :($Base.map(f, coll))
    @test_expr xmapreduce(:f, :op, :coll) == :($Base.mapreduce(f, op, coll))
    @test_expr xiterate(:it) == :($Base.iterate(it))
    @test_expr xiterate(:it, :st) == :($Base.iterate(it, st))
end


text:
n_features = (32, 32, 3)
n_targets = 1

@testset "trainset" begin
    d = SVHN2()

    @test d.split == :train
    @test extrema(d.features) == (0, 1)
    @test convert2image(d, 1) isa AbstractMatrix{<:RGB}
    @test convert2image(d, 1:2) isa AbstractArray{<:RGB,3}

    test_supervised_array_dataset(d;
        n_features, n_targets, n_obs=73257,
        Tx=Float32, Ty=Int)
end

@testset "testset" begin
    d = SVHN2(split=:test, Tx=UInt8)

    @test d.split == :test
    @test extrema(d.features) == (0, 255)
    @test convert2image(d, 1) isa AbstractMatrix{<:RGB}
    @test convert2image(d, 1:2) isa AbstractArray{<:RGB,3}

    test_supervised_array_dataset(d;
        n_features, n_targets, n_obs=26032,
        Tx=UInt8, Ty=Int)
end

# @testset "extraset" begin
#     d = SVHN2(split=:test)

#     @test d.split == :extra
#     @test extrema(d.features) == (0, 1)
#     @test convert2image(d, 1) isa AbstractMatrix{<:RGB}
#     @test convert2image(d, 1:2) isa AbstractArray{<:RGB,3}

#     test_supervised_array_dataset(d;
#         n_features, n_targets, n_obs=531131,
#         Tx=Float32, Ty=Int)
# end


text:
module TestPool

using Test
using ParallelRadixSort: Pool, initpool!, alloc!, free!

zeros3() = zeros(3)

@testset begin
    capacity = 4
    pool = Pool(zeros3, capacity)
    initpool!(pool)
    xs1 = alloc!(pool, 2)
    @test length(xs1) == 2
    xs2 = alloc!(pool, 4)
    @test length(xs2) == 4
    free!(pool, xs1)
    free!(pool, xs2)
    @test length(pool.buffers[Threads.threadid()]) == capacity
end

end  # module


text:
@testset "Random generation" begin

    # Make two random seeds, extremely low probability they will be the same
    seed1 = 1
    seed2 = 2

    volfrac = 0.1
    radius = 0.5

    region_shape = Sphere(2, 20.0)
    particle_shape = Circle(radius)
    medium = Acoustic(2; ρ = 0.2, c = 0.2)

    particles1  = random_particles(medium, particle_shape, region_shape, volfrac; seed=seed1)
    particles1a = random_particles(medium, particle_shape, region_shape, volfrac; seed=seed1)
    particles2  = random_particles(medium, particle_shape, region_shape, volfrac; seed=seed2)

    # Particles should be determined solely by the seed
    @test particles1 == particles1a
    @test particles1 != particles2

    @test_throws ErrorException random_particles(medium, particle_shape, region_shape, 0.9)
    @test_throws ErrorException random_particles(medium, particle_shape, region_shape, 0.5; seed=3)

    region_shape = Sphere(3, 20.0)
    particle_shape = Sphere(radius)
    medium = Acoustic(3; ρ = 0.2, c = 0.2)

    # If a seed is not provided, then a different seed will be used everytime
    particles2  = random_particles(medium, particle_shape, region_shape, volfrac)
end


text:
using LinearAlgebra
using Random
using PyPlot
using Interpolations
using SparseArrays
using Optim
using ProgressBars
using StatsBase
using AlphaStableDistributions
using AdvancedHMC
using Dates
using MAT
using MCMCDiagnostics

seed_shift = (@isdefined shift) ? shift : 0

function measurementmatrix(X,MX,kernel)
    s = X[2]-X[1];
    F = zeros(size(MX)[1],size(X)[1])
    for i = 1:size(MX)[1]
        F[i,:] = s.*kernel.(X,MX[i]);
    end
    return F
end

function smatrix(X0,a,b)
    N = length(X0);
    dx = X0[2]-X0[1];

    M = zeros(N,N);
    for i=2:N-1
        M[i,i-1] = -a/dx^2;
        M[i,i+1] = -a/dx^2;
        M[i,i] = 2*a/dx^2 +b;

    end
    #M[1,:] = [b+2*a/dx^2 -a/dx^2 zeros(1,N-3) -a/dx^2];
    #M[N,:] = [-a/dx^2 zeros(1,N-3) -a/dx^2 b+2*a/dx^2];

    M[1,1] = 1;
    #M(N,1) = 1/dx;
    #M(N,2) = -1/dx;

    M[N,N] = 1
    return M

end

function partialsum!(z, s0, sp, spc)
    N = size(spc)[1]
     for i = 1:N
        r = spc[i]
        z[r] = s0[r] + sp[i]
    end
    return nothing
end

function ravel(q)
    return dropdims(q, dims = tuple(findall(size(q) .== 1)...))
end

function difference1(X)
    N = length(X);
    M = zeros(N,N);

    for i  = 2:N
       M[i,i] = -1;
       M[i,i-1] = 1;
    end
    M[1,1] = 1;
    return M
end

function  difference2(X)
    N = length(X);
    M = zeros(N,N);

    for i  = 3:N
       M[i,i] = -1;
       M[i,i-1] = 2;
       M[i,i-2] = -1;
    end
    M[1,1] = 1; M[2,1] = 1; M[2,2] = -1;
    return M

end

function divide(N,NT)
    even = mod(N,NT) == 0
    k = Int64(floor(N/NT))
    t = 0
    if(even)
        q = Vector{Vector{Int64}}(undef,k)
        t = k

    else
        q = Vector{Vector{Int64}}(undef,k+1)
        t = k +1
    end

    for i = 1:t
        q[i] = Vector(NT*(i-1)+1:min(NT*i,N))
    end

    return q

end

@inbounds  function measurementmatrix2d(X,Y,Xt,Yt,kernel;constant=10.0)
    dy = abs(Y[2,1]-Y[1,1])
    dx = abs(X[1,2]-X[1,1])
    d = Float64(dy*dx)
    N = length(X);
    Nt = length(Xt)
    #F = zeros(Nt,N)

    minval = 0.0001*d*kernel(0,0,0,0,constant=constant)
    Nth = Threads.nthreads()
    rows = Vector{Vector{Int64}}(undef, Nth)
    cols = Vector{Vector{Int64}}(undef, Nth)
    vals = Vector{Vector{Float64}}(undef, Nth)
    for p = 1:Nth
        rows[p] = []
        cols[p] = []
        vals[p] = []
    end

    Threads.@threads  for i = 1:N
        for j = 1:Nt
            xi = div(i-1,N)+1;
            yi = rem(i,N)+1
            xj = div(j-1,Nt)+1;
            yj = rem(j,Nt)+1
            val = d*kernel(X[i],Xt[j],Y[i],Yt[j],constant=constant);
            if (val >= minval)
                #F[j,i] = val
                push!(rows[Threads.threadid()],j)
                push!(cols[Threads.threadid()],i)
                push!(vals[Threads.threadid()],val)
            end

            #F[i,j] = d*kernel(X[xi],X[xj],Y[yi],Y[yj],constant=constant);
            #F[j,i] = F[i,j]
        end
    end
    rows = vcat(rows...)
    cols = vcat(cols...)
    vals = vcat(vals...)
    F = sparse(rows,cols,vals,Nt,N)

    return F

end

function  s2dmatrix(x0,a,b)
    N = length(x0);
    dx = x0[2]- x0[1]

   M = zeros(N,N);
   for i=2:N-1
       M[i,i-1] = -a/dx^2;
       M[i,i+1] = -a/dx^2;
       M[i,i] = 2*a/dx^2;

   end
   M[1,:] = [2*a/dx^2 -a/dx^2 zeros(1,N-3) -a/dx^2];
   M[N,:] = [-a/dx^2 zeros(1,N-3) -a/dx^2 2*a/dx^2];

   M = kron(M,I(N)) + kron(I(N),M)
   M = M + I(size(M)[2])*b;

   return M
end

    function linspace(start,stop,length)
        return Vector(range(start,stop=stop,length=length))
    end

    function partialreplace(orig,replace,ix)
        a = copy(orig)
        a[ix] = replace
        return a
    end

   @inline function meshgrid(x,y)
        grid_a = [i for i in x, j in y]
        grid_b = [j for i in x, j in y]

        return grid_a,grid_b
    end

@inbounds @inline function storevector!(target,source)
    N = length(target)
    @assert N == length(source)
    @simd for i = 1:N
        target[i] = source[i]
    end
end

@inbounds @inline function storesubst!(target,a1,a2)
    N = length(target)
    @assert N == length(a1) == length(a2)
    @simd for i = 1:N
        target[i] = a1[i] - a2[i]
    end
end

@inbounds @inline function storeadd!(target,a1,a2)
    N = length(target)
    @assert N == length(a1) == length(a2)
    @simd for i = 1:N
        target[i] = a1[i] + a2[i]
    end
end

@inbounds @inline function slicemadd!(target,cols,vec,ix)
    N = length(ix)
    @assert N == length(vec)
    #mul!(target,cols[ix[1]],vec[1])
    for i = 1:N
        mul!(target,cols[ix[i]],vec[i],1,1)
    end
end

function  logpdiff_prior(x,args,cache)
    D = args.D;
    noisesigma = args.noisesigma
    scale = args.scale; #y  = args.y; F = args.F

    #Fxprop = cache.Fxprop
    Dprop = cache.Dprop
    #res = cache.residual

    #mul!(Fxprop,F,x)
    mul!(Dprop,D,x)

    #res .= Fxprop-y
    #storesubst!(res,Fxprop,y)
    #logp = -0.5/noisesigma^2*dotitself(res);
    logp = sum(log.(scale./((scale^2 .+ Dprop.^2))));

    return logp

end

function  logpdiffpartial(ixl,nth,x,args,cache,current)
    noisesigma = args.noisesigma
    ix = ixl[nth]
    scale = args.scale; y  = args.y; bscale = args.bscale
    # F =   args.F[:,ix]
    # Dx =  args.Dx[:,ix]; Dy =  args.Dy[:,ix]
    # Db =  args.Db[:,ix]
    F =   args.Fs[nth] ; D =   args.Ds[nth]
    Ds = args.Ds; Dnz = args.Dnz;
    Fs = args.Fs; Fnz = args.Fnz

    res = cache.residual
    Dprop = cache.Dprop
    Fxprop = cache.Fxprop
    xprop = cache.xprop

    currx = @view current.xcurr[ix]
    Fxcurr = current.Fxcurr
    Dcurr =  current.Dcurr

    xprop[ix] = x

    # Fnzall = union(vcat(Fnz[ix]...))
    Fnzall = Fnz[nth]
    # Dxnzall = union(vcat(Dxnz[ix]...))
    # Dynzall = union(vcat(Dynz[ix]...))
    Dnzall = Dnz[nth]
    # Dbnzall = union(vcat(Dbnz[ix]...))

    compx = x-currx
    #Fxprop .= Fxcurr
    #Fxprop .+= (F*compx)
    storevector!(view(Fxprop,Fnzall),view(Fxcurr,Fnzall))
    #storevector!(Fxprop,Fxcurr)
    mul!(Fxprop,F,compx,1,1)
    #slicemadd!(Fxprop,Fs,compx,ix)

    #Dxprop .= Dxcurr
    #Dxprop .+= Dx*compx
    storevector!(view(Dprop,Dnzall),view(Dcurr,Dnzall))
    #storevector!(Dprop,Dcurr)
    mul!(Dprop,D,compx,1,1)

    #res .= Fxprop
    storesubst!(res,Fxprop,y)
    #storevector!(res,Fxprop)
    #res .-=  y

    resv = @view res[Fnzall]
    Dv = @view Dprop[Dnzall]

    #return   -0.5/noisesigma^2*res'*res -3/2*sum(log.(scale^2 .+ Dxprop.^2 .+ Dyprop.^2)) - sum(log.(bscale^2 .+ Dbprop.^2))
    return   -0.5/noisesigma^2*dotitself(resv) + sum(log.(scale./((scale^2 .+ Dv.^2))));
end

function  logpdiffgradi_prior(x,arguments,cache;both=true)
    D = arguments.D;
    noisesigma = arguments.noisesigma
    scale = arguments.scale; #y  = arguments.y; F = arguments.F

    logp = 0.0

    #Fxprop = cache.Fxprop
    Dprop = cache.Dprop
    G = cache.gradiprop
    #res = cache.residual

    #mul!(Fxprop,F,x)
    mul!(Dprop,D,x)
    #storesubst!(res,Fxprop,y)

    den = (scale^2 .+ Dprop.^2)
    if both
        #logp = -0.5/noisesigma^2*dotitself(res);
        logp =  sum(log.(scale./(den)));
    end

     #mul!(G,F',(-((res)/noisesigma.^2)))
     #G= F'*(-((res)./noisesigma.^2))

    #Gd =   M'*(-2.0*Dprop./den);
    mul!(G,D',(-2.0*Dprop./den))
    #G = G + Gd;

    return logp,G

end

function  logpisodiff_prior(x,args,cache)
    noisesigma = args.noisesigma
    scale = args.scale; #y  = args.y;# F = args.F
    bscale = args.bscale
    Dx = args.Dx; Dy = args.Dy
    Db = args.Db

    #res = cache.residual
    Dxprop = cache.Dxprop
    Dyprop = cache.Dyprop
    Dbprop = cache.Dbprop
    #Fxprop = cache.Fxprop

    #mul!(Fxprop,F,x)
    #res .= Fxprop - y
    #storesubst!(res,Fxprop,y)
    #Dxprop .= Dx*x
    mul!(Dxprop,Dx,x)
    #Dyprop .= Dy*x
    mul!(Dyprop,Dy,x)
    #Dbprop .= Db*x
    mul!(Dbprop,Db,x)

    return   -sum(log.(bscale^2 .+ Dbprop.^2)) -3/2*sum(log.(scale^2 .+ Dxprop.^2 + Dyprop.^2))

end

function  logpisodiffpartial(ixl,nth,x,args,cache,current)
    noisesigma = args.noisesigma
    ix = ixl[nth]
    scale = args.scale; y  = args.y; bscale = args.bscale
    # F =   args.F[:,ix]
    # Dx =  args.Dx[:,ix]; Dy =  args.Dy[:,ix]
    # Db =  args.Db[:,ix]
    F =   args.Fs[nth]
    Dx =  args.Dxs[nth]; Dy =  args.Dys[nth]
    Db =  args.Dbs[nth]
    Dxs = args.Dxs; Dxnz = args.Dxnz;
    Dys = args.Dys; Dynz = args.Dynz
    Dbs = args.Dbs; Dbnz = args.Dbnz
    Dinz = args.Dinz
    Fs = args.Fs; Fnz = args.Fnz

    #res = cache.residual
    Dxprop = cache.Dxprop
    Dyprop = cache.Dyprop
    Dbprop = cache.Dbprop
    Fxprop = cache.Fxprop
    xprop = cache.xprop

    currx = @view current.xcurr[ix]
    Fxcurr = current.Fxcurr
    Dxcurr =  current.Dxcurr
    Dycurr = current.Dycurr
    Dbcurr = current.Dbcurr

    xprop[ix] = x

    # Fnzall = union(vcat(Fnz[ix]...))
    Fnzall = Fnz[nth]
    # Dxnzall = union(vcat(Dxnz[ix]...))
    Dxnzall = Dxnz[nth]
    # Dynzall = union(vcat(Dynz[ix]...))
    Dynzall = Dynz[nth]
    # Dbnzall = union(vcat(Dbnz[ix]...))
    Dbnzall = Dbnz[nth]
    Diall = Dinz[nth]
    #Diall = union(Dxnzall,Dynzall)

    compx = x-currx
    #Fxprop .= Fxcurr
    #Fxprop .+= (F*compx)
    storevector!(view(Fxprop,Fnzall),view(Fxcurr,Fnzall))
    #storevector!(Fxprop,Fxcurr)
    mul!(Fxprop,F,compx,1,1)
    #slicemadd!(Fxprop,Fs,compx,ix)

    #Dxprop .= Dxcurr
    #Dxprop .+= Dx*compx
    storevector!(view(Dxprop,Dxnzall),view(Dxcurr,Dxnzall))
    #storevector!(Dxprop,Dxcurr)
    mul!(Dxprop,Dx,compx,1,1)

    #Dyprop .= Dycurr
    #Dyprop .+= Dy*compx
    storevector!(view(Dyprop,Dynzall),view(Dycurr,Dynzall))
    #storevector!(Dyprop,Dycurr)
    mul!(Dyprop,Dy,compx,1,1)

    #Dbprop .= Dbcurr
    #Dbprop .+= Db*compx
    storevector!(view(Dbprop,Dbnzall),view(Dbcurr,Dbnzall))
    #storevector!(Dbprop,Dbcurr)
    mul!(Dbprop,Db,compx,1,1)

    #res .= Fxprop
    storesubst!(res,Fxprop,y)
    #storevector!(res,Fxprop)
    #res .-=  y

    resv = @view res[Fnzall]
    Dxv = @view Dxprop[Diall]
    Dyv = @view Dyprop[Diall]
    Dbv = @view Dbprop[Dbnzall]

    #return   -0.5/noisesigma^2*res'*res -3/2*sum(log.(scale^2 .+ Dxprop.^2 .+ Dyprop.^2)) - sum(log.(bscale^2 .+ Dbprop.^2))
    return   -0.5/noisesigma^2*dotitself(resv) - sum(log.(bscale^2 .+ Dbv.^2)) -3/2*sum(log.(scale^2 .+ Dxv.^2 + Dyv.^2))

end

function logpisodiffgradi_prior(x,args,cache;both=true)
    noisesigma = args.noisesigma
    scale = args.scale; #y  = args.y; #F = args.F
    bscale = args.bscale
    Dx = args.Dx; Dy = args.Dy
    Db = args.Db

    logp = 0.0

    #res = cache.residual
    Dxprop = cache.Dxprop
    Dyprop = cache.Dyprop
    Dbprop = cache.Dbprop
    #Fxprop = cache.Fxprop
    G = cache.gradiprop

    #Fxprop .= F*x
    #mul!(Fxprop,F,x)
    #storesubst!(res,Fxprop,y)
    #res .= Fxprop - y
    #Dxprop .= Dx*x
    mul!(Dxprop,Dx,x)
    #Dyprop .= Dy*x
    mul!(Dyprop,Dy,x)
    #Dbprop .= Db*x
    mul!(Dbprop,Db,x)

    #G .= F'*(-((res)./noisesigma.^2))
    #mul!(G,F',-((res)/noisesigma.^2))

    den = (scale^2 .+ Dyprop.^2 + Dxprop.^2)

    if both
        logp =  -3/2*sum(log.(den)) - sum(log.(bscale^2 .+ Dbprop.^2))
    end

    #Gd1 =  Dx'*(-3.0*Lxx./(scale^2 .+ Lxx.^2));
    #G .= G  -3.0*(Dx)'*(Dxprop./(scale^2 .+ Dyprop.^2 + Dxprop.^2))  - 3.0*(Dy)'*(Dyprop./(scale^2 .+ Dyprop.^2 + Dxprop.^2))   -Db'*(2.0*Dbprop./(bscale^2 .+ Dbprop.^2));

    mul!(G,-3.0*(Dx)',(Dxprop./den))
    mul!(G,-3.0*(Dy)',(Dyprop./den),1,1)
    mul!(G,-Db',(2.0*Dbprop./(bscale^2 .+ Dbprop.^2)),1,1)

    return logp, G
end

function regmatrices_first(dim)
    reg1d = spdiagm(Pair(0,-1*ones(dim))) + spdiagm(Pair(1,ones(dim-1))) + spdiagm(Pair(-dim+1,ones(1))) ;reg1d[dim,dim] = 0
    #reg1d = reg1d[1:dim-1,:]
    iden = I(dim)
    regx = kron(reg1d,iden)
    regy = kron(iden,reg1d)

    rmxix = sum(abs.(regx) ,dims=2) .< 2
    rmyix = sum(abs.(regy) ,dims=2) .< 2
    boundary = ((rmxix + rmyix)[:]) .!= 0
    q = findall(boundary .== 1)
    regx = regx[setdiff(1:dim^2,q), :]
    regy = regy[setdiff(1:dim^2,q), :]

    s = length(q)
    bmatrix = sparse(zeros(s,dim*dim))
    for i=1:s
        v = q[i]
        bmatrix[i,v] = 1
    end
    #bmatrix = bmatrix[i,i] .= 1

    return regx,regy,bmatrix
end

function regmatrices_second(dim)
    reg1d = spdiagm(Pair(0,2*ones(dim))) + spdiagm(Pair(1,-1*ones(dim-1))) + spdiagm(Pair(-1,-1*ones(dim-1))) #;reg1d[dim,dim] = 0
    #reg1d = reg1d[1:dim-1,:]
    iden = I(dim)
    regx = kron(reg1d,iden)
    regy = kron(iden,reg1d)

    rmxix = sum(abs.(regx), dims=2) .< 4
    rmyix = sum(abs.(regy), dims=2) .< 4
    boundary = ((rmxix + rmyix)[:]) .!= 0
    q = findall(boundary .== 1)
    regx = regx[setdiff(1:dim^2,q), :]
    regy = regy[setdiff(1:dim^2,q), :]

    #q = findall(boundary .== 1)
    #bmatrix = zeros(dim*dim,dim*dim)
    #for i in q
    #     bmatrix[i,i] = 1
    #end

    h2 = sparse(zeros(2,N))
    h2[1,1] = 1; h2[2,N] = 1
    h1 = hcat([sparse(zeros(dim-2)),I(dim-2),sparse(zeros(dim-2))]...)
    bmatrix = [ kron(h2,I(dim)); kron(h1,h2) ]

    return regx,regy,bmatrix
end

function spdematrix(xs,a,b)
    dx = abs(xs[2]-xs[1])
    N = length(xs)

    M = sparse(zeros(N,N))
    for i=2:N-1
        M[i,i-1] = -a/dx^2;
        M[i,i+1] = -a/dx^2;
        M[i,i] = 2*a/dx^2;# +b;

    end

    M[1,1] = 1 # Option 1
    M[N,N] = 1
    M[N,N-1] = -1

    #M[1,:] = [2*a/dx^2;-a/dx^2; zeros(N-3);-a/dx^2]; # Option 2
    #M[N,:] = [-a/dx^2; zeros(N-3); -a/dx^2;2*a/dx^2];

    M = kron(M,I(N)) + kron(I(N),M)
    M = M + I(size(M,2))*b;

    return M
end

function nzrows(M::SparseVector)
    return unique(M.nzind)
end

function nzrows(M::SparseMatrixCSC)
    return unique(M.rowval)
end

function splitmatrix(M::SparseMatrixCSC,ixlist::Array{Array{Int64,1},1})
    N = length(ixlist)
    zl = Vector{Vector{Int64}}(undef,N)
    Ml = Vector{AbstractSparseArray}(undef,N)
    for i =1:N
        Mi = M[:,ixlist[i]]
        nz = nzrows(Mi)
        zl[i] = nz
        Ml[i] = Mi
        #push!(zl,nz)
        #push!(Ml,Mi)
    end
    return Ml,zl
end

function splitmatrix(M::SparseMatrixCSC)
    N = size(M)[2]
    zl = Vector{Vector{Int64}}(undef,N)
    Ml = Vector{AbstractSparseArray}(undef,N)
    for i =1:N
        Mi = M[:,i]
        nz = nzrows(Mi)
        zl[i] = nz
        Ml[i] = Mi
        #push!(zl,nz)
        #push!(Ml,Mi)
    end
    return Ml,zl
end

mutable struct rram
    n::Int64
    acc::Int64
    Nadapt::Int64
    C::Array{Float64,2}
    Cho::Cholesky{Float64,Array{Float64,2}}
    xm::Vector{Float64}
end

function refreshcurrent!(cache,current;mode=1,arg=nothing,ixl=nothing,nth=nothing)
    if isnothing(arg)
        if (mode ==1)
            current.Mxcurr .= cache.Mxprop
            current.noisecurr .= cache.noiseprop
        elseif (mode==2)
            # current.xcurr .= cache.xprop
            # current.Fxcurr .= cache.Fxprop
            # current.Dcurr .= cache.Dprop
            storevector!(current.xcurr, cache.xprop)
            storevector!(current.Fxcurr, cache.Fxprop)
            storevector!(current.Dcurr, cache.Dprop)

        elseif (mode==3)
            # current.xcurr .= cache.xprop
            # current.Fxcurr .= cache.Fxprop
            # current.Dxcurr .= cache.Dxprop
            # current.Dycurr .= cache.Dyprop
            # current.Dbcurr .= cache.Dbprop
            storevector!(current.xcurr, cache.xprop)
            storevector!(current.Fxcurr, cache.Fxprop)
            storevector!(current.Dxcurr, cache.Dxprop)
            storevector!(current.Dycurr, cache.Dyprop)
            storevector!(current.Dbcurr, cache.Dbprop)
        else
            error("")
        end

    else
        if (mode ==1)
            current.Mxcurr .= cache.Mxprop
            current.noisecurr .= cache.noiseprop
        elseif (mode==2)
            storevector!(view(current.xcurr,ixl[nth]), view(cache.xprop,ixl[nth]))
            storevector!(view(current.Fxcurr,arg.Fnz[nth]), view(cache.Fxprop,arg.Fnz[nth]))
            storevector!(view(current.Dcurr,arg.Dnz[nth]), view(cache.Dprop,arg.Dnz[nth]))
            #current.xcurr .= cache.xprop
            #current.Fxcurr .= cache.Fxprop
            #current.Dcurr .= cache.Dprop

        elseif (mode==3)
            # current.xcurr .= cache.xprop
            # current.Fxcurr .= cache.Fxprop
            # current.Dxcurr .= cache.Dxprop
            # current.Dycurr .= cache.Dyprop
            # current.Dbcurr .= cache.Dbprop
            storevector!(view(current.xcurr,ixl[nth]), view(cache.xprop,ixl[nth]))
            storevector!(view(current.Fxcurr,arg.Fnz[nth]), view(cache.Fxprop,arg.Fnz[nth]))
            storevector!(view(current.Dxcurr,arg.Dxnz[nth]), view(cache.Dxprop,arg.Dxnz[nth]))
            storevector!(view(current.Dycurr,arg.Dynz[nth]), view(cache.Dyprop,arg.Dynz[nth]))
            storevector!(view(current.Dbcurr,arg.Dbnz[nth]), view(cache.Dbprop,arg.Dbnz[nth]))

        else
            error("")
        end
    end
end

function mwgpartial(lpdf,gradilpdf,ixl,nth,x00,args,cache,current,tuning;Nruns=5,Np=3,T=5.0,mode=1,extra=nothing)
    ix = ixl[nth]
    x0 = repeat(x00[ix],Np)
    x = x0# copy(x0)
    Nall = length(x)

    xnew = similar(x)

    logpdf_can(xp) = lpdf(ixl,nth,xp,args,cache,current)
    #logpdf_can(xp) = lpdf(partialreplace(x00,xp,ix),args,cache)

    logpdf(x) = logpdf_can(x)

    ll = logpdf(x)

    for _ = 1:Nruns

        tuning.n = tuning.n + 1

        V=randn(Nall)
        xnew .= x+ tuning.Cho.L*V

        lnew = logpdf(xnew)

        U = rand()
        apr = min(1.0,exp(-ll + lnew))

        if (U < apr)
            ll = lnew
            x .= xnew
            tuning.acc = tuning.acc +1
        end

        if (tuning.n <= tuning.Nadapt)
            xmp = copy(tuning.xm)
            tuning.xm .= tuning.xm + (x-tuning.xm)/tuning.n
            xm = tuning.xm
            tuning.C .= ((tuning.n-1)*tuning.C + (x-xm)*(x-xmp)')/tuning.n
            tuning.Cho = cholesky(Hermitian(2.38^2/Nall*tuning.C + 1e-11*I(Nall)))

            # atarg = 0.3
            # V = V/norm(V)
            # z = sqrt(tuning.n^(-0.7) *abs(tuning.acc/tuning.n - atarg)) *  tuning.Cho.L*V
            # #println(z)
            # if (tuning.acc/tuning.n >= atarg)
            #     lowrankupdate!(tuning.Cho, z)

            # else
            #     lowrankdowndate!(tuning.Cho, z)
            # end

        end

    end

    xret =  x#reweight(x,Np,logpdf_can,T)
    logpdf_can(xret)
    refreshcurrent!(cache,current;mode=mode,arg=args,ixl=ixl,nth=nth)

    return xret

end

function repatpartial(lpdf,placeholder,ixl,nth,x0,args,cache,current,tuning;mode=1,Np=1,Nruns=1,T=-Inf,extra=copy(x0))::Vector{Float64}
    ix = ixl[nth]
    x = copy(x0[ix])
    dim = length(x)

    logpdf(xp) = lpdf(ixl,nth,xp,args,cache,current)
    #logpdf(xp) = lpdf(partialreplace(x0,xp,ixl[nth]),args,cache)
    density = logpdf(x)

    z = copy(extra[ix])
    u = randn(dim)
    densityz = logpdf(z)

    for i = 1:Nruns
        R = tuning.Cho.L/sqrt(2)
		u .= randn(dim)
		xtilde = x + R * u
		densitytilde = logpdf(xtilde)
		while (log(rand()) >=  density - densitytilde)
			u .= randn(dim)
			xtilde .= x + R * u
			densitytilde = logpdf(xtilde)
		end

		u .= randn(dim)
		xstar = xtilde + R * u
		densitystar = logpdf(xstar)

		while (log(rand()) >=  densitystar - densitytilde)
			u .= randn(dim)
			xstar .= x + R * u
			densitystar = logpdf(xstar)
		end

		u .= randn(dim)
		zstar = xstar + R * u
		densityzstar = logpdf(zstar)

		while (log(rand()) >=  densitystar - densityzstar)
			u .= randn(dim)
			zstar .= xstar + R * u
			densityzstar = logpdf(zstar)
		end

        ratio = densitystar + min(0,density-densityz) - density - min(0,densitystar-densityzstar)
        if (log(rand()) <= ratio)
            x .= xstar
            z .= zstar
			extra[ix] = zstar
			density = densitystar
			densityz = densityzstar

        end
        tuning.n = tuning.n + 1

        if (tuning.n <= tuning.Nadapt)
            xmp = copy(tuning.xm)
            tuning.xm .= tuning.xm + (x-tuning.xm)/tuning.n
            xm = tuning.xm
            tuning.C .= ((tuning.n-1)*tuning.C + (x-xm)*(x-xmp)')/tuning.n
            tuning.Cho = cholesky(Hermitian(2.38^2/dim*tuning.C + 1e-11*I(dim)))
        end

    end
    logpdf(x)
    refreshcurrent!(cache,current;mode=mode,arg=args,nth=nth,ixl=ixl)
    return x
end

function totalsample(lpdf,glpdf,x0,meas,arg,ixl;N=1,Nadapt=1000,Nruns=5,thinning=10,algo=repatpartial,mode=2,Np=2,T=2.0)
    x = copy(x0)
    Npar = length(x0)
    chain = zeros(Npar, Int(floor(N / thinning)))
    Nix = length(ixl)

    if (mode==2)
        cache=(xprop=copy(x0),Fxprop=(arg.F*x0),Dprop=(arg.D*x0), residual=similar(arg.y),gradiprop=similar(x0))
        current = (Fxcurr=(arg.F*x0),Dcurr=(arg.D*x0), xcurr = copy(x0))
    elseif (mode==3)
        cache =(xprop=copy(x0),Fxprop=(arg.F*x0),Dxprop=(arg.Dx*x0),Dyprop=(arg.Dy*x0), Dbprop=(arg.Db*x0),  residual=arg.F*x0-arg.y,gradiprop=similar(x0))
        current = (xcurr = copy(x0), Dxcurr = arg.Dx*x0, Dycurr = arg.Dy*x0, Dbcurr = arg.Db*x0, Fxcurr = F*x0)
    else
        error("")
    end

    if algo == repatpartial
        extra = copy(x0)
        tuning = Vector{rram}(undef,Nix)
        for k = 1:Nix
            n = length(ixl[k])
            C = 0.1*Array(I(n))
            Cho = cholesky(C)
            xm = copy(x0[ixl[k]])
            tuning[k] = rram(1,0,Nadapt,C,Cho,xm)
        end

    elseif (algo==mwgpartial)
        extra = []
        tuning = Vector{rram}(undef,Nix)
        #tuning = Vector{dahmc2}(undef,Nix)
        #tuning = Vector{ahmc}(undef,Nix)
        for k = 1:Nix

            n = length(ixl[k])*Np
            C = 0.1*Array(I(n))
            Cho = cholesky(C)
            xm = repeat(copy(x0[ixl[k]]),Np)
            tuning[k] = rram(1,0,Nadapt,C,Cho,xm)

        end
    end

    pb = ProgressBar(1:N)

    for i in pb
        for k = 1:Nix
            x[ixl[k]] = algo(lpdf,glpdf,ixl,k,x,arg,cache,current,tuning[k];mode=mode,Nruns=Nruns,extra=extra,Np=Np,T=T)#lpdf,ix,x0,args,cache,current,tuning
            if (i % thinning == 0)
                chain[ixl[k], Int(i / thinning)] = x[ixl[k]]
            end

        end
    end

    return chain,tuning
end

function savechain(c,prefix)
    p = pwd()
    f = p*"/chains/"*prefix*"_"*Dates.format(Dates.now(), "yyyy_mm_dd_HH_MM_SS")*".mat"
    matwrite(f,Dict("chain"=>c))

end

cw = 150.0
kernel(xi,xj,yi,yj;constant=cw) = constant/pi*exp(-constant*((xi-xj)^2 + (yi-yj)^2) )
tf(x,y) =  1.0*exp.(-20*sqrt.((x .- 0.3).^2 .+ (y .- 0.3).^2)) + 1*((y-x) .< 0.7).*((y-x) .>= -0.7).*((-y-x) .<= 0.8).*((-y-x) .>= 0.4).*(((-x+y) .+ 0.7)/1.39)  + 1*(-x+y .< 1).*(-x+y .>= 0.8).*(abs.(x) .<= 1).*(abs.(y) .<= 1) + ( 4*0.25 .- 4*sqrt.((x.-0.5).^2+(y.+0.6).^2)).*(sqrt.((x.-0.5).^2+(y.+0.6).^2) .<= 0.25);

Random.seed!(1)

noisevar = 0.01^2

Nbig = 300
dimbig = Nbig
N = 180
dim = N
Nmeas = 100
dimmeas = Nmeas

xsbig = -1+1/dimbig:2/dimbig:1-1/(dimbig)
ysbig = -1+1/dimbig:2/dimbig:1-1/(dimbig)

xs = -1+1/dim:2/dim:1-1/(dim)
ys = -1+1/dim:2/dim:1-1/(dim)

xibig = linspace(-1,1,dimbig)
Ybig,Xbig = meshgrid(xibig,xibig)

ximeas = linspace(-1,1,dimmeas)
Ymeas,Xmeas = meshgrid(ximeas,ximeas)
Y,X = meshgrid(xs,xs)

Zbig = tf(Xbig,Ybig)
@time Fbig = measurementmatrix2d(Xbig,Ybig,Xmeas,Ymeas,kernel,constant=cw); # Theory matrix for measurements.
@time F = measurementmatrix2d(X,Y,Xmeas,Ymeas,kernel,constant=cw); # Theory matrix for inversion.

zx = Fbig*Zbig[:]
mebig = reshape(zx,(dimmeas,dimmeas))
y = mebig[:] + randn(Nmeas*Nmeas,)*sqrt(noisevar)

xp = randn(N*N)
ixl = divide(N*N,1)

Dx,Dy,Db = regmatrices_first(N);
Dxs, Dxnz = splitmatrix(Dx,ixl); Dys, Dynz = splitmatrix(Dy,ixl); Dbs, Dbnz = splitmatrix(Db,ixl); #Fs, Fnz = splitmatrix(F,ixl)
S = spdematrix(xs,0.001,0.1)
D = [Dx;Dy;Db]
Ds, Dnz = splitmatrix(D,ixl)
Dinz = [union(Dxnz[i],Dynz[i]) for i = 1:N^2]
argi  = (Dinz = Dinz, Dxs = Dxs, Dxnz = Dxnz , Dys = Dys, Dynz = Dynz, Dbs = Dbs, Dbnz = Dbnz, Dx = Dx, Dy = Dy, D=D,noisesigma = sqrt(noisevar), scale = 0.03, bscale = 0.03, Db = Db, Ds = Ds, Dnz = Dnz )
argis  = ( D=S, F = F,noisesigma = sqrt(noisevar), scale = 0.3, y = y, bscale = 1.0 )

cacheiso =(xprop=zeros(N*N),Dxprop=zeros(size(Dx)[1]),Dyprop=zeros(size(Dy)[1]),Dbprop=zeros(size(Db)[1]),gradiprop=zeros(N*N))
cached =(xprop=zeros(N*N),Dprop=zeros(size(D)[1]),gradiprop=zeros(N*N))
caches =(xprop=zeros(N*N),Fxprop=zeros(Nmeas^2),Dprop=zeros(size(S)[1]), residual=similar(y),gradiprop=zeros(N*N))

currentiso = (xcurr = copy(xp), Dxcurr = Dx*xp, Dycurr = Dy*xp, Dbcurr = Db*xp, Fxcurr = F*xp)
currentaniso = (xcurr = copy(xp), Dcurr = D*xp, Fxcurr = F*xp)

target1iso(x) = -logpisodiff(x,argi,cacheiso)# -logpdiff(w,diff1arg,cached)##
target1isograd(x) = -logpisodiffgradi(x,argi,cacheiso;both=false)[2]#  -logpdiffgradi(w,diff1arg,cached;both=false)[2]#
res = Optim.optimize(target1iso, target1isograd, 0*randn(N*N,),Optim.LBFGS(), Optim.Options(allow_f_increases=true,show_trace=true,iterations=700); inplace = false)
MAP_diff1iso = res.minimizer
imshow(reshape(MAP_diff1iso,(N,N))); colorbar(); clim(-0.1,1.1)

target1(x) = -logpdiff(x,argi,cached)# -logpdiff(w,diff1arg,cached)##
target1grad(x) = -logpdiffgradi(x,argi,cached;both=false)[2]#  -logpdiffgradi(w,diff1arg,cached;both=false)[2]#
res = Optim.optimize(target1, target1grad, 0*randn(N*N,),Optim.LBFGS(), Optim.Options(allow_f_increases=true,show_trace=true,iterations=700); inplace = false)
MAP_diff1 = res.minimizer
imshow(reshape(MAP_diff1,(N,N))); colorbar(); clim(-0.1,1.2)

targets(x) = -logpdiff(x,argis,caches)# -logpdiff(w,diff1arg,cached)##
targetsgrad(x) = -logpdiffgradi(x,argis,caches;both=false)[2]#  -logpdiffgradi(w,diff1arg,cached;both=false)[2]#
res = Optim.optimize(targets, targetsgrad, 0*randn(N*N,),Optim.LBFGS(), Optim.Options(allow_f_increases=true,show_trace=true,iterations=700); inplace = false)
MAP_s = res.minimizer
imshow(reshape(MAP_s,(N,N))); colorbar(); clim(-0.5,12)

function NUTS_sample(Niter,arg,cac)
    n_samples, n_adapts = 2000, 1000
    ss = Vector{Any}(undef,Niter)
    st = Vector{Any}(undef,Niter)
    Np = 1
    T = 1
    for i = (1:Niter) .+ seed_shift

        Ntot = N^2
        Random.seed!(i)
        println(i)
        initial_θ = copy(MAP_diff1)
        global c

        targetdiff(w) = -logpdiff(x,arg,cac)
        targetdiffgrad(w) = logpdiffgradi(w,arg,cac;both=true)

        metric = DiagEuclideanMetric(Ntot)# DenseEuclideanMetric(Ntot)# DiagEuclideanMetric(Ntot)
        hamiltonian = Hamiltonian(metric, targetdiff, targetdiffgrad)

        initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
        integrator =  Leapfrog(initial_ϵ)#  TemperedLeapfrog(initial_ϵ, 1.05)# Leapfrog(initial_ϵ)
        proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator,8,1000.0)
        adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
        c, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; progress=true)
        #ss[i] = hcat(c...)
        c = hcat(c...)
        # ss[i] = c
        #st[i] = stats
        pr="NUTS_"*string("anisodiff1d_prior")
        savechain(c,pr)

    end
    return ss, st
end
ch,b = NUTS_sample(1,argi,cached)

algo =  repatpartial#mwgpartial#
grfunk = nothing#
funk =  logpdiffpartial
par = argi
mode = 2# 1
#x000  = randn(N*N,)

for t=1:0
    Random.seed!(t)
    xz = copy(MAP_diff1)

    Nsamples = 100; Nadapt = 10; thinning = 10
    Np = 1; Nruns = 1; T = 1.0

    c,tun=totalsample(funk,grfunk,xz,y,par,ixl;N=Nsamples,Nadapt=Nadapt,thinning=thinning,algo=algo,mode=mode,Np=Np,T=T,Nruns=Nruns)
    savechain(c,"anisodiff1d_normal_"*string(nameof(algo))*"_"*string(nameof(funk))*"_N"*string(Nsamples)*"_Nadapt"*string(Nadapt)*"_thin"*string(thinning)*"_Np"*string(Np)*"_Nruns"*string(Nruns)*"_T"*string(T))
    println(effective_sample_size(c[1,:]))
    global c
    global tun
end

max_stars_repo_name: suurj/markovcauchy

max_issues_repo_name: suurj/markovcauchy

max_forks_repo_name: suurj/markovcauchy