text:
__precompile__()
module Primes

if VERSION >= v"0.5.0-dev+4340"
    if isdefined(Base,:isprime)
        import Base: isprime, primes, primesmask, factor
    else
        export isprime, primes, primesmask, factor
    end
end

# Primes generating functions
#     https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes
#     https://en.wikipedia.org/wiki/Wheel_factorization
#     http://primesieve.org
#     Jonathan Sorenson, "An analysis of two prime number sieves", Computer Science Technical Report Vol. 1028, 1991
const wheel         = [4,  2,  4,  2,  4,  6,  2,  6]
const wheel_primes  = [7, 11, 13, 17, 19, 23, 29, 31]
const wheel_indices = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 7, 7]

@inline function wheel_index(n)
    (d, r) = divrem(n - 1, 30)
    8d + wheel_indices[r+2]
end
@inline function wheel_prime(n)
    (d, r) = ((n - 1) >>> 3, (n - 1) & 7)
    30d + wheel_primes[r+1]
end

function _primesmask(limit::Int)
    limit < 7 && throw(ArgumentError("limit must be at least 7, got $limit"))
    n = wheel_index(limit)
    m = wheel_prime(n)
    sieve = ones(Bool, n)
    @inbounds for i = 1:wheel_index(isqrt(limit))
        if sieve[i]; p = wheel_prime(i)
            q = p * p
            j = (i - 1) & 7 + 1
            while q ≤ m
                sieve[wheel_index(q)] = false
                q = q + wheel[j] * p
                j = j & 7 + 1
            end
        end
    end
    return sieve
end

function _primesmask(lo::Int, hi::Int)
    7 ≤ lo ≤ hi || throw(ArgumentError("the condition 7 ≤ lo ≤ hi must be met"))
    lo == 7 && return _primesmask(hi)
    wlo, whi = wheel_index(lo - 1), wheel_index(hi)
    m = wheel_prime(whi)
    sieve = ones(Bool, whi - wlo)
    hi < 49 && return sieve
    small_sieve = _primesmask(isqrt(hi))
    @inbounds for i = 1:length(small_sieve)  # don't use eachindex here
        if small_sieve[i]; p = wheel_prime(i)
            j = wheel_index(2 * div(lo - p - 1, 2p) + 1)
            q = p * wheel_prime(j + 1); j = j & 7 + 1
            while q ≤ m
                sieve[wheel_index(q)-wlo] = false
                q = q + wheel[j] * p
                j = j & 7 + 1
            end
        end
    end
    return sieve
end

"""
    primesmask([lo,] hi)

Returns a prime sieve, as a `BitArray`, of the positive integers (from `lo`, if specified)
up to `hi`. Useful when working with either primes or composite numbers.
"""
function primesmask(lo::Int, hi::Int)
    0 < lo ≤ hi || throw(ArgumentError("the condition 0 < lo ≤ hi must be met"))
    sieve = falses(hi - lo + 1)
    lo ≤ 2 ≤ hi && (sieve[3-lo] = true)
    lo ≤ 3 ≤ hi && (sieve[4-lo] = true)
    lo ≤ 5 ≤ hi && (sieve[6-lo] = true)
    hi < 7 && return sieve
    wheel_sieve = _primesmask(max(7, lo), hi)
    lsi = lo - 1
    lwi = wheel_index(lsi)
    @inbounds for i = 1:length(wheel_sieve)   # don't use eachindex here
        sieve[wheel_prime(i + lwi) - lsi] = wheel_sieve[i]
    end
    return sieve
end
primesmask{T<:Integer}(lo::T, hi::T) = lo <= hi <= typemax(Int) ? primesmask(Int(lo), Int(hi)) :
    throw(ArgumentError("both endpoints of the interval to sieve must be ≤ $(typemax(Int)), got $lo and $hi"))

primesmask(limit::Int) = primesmask(1, limit)
primesmask(n::Integer) = n <= typemax(Int) ? primesmask(Int(n)) :
    throw(ArgumentError("requested number of primes must be ≤ $(typemax(Int)), got $n"))

"""
    primes([lo,] hi)

Returns a collection of the prime numbers (from `lo`, if specified) up to `hi`.
"""
function primes(lo::Int, hi::Int)
    lo ≤ hi || throw(ArgumentError("the condition lo ≤ hi must be met"))
    list = Int[]
    lo ≤ 2 ≤ hi && push!(list, 2)
    lo ≤ 3 ≤ hi && push!(list, 3)
    lo ≤ 5 ≤ hi && push!(list, 5)
    hi < 7 && return list
    sizehint!(list, floor(Int, hi / log(hi)))
    sieve = _primesmask(max(7, lo), hi)
    lwi = wheel_index(lo - 1)
    @inbounds for i = 1:length(sieve)   # don't use eachindex here
        sieve[i] && push!(list, wheel_prime(i + lwi))
    end
    return list
end
primes(n::Int) = primes(1, n)

const PRIMES = primes(2^16)

"""
    isprime(x::Integer) -> Bool

Returns `true` if `x` is prime, and `false` otherwise.

```jldoctest
julia> isprime(3)
true
```
"""
function isprime(n::Integer)
    # Small precomputed primes + Miller-Rabin for primality testing:
    #     https://en.wikipedia.org/wiki/Miller–Rabin_primality_test
    #
    (n < 3 || iseven(n)) && return n == 2
    n <= 2^16 && return PRIMES[searchsortedlast(PRIMES,n)] == n
    s = trailing_zeros(n-1)
    d = (n-1) >>> s
    for a in witnesses(n)
        x = powermod(a,d,n)
        x == 1 && continue
        t = s
        while x != n-1
            (t-=1) <= 0 && return false
            x = oftype(n, widemul(x,x) % n)
            x == 1 && return false
        end
    end
    return true
end

"""
    isprime(x::BigInt, [reps = 25]) -> Bool

Probabilistic primality test. Returns `true` if `x` is prime with high probability (pseudoprime);
and `false` if `x` is composite (not prime). The false positive rate is about `0.25^reps`.
`reps = 25` is considered safe for cryptographic applications (Knuth, Seminumerical Algorithms).

```jldoctest
julia> isprime(big(3))
true
```
"""
isprime(x::BigInt, reps=25) = ccall((:__gmpz_probab_prime_p,:libgmp), Cint, (Ptr{BigInt}, Cint), &x, reps) > 0

# Miller-Rabin witness choices based on:
#     http://mathoverflow.net/questions/101922/smallest-collection-of-bases-for-prime-testing-of-64-bit-numbers
#     http://primes.utm.edu/prove/merged.html
#     http://miller-rabin.appspot.com
#
witnesses(n::Union{UInt8,Int8,UInt16,Int16}) = (2,3)
witnesses(n::Union{UInt32,Int32}) = n < 1373653 ? (2,3) : (2,7,61)
witnesses(n::Union{UInt64,Int64}) =
        n < 1373653         ? (2,3) :
        n < 4759123141      ? (2,7,61) :
        n < 2152302898747   ? (2,3,5,7,11) :
        n < 3474749660383   ? (2,3,5,7,11,13) :
                              (2,325,9375,28178,450775,9780504,1795265022)

isprime(n::UInt128) =
    n <= typemax(UInt64) ? isprime(UInt64(n)) : isprime(BigInt(n))
isprime(n::Int128) = n < 2 ? false :
    n <= typemax(Int64)  ? isprime(Int64(n))  : isprime(BigInt(n))

# Trial division of small (< 2^16) precomputed primes +
# Pollard rho's algorithm with Richard P. Brent optimizations
#     https://en.wikipedia.org/wiki/Trial_division
#     https://en.wikipedia.org/wiki/Pollard%27s_rho_algorithm
#     http://maths-people.anu.edu.au/~brent/pub/pub051.html
#
"""
    factor(n) -> Dict

Compute the prime factorization of an integer `n`. Returns a dictionary. The keys of the
dictionary correspond to the factors, and hence are of the same type as `n`. The value
associated with each key indicates the number of times the factor appears in the
factorization.

```jldoctest
julia> factor(100) # == 2*2*5*5
Dict{Int64,Int64} with 2 entries:
  2 => 2
  5 => 2
```
"""
function factor{T<:Integer}(n::T)
    0 < n || throw(ArgumentError("number to be factored must be ≥ 0, got $n"))
    h = Dict{T,Int}()
    n == 1 && return h
    isprime(n) && (h[n] = 1; return h)
    local p::T
    for p in PRIMES
        if n % p == 0
            h[p] = get(h,p,0)+1
            n = div(n,p)
            while n % p == 0
                h[p] = get(h,p,0)+1
                n = div(n,p)
            end
            n == 1 && return h
            isprime(n) && (h[n] = 1; return h)
        end
    end
    T <: BigInt || widemul(n-1,n-1) <= typemax(n) ? pollardfactors!(n, h) : pollardfactors!(widen(n), h)
end

function pollardfactors!{T<:Integer,K<:Integer}(n::T, h::Dict{K,Int})
    while true
        local c::T = rand(1:(n-1)), G::T = 1, r::K = 1, y::T = rand(0:(n-1)), m::K = 1900
        local ys::T, q::T = 1, x::T
        while c == n - 2
            c = rand(1:(n-1))
        end
        while G == 1
            x = y
            for i in 1:r
                y = (y*y)%n
                y = (y+c)%n
            end
            local k::K = 0
            G = 1
            while k < r && G == 1
                for i in 1:(m>(r-k)?(r-k):m)
                    ys = y
                    y = (y*y)%n
                    y = (y+c)%n
                    q = (q*(x>y?x-y:y-x))%n
                end
                G = gcd(q,n)
                k = k + m
            end
            r = 2 * r
        end
        G == n && (G = 1)
        while G == 1
            ys = (ys*ys)%n
            ys = (ys+c)%n
            G = gcd(x>ys?x-ys:ys-x,n)
        end
        if G != n
            isprime(G) ? h[G] = get(h,G,0) + 1 : pollardfactors!(G,h)
            G2 = div(n,G)
            isprime(G2) ? h[G2] = get(h,G2,0) + 1 : pollardfactors!(G2,h)
            return h
        end
    end
end

end # module


text:
# routines for assembling a coarse model

"""
`build_coarse_transition_matrix`: Contruct a transition matrix amongst the bins (serial version).

### Arguments
* `mutation!` - an in place mutation function
* `bin_id` - bin identification function
* `x0_vals` - an array of starting values
* `bin0_vals` - an array of the bins corresponding to `x0_vals`
* `n_bins` - total number of bins
* `n_trials` - number of independent trials for each x0 starting value
"""
function build_coarse_transition_matrix(mutation!, bin_id, x0_vals, n_bins, n_trials)

   n_x0 = length(x0_vals);
   X = similar(x0_vals[1]);
   row_vals = zeros(Int, n_x0*n_trials);
   col_vals = zeros(Int,n_x0*n_trials);

   for k in 1:n_trials
      for l in 1:n_x0
         X .= deepcopy(x0_vals[l]);
         row_vals[n_x0*(k-1) + l] = bin_id(X);
         mutation!(X);
         col_vals[n_x0*(k-1) + l] = bin_id(X);
      end
   end

   K = sparse(row_vals,col_vals,ones(size(row_vals)),n_bins, n_bins);
   for i in 1:n_bins
      rowZ = sum(K[i,:]);
      @. K[i,:] = K[i,:]/rowZ;
   end

   return K
end

"""
`build_coarse_vectors`: Assemble the conditional expectation and 1- step
variance approximations on a coarser model, given the transition matrix, `K̃`,
and a coarse scale QoI function, `f̃`.

### Arguments
* `n_we_steps` - number of WE steps
* `K̃` - coarse scale transition matrix
* `f̃` - quantity of interest vector on the bin space
"""
function build_coarse_vectors(n_we_steps, K̃, f̃)
   n_bins = length(f̃);
   ṽ²_vals = [zeros(n_bins) for j in 1:n_we_steps];
   h̃_vals = [zeros(n_bins) for j in 1:n_we_steps+1];
   h̃ = deepcopy(f̃);
   K̃h̃ =similar(f̃);
   # vector of 1's
   e = ones(n_bins);
   @. h̃_vals[end] = h̃;

   # Use t+1 in indexing since Julia arrays start at 1
   for t in n_we_steps-1:-1:0

      K̃h̃ .= K̃ * h̃;
      # compute variance row by row
      for p in 1:n_bins
           ṽ²_vals[t+1][p] = K̃[p,:]⋅((h̃ .- e * (K̃h̃[p])).^2)
      end
      # update for next iterate
      @. h̃ = K̃h̃;
      @. h̃_vals[t+1] = h̃;
   end

   return h̃_vals, ṽ²_vals
end

"""
`build_coarse_poisson`: Construct the solution to the Poisson problem and the
1-step variance approximation on the coarser model given the transition matrix,
`K̃`, and a coarse scale QoI function, `f̃`.

### Arguments
* `K̃` - coarse scale transition matrix
* `f̃` - quantity of interest vector on the bin space
"""
function build_coarse_poisson(K̃, f̃)

   n_bins = length(f̃);
   _, evecs  = eigs(transpose(K̃),nev=1);
   μ̃ = real.(evecs[:,1]);μ̃ .= μ̃/ sum(μ̃); #normalize the invariant measure

   # compute solution to Poisson
   h̃ = gmres(I - K̃, f̃ .- μ̃ ⋅ f̃);
   # normalize
   h̃ = h̃ .- μ̃⋅h̃;
   K̃h̃ = K̃*h̃;
   e = ones(n_bins);
   ṽ² =zeros(n_bins);
   for p in 1:n_bins
      ṽ²[p] = K̃[p,:]⋅((h̃ .- e * (K̃h̃[p])).^2)
   end

   return h̃, ṽ²
end

ext: jl


text:
begin
	# Exponential Search

# Julia code to Exponential search x in arr[].
# If x is present then return its location,
# otherwise return -1
import Base:length,show,AbstractArray
arr = [2, 3, 4, 10, 40]
x = 48
n = length(arr)
function binarysearch(arr::AbstractArray{},l,r,x)
    if r >= l
		value=r-l//2
        mid = round(Int,l + (r - l) // 2)
        # If element is present at the middle itself
        if arr[mid] == x
            return mid
        # If element is smaller than mid, then it
        # can only be present in left subarray
		elseif arr[mid] > x
            return binarysearch(arr, l, mid-1, x)
        # Else the element can only be present
        # in right subarray
        else
            return binarysearch(arr, mid + 1, r, x)
		end
    else
        # Element is not present in the array
        return -1
	end
end

# Returns the position of first
# occurrence of x in array
function exponentialsearch(arr::AbstractArray{}, n, x)
    # IF x is present at first
    # location itself
    if arr[1] == x
        return 1
	end
    # Find range for binary search
    # j by repeated doubling
    i = 2
    while i < n+1 && arr[i] <= x
        i = i * 2
	end
	val=round(Int,i/2)
    # Call binary search for the found range
    return binarysearch( arr,val, min(i, n), x)
end
global result=-1

result=exponentialsearch(arr, n, x)

println("Result after function execution:",result)
if result == -1
    println("Element:",x," is not present in array",arr)
else
    println("Element:",x," is present at index:", result," in array",arr)
end
end

max_stars_repo_name: Learning-Julia/Algorithms

max_issues_repo_name: Learning-Julia/Algorithms

max_forks_repo_name: Learning-Julia/Algorithms


text:
### A Pluto.jl notebook ###
# v0.19.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 84388088-c42d-11ec-249b-4db2ebc98d5b
begin
	using  CairoMakie, ForwardDiff,PlutoUI,ForwardDiff
	import PlutoUI:Slider

end

# ╔═╡ 67be87c3-ccaa-4c54-94c0-be6a24471043
function phaseline(xr,yr,us,vs)
	     strength = vec(sqrt.(us .^2 .+ vs .^2))
	     cmap = :gnuplot
	    fig = Figure(resolution = (600, 400))
	    ax = Axis(fig[1,1], xlabel = "x", ylabel = "y", aspect = DataAspect())
	    arrows!(ax, xr, yr, us, vs, arrowsize = 10, lengthscale = 0.1,
	        arrowcolor = strength, linecolor = strength, colormap = cmap)
	    Colorbar(fig[1,2], limits =(minimum(strength), maximum(strength)),
	        nsteps =100, colormap = cmap, ticksize=15, width = 15, tickalign=1)
	    limits!(ax, -4,4,-4,4)
	    colsize!(fig.layout, 1, Aspect(1, 1.0))
   return fig
end

# ╔═╡ 3f51d43c-921b-407a-96f6-a1984c88ab79
function vectorfileds(xr,yr,us,vs) # 和上面的函数其实一样
	     strength = vec(sqrt.(us .^2 .+ vs .^2))
	     cmap = :gnuplot
	    fig = Figure(resolution = (600, 400))
	    ax = Axis(fig[1,1], xlabel = "x", ylabel = "y", aspect = DataAspect())
	    arrows!(ax, xr, yr, us, vs, arrowsize = 10, lengthscale = 0.5,
	        arrowcolor = strength, linecolor = strength, colormap = cmap)
	    Colorbar(fig[1,2], limits =(minimum(strength), maximum(strength)),
	        nsteps =100, colormap = cmap, ticksize=15, width = 15, tickalign=1)
	    limits!(ax, -4,4,-4,4)
	    colsize!(fig.layout, 1, Aspect(1, 1.0))
   return fig
end

# ╔═╡ 923f3a66-fc41-4672-b0da-40b7c720c593
function vectorfileds2(xr,yr,us,vs) # 去掉强度

	     cmap = :gnuplot
	    fig = Figure(resolution = (600, 400))
	    ax = Axis(fig[1,1], xlabel = "x", ylabel = "y", aspect = DataAspect())
	    arrows!(ax, xr, yr, us, vs, arrowsize = 10, lengthscale = 0.5, colormap = cmap)

	    limits!(ax, -4,4,-4,4)
	    colsize!(fig.layout, 1, Aspect(1, 1.0))
   return fig
end

# ╔═╡ 36d1a402-de17-432e-b2dd-ec8880eb15c7
begin
range=-3:3
@bind p Slider(range; default=0, show_value=true)
end

# ╔═╡ 98767b85-6371-45c6-8822-48cdb25234a1
let
	f(p)=p[1]^2+2*p[2]^2
	#f1(p)=p[1]  #返回数组第一个元素,如果直接返回 0,就得到 phaseline
	fx(p)=0
	fy(p)=p[2]  #返回数组第二个元素
    xr = LinRange(p+0.0001, p+0.00011, 2)  #取极小区间
    yr = LinRange(-3, 3, 20)
    ps = [[x,y] for x in xr, y in yr]
    temp = [ForwardDiff.gradient(f, p) for p in ps] # 中间变量用于计算梯度向量
	us=fx.(temp)
	vs=fy.(temp)
    fig=phaseline(xr,yr,us,vs)

end

# ╔═╡ b6efecfe-5daa-419e-8a76-1bc0e77bc0e5
let
	f(p)=p[1]^2+2*p[2]^2
	#f1(p)=p[1]  #返回数组第一个元素,如果直接返回 0,就得到 phaseline
	fx(p)=0
	fy(p)=p[2]  #返回数组第二个元素
    xr = LinRange(-5, 5, 20)
    yr = LinRange(-5, 5, 20)
    # 中间变量用于计算梯度向量
	 ps = [[x,y] for x in xr, y in yr]
	temp = [ForwardDiff.gradient(f, p) for p in ps]
	us=fx.(temp)
	vs=fy.(temp)
    fig=phaseline(xr,yr,us,vs)

end

# ╔═╡ 3a89db02-cf60-40d2-9066-5c399c11899e
let

    xspan = LinRange(-4, 4, 21)  #取极小区间
    yspan = LinRange(-4, 4, 21)
    us = [0  for x in xspan, y in yspan]  # 梯度要求导数
    vs = [cos(y)  for x in xspan, y in yspan]
    fig=vectorfileds(xspan,yspan,us,vs)
end

# ╔═╡ a9576e89-4680-4c67-bb46-ca14de29d17e
let
	f(p)=sin(p[2])
	#f1(p)=p[1]  #返回数组第一个元素,如果直接返回 0,就得到 phaseline
	fx(p)=p[1]
	fy(p)=p[2]  #返回数组第二个元素
    xr=yr = LinRange(-4,4, 21)  #取极小区间
    ps = [[x,y] for x in xr, y in yr]
    temp = [ForwardDiff.gradient(f, p) for p in ps] # 中间变量用于计算梯度向量
	us=fx.(temp)
	vs=fy.(temp)
    fig=vectorfileds2(xr,yr,us,vs)

end

# ╔═╡ Cell order:
# ╟─84388088-c42d-11ec-249b-4db2ebc98d5b
# ╠═67be87c3-ccaa-4c54-94c0-be6a24471043
# ╠═3f51d43c-921b-407a-96f6-a1984c88ab79
# ╠═923f3a66-fc41-4672-b0da-40b7c720c593
# ╟─36d1a402-de17-432e-b2dd-ec8880eb15c7
# ╠═98767b85-6371-45c6-8822-48cdb25234a1
# ╠═b6efecfe-5daa-419e-8a76-1bc0e77bc0e5
# ╠═3a89db02-cf60-40d2-9066-5c399c11899e
# ╠═a9576e89-4680-4c67-bb46-ca14de29d17e

max_stars_repo_name: phpsmarter/plutonote

max_issues_repo_name: phpsmarter/plutonote

max_forks_repo_name: phpsmarter/plutonote


text:
"""
	friction
"""
struct FrictionConstraints <: Constraints
    n
    ineq
    n_stage
end

function friction!(c, x⁺, u, model)

    v = x⁺[3]
    β = u[2:3]
    ψ = u[4]
    η = u[5:6]
    s = u[7]

    n = (model.mc + model.mp) * model.g

    c[1] = v + ψ - η[1]
    c[2] = -v + ψ - η[2]
    c[3] = model.μ * n - sum(β)
    c[4] = s - ψ * (model.μ * n - sum(β))
    c[5] = s - β' * η

    return nothing
end

function constraints!(c, Z, con::FrictionConstraints, model, idx, h, T)
    for t = 1:T-1
        x⁺ = view(Z, idx.x[t + 1])
        u = view(Z, idx.u[t])
        friction!(view(c, (t - 1) * con.n_stage .+ (1:con.n_stage)), x⁺, u, model)
    end

    nothing
end

function constraints_jacobian!(∇c, Z, con::FrictionConstraints, model, idx, h, T)
    n = model.n
    m = model.m

    c_tmp = zeros(con.n_stage)
    shift = 0

    for t = 1:T-1
        x⁺ = view(Z, idx.x[t + 1])
        u = view(Z, idx.u[t])

        cx!(c, z) = friction!(c, z, u, model)
        cu!(c, z) = friction!(c, x⁺, z, model)

        r_idx = (t - 1) * con.n_stage .+ (1:con.n_stage)

        c_idx = idx.x[t + 1]
        len = length(r_idx) * length(c_idx)
        ForwardDiff.jacobian!(reshape(view(∇c, shift .+ (1:len)), con.n_stage, n),
            cx!, c_tmp, x⁺)
        shift += len

        c_idx = idx.u[t]
        len = length(r_idx) * length(c_idx)
        ForwardDiff.jacobian!(reshape(view(∇c, shift .+ (1:len)), con.n_stage, m),
            cu!, c_tmp, u)
        shift += len
    end

    return nothing
end

function constraints_sparsity(con::FrictionConstraints, model, idx, T;
    shift_row = 0, shift_col = 0)
    row = []
    col = []

    for t = 1:T-1
        r_idx = shift_row + (t - 1) * con.n_stage .+ (1:con.n_stage)

        c_idx = shift_col .+ idx.x[t + 1]
        row_col!(row, col, r_idx, c_idx)

        c_idx = shift_col .+ idx.u[t]
        row_col!(row, col, r_idx, c_idx)
    end

    return collect(zip(row, col))
end

max_issues_repo_name: zacmanchester/DirectMotionPlanning

max_forks_repo_name: zacmanchester/DirectMotionPlanning


text:
mutable struct SummaryStat
    step::Int64
    totalStep::Int64
    visited::Vector{Float64}
    reweight::Vector{Float64} # only record the reweight factor of a typical configuration
    propose::Array{Float64,3}
    accept::Array{Float64,3}
end

function addStat(config, summary = nothing)
    if isnothing(summary)
        return SummaryStat(config.step, config.totalStep, config.visited, config.reweight, config.propose, config.accept)
    else
        summary = deepcopy(summary)
        summary.step += config.step
        summary.totalStep += config.totalStep
        summary.visited += config.visited
        summary.reweight = config.reweight # only record the newest reweight factor
        summary.propose += config.propose
        summary.accept += config.accept
        return summary
    end
end

# sum SummaryStat from all workers to the root
function reduceStat(summary, root, comm)
    step = MPI.Reduce(summary.step, MPI.SUM, root, comm)
    totalStep = MPI.Reduce(summary.totalStep, MPI.SUM, root, comm)
    visited = MPI.Reduce(summary.visited, MPI.SUM, root, comm)
    reweight = MPI.Reduce(summary.reweight, MPI.SUM, root, comm)
    propose = MPI.Reduce(summary.propose, MPI.SUM, root, comm)
    accept = MPI.Reduce(summary.accept, MPI.SUM, root, comm)
    if MPI.Comm_rank(comm) == root
        reweight ./= MPI.Comm_size(comm)
        return SummaryStat(step, totalStep, visited, reweight, propose, accept)
    else
        return summary
    end
end

function MPIreduce(data)
    comm = MPI.COMM_WORLD
    Nworker = MPI.Comm_size(comm)  # number of MPI workers
    rank = MPI.Comm_rank(comm)  # rank of current MPI worker
    root = 0 # rank of the root worker

    if Nworker == 1 #no parallelization
        return data
    end
    if typeof(data) <: AbstractArray
        MPI.Reduce!(data, MPI.SUM, root, comm) # root node gets the sum of observables from all blocks
        return data
    else
        result = [data,]  # MPI.Reduce works for array only
        MPI.Reduce!(result, MPI.SUM, root, comm) # root node gets the sum of observables from all blocks
        return result[1]
    end
end


text:
"""
This file is for Gym mountaincar environment https://github.com/openai/gym/tree/master/gym/envs/classic_control.
These are based on https://github.com/WilsonWangTHU/mbbl
"""

using LyceumBase
using Shapes
using Distributions
using Base: @propagate_inbounds

struct MountainCar{D,SP,OP,AP,RP} <: AbstractEnvironment
    state::Vector{D}
    obs::Vector{D}
    action::Vector{D}
    time::Vector{D}
    nextobs::Vector{D}

    ctrlrange::Vector{D}
    maxspeed::D
    power::D
    posrange::Vector{D}

    sp::SP
    op::OP
    ap::AP
    rp::RP
    function MountainCar{T}() where T<:AbstractFloat
        nobs = 2
        naction = 1

        ctrlrange    = zeros(T, 2)
        ctrlrange[1] = -1.;  ctrlrange[2] = 1.
        maxspeed     = T(0.07); power        = T(0.0015)
        posrange     = zeros(T, 2)
        posrange[1]  = -1.2; posrange[2]  = 0.6

        sp = VectorShape(T, nobs)
        op = VectorShape(T, nobs)
        ap = VectorShape(T, naction)
        rp = ScalarShape(T)
        s = zeros(T, nobs)
        o = zeros(T, nobs)
        a = zeros(T, naction)
        env = new{T,
        typeof(sp), typeof(op), typeof(ap), typeof(rp)}(s,o,a,zeros(T,1),
                                                        zeros(T, nobs),
                                                        ctrlrange, maxspeed, power, posrange,
                                                        sp,op,ap,rp)
        reset!(env)
    end
end

function LyceumBase.tconstruct(::Type{MountainCar}, n::Integer)
    Tuple(MountainCar{Float64}() for _ in 1:n)
end
MountainCar() = first(tconstruct(MountainCar, 1))

@inline LyceumBase.statespace(env::MountainCar) = env.sp
@inline LyceumBase.obsspace(env::MountainCar) = env.op
@inline LyceumBase.actionspace(env::MountainCar) = env.ap
@inline LyceumBase.rewardspace(env::MountainCar) = env.rp

@propagate_inbounds @inline function LyceumBase.getstate!(state, env::MountainCar)
    #copyto!(state, env.state)
    state[1] = env.state[1] # unroll for speed
    state[2] = env.state[2]
end
@propagate_inbounds @inline function LyceumBase.setstate!(env::MountainCar, state)
    #env.state .= state
    env.state[1] = state[1] # unroll for speed
    env.state[2] = state[2]
end

@inline LyceumBase.getaction!(action, env::MountainCar) = action .= env.action
@inline LyceumBase.setaction!(env::MountainCar, action) = env.action .= action

@propagate_inbounds function LyceumBase.isdone(state, action, obs, env::MountainCar)
    false
end
@propagate_inbounds function LyceumBase.getobs!(obs, env::MountainCar)
    #copyto!(env.obs, env.state)
    #copyto!(obs, env.state)
    env.obs[1] = env.state[1] # unroll for speed
    env.obs[2] = env.state[2]
    obs[1] = env.state[1]
    obs[2] = env.state[2]
end
@propagate_inbounds function LyceumBase.getreward(s, a, o, env::MountainCar)
    # In benchmark paper
    return o[1]
end
@propagate_inbounds function LyceumBase.geteval(s, a, o, env::MountainCar)
    return o[1]
end
@propagate_inbounds function LyceumBase.reset!(env::MountainCar)
    randreset!(env)
end
@propagate_inbounds function LyceumBase.randreset!(env::MountainCar)
    env.state[1] = rand(Uniform(-0.6, -0.4))
    env.state[2] = 0.0
    getobs!(env.obs, env)
    env.time[1] = 0.0
    env
end

@propagate_inbounds function LyceumBase.step!(env::MountainCar)

    state, act, ctrlrange, maxspeed, power, posrange = env.state, env.action, env.ctrlrange, env.maxspeed, env.power, env.posrange

    position = state[1]
    velocity = state[2]
    """
        Note this is from Benchmark paper, not original in Gym env
        corresponding to action = np.clip(action, -1., 1.)
        also in continuous MC Gym env force = min(max(action[0], self.min_action), self.max_action)
        where min_action -1 max_action +1
    """
    force = clamp(act[1], ctrlrange[1], ctrlrange[2])

    velocity += force * power - 0.0025 * cos(3.0 * position)
    velocity = clamp(velocity, -maxspeed, maxspeed)

    position += velocity
    position = clamp(position, posrange[1], posrange[2])
    if (position == posrange[1]) & (velocity < 0.0);  velocity = 0.0;  end

    env.state[1] = position
    env.state[2] = velocity
    env.time[1] += timestep(env)
    env
end

@inline LyceumBase.timestep(env::MountainCar) = 1.0
@inline Base.time(env::MountainCar) = env.time[1]


text:
"""
    write_wat(filename, m::Module)

Write the WebAssembly module `m` to WebAssembly text format in `filename`.
"""
function write_wat(filename, m::Module)
  open(filename, "w") do f
    show(f, m)
  end
end

"""
    write_wat(filename, m::Module)

Write the WebAssembly module `m` to WebAssembly binary format in `filename`.
"""
function binary(m::Module, file; optimise = true)
  wat = tempname() * ".wat"
  write_wat(wat, m)
  run(`$(WABT.wat2wasm) --enable-all $wat -o $file`)
  optimise && run(`$(Binaryen.wasm_opt) $file -O4 -o $file`)
  rm(wat)
  return
end


text:
using DataValueArrays
using Base.Test

my_tests = [
    "typedefs.jl",
    "constructors.jl",
    "primitives.jl",
    "indexing.jl",
    "map.jl",
    "broadcast.jl",
    "nullablevector.jl",
    "nullablematrix.jl",
    "reduce.jl",
    "subarray.jl",
    "show.jl",
]

immutable SurvEvent
    time::Float64
    censored::Bool
end

@testset "DataValueArrays" begin

for my_test in my_tests
    include(my_test)
end

end


text:
using Zlib
zliblength(str) = length(Zlib.compress(str,9,false,true))

isinstalled(pkg) = try
    isa(Pkg.installed(pkg), VersionNumber)
  catch
    return false
end

# If Blosc is installed we can also use it for compression.
# If you expect strings involved in NCD calculations to be long
# you should definetely use qlz4length, lz4length or lz4hclength
# instead of zliblength as compressiors since they are typically
# an order of magnitude faster.
if isinstalled("Blosc")
  using Blosc
  blosc_compress(s, level = 9, shuffle = false, compressor = "blosclz") = begin
    Blosc.set_compressor(compressor)
    compress(s; level = level, shuffle = shuffle)
  end
  blosclzlength(s) = length(blosc_compress(s, 9, false, "blosclz"))
  lz4length(s) = length(blosc_compress(s, 9, false, "lz4"))
  lz4hclength(s) = length(blosc_compress(s, 9, false, "lz4hc"))
  snappylength(s) = length(blosc_compress(s, 9, false, "snappy"))
  blosczliblength(s) = length(blosc_compress(s, 9, false, "zlib"))

  qlz4length(s) = length(blosc_compress(s, 1, false, "lz4"))
  qsnappylength(s) = length(blosc_compress(s, 1, false, "snappy"))
  qblosczliblength(s) = length(blosc_compress(s, 1, false, "zlib"))
end

# If we are using GZip we need to save to file:
#using GZip
#
#function gzip_length(str) # We should not have to write to disk but couldn't find a way to do it in mem with GZip.jl. Investigate alternatives!
#  fn = tempfilename()
#  gfh = GZip.open(fn, "w")
#  println(gfh, str)
#  close(gfh)
#  len = filesize(fn)
#  rm(fn)
#  len
#end
#function tempfilename()
#  "tempfile_" * strftime("%Y%m%d_%H%M%S_", time()) * string(rand(1:1000000000)) * ".tmp"
#end

ncd(x, y, clen = lz4length) = ncd(string(x), string(y), clen)

function ncd(x::ASCIIString, y::ASCIIString, clen = lz4length)
  cxy = clen(x * y)
  cx = clen(x)
  cy = clen(y)
  (cxy - min(cx, cy)) / max(cx, cy)
end

# NCD for lists is here called NCDL. It is an implementation of the "NCD for Multisets"
# as described in the paper:
#  A. Cohen & P. Vitanyi, "", 2012, http://arxiv.org/pdf/1212.5711.pdf
# This is a heuristic to calc the NCDL in O(n^2) operations since the def of NCDL implies
# a O(2^n) algorithm.
function ncdl{T <: Any}(X::Array{T, 1}, clen = lz4length)
  ncdl(map(string, X), clen)
end

function ncdl(Xstrs::Array{ASCIIString, 1}, clen = lz4length, return_indexset = false)
  Gxs = map(clen, Xstrs)
  n = length(Xstrs)

  maxncd1 = -10.0 # Start with value < 0.0 to ensure first one is larger (NCD values are between 0.0 and 1+epsilon).

  # We will deselect one element in each step but start with all selected:
  selected = collect(1:n) .> 0
  selectedindices = Set(collect(1:n))
  numselected = n
  best = largest = 0

  while numselected >= 2

    # Calc max G(X \ {x}) for the still selected ones
    maxGxminus1 = -Inf
    for i in selectedindices
      selected[i] = false
      newGxminus1 = clen(join(Xstrs[selected], ""))
      selected[i] = true
      if newGxminus1 > maxGxminus1
        maxGxminus1 = newGxminus1
        largest = i
      end
    end

    # NCD1 = (G(X) - minimum(G(X))) / maxGxminus1
    ncd1 = (clen(join(Xstrs[selected], "")) - minimum(Gxs[selected])) / maxGxminus1
    #println("$numselected: $ncd1")
    if ncd1 > maxncd1
      maxncd1 = ncd1
      best = copy(selectedindices)
    end
    selected[largest] = false
    selectedindices = setdiff(selectedindices, Set(largest))
    numselected -= 1

  end

  if return_indexset
    return (maxncd1, best)
  else
    return maxncd1
  end
end

# We can then define the Test Set Diameter (TSD) as the NCDL of the test inputs or the test cases or the traces when executing these tests.
# From TSD we can then define the Test Strategy Diameter for size N, TSD(N), as the average
# TSD for a test set of size generated with that strategy. Strategy here is taken in its most
# general sense, i.e. can be a manual strategy for how to design tests (executed independently many times by different testers)
# or it can be any automated test generation strategy. Since the TSD(N) is a stochastic variable
# we should also use the standard deviation to characterize it. We can also compare different
# strategies (generators) by applying them multiple times and consider which one has the largest
# diameter. We hypothesise that if no further information about the SUT is known a strategy S1
# with a statistically significantly higher TSD(N) than another strategy S2 (TSD(N)_S1 >~ TSD(N)_S2)
# (evaluated with something like a Wilcoxon test, for example) is likely to yield higher fault
# finding ability.

ncdl(ary) = ncdm(ary, Lz4hc)

ncdl([1,2,3,4,5])

# We expect a random uniform sampling to give a larger ncd than a normal sampled one if the latter
# is confined to smaller range.
N = 100
ncdl(int(100.0 * rand(N)))
ncdl(int(50 + 5*randn(N)))

# We expect a generator of arrays of integers that can generate longer arrays to have larger NCDL
# than one which generates shorter arrays.
genintarrays(maxsize, minInt = -100, maxInt = 100) = map((i) -> rand(minInt:maxInt), 1:rand(0:maxsize))
ncdl(map((i) -> genintarrays(3), 1:N))
ncdl(map((i) -> genintarrays(8), 1:N))

# We expect generators generating arrays of same length to have larger ncd if they sample a
# larger span of the ints:
ncdl(map((i) -> genintarrays(5, -5, 5), 1:N))
ncdl(map((i) -> genintarrays(5, -50, 50), 1:N))
ncdl(map((i) -> genintarrays(5, -5000, 5000), 1:N))

# Calculations are quite slow though:
input = map((i) -> genintarrays(5, -5000, 5000), 1:N);
@time ncdl(input, zliblength)
@time ncdl(input, lz4length)
@time ncdl(input, lz4hclength)
@time ncdl(input, qlz4length)
@time ncdl(input, blosclzlength)

# although I imagine there are many ways we can speed this up if it is really useful.
# A problem with these types of InfoTheory results though is that it is less clear how one
# can incorporate biases in which test sets one prefer. But maybe we solve that with other
# aspects of our tools.

# Benchmark two generators and count the number of times that the first give higher
# NCD than the second one. The generators generate one instance only, they will be called
# repeatedly to create the multisets used for comparison.
function ncd_compare_generators(gen1, gen2;
  lengths = [2, 5, 10, 20, 50, 100],
  reps = 30,
  clen = lz4length)

  counts_higher_ncd = map(lengths) do len
  println("Length = $len")
    count = 0
    for r in 1:reps
      list1 = map((i) -> gen1(), 1:len)
      list2 = map((i) -> gen2(), 1:len)
      ncd1 = ncdl(list1, clen)
      ncd2 = ncdl(list2, clen)
      count += (ncd1 > ncd2 ? 1 : 0)
    end
    count
  end

  100.0 * counts_higher_ncd / reps
end

ncd_compare_generators(() -> genintarrays(8), () -> genintarrays(6))
ncd_compare_generators(() -> genintarrays(5, -50, 50), () -> genintarrays(5, -5, 5); clen = lz4length)


text:
module ExceptionUnwrapping

# Document this Module via the README.md file.
@doc let path = joinpath(dirname(@__DIR__), "README.md")
    include_dependency(path)
    read(path, String)
end ExceptionUnwrapping

export unwrap_exception, has_wrapped_exception, is_wrapped_exception,
    unwrap_exception_until, unwrap_exception_to_root, @test_throws_wrapped

include("test_throws_wrapped.jl")

"""
    has_wrapped_exception(e, ExceptionType)::Bool

Returns true if the given exception instance, `e`, contains an exception of type `T`
anywhere in its chain of unwrapped exceptions.

Application code should prefer to use `has_wrapped_exception(e, T)` instead of `e isa T` in
catch-blocks, to keep code from breaking when libraries wrap user's exceptions.

This makes application code resilient to library changes that may cause wrapped exceptions,
such as e.g. changes to underlying concurrency decisions (thus maintaining concurrency's
cooperative benefits).

# Example
```julia
try
    # If this becomes concurrent in the future, the catch-block doesn't need to change.
    library_function(args...)
catch e
    if has_wrapped_exception(e, MyExceptionType)
        unwrapped = unwrap_exception_until(e, MyExceptionType)
        handle_my_exception(unwrapped, caught=e)
    else
        rethrow()
    end
end
```
"""
function has_wrapped_exception end

"""
    is_wrapped_exception(e)::Bool

Returns true if the given exception instance, `e` is a wrapped exception, such
that `unwrap_exception(e)` would return something different than `e`.
"""
function is_wrapped_exception end

"""
    unwrap_exception(exception_wrapper) -> wrapped_exception
    unwrap_exception(normal_exception) -> normal_exception

    # Add overrides for custom exception types
    ExceptionUnwrapping.unwrap_exception(e::MyWrappedException) = e.wrapped_exception

Unwraps a wrapped exception by one level. *New wrapped exception types should add a method
to this function.*

One example of a wrapped exception is the `TaskFailedException`, which wraps an exception
thrown by a `Task` with a new `Exception` describing the task failure.

It is useful to unwrap the exception to test what kind of exception was thrown in the first
place, which is useful in case you need different exception handling behavior for different
types of exceptions.

Authors of new wrapped exception types can overload this to indicate what field their
exception is wrapping, by adding an overload, e.g.:
```julia
ExceptionUnwrapping.unwrap_exception(e::MyWrappedException) = e.wrapped_exception
```

This is used in the implementations of the other functions in the module:
- [`has_wrapped_exception(e, ::Type)`](@ref)
- [`unwrap_exception_to_root(e)`](@ref)
"""
function unwrap_exception end

"""
    unwrap_exception_until(e, ExceptionType)::ExceptionType

Recursively unwrap a wrapped exception `e` until reaching an instance of `ExceptionType`.
"""
function unwrap_exception_until end

"""
    unwrap_exception_to_root(exception_wrapper) -> wrapped_exception
    unwrap_exception_to_root(normal_exception) -> normal_exception

Unwrap a wrapped exception to its bottom layer.
"""
function unwrap_exception_to_root end

struct UnwrappedExceptionNotFound{RequestedType, ExceptionType} <: Base.Exception
    exception::ExceptionType
end
UnwrappedExceptionNotFound{R}(e::E) where {R,E} = UnwrappedExceptionNotFound{R,E}(e)

# Base case is that e -> e
unwrap_exception(e) = e
# Add overloads for wrapped exception types to unwrap the exception.
if VERSION >= v"1.3.0-"
    # TaskFailedExceptions wrap a failed task, which contains the exception that caused it
    # to fail. You can unwrap the exception to discover the root cause of the failure.
    unwrap_exception(e::Base.TaskFailedException) = e.task.exception
end

function has_wrapped_exception(e, ::Type{T}) where T
    if e isa T
        true
    else
        is_wrapped_exception(e) ? has_wrapped_exception(unwrap_exception(e), T) : false
    end
end

function is_wrapped_exception(e)
    return e !== unwrap_exception(e)
end

function unwrap_exception_until(e, ::Type{T}) where T
    if e isa T
        e
    else
        if is_wrapped_exception(e)
            unwrap_exception_until(unwrap_exception(e), T)
        else
            throw(UnwrappedExceptionNotFound{T}(e))
        end
    end
end

function unwrap_exception_to_root(e)
    is_wrapped_exception(e) ? unwrap_exception_to_root(unwrap_exception(e)) : e
end

end # module


text:
using QuantumInformation.ShortNames
using Base.Test

# write your own tests here
include("gates.jl")
include("printing.jl")
include("basic.jl")


text:
# Util functions.
module Util
using Random

export weighted_median,
       err_must_be_overriden,
       holdout

# Weighted median.
#
# @param weights Weights of values.
# @param values Values.
# @return Weighted median.
function weighted_median( weights::AbstractVector{U}, values::AbstractVector{V}) where {U,V<:Real}

  k = 1
  sorted_ind = sortperm(values)
  weight_sum = sum(weights)

  remaining_sum = weight_sum - weights[sorted_ind[k]]
  while remaining_sum > weight_sum / 2.0
    k += 1
    remaining_sum -= weights[sorted_ind[k]]
  end

  return values[sorted_ind[k]]
end

function err_must_be_overriden(s = "")
  error("Function must be overriden.")
end

# Holdout method that partitions a collection
# into two partitions.
#
# @param n Size of collection to partition.
# @param right_prop Percentage of collection placed in right partition.
# @return Two partitions of indices, left and right.
function holdout(n, right_prop)
  shuffled_indices = randperm(n)
  partition_pivot = round(Int, right_prop * n)
  right = shuffled_indices[1:partition_pivot]
  left = shuffled_indices[partition_pivot+1:end]
  return (left, right)
end

end # module


text:
################################################################################
#
#  Field access
#
################################################################################

function elem_type(::Type{RelOrdQuoRing{T1, T2, T3}}) where { T1, T2, T3 }
  S = elem_type(T1)
  return RelOrdQuoRingElem{T1, T2, T3, S}
end

function elem_type(::RelOrdQuoRing{T1, T2, T3}) where { T1, T2, T3 }
  S = elem_type(T1)
  return RelOrdQuoRingElem{T1, T2, T3, S}
end

base_ring(Q::RelOrdQuoRing) = Q.base_ring

ideal(Q::RelOrdQuoRing) = Q.ideal

basis_pmatrix(Q::RelOrdQuoRing) = Q.basis_pmatrix

parent(x::RelOrdQuoRingElem) = x.parent

parent_type(::Type{RelOrdQuoRingElem{T1, T2, T3, S}}) where { T1, T2, T3, S } = RelOrdQuoRing{T1, T2, T3}

################################################################################
#
#  Hashing
#
################################################################################

hash(x::RelOrdQuoRingElem, h::UInt) = hash(x.elem, h)

################################################################################
#
#  Copying
#
################################################################################

Base.deepcopy_internal(x::RelOrdQuoRingElem, dict::IdDict) =
        RelOrdQuoRingElem(parent(x), Base.deepcopy_internal(x.elem, dict))

################################################################################
#
#  I/O
#
################################################################################

function show(io::IO, Q::RelOrdQuoRing)
  print(io, "Quotient of $(base_ring(Q))")
end

function show(io::IO, x::RelOrdQuoRingElem)
  print(io, "$(x.elem)")
end

################################################################################
#
#  Parent object overloading
#
################################################################################

function (Q::RelOrdQuoRing{T1, T2, T3})(x::RelOrdQuoRingElem{T1, T2, T3, S}) where { T1, T2, T3, S }
  parent(x) !== Q && error("Cannot coerce element into quotient ring")
  return x
end

function (Q::RelOrdQuoRing{T1, T2, T3})(x::S) where { T1, T2, T3, S }
  parent(x) !== base_ring(Q) && error("Cannot coerce element into the quotient ring")
  return RelOrdQuoRingElem(Q, x)
end

function (Q::RelOrdQuoRing)(x::Integer)
  return Q(base_ring(Q)(x))
end

function (Q::RelOrdQuoRing)(x::fmpz)
  return Q(base_ring(Q)(x))
end

################################################################################
#
#  Quotient function
#
################################################################################

function quo(O::Union{NfRelOrd, AlgAssRelOrd}, I::Union{NfRelOrdIdl, AlgAssRelOrdIdl})
  @assert order(I) === O
  # We should check that I is not zero
  Q = RelOrdQuoRing(O, I)
  f = RelOrdQuoMap(O, Q)
  return Q, f
end

Nemo.ResidueRing(O::Union{NfRelOrd, AlgAssRelOrd}, I::Union{NfRelOrdIdl, AlgAssRelOrdIdl}) = RelOrdQuoRing(O, I)

################################################################################
#
#  Parent check
#
################################################################################

function check_parent(x::RelOrdQuoRingElem, y::RelOrdQuoRingElem)
  if parent(x) !== parent(y)
    error("Elements must have same parents")
  end
  return true
end

################################################################################
#
#  Arithmetic
#
################################################################################

function +(x::RelOrdQuoRingElem, y::RelOrdQuoRingElem)
  check_parent(x, y)
  return parent(x)(x.elem + y.elem)
end

function -(x::RelOrdQuoRingElem, y::RelOrdQuoRingElem)
  check_parent(x, y)
  return parent(x)(x.elem - y.elem)
end

function -(x::RelOrdQuoRingElem)
  return parent(x)(-x.elem)
end

function *(x::RelOrdQuoRingElem, y::RelOrdQuoRingElem)
  check_parent(x, y)
  return parent(x)(x.elem * y.elem)
end

function mul!(z::RelOrdQuoRingElem, x::RelOrdQuoRingElem, y::RelOrdQuoRingElem)
  z.elem = mul!(z.elem, x.elem, y.elem)
  z.elem = mod!(z.elem, parent(z))
  return z
end

function add!(z::RelOrdQuoRingElem, x::RelOrdQuoRingElem, y::RelOrdQuoRingElem)
  z.elem = add!(z.elem, x.elem, y.elem)
  z.elem = mod!(z.elem, parent(z))
  return z
end

function *(x::Integer, y::RelOrdQuoRingElem)
  return parent(y)(x * y.elem)
end

*(x::RelOrdQuoRingElem, y::Integer) = y*x

function *(x::fmpz, y::RelOrdQuoRingElem)
  return parent(y)(x * y.elem)
end

*(x::RelOrdQuoRingElem, y::fmpz) = y*x

function ^(a::RelOrdQuoRingElem, f::fmpz)
  if fits(Int, f)
    return a^Int(f)
  end
  f == 0 && return one(parent(a))
  f == 1 && return deepcopy(a)
  if f < 0
    f = -f
    a = inv(a)
  end
  b = a^(div(f, 2))
  b *= b
  if isodd(f)
    b *= a
  end
  return b
end

function ^(a::RelOrdQuoRingElem, b::Int)
  if b == 0
    return one(parent(a))
  elseif b == 1
    return deepcopy(a)
  else
    if b < 0
      a = inv(a)
      b = -b
    end
    bit = ~((~UInt(0)) >> 1)
    while (UInt(bit) & b) == 0
      bit >>= 1
    end
    z = deepcopy(a)
    bit >>= 1
    while bit != 0
      z = mul!(z, z, z)
      if (UInt(bit) & b) != 0
        z = mul!(z, z, a)
      end
      bit >>= 1
    end
    return z
  end
end

################################################################################
#
#  Special elements
#
################################################################################

iszero(x::RelOrdQuoRingElem) = iszero(x.elem)

isone(x::RelOrdQuoRingElem) = isone(x.elem)

function one(Q::RelOrdQuoRing)
  return Q(one(base_ring(Q)))
end

function zero(Q::RelOrdQuoRing)
  return Q(zero(base_ring(Q)))
end

################################################################################
#
#  Equality
#
################################################################################

function ==(x::RelOrdQuoRing, y::RelOrdQuoRing)
  return base_ring(x) === base_ring(y) && ideal(x) == ideal(y)
end

==(x::RelOrdQuoRingElem, y::RelOrdQuoRingElem) = x.elem == y.elem


text:
function mapcolumns(f, cols)
    columnnames = Symbol[]
    columndict = Dict{Symbol,Vector}()
    nums = Int[]
    for (k, v) in pairs(cols)
        v = f(v)
        if v !== nothing
            columndict[k] = v
            push!(columnnames, k)
            push!(nums, length(v))
        end
    end
    isempty(columnnames) && return NamedTuple()
    @assert all(==(nums[1]), nums)
    return (; (k => columndict[k] for k in columnnames)...)
end

function ishomogenous(xs)
    y = iterate(xs)
    y === nothing && return true
    x1, s = y
    while true
        y = iterate(xs, s)
        y === nothing && return true
        x, s = y
        isequal(x, x1) || return false
    end
end


text:
# julia --project -O3 --check-bounds=no -t 4 ./scripts-part1/diffusion3D_visualize.jl

# Initialize Globals
const USE_GPU = true
const BENCHMARK = false
const VISUALIZE = true

# Include main script
include("./diffusion3D_multixpu.jl")

# Run the simulation
diffusion_3D(256)


text:
using Test

using GCMMesh
using GCMMesh.Mesh
using GCMMesh.Partition

FT = Float64

@testset "test rsb_partition" begin
    x1min, x1max = FT(0), FT(1)
    x2min, x2max = FT(0), FT(1)
    limits = (x1min, x2min, x1max, x2max)
    n1, n2 = 4, 4 # # of elements along x1 and x2 axis respectively
    mesh = equispaced_rectangular_mesh(limits..., n1, n2, (false, false))

    @testset "test rsb_partition with 5 partitions" begin
        npart = 5 # # of partitions
        partition = rsb_partition(mesh, npart)

        partlen = Vector{Int}(undef, npart)
        for i = 1:npart
            partlen[i] = count(partition .== i - 1)
        end
        @test extrema(partition) == (0, npart - 1) # verify if process numbers are assigned to all elems
        @test sum(partlen) == mesh.nelems             # verify sum(nlocalelems) == nelems
        @test maximum(partlen) - minimum(partlen) ≤ 1 # check if partion lengths are even
    end

    @testset "test rsb_partition with single partition" begin
        npart = 1
        partition = rsb_partition(mesh, npart)
        @test partition == zeros(n1 * n2)
    end
end

max_stars_repo_name: CliMA/GCMMesh

max_issues_repo_name: CliMA/GCMMesh

max_forks_repo_name: CliMA/GCMMesh


text:
module MieScatter

export compute_mie, size_parameter

size_parameter(radius::Real, wavelength::Real) = 2π * radius / wavelength

function compute_mie(size_param::Real, ref_idx::Number, N_angles::Integer)
    delta_angle = π / (N_angles-1)
    compute_mie(size_param, ref_idx, [(i-1)*delta_angle for i = 1:N_angles])
end

function compute_mie(size_param::Real, ref_idx::Number, angles::Vector{Float64})
    y = size_param * ref_idx
    x_stop = size_param + 4 * size_param^0.3333 + 2.0

    y_modulus = abs(y)
    n_max = round(Int, max(x_stop, y_modulus) + 15)

    mu_table = [cos(theta) for theta in angles]
    N_angles = length(angles)
    idx_forward_scatter = -1
    idx_backscatter = -1
    for i = 1:N_angles
        if angles[i] ≈ 0.0
            idx_forward_scatter = i
        end
        if angles[i] - π ≈ 0.0
            idx_backscatter = i
        end
    end

    d = zeros(ComplexF64, n_max+1)

    for n = n_max-1:-1:1
        rn = n+1
        d[n] = (rn/y) - (1 / (d[n+1] + rn/y))
    end

    pi0 = zeros(Float64, N_angles)
    pi1 = ones(Float64, N_angles)
    piX = zeros(Float64, N_angles)

    s1 = zeros(ComplexF64, N_angles)
    s2 = zeros(ComplexF64, N_angles)
    tau = zeros(Float64, N_angles)

    psi0 = cos(size_param)
    psi1 = sin(size_param)
    chi0 = -sin(size_param)
    chi1 = cos(size_param)

    S = zeros(Float64, N_angles, 4)

    xi1 = ComplexF64(psi1, -chi1)
    Qsca = 0.0

    for n = 1:round(Int, x_stop)
        fn = (2n+1) / (n*(n+1))
        psi = (2n-1) * psi1/size_param - psi0
        chi = (2n-1) * chi1/size_param - chi0
        xi = ComplexF64(psi, -chi)
        t_a = d[n] / ref_idx + n/size_param
        t_b = d[n] * ref_idx + n/size_param
        an = (t_a * psi - psi1) / (t_a * xi - xi1)
        bn = (t_b * psi - psi1) / (t_b * xi - xi1)
        Qsca += (2n+1) * (abs(an)^2 + abs(bn)^2)

        for j = 1:N_angles
            piX[j] = pi1[j]
            tau[j] = n * mu_table[j] * piX[j] - (n+1)*pi0[j]
            t = (-1)^n
            p = (-1)^(n-1)
            s1[j] = s1[j] + fn * (an*piX[j] + bn*tau[j])
            s2[j] = s2[j] + fn * (an*tau[j] + bn*piX[j])
        end

        psi0 = psi1
        psi1 = psi
        chi0 = chi1
        chi1 = chi
        xi1 = ComplexF64(psi1, -chi1)

        for i = 1:N_angles
            pi1[i] = (2n+1)/n * mu_table[i] * pi1[i] - (n+1)*pi0[i] / n
            pi0[i] = piX[i]
        end
    end

    Qsca *= 2/size_param^2
    Qext = (idx_forward_scatter > 0) ? (4 / size_param^2) * real(s1[1]) : NaN
    Qback = (idx_backscatter > 0) ? (4 / size_param^2) * abs(s1[end])^2 : NaN

    for i = 1:N_angles
        S[i,1] = 0.5 * (abs(s1[i])^2 + abs(s2[i])^2)
        S[i,2] = -0.5 * (abs(s1[i])^2 - abs(s2[i])^2)
        S[i,3] = real(s2[i] * conj(s1[i]))
        S[i,4] = imag(s2[i] * conj(s1[i]))
    end

    return S, Qsca, Qext, Qback

end

end # module


text:
function spectralClustering(A::Array{Float64,2}, nEigen::Int)

	D = Diagonal(sum(A, dims=1)[:])
	L = D-A
	NL = D^(-1/2) * L * D^(-1/2)
	vals, vects = eigen(NL)
	X = vects[:,1:nEigen]
	XX = fill(0.0,size(X))

	for i in 1:size(X,1)
		XX[i,:] = X[i,:] ./ sqrt(sum(X[i,:].^2))
	end

	return vals, XX
end

max_stars_repo_name: neuro-myoung/CellSegmenter

max_issues_repo_name: neuro-myoung/CellSegmenter

max_forks_repo_name: neuro-myoung/CellSegmenter


text:
#### EDMF model kernels

using CLIMAParameters.Planet: e_int_v0, grav, day, R_d, R_v, molmass_ratio
using Printf
using ClimateMachine.Atmos: nodal_update_auxiliary_state!, Advect

using ClimateMachine.BalanceLaws

using ClimateMachine.MPIStateArrays: MPIStateArray
using ClimateMachine.DGMethods: LocalGeometry, DGModel

import ClimateMachine.BalanceLaws:
    vars_state,
    prognostic_vars,
    prognostic_to_primitive!,
    primitive_to_prognostic!,
    get_prog_state,
    flux,
    precompute,
    source,
    eq_tends,
    update_auxiliary_state!,
    init_state_prognostic!,
    compute_gradient_argument!,
    compute_gradient_flux!

import ClimateMachine.TurbulenceConvection:
    init_aux_turbconv!,
    turbconv_nodal_update_auxiliary_state!,
    turbconv_boundary_state!,
    turbconv_normal_boundary_flux_second_order!

using ClimateMachine.Thermodynamics: air_pressure, air_density

include(joinpath("helper_funcs", "nondimensional_exchange_functions.jl"))
include(joinpath("helper_funcs", "lamb_smooth_minimum.jl"))
include(joinpath("helper_funcs", "utility_funcs.jl"))
include(joinpath("helper_funcs", "subdomain_statistics.jl"))
include(joinpath("helper_funcs", "diagnose_environment.jl"))
include(joinpath("helper_funcs", "subdomain_thermo_states.jl"))
include(joinpath("helper_funcs", "save_subdomain_temperature.jl"))
include(joinpath("closures", "entr_detr.jl"))
include(joinpath("closures", "pressure.jl"))
include(joinpath("closures", "mixing_length.jl"))
include(joinpath("closures", "turbulence_functions.jl"))
include(joinpath("closures", "surface_functions.jl"))

function vars_state(m::NTuple{N, Updraft}, st::Auxiliary, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

vars_state(::Updraft, ::Auxiliary, FT) = @vars(T::FT)
vars_state(::Environment, ::Auxiliary, FT) = @vars(T::FT)

function vars_state(m::EDMF, st::Auxiliary, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Prognostic, FT)
    @vars(ρa::FT, ρaw::FT, ρaθ_liq::FT, ρaq_tot::FT,)
end

function vars_state(::Environment, ::Prognostic, FT)
    @vars(ρatke::FT, ρaθ_liq_cv::FT, ρaq_tot_cv::FT, ρaθ_liq_q_tot_cv::FT,)
end

function vars_state(::Updraft, ::Primitive, FT)
    @vars(a::FT, aw::FT, aθ_liq::FT, aq_tot::FT,)
end

function vars_state(::Environment, ::Primitive, FT)
    @vars(atke::FT, aθ_liq_cv::FT, aq_tot_cv::FT, aθ_liq_q_tot_cv::FT,)
end

function vars_state(
    m::NTuple{N, Updraft},
    st::Union{Prognostic, Primitive},
    FT,
) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Union{Prognostic, Primitive}, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Gradient, FT)
    @vars(w::FT,)
end

function vars_state(::Environment, ::Gradient, FT)
    @vars(
        θ_liq::FT,
        q_tot::FT,
        w::FT,
        tke::FT,
        θ_liq_cv::FT,
        q_tot_cv::FT,
        θ_liq_q_tot_cv::FT,
        θv::FT,
        h_tot::FT,
    )
end

function vars_state(m::NTuple{N, Updraft}, st::Gradient, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Gradient, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT),
        u::FT,
        v::FT
    )
end

function vars_state(m::NTuple{N, Updraft}, st::GradientFlux, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, st::GradientFlux, FT)
    @vars(∇w::SVector{3, FT},)
end

function vars_state(::Environment, ::GradientFlux, FT)
    @vars(
        ∇θ_liq::SVector{3, FT},
        ∇q_tot::SVector{3, FT},
        ∇w::SVector{3, FT},
        ∇tke::SVector{3, FT},
        ∇θ_liq_cv::SVector{3, FT},
        ∇q_tot_cv::SVector{3, FT},
        ∇θ_liq_q_tot_cv::SVector{3, FT},
        ∇θv::SVector{3, FT},
        ∇h_tot::SVector{3, FT},
    )

end

function vars_state(m::EDMF, st::GradientFlux, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT),
        ∇u::SVector{3, FT},
        ∇v::SVector{3, FT}
    )
end

abstract type EDMFPrognosticVariable <: AbstractPrognosticVariable end

abstract type EnvironmentPrognosticVariable <: EDMFPrognosticVariable end
struct en_ρatke <: EnvironmentPrognosticVariable end
struct en_ρaθ_liq_cv <: EnvironmentPrognosticVariable end
struct en_ρaq_tot_cv <: EnvironmentPrognosticVariable end
struct en_ρaθ_liq_q_tot_cv <: EnvironmentPrognosticVariable end

abstract type UpdraftPrognosticVariable{i} <: EDMFPrognosticVariable end
struct up_ρa{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaw{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaθ_liq{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaq_tot{i} <: UpdraftPrognosticVariable{i} end

prognostic_vars(m::EDMF) =
    (prognostic_vars(m.environment)..., prognostic_vars(m.updraft)...)
prognostic_vars(m::Environment) =
    (en_ρatke(), en_ρaθ_liq_cv(), en_ρaq_tot_cv(), en_ρaθ_liq_q_tot_cv())

function prognostic_vars(m::NTuple{N, Updraft}) where {N}
    t_ρa = vuntuple(i -> up_ρa{i}(), N)
    t_ρaw = vuntuple(i -> up_ρaw{i}(), N)
    t_ρaθ_liq = vuntuple(i -> up_ρaθ_liq{i}(), N)
    t_ρaq_tot = vuntuple(i -> up_ρaq_tot{i}(), N)
    t = (t_ρa..., t_ρaw..., t_ρaθ_liq..., t_ρaq_tot...)
    return t
end

get_prog_state(state, ::en_ρatke) = (state.turbconv.environment, :ρatke)
get_prog_state(state, ::en_ρaθ_liq_cv) =
    (state.turbconv.environment, :ρaθ_liq_cv)
get_prog_state(state, ::en_ρaq_tot_cv) =
    (state.turbconv.environment, :ρaq_tot_cv)
get_prog_state(state, ::en_ρaθ_liq_q_tot_cv) =
    (state.turbconv.environment, :ρaθ_liq_q_tot_cv)

get_prog_state(state, ::up_ρa{i}) where {i} = (state.turbconv.updraft[i], :ρa)
get_prog_state(state, ::up_ρaw{i}) where {i} = (state.turbconv.updraft[i], :ρaw)
get_prog_state(state, ::up_ρaθ_liq{i}) where {i} =
    (state.turbconv.updraft[i], :ρaθ_liq)
get_prog_state(state, ::up_ρaq_tot{i}) where {i} =
    (state.turbconv.updraft[i], :ρaq_tot)

struct EntrDetr{N_up} <: TendencyDef{Source} end
struct PressSource{N_up} <: TendencyDef{Source} end
struct BuoySource{N_up} <: TendencyDef{Source} end
struct ShearSource <: TendencyDef{Source} end
struct DissSource <: TendencyDef{Source} end
struct GradProdSource <: TendencyDef{Source} end

prognostic_vars(::EntrDetr{N_up}) where {N_up} = (
    vuntuple(i -> up_ρa{i}, N_up)...,
    vuntuple(i -> up_ρaw{i}, N_up)...,
    vuntuple(i -> up_ρaθ_liq{i}, N_up)...,
    vuntuple(i -> up_ρaq_tot{i}, N_up)...,
    en_ρatke(),
    en_ρaθ_liq_cv(),
    en_ρaq_tot_cv(),
    en_ρaθ_liq_q_tot_cv(),
)
prognostic_vars(::PressSource{N_up}) where {N_up} =
    vuntuple(i -> up_ρaw{i}(), N_up)

prognostic_vars(::BuoySource{N_up}) where {N_up} =
    vuntuple(i -> up_ρaw{i}(), N_up)

EntrDetr(m::EDMF) = EntrDetr{n_updrafts(m)}()
BuoySource(m::EDMF) = BuoySource{n_updrafts(m)}()
PressSource(m::EDMF) = PressSource{n_updrafts(m)}()

# Dycore tendencies
eq_tends(
    pv::Union{Momentum, Energy, TotalMoisture},
    m::EDMF,
    flux::Flux{SecondOrder},
) = eq_tends(pv, m.coupling, flux)

eq_tends(
    pv::Union{Momentum, Energy, TotalMoisture},
    ::Decoupled,
    ::Flux{SecondOrder},
) = ()

eq_tends(
    pv::Union{Momentum, Energy, TotalMoisture},
    ::Coupled,
    ::Flux{SecondOrder},
) = (SGSFlux(),)

# Turbconv tendencies
eq_tends(pv::EDMFPrognosticVariable, m::AtmosModel, tt::Flux{O}) where {O} =
    eq_tends(pv, turbconv_model(m), tt)

eq_tends(::EDMFPrognosticVariable, m::EDMF, ::Flux{O}) where {O} = ()

eq_tends(::EnvironmentPrognosticVariable, m::EDMF, ::Flux{SecondOrder}) =
    (Diffusion(),)

eq_tends(pv::EDMFPrognosticVariable, m::EDMF, ::Flux{FirstOrder}) = (Advect(),)

eq_tends(pv::PV, m::EDMF, ::Source) where {PV} = ()

eq_tends(::EDMFPrognosticVariable, m::EDMF, ::Source) = (EntrDetr(m),)

eq_tends(pv::en_ρatke, m::EDMF, ::Source) =
    (EntrDetr(m), PressSource(m), BuoySource(m), ShearSource(), DissSource())

eq_tends(
    ::Union{en_ρaθ_liq_cv, en_ρaq_tot_cv, en_ρaθ_liq_q_tot_cv},
    m::EDMF,
    ::Source,
) = (EntrDetr(m), DissSource(), GradProdSource())

eq_tends(::up_ρaw, m::EDMF, ::Source) =
    (EntrDetr(m), PressSource(m), BuoySource(m))

struct SGSFlux <: TendencyDef{Flux{SecondOrder}} end

"""
    init_aux_turbconv!(
        turbconv::EDMF{FT},
        m::AtmosModel{FT},
        aux::Vars,
        geom::LocalGeometry,
    ) where {FT}

Initialize EDMF auxiliary variables.
"""
function init_aux_turbconv!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    aux::Vars,
    geom::LocalGeometry,
) where {FT} end;

function turbconv_nodal_update_auxiliary_state!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    save_subdomain_temperature!(m, state, aux)
end;

function prognostic_to_primitive!(
    turbconv::EDMF,
    atmos,
    moist::DryModel,
    prim::Vars,
    prog::Vars,
)
    N_up = n_updrafts(turbconv)
    prim_en = prim.turbconv.environment
    prog_en = prog.turbconv.environment
    prim_up = prim.turbconv.updraft
    prog_up = prog.turbconv.updraft

    ρ_inv = 1 / prog.ρ
    prim_en.atke = prog_en.ρatke * ρ_inv
    prim_en.aθ_liq_cv = prog_en.ρaθ_liq_cv * ρ_inv
    prim_en.aθ_liq_q_tot_cv = prog_en.ρaθ_liq_q_tot_cv * ρ_inv

    if moist isa DryModel
        prim_en.aq_tot_cv = 0
    else
        prim_en.aq_tot_cv = prog_en.ρaq_tot_cv * ρ_inv
    end
    @unroll_map(N_up) do i
        prim_up[i].a = prog_up[i].ρa * ρ_inv
        prim_up[i].aw = prog_up[i].ρaw * ρ_inv
        prim_up[i].aθ_liq = prog_up[i].ρaθ_liq * ρ_inv
        if moist isa DryModel
            prim_up[i].aq_tot = 0
        else
            prim_up[i].aq_tot = prog_up[i].ρaq_tot * ρ_inv
        end
    end
end

function primitive_to_prognostic!(
    turbconv::EDMF,
    atmos,
    moist::DryModel,
    prog::Vars,
    prim::Vars,
)

    N_up = n_updrafts(turbconv)
    prim_en = prim.turbconv.environment
    prog_en = prog.turbconv.environment
    prim_up = prim.turbconv.updraft
    prog_up = prog.turbconv.updraft

    ρ_gm = prog.ρ
    prog_en.ρatke = prim_en.atke * ρ_gm
    prog_en.ρaθ_liq_cv = prim_en.aθ_liq_cv * ρ_gm
    prog_en.ρaθ_liq_q_tot_cv = prim_en.aθ_liq_q_tot_cv * ρ_gm

    if moist isa DryModel
        prog_en.ρaq_tot_cv = 0
    else
        prog_en.ρaq_tot_cv = prim_en.aq_tot_cv * ρ_gm
    end
    @unroll_map(N_up) do i
        prog_up[i].ρa = prim_up[i].a * ρ_gm
        prog_up[i].ρaw = prim_up[i].aw * ρ_gm
        prog_up[i].ρaθ_liq = prim_up[i].aθ_liq * ρ_gm
        if moist isa DryModel
            prog_up[i].ρaq_tot = 0
        else
            prog_up[i].ρaq_tot = prim_up[i].aq_tot * ρ_gm
        end
    end

end

function compute_gradient_argument!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    z = altitude(m, aux)

    # Aliases:
    gm_tf = transform.turbconv
    up_tf = transform.turbconv.updraft
    en_tf = transform.turbconv.environment
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, N_up)
    param_set = parameter_set(m)
    @unroll_map(N_up) do i
        up_tf[i].w = fix_void_up(up[i].ρa, up[i].ρaw / up[i].ρa)
    end
    _grav::FT = grav(param_set)

    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)

    if moisture_model(m) isa DryModel
        q_tot_en = FT(0)
    else
        q_tot_en = total_specific_humidity(ts.en)
    end

    # populate gradient arguments
    en_tf.θ_liq = θ_liq_en
    en_tf.q_tot = q_tot_en
    en_tf.w = env.w

    en_tf.tke = en.ρatke / (env.a * gm.ρ)
    en_tf.θ_liq_cv = en.ρaθ_liq_cv / (env.a * gm.ρ)

    if moisture_model(m) isa DryModel
        en_tf.q_tot_cv = FT(0)
        en_tf.θ_liq_q_tot_cv = FT(0)
    else
        en_tf.q_tot_cv = en.ρaq_tot_cv / (env.a * gm.ρ)
        en_tf.θ_liq_q_tot_cv = en.ρaθ_liq_q_tot_cv / (env.a * gm.ρ)
    end

    en_tf.θv = virtual_pottemp(ts.en)
    en_e_kin =
        FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w^2) # TBD: Check
    en_e_tot = total_energy(en_e_kin, _grav * z, ts.en)
    en_tf.h_tot = total_specific_enthalpy(ts.en, en_e_tot)

    gm_tf.u = gm.ρu[1] * ρ_inv
    gm_tf.v = gm.ρu[2] * ρ_inv
end;

function compute_gradient_flux!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    args = (; diffusive, state, aux, t)
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    gm_dif = diffusive.turbconv
    gm_∇tf = ∇transform.turbconv
    up_dif = diffusive.turbconv.updraft
    up_∇tf = ∇transform.turbconv.updraft
    en = state.turbconv.environment
    en_dif = diffusive.turbconv.environment
    en_∇tf = ∇transform.turbconv.environment

    @unroll_map(N_up) do i
        up_dif[i].∇w = up_∇tf[i].w
    end

    env = environment_vars(state, N_up)
    ρa₀ = gm.ρ * env.a
    # first moment grid mean coming from environment gradients only
    en_dif.∇θ_liq = en_∇tf.θ_liq
    en_dif.∇q_tot = en_∇tf.q_tot
    en_dif.∇w = en_∇tf.w
    # second moment env cov
    en_dif.∇tke = en_∇tf.tke
    en_dif.∇θ_liq_cv = en_∇tf.θ_liq_cv
    en_dif.∇q_tot_cv = en_∇tf.q_tot_cv
    en_dif.∇θ_liq_q_tot_cv = en_∇tf.θ_liq_q_tot_cv

    en_dif.∇θv = en_∇tf.θv
    en_dif.∇h_tot = en_∇tf.h_tot

    gm_dif.∇u = gm_∇tf.u
    gm_dif.∇v = gm_∇tf.v
end;

function source(::up_ρa{i}, ::EntrDetr, atmos, args) where {i}
    @unpack E_dyn, Δ_dyn, ρa_up = args.precomputed.turbconv
    return fix_void_up(ρa_up[i], E_dyn[i] - Δ_dyn[i])
end

function source(::up_ρaw{i}, ::EntrDetr, atmos, args) where {i}
    @unpack E_dyn, Δ_dyn, E_trb, env, ρa_up, w_up = args.precomputed.turbconv
    up = args.state.turbconv.updraft
    entr = fix_void_up(ρa_up[i], (E_dyn[i] + E_trb[i]) * env.w)
    detr = fix_void_up(ρa_up[i], (Δ_dyn[i] + E_trb[i]) * w_up[i])

    return entr - detr
end

function source(::up_ρaθ_liq{i}, ::EntrDetr, atmos, args) where {i}
    @unpack E_dyn, Δ_dyn, E_trb, env, ρa_up, ts_en = args.precomputed.turbconv
    up = args.state.turbconv.updraft
    θ_liq_en = liquid_ice_pottemp(ts_en)
    entr = fix_void_up(ρa_up[i], (E_dyn[i] + E_trb[i]) * θ_liq_en)
    detr =
        fix_void_up(ρa_up[i], (Δ_dyn[i] + E_trb[i]) * up[i].ρaθ_liq / ρa_up[i])

    return entr - detr
end

function source(::up_ρaq_tot{i}, ::EntrDetr, atmos, args) where {i}
    @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv
    up = args.state.turbconv.updraft
    q_tot_en = total_specific_humidity(ts_en)
    entr = fix_void_up(ρa_up[i], (E_dyn[i] + E_trb[i]) * q_tot_en)
    detr =
        fix_void_up(ρa_up[i], (Δ_dyn[i] + E_trb[i]) * up[i].ρaq_tot / ρa_up[i])

    return entr - detr
end

function source(::en_ρatke, ::EntrDetr, atmos, args)
    @unpack E_dyn, Δ_dyn, E_trb, env, ρa_up, w_up = args.precomputed.turbconv
    @unpack state = args
    up = state.turbconv.updraft
    en = state.turbconv.environment
    gm = state
    N_up = n_updrafts(turbconv_model(atmos))
    ρ_inv = 1 / gm.ρ
    tke_en = enforce_positivity(en.ρatke) * ρ_inv / env.a

    entr_detr = vuntuple(N_up) do i
        fix_void_up(
            ρa_up[i],
            E_trb[i] * (env.w - gm.ρu[3] * ρ_inv) * (env.w - w_up[i]) -
            (E_dyn[i] + E_trb[i]) * tke_en +
            Δ_dyn[i] * (w_up[i] - env.w) * (w_up[i] - env.w) / 2,
        )
    end
    return sum(entr_detr)
end

function source(::en_ρaθ_liq_cv, ::EntrDetr, atmos, args)
    @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv
    @unpack state = args
    ts_gm = args.precomputed.ts
    up = state.turbconv.updraft
    en = state.turbconv.environment
    N_up = n_updrafts(turbconv_model(atmos))
    θ_liq = liquid_ice_pottemp(ts_gm)
    θ_liq_en = liquid_ice_pottemp(ts_en)

    entr_detr = vuntuple(N_up) do i
        fix_void_up(
            ρa_up[i],
            Δ_dyn[i] *
            (up[i].ρaθ_liq / ρa_up[i] - θ_liq_en) *
            (up[i].ρaθ_liq / ρa_up[i] - θ_liq_en) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq / ρa_up[i]) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq / ρa_up[i]) -
            (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_cv,
        )
    end
    return sum(entr_detr)
end

function source(::en_ρaq_tot_cv, ::EntrDetr, atmos, args)
    @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv
    @unpack state = args
    FT = eltype(state)
    up = state.turbconv.updraft
    en = state.turbconv.environment
    gm = state
    N_up = n_updrafts(turbconv_model(atmos))
    q_tot_en = total_specific_humidity(ts_en)
    ρ_inv = 1 / gm.ρ
    ρq_tot = moisture_model(atmos) isa DryModel ? FT(0) : gm.moisture.ρq_tot

    entr_detr = vuntuple(N_up) do i
        fix_void_up(
            ρa_up[i],
            Δ_dyn[i] *
            (up[i].ρaq_tot / ρa_up[i] - q_tot_en) *
            (up[i].ρaq_tot / ρa_up[i] - q_tot_en) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot / ρa_up[i]) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot / ρa_up[i]) -
            (E_dyn[i] + E_trb[i]) * en.ρaq_tot_cv,
        )
    end
    return sum(entr_detr)
end

function source(::en_ρaθ_liq_q_tot_cv, ::EntrDetr, atmos, args)
    @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv
    @unpack state = args
    FT = eltype(state)
    ts_gm = args.precomputed.ts
    up = state.turbconv.updraft
    en = state.turbconv.environment
    gm = state
    N_up = n_updrafts(turbconv_model(atmos))
    q_tot_en = total_specific_humidity(ts_en)
    θ_liq = liquid_ice_pottemp(ts_gm)
    θ_liq_en = liquid_ice_pottemp(ts_en)
    ρ_inv = 1 / gm.ρ
    ρq_tot = moisture_model(atmos) isa DryModel ? FT(0) : gm.moisture.ρq_tot

    entr_detr = vuntuple(N_up) do i
        fix_void_up(
            ρa_up[i],
            Δ_dyn[i] *
            (up[i].ρaθ_liq / ρa_up[i] - θ_liq_en) *
            (up[i].ρaq_tot / ρa_up[i] - q_tot_en) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (q_tot_en - up[i].ρaq_tot / ρa_up[i]) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (θ_liq_en - up[i].ρaθ_liq / ρa_up[i]) -
            (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_q_tot_cv,
        )
    end
    return sum(entr_detr)
end

function source(::en_ρatke, ::PressSource, atmos, args)
    @unpack env, ρa_up, dpdz, w_up = args.precomputed.turbconv
    up = args.state.turbconv.updraft
    N_up = n_updrafts(turbconv_model(atmos))
    press_tke = vuntuple(N_up) do i
        fix_void_up(ρa_up[i], ρa_up[i] * (w_up[i] - env.w) * dpdz[i])
    end
    return sum(press_tke)
end

function source(::en_ρatke, ::ShearSource, atmos, args)
    @unpack env, K_m, Shear² = args.precomputed.turbconv
    gm = args.state
    ρa₀ = gm.ρ * env.a
    # production from mean gradient and Dissipation
    return ρa₀ * K_m * Shear² # tke Shear source
end

function source(::en_ρatke, ::BuoySource, atmos, args)
    @unpack env, K_h, ∂b∂z_env = args.precomputed.turbconv
    gm = args.state
    ρa₀ = gm.ρ * env.a
    return -ρa₀ * K_h * ∂b∂z_env   # tke Buoyancy source
end

function source(::en_ρatke, ::DissSource, atmos, args)
    @unpack Diss₀ = args.precomputed.turbconv
    en = args.state.turbconv.environment
    return -Diss₀ * en.ρatke  # tke Dissipation
end

function source(::en_ρaθ_liq_cv, ::DissSource, atmos, args)
    @unpack Diss₀ = args.precomputed.turbconv
    en = args.state.turbconv.environment
    return -Diss₀ * en.ρaθ_liq_cv
end

function source(::en_ρaq_tot_cv, ::DissSource, atmos, args)
    @unpack Diss₀ = args.precomputed.turbconv
    en = args.state.turbconv.environment
    return -Diss₀ * en.ρaq_tot_cv
end

function source(::en_ρaθ_liq_q_tot_cv, ::DissSource, atmos, args)
    @unpack Diss₀ = args.precomputed.turbconv
    en = args.state.turbconv.environment
    return -Diss₀ * en.ρaθ_liq_q_tot_cv
end

function source(::en_ρaθ_liq_cv, ::GradProdSource, atmos, args)
    @unpack env, K_h = args.precomputed.turbconv
    gm = args.state
    en_dif = args.diffusive.turbconv.environment
    ρa₀ = gm.ρ * env.a
    return ρa₀ * (2 * K_h * en_dif.∇θ_liq[3] * en_dif.∇θ_liq[3])
end

function source(::en_ρaq_tot_cv, ::GradProdSource, atmos, args)
    @unpack env, K_h = args.precomputed.turbconv
    gm = args.state
    en_dif = args.diffusive.turbconv.environment
    ρa₀ = gm.ρ * env.a
    return ρa₀ * (2 * K_h * en_dif.∇q_tot[3] * en_dif.∇q_tot[3])
end

function source(::en_ρaθ_liq_q_tot_cv, ::GradProdSource, atmos, args)
    @unpack env, K_h = args.precomputed.turbconv
    gm = args.state
    en_dif = args.diffusive.turbconv.environment
    ρa₀ = gm.ρ * env.a
    return ρa₀ * (2 * K_h * en_dif.∇θ_liq[3] * en_dif.∇q_tot[3])
end

function source(::up_ρaw{i}, ::BuoySource, atmos, args) where {i}
    @unpack buoy = args.precomputed.turbconv
    up = args.state.turbconv.updraft
    return up[i].ρa * buoy.up[i]
end

function source(::up_ρaw{i}, ::PressSource, atmos, args) where {i}
    @unpack dpdz = args.precomputed.turbconv
    up = args.state.turbconv.updraft
    return -up[i].ρa * dpdz[i]
end

function compute_ρa_up(atmos, state, aux)
    # Aliases:
    turbconv = turbconv_model(atmos)
    gm = state
    up = state.turbconv.updraft
    N_up = n_updrafts(turbconv)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max
    # in future GCM implementations we need to think about grid mean advection
    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa / gm.ρ, a_min, a_max)
    end
    return ρa_up
end

function flux(::up_ρa{i}, ::Advect, atmos, args) where {i}
    @unpack state, aux = args
    @unpack ρa_up = args.precomputed.turbconv
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(atmos, aux)
    return fix_void_up(ρa_up[i], up[i].ρaw) * ẑ
end
function flux(::up_ρaw{i}, ::Advect, atmos, args) where {i}
    @unpack state, aux = args
    @unpack ρa_up, w_up = args.precomputed.turbconv
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(atmos, aux)
    return fix_void_up(ρa_up[i], up[i].ρaw * w_up[i]) * ẑ

end
function flux(::up_ρaθ_liq{i}, ::Advect, atmos, args) where {i}
    @unpack state, aux = args
    @unpack ρa_up, w_up = args.precomputed.turbconv
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(atmos, aux)
    return fix_void_up(ρa_up[i], w_up[i] * up[i].ρaθ_liq) * ẑ

end
function flux(::up_ρaq_tot{i}, ::Advect, atmos, args) where {i}
    @unpack state, aux = args
    @unpack ρa_up, w_up = args.precomputed.turbconv
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(atmos, aux)
    return fix_void_up(ρa_up[i], w_up[i] * up[i].ρaq_tot) * ẑ

end

function flux(::en_ρatke, ::Advect, atmos, args)
    @unpack state, aux = args
    @unpack env = args.precomputed.turbconv
    en = state.turbconv.environment
    ẑ = vertical_unit_vector(atmos, aux)
    return en.ρatke * env.w * ẑ
end
function flux(::en_ρaθ_liq_cv, ::Advect, atmos, args)
    @unpack state, aux = args
    @unpack env = args.precomputed.turbconv
    en = state.turbconv.environment
    ẑ = vertical_unit_vector(atmos, aux)
    return en.ρaθ_liq_cv * env.w * ẑ
end
function flux(::en_ρaq_tot_cv, ::Advect, atmos, args)
    @unpack state, aux = args
    @unpack env = args.precomputed.turbconv
    en = state.turbconv.environment
    ẑ = vertical_unit_vector(atmos, aux)
    return en.ρaq_tot_cv * env.w * ẑ
end
function flux(::en_ρaθ_liq_q_tot_cv, ::Advect, atmos, args)
    @unpack state, aux = args
    @unpack env = args.precomputed.turbconv
    en = state.turbconv.environment
    ẑ = vertical_unit_vector(atmos, aux)
    return en.ρaθ_liq_q_tot_cv * env.w * ẑ
end

function precompute(::EDMF, bl, args, ts, ::Flux{FirstOrder})
    @unpack state, aux = args
    FT = eltype(state)
    turbconv = turbconv_model(bl)
    env = environment_vars(state, n_updrafts(turbconv))
    ρa_up = compute_ρa_up(bl, state, aux)
    up = state.turbconv.updraft
    gm = state
    N_up = n_updrafts(turbconv)
    ρ_inv = 1 / gm.ρ
    w_up = vuntuple(N_up) do i
        fix_void_up(ρa_up[i], up[i].ρaw / ρa_up[i])
    end

    θ_liq_up = vuntuple(N_up) do i
        fix_void_up(up[i].ρa, up[i].ρaθ_liq / up[i].ρa, liquid_ice_pottemp(ts))
    end
    a_up = vuntuple(N_up) do i
        fix_void_up(up[i].ρa, up[i].ρa * ρ_inv)
    end
    if !(moisture_model(bl) isa DryModel)
        q_tot_up = vuntuple(N_up) do i
            fix_void_up(up[i].ρa, up[i].ρaq_tot / up[i].ρa, gm.moisture.ρq_tot)
        end
    else
        q_tot_up = vuntuple(i -> FT(0), N_up)
    end

    return (; env, a_up, q_tot_up, ρa_up, θ_liq_up, w_up)
end

function precompute(::EDMF, bl, args, ts, ::Flux{SecondOrder})
    @unpack state, aux, diffusive, t = args
    ts_gm = ts
    up = state.turbconv.updraft
    turbconv = turbconv_model(bl)
    N_up = n_updrafts(turbconv)
    env = environment_vars(state, N_up)
    ts_en = new_thermo_state_en(bl, moisture_model(bl), state, aux, ts_gm)
    ts_up = new_thermo_state_up(bl, moisture_model(bl), state, aux, ts_gm)

    buoy = compute_buoyancy(bl, state, env, ts_en, ts_up, aux.ref_state)

    E_dyn, Δ_dyn, E_trb = entr_detr(bl, state, aux, ts_up, ts_en, env, buoy)

    Shear² =
        diffusive.turbconv.∇u[3]^2 +
        diffusive.turbconv.∇v[3]^2 +
        diffusive.turbconv.environment.∇w[3]^2
    l_mix, ∂b∂z_env, Pr_t = mixing_length(
        bl,
        turbconv.mix_len,
        args,
        Δ_dyn,
        E_trb,
        Shear²,
        ts_gm,
        ts_en,
        env,
    )
    ρa_up = compute_ρa_up(bl, state, aux)

    en = state.turbconv.environment
    tke_en = enforce_positivity(en.ρatke) / env.a / state.ρ
    K_m = turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    K_h = K_m / Pr_t
    ρaw_up = vuntuple(i -> up[i].ρaw, N_up)

    return (;
        env,
        ρa_up,
        ρaw_up,
        ts_en,
        ts_up,
        E_dyn,
        Δ_dyn,
        E_trb,
        l_mix,
        ∂b∂z_env,
        K_h,
        K_m,
        Pr_t,
    )
end

"""
    compute_buoyancy(
        bl::BalanceLaw,
        state::Vars,
        env::NamedTuple,
        ts_en::ThermodynamicState,
        ts_up,
        ref_state::Vars
    )

Compute buoyancies of subdomains
"""
function compute_buoyancy(
    bl::BalanceLaw,
    state::Vars,
    env::NamedTuple,
    ts_en::ThermodynamicState,
    ts_up,
    ref_state::Vars,
)
    FT = eltype(state)
    param_set = parameter_set(bl)
    N_up = n_updrafts(turbconv_model(bl))
    _grav::FT = grav(param_set)
    gm = state
    ρ_inv = 1 / gm.ρ
    buoyancy_en = -_grav * (air_density(ts_en) - ref_state.ρ) * ρ_inv
    up = state.turbconv.updraft

    a_up = vuntuple(N_up) do i
        fix_void_up(up[i].ρa, up[i].ρa * ρ_inv)
    end

    abs_buoyancy_up = vuntuple(N_up) do i
        -_grav * (air_density(ts_up[i]) - ref_state.ρ) * ρ_inv
    end
    b_gm = grid_mean_b(env, a_up, N_up, abs_buoyancy_up, buoyancy_en)

    # remove the gm_b from all subdomains
    buoyancy_up = vuntuple(N_up) do i
        abs_buoyancy_up[i] - b_gm
    end

    buoyancy_en -= b_gm
    return (; up = buoyancy_up, en = buoyancy_en)
end

function precompute(::EDMF, bl, args, ts, ::Source)
    @unpack state, aux, diffusive, t = args
    ts_gm = ts
    gm = state
    up = state.turbconv.updraft
    turbconv = turbconv_model(bl)
    N_up = n_updrafts(turbconv)
    # Get environment variables
    env = environment_vars(state, N_up)
    # Recover thermo states
    ts_en = new_thermo_state_en(bl, moisture_model(bl), state, aux, ts_gm)
    ts_up = new_thermo_state_up(bl, moisture_model(bl), state, aux, ts_gm)
    ρa_up = compute_ρa_up(bl, state, aux)

    Shear² =
        diffusive.turbconv.∇u[3]^2 +
        diffusive.turbconv.∇v[3]^2 +
        diffusive.turbconv.environment.∇w[3]^2

    buoy = compute_buoyancy(bl, state, env, ts_en, ts_up, aux.ref_state)
    E_dyn, Δ_dyn, E_trb = entr_detr(bl, state, aux, ts_up, ts_en, env, buoy)

    dpdz = perturbation_pressure(bl, args, env, buoy)

    l_mix, ∂b∂z_env, Pr_t = mixing_length(
        bl,
        turbconv.mix_len,
        args,
        Δ_dyn,
        E_trb,
        Shear²,
        ts_gm,
        ts_en,
        env,
    )

    w_up = vuntuple(N_up) do i
        fix_void_up(ρa_up[i], up[i].ρaw / ρa_up[i])
    end

    en = state.turbconv.environment
    tke_en = enforce_positivity(en.ρatke) / env.a / state.ρ
    K_m = turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    K_h = K_m / Pr_t
    Diss₀ = turbconv.mix_len.c_d * sqrt(tke_en) / l_mix
    tke_buoy_prod = -gm.ρ * env.a * K_h * ∂b∂z_env   # tke Buoyancy source

    return (;
        env,
        Diss₀,
        buoy,
        K_m,
        K_h,
        ρa_up,
        w_up,
        ts_en,
        ts_up,
        E_dyn,
        Δ_dyn,
        E_trb,
        dpdz,
        l_mix,
        ∂b∂z_env,
        Pr_t,
        Shear²,
        tke_buoy_prod,
    )
end

function flux(::Energy, ::SGSFlux, atmos, args)
    @unpack state, aux, diffusive = args
    @unpack ts = args.precomputed
    @unpack env, K_h, ρa_up, ρaw_up, ts_up, ts_en = args.precomputed.turbconv
    FT = eltype(state)
    param_set = parameter_set(atmos)
    _grav::FT = grav(param_set)
    z = altitude(atmos, aux)
    en_dif = diffusive.turbconv.environment
    up = state.turbconv.updraft
    gm = state
    ρ_inv = 1 / gm.ρ
    N_up = n_updrafts(turbconv_model(atmos))
    ρu_gm_tup = Tuple(gm.ρu)
    ρa_en = gm.ρ * env.a
    # TODO: Consider turbulent contribution:

    e_kin_up = vuntuple(N_up) do i
        FT(1 // 2) * (
            (gm.ρu[1] * ρ_inv)^2 +
            (gm.ρu[2] * ρ_inv)^2 +
            (fix_void_up(up[i].ρa, up[i].ρaw / up[i].ρa))^2
        )
    end
    e_kin_en =
        FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w)^2

    e_tot_up = ntuple(i -> total_energy(e_kin_up[i], _grav * z, ts_up[i]), N_up)
    h_tot_up = ntuple(i -> total_specific_enthalpy(ts_up[i], e_tot_up[i]), N_up)

    e_tot_en = total_energy(e_kin_en, _grav * z, ts_en)
    h_tot_en = total_specific_enthalpy(ts_en, e_tot_en)
    h_tot_gm = total_specific_enthalpy(ts, gm.energy.ρe * ρ_inv)

    massflux_h_tot = sum(
        ntuple(N_up) do i
            fix_void_up(
                ρa_up[i],
                ρa_up[i] *
                (h_tot_gm - h_tot_up[i]) *
                (gm.ρu[3] * ρ_inv - ρaw_up[i] / ρa_up[i]),
            )
        end,
    )
    massflux_h_tot +=
        (ρa_en * (h_tot_gm - h_tot_en) * (ρu_gm_tup[3] * ρ_inv - env.w))
    ρh_sgs_flux = -gm.ρ * env.a * K_h * en_dif.∇h_tot[3] + massflux_h_tot
    return SVector{3, FT}(0, 0, ρh_sgs_flux)
end

function flux(::TotalMoisture, ::SGSFlux, atmos, args)
    @unpack state, diffusive = args
    @unpack env, K_h, ρa_up, ρaw_up, ts_en = args.precomputed.turbconv
    FT = eltype(state)
    en_dif = diffusive.turbconv.environment
    up = state.turbconv.updraft
    gm = state
    ρ_inv = 1 / gm.ρ
    N_up = n_updrafts(turbconv_model(atmos))
    ρq_tot = moisture_model(atmos) isa DryModel ? FT(0) : gm.moisture.ρq_tot
    ρaq_tot_up = vuntuple(i -> up[i].ρaq_tot, N_up)
    ρa_en = gm.ρ * env.a
    q_tot_en = total_specific_humidity(ts_en)

    ρu_gm_tup = Tuple(gm.ρu)

    massflux_q_tot = sum(
        ntuple(N_up) do i
            fix_void_up(
                ρa_up[i],
                ρa_up[i] *
                (ρq_tot * ρ_inv - ρaq_tot_up[i] / ρa_up[i]) *
                (ρu_gm_tup[3] * ρ_inv - ρaw_up[i] / ρa_up[i]),
            )
        end,
    )
    massflux_q_tot +=
        (ρa_en * (ρq_tot * ρ_inv - q_tot_en) * (ρu_gm_tup[3] * ρ_inv - env.w))
    ρq_tot_sgs_flux = -gm.ρ * env.a * K_h * en_dif.∇q_tot[3] + massflux_q_tot
    return SVector{3, FT}(0, 0, ρq_tot_sgs_flux)
end

function flux(::Momentum, ::SGSFlux, atmos, args)
    @unpack state, diffusive = args
    @unpack env, K_m, ρa_up, ρaw_up = args.precomputed.turbconv
    FT = eltype(state)
    en_dif = diffusive.turbconv.environment
    gm_dif = diffusive.turbconv
    up = state.turbconv.updraft
    gm = state
    ρ_inv = 1 / gm.ρ
    N_up = n_updrafts(turbconv_model(atmos))
    ρa_en = gm.ρ * env.a

    ρu_gm_tup = Tuple(gm.ρu)

    massflux_w = sum(
        ntuple(N_up) do i
            fix_void_up(
                ρa_up[i],
                ρa_up[i] *
                (ρu_gm_tup[3] * ρ_inv - ρaw_up[i] / ρa_up[i]) *
                (ρu_gm_tup[3] * ρ_inv - ρaw_up[i] / ρa_up[i]),
            )
        end,
    )
    massflux_w += (
        ρa_en * (ρu_gm_tup[3] * ρ_inv - env.w) * (ρu_gm_tup[3] * ρ_inv - env.w)
    )
    ρw_sgs_flux = -gm.ρ * env.a * K_m * en_dif.∇w[3] + massflux_w
    ρu_sgs_flux = -gm.ρ * env.a * K_m * gm_dif.∇u[3]
    ρv_sgs_flux = -gm.ρ * env.a * K_m * gm_dif.∇v[3]
    return SMatrix{3, 3, FT, 9}(
        0,
        0,
        ρu_sgs_flux,
        0,
        0,
        ρv_sgs_flux,
        0,
        0,
        ρw_sgs_flux,
    )
end

function flux(::en_ρaθ_liq_cv, ::Diffusion, atmos, args)
    @unpack state, aux, diffusive = args
    @unpack env, l_mix, Pr_t, K_h = args.precomputed.turbconv
    en_dif = diffusive.turbconv.environment
    gm = state
    ẑ = vertical_unit_vector(atmos, aux)
    return -gm.ρ * env.a * K_h * en_dif.∇θ_liq_cv[3] * ẑ
end
function flux(::en_ρaq_tot_cv, ::Diffusion, atmos, args)
    @unpack state, aux, diffusive = args
    @unpack env, l_mix, Pr_t, K_h = args.precomputed.turbconv
    en_dif = diffusive.turbconv.environment
    gm = state
    ẑ = vertical_unit_vector(atmos, aux)
    return -gm.ρ * env.a * K_h * en_dif.∇q_tot_cv[3] * ẑ
end
function flux(::en_ρaθ_liq_q_tot_cv, ::Diffusion, atmos, args)
    @unpack state, aux, diffusive = args
    @unpack env, l_mix, Pr_t, K_h = args.precomputed.turbconv
    en_dif = diffusive.turbconv.environment
    gm = state
    ẑ = vertical_unit_vector(atmos, aux)
    return -gm.ρ * env.a * K_h * en_dif.∇θ_liq_q_tot_cv[3] * ẑ
end
function flux(::en_ρatke, ::Diffusion, atmos, args)
    @unpack state, aux, diffusive = args
    @unpack env, K_m = args.precomputed.turbconv
    gm = state
    en_dif = diffusive.turbconv.environment
    ẑ = vertical_unit_vector(atmos, aux)
    return -gm.ρ * env.a * K_m * en_dif.∇tke[3] * ẑ
end

# First order boundary conditions
function turbconv_boundary_state!(
    nf,
    bc::EDMFBottomBC,
    atmos::AtmosModel{FT},
    state⁺::Vars,
    args,
) where {FT}
    @unpack state⁻, aux⁻, aux_int⁻ = args
    turbconv = turbconv_model(atmos)
    N_up = n_updrafts(turbconv)
    up⁺ = state⁺.turbconv.updraft
    en⁺ = state⁺.turbconv.environment
    gm⁻ = state⁻
    gm_a⁻ = aux⁻

    zLL = altitude(atmos, aux_int⁻)
    surf_vals = subdomain_surface_values(atmos, gm⁻, gm_a⁻, zLL)
    a_up_surf = surf_vals.a_up_surf

    @unroll_map(N_up) do i
        up⁺[i].ρaw = FT(0)
        up⁺[i].ρa = gm⁻.ρ * a_up_surf[i]
        up⁺[i].ρaθ_liq = gm⁻.ρ * a_up_surf[i] * surf_vals.θ_liq_up_surf[i]
        if !(moisture_model(atmos) isa DryModel)
            up⁺[i].ρaq_tot = gm⁻.ρ * a_up_surf[i] * surf_vals.q_tot_up_surf[i]
        else
            up⁺[i].ρaq_tot = FT(0)
        end
    end

    a_en = environment_area(gm⁻, N_up)
    en⁺.ρatke = gm⁻.ρ * a_en * surf_vals.tke
    en⁺.ρaθ_liq_cv = gm⁻.ρ * a_en * surf_vals.θ_liq_cv
    if !(moisture_model(atmos) isa DryModel)
        en⁺.ρaq_tot_cv = gm⁻.ρ * a_en * surf_vals.q_tot_cv
        en⁺.ρaθ_liq_q_tot_cv = gm⁻.ρ * a_en * surf_vals.θ_liq_q_tot_cv
    else
        en⁺.ρaq_tot_cv = FT(0)
        en⁺.ρaθ_liq_q_tot_cv = FT(0)
    end
end;
function turbconv_boundary_state!(
    nf,
    bc::EDMFTopBC,
    atmos::AtmosModel{FT},
    state⁺::Vars,
    args,
) where {FT}
    N_up = n_updrafts(turbconv_model(atmos))
    up⁺ = state⁺.turbconv.updraft
    @unroll_map(N_up) do i
        up⁺[i].ρaw = FT(0)
    end
end;

# The boundary conditions for second-order unknowns
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFBottomBC,
    atmos::AtmosModel,
    fluxᵀn::Vars,
    _...,
)
    nothing
end;
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFTopBC,
    atmos::AtmosModel{FT},
    fluxᵀn::Vars,
    args,
) where {FT}

    turbconv = turbconv_model(atmos)
    N_up = n_updrafts(turbconv)
    up_flx = fluxᵀn.turbconv.updraft
    en_flx = fluxᵀn.turbconv.environment
    @unroll_map(N_up) do i
        up_flx[i].ρa = FT(0)
        up_flx[i].ρaθ_liq = FT(0)
        up_flx[i].ρaq_tot = FT(0)
    end
    en_flx.ρatke = FT(0)
    en_flx.ρaθ_liq_cv = FT(0)
    en_flx.ρaq_tot_cv = FT(0)
    en_flx.ρaθ_liq_q_tot_cv = FT(0)

end;


text:
using ForwardDiff

include(string(pwd(),"/src/likelihood_and_diff_old_kalman.jl"))

to_unit(x) = 1/(1+exp(-x))
to_positive(x) = exp(x)
to_one_inf(x) = exp(x) + 1

d_to_unit(x) = exp(-x)/(exp(-x)+1)^2
d_to_positive(x) = exp(x)
d_to_one_inf(x) = exp(x)

d2_to_unit(x) = ForwardDiff.derivative(xx->abs(d_to_unit(xx)),x)
d2_to_positive(x) = exp(x)
d2_to_one_inf(x) = exp(x)

function renormalization(par)

    J_1 = d_to_unit(par[2])
    J_2 = d_to_positive(par[3])
    J_3 = d_to_unit(par[4])
    J_4 = d_to_positive(par[5])
    J_5 = d_to_positive(par[6])
    J_6 = d_to_positive(par[7])
    J_7 = d_to_one_inf(par[8])
    J_8 = d_to_positive(par[9])
    J_9 = d_to_unit(par[10])

    return([J_1;J_2;J_3;J_4;J_5;J_6;J_7;J_8;J_9])
end

function dens(par,data)
    alfa = par[1]
    bet = to_unit(par[2])
    epsilon = to_positive(par[3])
    theta = to_unit(par[4])
    sig = to_positive(par[5])
    s2 = to_positive(par[6])
    phi = to_positive(par[7])
    phi_pi = to_one_inf(par[8])
    phi_y = to_positive(par[9])
    rho_v = to_unit(par[10])

    renorm_par = [alfa;bet;epsilon;theta;sig;s2;phi;phi_pi;phi_y;rho_v]

    renorm = renormalization(par)
    jacob = abs(prod(renorm))

    log_p_bet(bet) = logpdf(prior_bet,bet)
    log_p_epsilon(epsilon) = logpdf(prior_epsilon,epsilon)
    log_p_theta(theta) = logpdf(prior_theta,theta)
    log_p_sig(sig) = logpdf(prior_sig,sig)
    log_p_s2(s2) = logpdf(prior_s2,s2)
    log_p_phi(phi) = logpdf(prior_phi,phi)
    log_p_phi_pi(phi_pi) = logpdf(prior_phi_pi,phi_pi)
    log_p_phi_y(phi_y) = logpdf(prior_phi_y,phi_y)
    log_p_rho_v(rho_v) = logpdf(prior_rho_v,rho_v)
    llh,dll = log_like_dsge(renorm_par,data)
    llh = llh + log(jacob) + log_p_bet(bet) + log_p_epsilon(epsilon) + log_p_theta(theta) + log_p_sig(sig) + log_p_s2(s2) + log_p_phi(phi) + log_p_phi_pi(phi_pi) + log_p_phi_y(phi_y) + log_p_rho_v(rho_v) + sum(log.(renorm))
    return llh
end

function dens_and_grad(par,data)
    alfa = par[1]
    bet = to_unit(par[2])
    epsilon = to_positive(par[3])
    theta = to_unit(par[4])
    sig = to_positive(par[5])
    s2 = to_positive(par[6])
    phi = to_positive(par[7])
    phi_pi = to_one_inf(par[8])
    phi_y = to_positive(par[9])
    rho_v = to_unit(par[10])

    nobs = size(data,1)

    renorm_par = [alfa;bet;epsilon;theta;sig;s2;phi;phi_pi;phi_y;rho_v]

    renorm = renormalization(par)
    jacob = abs(prod(renorm))

    renorm_diff = ForwardDiff.gradient(x->abs(prod(renormalization(x))),par) #derivate of abs det of Jacobian matrix

    diff_correction = 1/jacob*renorm_diff
    diff_correction = diff_correction[2:10]

    log_p_bet(bet) = logpdf(prior_bet,bet)
    log_p_epsilon(epsilon) = logpdf(prior_epsilon,epsilon)
    log_p_theta(theta) = logpdf(prior_theta,theta)
    log_p_sig(sig) = logpdf(prior_sig,sig)
    log_p_s2(s2) = logpdf(prior_s2,s2)
    log_p_phi(phi) = logpdf(prior_phi,phi)
    log_p_phi_pi(phi_pi) = logpdf(prior_phi_pi,phi_pi)
    log_p_phi_y(phi_y) = logpdf(prior_phi_y,phi_y)
    log_p_rho_v(rho_v) = logpdf(prior_rho_v,rho_v)

    ## Derivates for the priors

    diff_prior = [ForwardDiff.derivative(log_p_bet,bet); ForwardDiff.derivative(log_p_epsilon,epsilon); ForwardDiff.derivative(log_p_theta,theta); ForwardDiff.derivative(log_p_sig,sig); ForwardDiff.derivative(log_p_s2,s2); ForwardDiff.derivative(log_p_phi,phi); ForwardDiff.derivative(log_p_phi_pi,phi_pi); ForwardDiff.derivative(log_p_phi_y,phi_y); ForwardDiff.derivative(log_p_rho_v,rho_v)]

    #diff_prior_cor = [d2_to_unit(par[2]);d2_to_positive(par[3]);d2_to_unit(par[4]);d2_to_positive(par[5]);d2_to_positive(par[6]);d2_to_positive(par[7]);d2_to_one_inf(par[8]);d2_to_positive(par[9]);d2_to_unit(par[10])]

    #pdf_prior = [pdf(prior_bet,bet);pdf(prior_epsilon,epsilon);pdf(prior_theta,theta);pdf(prior_sig,sig);pdf(prior_s2,s2);pdf(prior_phi,phi);pdf(prior_phi_pi,phi_pi);pdf(prior_phi_y,phi_y);pdf(prior_rho_v,rho_v)]

    llh,dll = log_like_dsge(renorm_par,data)
    llh = llh +log(jacob) + log_p_bet(bet) + log_p_epsilon(epsilon) + log_p_theta(theta) + log_p_sig(sig) + log_p_s2(s2) + log_p_phi(phi) + log_p_phi_pi(phi_pi) + log_p_phi_y(phi_y) + log_p_rho_v(rho_v) + sum(log.(abs.(renorm)))

    dll = dll[2:10]

    dll = dll .*renorm + diff_correction + diff_prior.*renorm

    #dll = dll/nobs

    return llh,dll
end

max_stars_repo_name: danmrc/DSGE-HMC

max_issues_repo_name: danmrc/DSGE-HMC

max_forks_repo_name: danmrc/DSGE-HMC


text:
"""
    G

Gravitational constant in the code units:
- Mass in solar masses (MSun),
- Length in milliparsecs (mpc),
- Time in millions of years (Myr).

Set as `4498502.15157528`.
This value comes from `astropy`:
```python-repl
>>> astropy.__version__
'4.0.1.post1'
```

```python-repl
>>> from astropy import units as u
>>> from astropy.constants import G
>>> uG = (u.mpc ** 3)/(u.M_sun * u.Myr ** 2)
>>> G.to(uG).value
4498502.151575286
```
"""
const G = 4498502.151575286

"""
    cvel

Speed of light in the code units:
- Mass in solar masses (MSun),
- Length in milliparsecs (mpc),
- Time in millions of years (Myr).

Set as `306601393.7879527`.
This value comes from `astropy`:
```python-repl
>>> astropy.__version__
'4.0.1.post1'
```

```python-repl
>>> from astropy import units as u
>>> from astropy.constants import c
>>> uc = (u.mpc/u.Myr)
>>> c.to(uc).value
306601393.7879527
```
"""
const cvel = 306601393.7879527

"""
    d_BH

Distance to SMBH Sagittarius A* (in mpc).
Taken from Gillessen et al. (2017): `An Update on Monitoring Stellar Orbits in the Galactic Center`
"""
const d_BH = 8.32*10^6 # Distance to Sgr A* (in mpc)


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "cmark_gfm"
version = v"0.29.0" # 0.29.0.gfm.3

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/github/cmark-gfm.git", "cf7577d2f74289cb83de0a652afc1a8b08a37036"),
    DirectorySource("./bundled"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/cmark-gfm/
atomic_patch -p1 ${WORKSPACE}/srcdir/patches/extension_type_exports.patch
mkdir build
cd build/
cmake -DCMAKE_INSTALL_PREFIX=$prefix -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} -DCMAKE_BUILD_TYPE=Release ..
make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

# The products that we will ensure are always built
products = [
    LibraryProduct("libcmark-gfm", :libcmark_gfm),
    LibraryProduct("libcmark-gfm-extensions", :libcmark_gfm_extensions),
    ExecutableProduct("cmark-gfm", :cmark_gfm)
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6")