text:
################################################################################
#
#   IMPLEMENTATIONS OF DIFFERENT UNITCELL CONSTRUCTION RELATED STUFF
#
#   STRUCTURE OF THE FILE
#
#   1) UNITCELL FROM SITE COLLECTION (automatically determine the connections)
#       - 2D (not exported)
#       - 3D (not exported)
#       - independent of the dimension
#
#   2) UNITCELL GENERATING CODE
#
################################################################################

#-----------------------------------------------------------------------------------------------------------------------------
#
#   GET A UNITCELL OBJECT BY ONLY PASSING SITES AND LATTICE VECTORS
#
#   Parameters (required)
#   - sites: an array indicating the position of sites
#   - lattice_vectors: an array indicating the lattice vectors of the lattice
#
#   Parameters (optional)
#   - max_ij(k): the maximum number of unitcells in all direction that is beeing considered to build geometric neighbors
#   - epsilon: The precision up to which distances are considered to be equal
#   - min_NN: minimal number of nearest neighbors that ALL sites have to have (-1: Auto)
#   - max_NN: maximal number of nearest neighbors that ALL sites have to have (-1: Auto)
#   - strength_NN: interaction strength of the new connections
#   - name: name of the lattice (if "AUTO" is selected, a name will be build automatically)
#   - save: if to save the newly constructed unitcell
#
#-----------------------------------------------------------------------------------------------------------------------------
function getUnitcellFromSites2D(
        sites::Array{Array{Float64, 1}, 1},
        lattice_vectors::Array{Array{Float64, 1}, 1};
        max_ij::Int64=3,
        epsilon::Float64=1e-8,
        min_NN::Int64=-1,
        max_NN::Int64=-1,
        strength_NN=1.0,
        name::String="AUTO",
        save::Bool=false
    )
    # make a new list with connections
    connections = Array{Any, 1}[]
    # define what "close to" means
    function distance(p1, p2)
        return sqrt(sum((p1.-p2).*(p1.-p2)))
    end
    function closeto(p1, p2)
        return distance(p1, p2) < epsilon
    end
    # function to sort the checklist
    function sortfunction(checklistentry)
        return checklistentry[3]
    end
    # iterate over all sites
    for s in 1:size(sites, 1)
        # build a new checklist for sites
        checklist = Array[]
        # check all other sites in a lattice around this one
        for i in -max_ij:max_ij
        for j in -max_ij:max_ij
        for alpha in 1:size(sites, 1)
            # site at i*a1 + j*a2 + basis[alpha]
            position = i.*lattice_vectors[1] .+ j.*lattice_vectors[2] .+ sites[alpha]
            # check if the position is close to the original position
            if closeto(position, sites[s])
                # just the original site
                continue
            end
            # put the site into the list
            # format: [site_index, (wrap), distance]
            push!(checklist, [alpha, (i,j), distance(position, sites[s])])
        end
        end
        end
        # sort the checklist so that now the first element is the closest
        sort!(checklist, by=sortfunction)
        # evaluate the checklist
        # first: find out the number of nearest neighbors
        number_NN = 0
        # if unrestricted: simply go by distance and search for all sites
        # that are as close as the first (closest) neighbor
        if min_NN <= 0 && max_NN <= 0
            for c in checklist
                if closeto(c[3], checklist[1][3])
                    number_NN = number_NN + 1
                else
                    break
                end
            end
        # if restricted, search only between the min_NN^th and max_NN^th (closest) neighbors
        elseif min_NN > 0 && max_NN <=0
            number_NN = min_NN-1
            for c in checklist[min_NN:end]
                if closeto(c[3], checklist[min_NN][3])
                    number_NN = number_NN + 1
                else
                    break
                end
            end
        else
            if min_NN <= 0
                min_NN = 1
            end
            number_NN = min_NN-1
            for c in checklist[min_NN:max_NN]
                if closeto(c[3], checklist[min_NN][3])
                    number_NN = number_NN + 1
                else
                    break
                end
            end
        end
        # print the number of nearest neighbors
        println("number of nearest neighbors of site $(s): $(number_NN)")
        # check if maybe cut at the wrong element, i.e. if there are more nearest neighbors
        if closeto(checklist[number_NN][3],checklist[number_NN+1][3])
            println("cut at the wrong number of nearest neighbors, there are more that are equally close!")
        end
        # restrict the checklist to that length
        checklist = checklist[1:number_NN]
        # parse into connections
        for c in checklist
            # define a new connection of interaction strength type float or int
            if typeof(strength_NN) == Float64 || typeof(strength_NN) == Int64
                push!(connections, [s, c[1], strength_NN, c[2]])
            elseif typeof(strength_NN) == String
                # automatic connection geometry based strength
                if typeof(strength_NN) == String && strength_NN == "AUTO"
                    push!(connections, [s, c[1], 1/c[3], c[2]])
                end
            else
                println("unknown connection strength type \"$(typeof(strength_NN))\" of connection strength $(strength_NN)")
                return
            end
        end
    end
    # correct all connections - i.e. find out if the returning connection is present, if not: add it
    connections_additional = Array{Any, 1}[]
    for c in connections
        # check if "-c" is present
        found = false
        for mc in connections
            if c[1] == mc[2] && c[2] == mc[1] && c[3] == mc[3] && -c[4][1] == mc[4][1] && -c[4][2] == mc[4][2]
                found = true
            end
        end
        if !found
            push!(connections_additional, [c[2], c[1], c[3], (-c[4][1], -c[4][2])])
        end
    end
    # add all additional connections
    for c in connections_additional
        push!(connections, c)
    end
    # build a filename
    if name == "AUTO"
        name = "from_$(size(sites, 1))_sites"
    end
    filename = "$(FOLDER_UNITCELLS)2d_$(name)_unitcell.jld"
    # build the finished unitcell object
    unitcell = Unitcell(
        lattice_vectors,
        sites,
        connections,
        filename
    )
    # maybe save
    if save
        saveUnitcell(unitcell)
    end
    # return the unitcell
    return unitcell
end
function getUnitcellFromSites3D(
        sites::Array{Array{Float64, 1}, 1},
        lattice_vectors::Array{Array{Float64, 1}, 1};
        max_ijk::Int64=3,
        epsilon::Float64=1e-8,
        min_NN::Int64=-1,
        max_NN::Int64=-1,
        strength_NN=1.0,
        name::String="AUTO",
        save::Bool=false
    )
    # make a new list with connections
    connections = Array{Any, 1}[]
    # define what "close to" means
    function distance(p1, p2)
        return sqrt(sum((p1.-p2).*(p1.-p2)))
    end
    function closeto(p1, p2)
        return distance(p1, p2) < epsilon
    end
    # function to sort the checklist
    function sortfunction(checklistentry)
        return checklistentry[3]
    end
    # iterate over all sites
    for s in 1:size(sites, 1)
        # build a new checklist for sites
        checklist = Array[]
        # check all other sites in a lattice around this one
        for i in -max_ijk:max_ijk
        for j in -max_ijk:max_ijk
        for k in -max_ijk:max_ijk
        for alpha in 1:size(sites, 1)
            # site at i*a1 + j*a2 + k*a3 + basis[alpha]
            position = i.*lattice_vectors[1] .+ j.*lattice_vectors[2] .+ k.*lattice_vectors[3] .+ sites[alpha]
            # check if the position is close to the original position
            if closeto(position, sites[s])
                # just the original site
                continue
            end
            # put the site into the list
            # format: [site_index, (wrap), distance]
            push!(checklist, [alpha, (i,j,k), distance(position, sites[s])])
        end
        end
        end
        end
        # sort the checklist
        sort!(checklist, by=sortfunction)
        # evaluate the checklist
        # first: find out the number of nearest neighbors
        number_NN = 0
        # if unrestricted: simply go by distance and search for all sites
        # that are as close as the first (closest) neighbor
        if min_NN <= 0 && max_NN <= 0
            for c in checklist
                if closeto(c[3], checklist[1][3])
                    number_NN = number_NN + 1
                else
                    break
                end
            end
        # if restricted, search only between the min_NN^th and max_NN^th (closest) neighbors
        elseif min_NN > 0 && max_NN <=0
            number_NN = min_NN-1
            for c in checklist[min_NN:end]
                if closeto(c[3], checklist[min_NN][3])
                    number_NN = number_NN + 1
                else
                    break
                end
            end
        else
            if min_NN <= 0
                min_NN = 1
            end
            number_NN = min_NN-1
            for c in checklist[min_NN:max_NN]
                if closeto(c[3], checklist[min_NN][3])
                    number_NN = number_NN + 1
                else
                    break
                end
            end
        end
        # print the number of nearest neighbors
        println("number of nearest neighbors of site $(s): $(number_NN)")
        # check if maybe cut at the wrong element, i.e. if there are more nearest neighbors
        if closeto(checklist[number_NN][3],checklist[number_NN+1][3])
            println("cut at the wrong number of nearest neighbors, there are more that are equally close!")
        end
        # restrict the checklist to that length
        checklist = checklist[1:number_NN]
        # parse into connections
        for c in checklist
            # define a new connection of interaction strength type float or int
            if typeof(strength_NN) == Float64 || typeof(strength_NN) == Int64
                push!(connections, [s, c[1], strength_NN, c[2]])
            elseif typeof(strength_NN) == String
                # automatic connection geometry based strength
                if typeof(strength_NN) == String && strength_NN == "AUTO"
                    push!(connections, [s, c[1], 1/c[3], c[2]])
                end
            else
                println("unknown connection strength type \"$(typeof(strength_NN))\" of connection strength $(strength_NN)")
                return
            end
        end
    end
    # correct all connections - i.e. find out if the returning connection is present, if not: add it
    connections_additional = Array{Any, 1}[]
    for c in connections
        # check if "-c" is present
        found = false
        for mc in connections
            if c[1] == mc[2] && c[2] == mc[1] && c[3] == mc[3] && -c[4][1] == mc[4][1] && -c[4][2] == mc[4][2] && -c[4][3] == mc[4][3]
                found = true
            end
        end
        if !found
            push!(connections_additional, [c[2], c[1], c[3], (-c[4][1], -c[4][2], -c[4][3])])
        end
    end
    # add all additional connections
    for c in connections_additional
        push!(connections, c)
    end
    # build a filename
    if name == "AUTO"
        name = "from_$(size(sites, 1))_sites"
    end
    filename = "$(FOLDER_UNITCELLS)3d_$(name)_unitcell.jld"
    # build the finished unitcell object
    unitcell = Unitcell(
        lattice_vectors,
        sites,
        connections,
        filename
    )
    # maybe save
    if save
        saveUnitcell(unitcell)
    end
    # return the unitcell
    return unitcell
end

"""
    getUnitcellFromSites(
            sites::Array{Array{Float64, 1}, 1},
            lattice_vectors::Array{Array{Float64, 1}, 1}
            ;
            max_ij::Int64=3,
            epsilon::Float64=1e-8,
            min_NN::Int64=-1,
            max_NN::Int64=-1,
            strength_NN=1.0,
            name::String="AUTO",
            save::Bool=false
        )

Function that generates a fully periodic `Unitcell` object in 2D or 3D which is only based on a collection of `sites`
and a collection of `lattice_vectors` that is passed.
Connections will be based on nearest neighbors and will have a strength of `strength_NN`

The search assigns a number of nearest neighbors (between `min_NN` and `max_NN`)
to every site which are at most `epsilon` different in distance apart.
Then in a subsequent step, it will also register all returning connections.

Since the periodicity is determined by the number of `lattice_vectors` passed, the sites should have the same
dimension as this number. I.e. it is only possible to build fully periodic lattices.

# Examples

```julia-repl
julia> getUnitcellFromSites(sites, lattice_vectors)
LatticePhysics.Unitcell(...)

julia> getUnitcellFromSites(sites, lattice_vectors, name="myfancylattice")
LatticePhysics.Unitcell(...)
```
"""
function getUnitcellFromSites(
        sites::Array{Array{Float64, 1}, 1},
        lattice_vectors::Array{Array{Float64, 1}, 1};
        max_ijk::Int64=3,
        epsilon::Float64=1e-8,
        min_NN::Int64=-1,
        max_NN::Int64=-1,
        strength_NN=1.0,
        name::String="AUTO",
        save::Bool=false
    )
    # check the dimension and call the respective function
    if size(lattice_vectors, 1) == 1
        println("building from sites with one lattice vector not supported")
    elseif size(lattice_vectors, 1) == 2
        return getUnitcellFromSites2D(
            sites,
            lattice_vectors;
            max_ij=max_ijk,
            epsilon=epsilon,
            min_NN=min_NN,
            max_NN=max_NN,
            strength_NN=strength_NN,
            name=name,
            save=save
        )
    elseif size(lattice_vectors, 1) == 3
        return getUnitcellFromSites3D(
            sites,
            lattice_vectors;
            max_ijk=max_ijk,
            epsilon=epsilon,
            min_NN=min_NN,
            max_NN=max_NN,
            strength_NN=strength_NN,
            name=name,
            save=save
        )
    elseif size(lattice_vectors, 1) > 3
        println("building from sites with more than 3 lattice vector not supported")
    end
end
export getUnitcellFromSites

"""
    printUnitcellGeneratingCode(unitcell::Unitcell [; version::Int64=1])

Function that prints julia code into the console which can be copied and used from there on.
The printed code is a function definiton for construction of the `Unitcell` object passed.

If the additional `version` parameter is passed as well, the code will be versionized in
the convention that is also used through the library.

# Examples

```julia-repl
julia> printUnitcellGeneratingCode(unitcell)
...

julia> printUnitcellGeneratingCode(unitcell, version=4)
...
```
"""
function printUnitcellGeneratingCode(unitcell::Unitcell; version::Int64=1)
    # distinguish if versioning yes or no
    if version != -1
        # header printing
        println("function getUnitcellCustom(version::Int64=$(version); save::Bool=false)")
        # version distinguish
        println("   # distinguish by version")
        println("   if version == $(version)")
        # lattice vectors
        println("       # the lattice vectors")
        index = 0
        for l in unitcell.lattice_vectors
            index = index+1
            println("       a$(index) = $(l)")
        end
        println("       lattice_vectors = Array{Float64, 1}[]")
        for index in 1:size(unitcell.lattice_vectors, 1)
            println("       push!(lattice_vectors, a$(index))")
        end
        # Basis
        println("       # Basis Definition")
        println("       basis = Array{Float64, 1}[")
        for b in unitcell.basis[1:end-1]
            println("           $(b),")
        end
        println("           $(unitcell.basis[end])")
        println("       ]")
        # Connections
        println("       # Connection Definition")
        println("       # [<from index>; <to index>; <strength>; (<lattice displaced by lattice vector j>)]")
        println("       connections = Array{Any, 1}[")
        for c in unitcell.connections[1:end-1]
            if typeof(c[3]) == String
                println("           [$(c[1]); $(c[2]); \"$(c[3])\"; $(c[4])],")
            else
                println("           [$(c[1]); $(c[2]); $(c[3]); $(c[4])],")
            end
        end
        c = unitcell.connections[end]
        if typeof(c[3]) == String
            println("           [$(c[1]); $(c[2]); \"$(c[3])\"; $(c[4])],")
        else
            println("           [$(c[1]); $(c[2]); $(c[3]); $(c[4])],")
        end
        println("       ]")
        # Filename
        println("       # filename")
        if size(unitcell.lattice_vectors,1) == 2
            println("       filename = \"\$(FOLDER_UNITCELLS)2d_custom_unitcell.jld\"")
        elseif size(unitcell.lattice_vectors,1) == 3
            println("       filename = \"\$(FOLDER_UNITCELLS)3d_custom_unitcell.jld\"")
        else
            println("       filename = \"\$(FOLDER_UNITCELLS)Xd_custom_unitcell.jld\"")
        end
        # end
        println("   end")
    else
        # header printing
        println("function getUnitcellCustom(version::Int64=1; save::Bool=false)")
        # lattice vectors
        println("   # the lattice vectors")
        index = 0
        for l in unitcell.lattice_vectors
            index = index+1
            println("   a$(index) = $(l)")
        end
        println("   lattice_vectors = Array{Float64, 1}[]")
        for index in 1:size(unitcell.lattice_vectors, 1)
            println("   push!(lattice_vectors, a$(index))")
        end
        # Basis
        println("   # Basis Definition")
        println("   basis = Array{Float64, 1}[")
        for b in unitcell.basis[1:end-1]
            println("       $(b),")
        end
        println("       $(unitcell.basis[end])")
        println("   ]")
        # Connections
        println("   # Connection Definition")
        println("   # [<from index>; <to index>; <strength>; (<lattice displaced by lattice vector j>)]")
        println("   connections = Array{Any, 1}[")
        for c in unitcell.connections[1:end-1]
            if typeof(c[3]) == String
                println("       [$(c[1]); $(c[2]); \"$(c[3])\"; $(c[4])],")
            else
                println("       [$(c[1]); $(c[2]); $(c[3]); $(c[4])],")
            end
        end
        c = unitcell.connections[end]
        if typeof(c[3]) == String
            println("       [$(c[1]); $(c[2]); \"$(c[3])\"; $(c[4])],")
        else
            println("       [$(c[1]); $(c[2]); $(c[3]); $(c[4])],")
        end
        println("   ]")
        # Filename
        println("   # filename")
        if size(unitcell.lattice_vectors,1) == 2
            println("   filename = \"\$(FOLDER_UNITCELLS)2d_custom_unitcell.jld\"")
        elseif size(unitcell.lattice_vectors,1) == 3
            println("   filename = \"\$(FOLDER_UNITCELLS)3d_custom_unitcell.jld\"")
        else
            println("   filename = \"\$(FOLDER_UNITCELLS)Xd_custom_unitcell.jld\"")
        end
    end
    # part which is not version dependent
    println("   # generate unitcell")
    println("   uc = Unitcell(lattice_vectors, basis, connections, filename)")
    println("   if save")
    println("       saveUnitcell(uc)")
    println("   end")
    println("   # return the unitcell")
    println("   return uc")
    println("end")
end
export printUnitcellGeneratingCode


text:
using MAT

include("estimator.jl")

# Load raw data
data = matread("data.mat")

function main()

    # Create Batch Estimator from data
    batch_estimator = BatchEstimator(data)

    # Perform state estimation
    states1 = dead_reconing(batch_estimator)
    states2 = gauss_newton(batch_estimator, states1)
    plot_trajectory(batch_estimator, states1, states2)
end

main()

ext: jl

max_stars_repo_name: jandaa/state-estimation

max_issues_repo_name: jandaa/state-estimation

max_forks_repo_name: jandaa/state-estimation


text:
function random_similar(demands::Vector{Edge}, avg::Float64, stdev::Float64, rng)
    # rng is assumed to have a normal distribution (i.e. both positive and
    # negative values can be generated from it).
    values = avg .+ stdev .* randn(rng, length(demands))
    values[values .< 0.0] .= 1.0 # Remove negative traffic value, they make no sense.
    return Dict(demands[i] => values[i] for i in length(demands))
end

random_similar(demands::Dict{Edge{Int}, Float64}, avg::Float64, stdev::Float64, rng) =
    random_similar(collect(keys(demands)), avg, stdev, rng)

random_similar(rd::RoutingData, avg::Float64, stdev::Float64, rng) =
    random_similar(collect(demands(rd)), avg, stdev, rng)

function random_experiment(rd::RoutingData, solver, oblivious_routing::Routing,
                           avg::Float64, stdev::Float64;
                           rng=MersenneTwister())
    traffic = random_similar(rd, avg, stdev, rng)

    # Compute the optimum load for this traffic matrix, with a routing that
    # is perfectly adapted to it.
    m = _create_model(rd)
    capacitated_demand_routing_model(m, rd, traffic)
    optimize!(m)
    optimum_load = objective_value(m)

    # Maybe we generated an infeasible matrix, so the results would make no sense. Generate a new matrix.
    if optimum_load > 1.0
        return random_experiment(rd, solver, oblivious_routing, avg, stdev, rng=rng)
    end

    oblivious_load = compute_max_load(oblivious_routing, traffic)
    return oblivious_load, optimum_load
end

function random_experiments(rd::RoutingData, solver, oblivious_routing,
                           avg::Float64, stdev::Float64, n::Int=1; kwargs...)
    oblivious_loads = Float64[]
    optimum_loads = Float64[]
    for _ in 1:n
        experiment = random_experiment(rd, solver, oblivious_routing, avg, stdev; kwargs...)
        push!(oblivious_loads, experiment[1])
        push!(optimum_loads, experiment[2])
    end

    return oblivious_loads, optimum_loads, oblivious_loads ./ optimum_loads
end


text:
module QJuliaBlas

import QJuliaRegisters
import QJuliaIntrinsics

using Base.Threads

@inline function qjulia_blas_info()
  println("Experimental: executing blas routines in ", Threads.nthreads(), " threads.")
end

#create function/type alias
#SSE
m128d   = QJuliaRegisters.double2
m128    = QJuliaRegisters.float4
#AVX/AVX2
m256d   = QJuliaRegisters.double4
m256    = QJuliaRegisters.float8
#AVX3
m512d   = QJuliaRegisters.double8
m512    = QJuliaRegisters.float16

mm_mul  = QJuliaIntrinsics.mm_mul
mm_add  = QJuliaIntrinsics.mm_add
mm_sub  = QJuliaIntrinsics.mm_sub
mm_mad  = QJuliaIntrinsics.mm_mad

#collection of generic blas operations (e.g., on vector registers etc.)

@inline function convert_c2r(y::Vector{m256d}, x::Vector{Complex{T}}) where T <: AbstractFloat
  for i in 1:length(y)
    c=[real(x[2i-1]), imag(x[2i-1]), real(x[2i]), imag(x[2i]) ]
    y[i] = m256d(ntuple(i->c[i],4))
  end
end #convert_c2r

@inline function convert_c2r(y::Vector{m256}, x::Vector{Complex{T}}) where T <: AbstractFloat
  for i in 1:length(y)
    c=[real(x[4i-3]), imag(x[4i-3]), real(x[4i-2]), imag(x[4i-2]), real(x[4i-1]), imag(x[4i-1]), real(x[4i]), imag(x[4i]) ]
    y[i] = m256(ntuple(i->c[i],8))
  end
end #convert_c2r

@inline function convert_c2r(y::Vector{T}, x::Vector{Complex{T}}) where T <: AbstractFloat
@threads  for i in 1:length(x); y[2i-1] = real(x[i]); y[2i] = imag(x[i]); end
end #convert_c2r

@inline function convert_c2r(y::AbstractArray, x::Vector{Complex{T}}) where T <: AbstractFloat
@threads  for i in 1:length(x); y[2i-1] = real(x[i]); y[2i] = imag(x[i]); end
end #convert_c2r

@inline function convert_r2c(y::Vector{Complex{T}}, x::Vector{T}) where T <: AbstractFloat
@threads  for i in 1:length(y); y[i] = x[2i-1] + x[2i]*im; end
end #convert_r2c

@inline function convert_r2c(y::Vector{Complex{T}}, x::AbstractArray) where T <: AbstractFloat
@threads  for i in 1:length(y); y[i] = x[2i-1] + x[2i]*im; end
end #convert_r2c

@inline function cpy(y::AbstractArray, x::AbstractArray)
  if pointer_from_objref(y) == pointer_from_objref(x); return; end
@threads for i in 1:length(x);  y[i] = x[i]; end
end #cpy

@inline function axpy(a, x::AbstractArray, y::AbstractArray)
@threads for i in 1:length(x);  y[i] = mm_mad(a, x[i], y[i]); end
end #axpy

@inline function xpay(x::AbstractArray, a, y::AbstractArray)
@threads for i in 1:length(x);  y[i] = mm_mad(a, y[i], x[i]); end
end #xpay

@inline function xmy(x::AbstractArray, y::AbstractArray)
@threads for i in 1:length(x);  y[i] = mm_sub(x[i], y[i]); end
end #xmy

@inline function xpy(x::AbstractArray, y::AbstractArray)
@threads for i in 1:length(x);  y[i] = mm_add(x[i], y[i]); end
end #xpy

#First performs the operation y[i] += a*x[i]
#Second performs the operator x[i] -= a*z[i]

@inline function caxpyXmaz(a, x::Vector{T}, y::Vector{T}, z::Vector{T})  where T <: AbstractFloat

         cx = view(reinterpret(Complex{T}, x), :)
         cy = view(reinterpret(Complex{T}, y), :)
         cz = view(reinterpret(Complex{T}, z), :)

@threads for i in 1:length(cx)
           cy[i]  += a*cx[i]
           cx[i]  -= a*cz[i]
         end
end #caxpyXmaz

@inline function axpyXmaz(a, x::Vector{T}, y::Vector{T}, z::Vector{T})  where T <: AbstractFloat

@threads for i in 1:length(x)
		y[i]  += a*x[i]
		x[i]  -= a*z[i]
	end

end #axpyXmaz

#First performs the operation x[i] = x[i] + a*p[i]
#Second performs the operator p[i] = u[i] + b*p[i]

@inline function axpyZpbx(a, p::Vector{T}, x::Vector{T}, u::Vector{T}, b)  where T <: AbstractFloat

@threads for i in 1:length(x)
           x[i] = x[i]+a*p[i]
           p[i] = u[i]+b*p[i]
         end
end #axpyZpbx

end #QJuliaBlas

max_stars_repo_name: alexstrel/QJulia

max_issues_repo_name: alexstrel/QJulia

max_forks_repo_name: alexstrel/QJulia


text:
### RigidBodyDynamics ###
function RigidBodyDynamics.principal_value!(mech_scen::MechanismScenario, x::Vector{Float64})
    copyto!(mech_scen.float, x)
    principal_value!(mech_scen.float.state)
    copyto!(x, mech_scen.float)
    return nothing
end

function RigidBodyDynamics.set_configuration!(mech_scen::MechanismScenario, joint::Joint, config)
    return set_configuration!(mech_scen.float.state, joint, config)
end

### Binary_BB_Trees
Binary_BB_Trees.fit_tri_obb(eM::eMesh{Tri,Nothing}, k::Int64) = fit_tri_obb(eM.point[eM.tri[k]])
function Binary_BB_Trees.fit_tet_obb(eM::eMesh{Nothing,Tet}, k::Int64)
    i = eM.tet[k]
    return fit_tet_obb(eM.point[i], eM.ϵ[i])
end

### Base ###
function Base.copyto!(dest::TypedMechanismScenario{T}, src::AbstractVector{T}) where {N,T}
    nq = num_positions(dest.state)
    nv = num_velocities(dest.state)
    ns = length(dest.s)
    @boundscheck length(src) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(parent(dest.state.q), 1, src,           1, nq)
    @inbounds copyto!(parent(dest.state.v), 1, src, nq +      1, nv)
    @inbounds copyto!(dest.s,               1, src, nq + nv + 1, ns)
    setdirty!(dest.state)
    dest
end
function Base.copyto!(dest::AbstractVector{T}, src::TypedMechanismScenario{T}) where {N,T}
    nq = num_positions(src.state)
    nv = num_velocities(src.state)
    ns = length(src.s)
    length(dest) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(dest,           1, src.state.q, 1, nq)
    @inbounds copyto!(dest, nq +      1, src.state.v, 1, nv)
    @inbounds copyto!(dest, nq + nv + 1, src.s,       1, ns)
    dest
end
function Base.copyto!(ẋ::AbstractVector, src::TypedMechanismScenario{T}, result::DynamicsResult) where {N,T}
    nq = length(result.q̇)
    nv = length(result.v̇)
    ns = length(src.ṡ)
    @boundscheck length(ẋ) == nq + nv + ns || throw(DimensionMismatch())
    @inbounds copyto!(ẋ,           1, result.q̇, 1, nq)
    @inbounds copyto!(ẋ, nq +      1, result.v̇, 1, nv)
    @inbounds copyto!(ẋ, nq + nv + 1, src.ṡ,    1, ns)
    ẋ
end


text:
import OhMyREPL: OUTPUT_PROMPT, update_interface

function Base.REPL.display(d::Base.REPL.REPLDisplay, mime::MIME"text/plain", x)
    global OUTPUT_PROMPT
    io = Base.REPL.outstream(d.repl)
    write(io, OUTPUT_PROMPT)
    Base.have_color && write(io, Base.REPL.answer_color(d.repl))
    show(IOContext(io, :limit => true), mime, x)
    println(io)
end


text:
# === HowTo ===
# > julia -p <P> covid_run_exp.jl path/to/config-folder/ <N> [--append]
# P: Number of processes run in parallel
# N: Number of runs
# config-folder has to contain covid_par_ini_.jl-file
# if "--append", existing baseline.dat will be extended.

using Distributed
using Serialization

@everywhere using Agents, Serialization, Statistics,Random

filename_prefix = ARGS[1]
no_runs = parse(Int64, ARGS[2])
append = false
if length(ARGS)[1] > 2
	if ARGS[3] == "--append"
		append = true
	end
end

@everywhere function launch_model(filename_prefix, number)
	println("starting run $(number)")
	flush(stdout)

	ini_file = "$(filename_prefix)covid_par_ini.jl"

	include("covid_model.jl")
	include("covid_snapshot_loader.jl")
	include(ini_file)
	include("covid_main_p.jl")

	println("finished run $(number)")
	flush(stdout)
end

all_results = []
if append
	all_results = deserialize(open("$(filename_prefix)batchdata.dat"))
end

for idx in 1:(Int(floor(no_runs / nworkers())) + 1)
	from = (idx-1)*nworkers()+1
	to = min((idx-1)*nworkers() + nworkers(), no_runs)
	if from <= to
		println("Spawning workers for runs $(from) to $(to)...")
		worker_futures = []
		lastactiveworker=0
		for idxx = 1:nworkers()
			if (idx-1)*nworkers()+idxx <= no_runs
				push!(worker_futures, @spawnat (idxx+1) launch_model("$(filename_prefix)", (idx-1)*nworkers()+idxx))
				sleep(1)
				lastactiveworker = (idxx)
			end
		end
		fetch.(worker_futures)

		result_futures = []
		for idx = 1:lastactiveworker
			push!(result_futures, @spawnat (idx+1) (() -> Dict(:RKIR0traj => RKIR0traj,
		                :castraj => castraj,:totcastraj => totcastraj, :unempltraj => unempltraj,
		                :pubacctraj => pubacctraj,:constraj => constraj,
		                :unempsectraj => unempsectraj,:tautraj => tautraj,
		                :conspercaptraj => conspercaptraj, :R0counttraj => R0counttraj,
		                :gdppercaptraj => gdppercaptraj, :f_av_savings => f_av_savings,
		                :f_av_workers => f_av_workers,:f_var_workers => f_var_workers,
		                :f_av_stock => f_av_stock,:f_av_demandexp => f_av_demandexp,
		                :f_var_demandexp => f_var_demandexp,:totinftraj => totinftraj,
						:inftraj => inftraj,:curinfregtrajy=>curinfregtrajy,:curinfregtrajo=>curinfregtrajo,
						:inactivefirmstraj=>inactivefirmstraj,:bailouttraj=>bailouttraj,
						:totcas=>totcas,:togdploss=>totgdploss,:totfirmtraj=>totfirmtraj,:gdplosstraj=>gdplosstraj,
						:bankrupttraj=>bankrupttraj,:contact_work_traj=>contact_work_traj,:contact_count_traj=>contact_count_traj,
						:contact_social_traj=>contact_social_traj,:contact_shop_traj=>contact_shop_traj,:RKIR0smtraj=>RKIR0smtraj,
						:polswitchcount=>polswitchcount,:badpoltime=>badpoltime))())
		end
		worker_results = fetch.(result_futures)
		global all_results = vcat(all_results, worker_results)
	end
end

println("All runs finished!")

open("$(filename_prefix)batchdata.dat", "w") do outfile
	serialize(outfile,all_results)
end


text:
# syntax: proto3
using ProtoBuf
import ProtoBuf.meta

mutable struct GraphDef <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function GraphDef(; kwargs...)
        obj = new(meta(GraphDef), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
        end
        obj
    end
end # mutable struct GraphDef
const __meta_GraphDef = Ref{ProtoMeta}()
function meta(::Type{GraphDef})
    ProtoBuf.metalock() do
        if !isassigned(__meta_GraphDef)
            __meta_GraphDef[] = target = ProtoMeta(GraphDef)
            fnum = Int[1,4,3,2]
            allflds = Pair{Symbol,Union{Type,String}}[:node => Base.Vector{NodeDef}, :versions => VersionDef, :version => Int32, :library => FunctionDefLibrary]
            meta(target, GraphDef, allflds, ProtoBuf.DEF_REQ, fnum, ProtoBuf.DEF_VAL, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES)
        end
        __meta_GraphDef[]
    end
end
function Base.getproperty(obj::GraphDef, name::Symbol)
    if name === :node
        return (obj.__protobuf_jl_internal_values[name])::Base.Vector{NodeDef}
    elseif name === :versions
        return (obj.__protobuf_jl_internal_values[name])::VersionDef
    elseif name === :version
        return (obj.__protobuf_jl_internal_values[name])::Int32
    elseif name === :library
        return (obj.__protobuf_jl_internal_values[name])::FunctionDefLibrary
    else
        getfield(obj, name)
    end
end

export GraphDef


text:
using RandomGeometricGraphs, LightGraphs, Test, Statistics

@testset "Boundaries" begin
    nodes = 100
    @test random_geometric_graph(nodes,0) |> edges |> length == 0
    @test random_geometric_graph(nodes,2) |> edges |> length ≈ nodes * (nodes - 1) / 2 atol=0.6
end

@testset "Equivalence to LightGraphs" begin
    nodes = 100
    dimensions = 2
    radius = 0.2
    lg_edges = [euclidean_graph(nodes,dimensions,cutoff=radius)[1] |> edges |> length for _ in 1:1000]
    rgg_edges = [random_geometric_graph(nodes,radius;dim=dimensions) |> edges |> length for _ in 1:1000]

    # This is an abuse of statistics
    @test mean(lg_edges) ≈ mean(rgg_edges) atol=std(lg_edges)
end


text:
#Goodness-of-fit output structure
struct UmGOF
  test_name::String
  t0::Float64
  tstar::Array{Float64}
end

#Parametric bootstrap
function parboot(fit::UnmarkedModel, nsims::Int, statistic::Any)

  out = Array{Float64}(undef, nsims)
  p = Progress(nsims, desc=string("Bootstrap (",nsims," sims) "))
  Threads.@threads for i = 1:nsims
    new_data = simulate(fit)
    new_fit = update(fit, new_data)
    out[i] = statistic(new_fit)
    next!(p)
  end

  return out
end

max_stars_repo_name: kenkellner/julia-hme

max_issues_repo_name: kenkellner/julia-hme

max_forks_repo_name: kenkellner/julia-hme


text:
module MPIReco
  using RegularizedLeastSquares
  @reexport using RegularizedLeastSquares
  @reexport using ImageUtils
  @reexport using MPIFiles
  const shape = MPIFiles.shape
  @reexport using DSP

  using LinearAlgebra
  using SparseArrays
  using Distributed
  using DistributedArrays
  # using TensorDecompositions
  using IniFile
  import LinearAlgebra: ldiv!, \

  include("Utils.jl")
  include("MultiContrast.jl")
  include("RecoParameters.jl")
  include("SystemMatrixCenter.jl")
  include("SystemMatrix.jl")
  include("Weighting.jl")
  include("Background.jl")
  include("Reconstruction.jl")
  include("MultiPatch.jl")
  include("SystemMatrixRecovery.jl")
  include("MotionCompensation/MotionCompensation.jl")
end # module


text:
# This file was generated, do not modify it. # hide
v = [1, 2, 3, 4]
stand_model = UnivariateStandardizer()
stand = machine(stand_model, v)

max_stars_repo_name: tlienart/MLJTutorials

max_issues_repo_name: tlienart/MLJTutorials

max_forks_repo_name: tlienart/MLJTutorials


text:
using SailRoute, Test

function circular_performance()
    twa = [a for a in LinRange(40.0, 180.0, 10)]
    tws = [10.0, 15.0]
    perf = fill(5.0, 10, 2)
    perf[:, 2] .= 7.5
    polar = SailRoute.setup_perf_interpolation(tws, twa, perf)
    performance = SailRoute.Performance(polar, 1.0, 1.0, nothing)
    return performance
end

@testset "Performance interpolation" begin
    test_performance = circular_performance()
    @test SailRoute.perf_interp(test_performance, 60.0, 10.0) ≈ 5.0
    @test SailRoute.perf_interp(test_performance, 60.0, 15.0) ≈ 7.5
    @test SailRoute.perf_interp(test_performance, 30.0, 10.0) ≈ 0.0
end

@testset "Cost function" begin
    @test SailRoute.cost_func(90.0, 10.0, 0.0, 0.0, 0.0, circular_performance()) ≈ 5.0
    @test SailRoute.cost_func(90.0, 10.0, 0.0, 1.0, 0.0, circular_performance()) ≈ 5.5
end

@testset "Speed given current solution" begin
    @test SailRoute.solve_speed_given_current(10.0, 0.0, 0.0, 0.0, 90.0, circular_performance()) ≈ 5.0
    @test SailRoute.solve_speed_given_current(10.0, 0.0, 1.0, 180.0, 90.0, circular_performance()) ≈ 5.408 atol=0.01
end


text:
import DataFrames
import FITSIO
import WCS

const band_letters = ['u', 'g', 'r', 'i', 'z']

"Test that the identity WCSTransform works as expected."
function test_id_wcs()
    rand_coord = rand(2, 10)
    @test WCS.pix_to_world(SampleData.wcs_id, rand_coord) == rand_coord
    @test WCS.world_to_pix(SampleData.wcs_id, rand_coord) == rand_coord
end

function test_linear_world_to_pix()
    rcf = SDSSIO.RunCamcolField(3900, 6, 269)
    image = SDSSIO.load_field_images(rcf, datadir)[1]
    wcs = image.wcs

    pix_center = Float64[0.5 * image.H, 0.5 * image.W]
    pix_loc = pix_center + [5., 3.]
    world_center = WCS.pix_to_world(wcs, pix_center)
    world_loc = WCS.pix_to_world(wcs, pix_loc)

    function test_jacobian(wcs, pix_center, world_center)
        wcs_jacobian = Model.pixel_world_jacobian(wcs, pix_center);

        pix_loc_test1 = WCS.world_to_pix(wcs, world_loc)
        pix_loc_test2 = Model.linear_world_to_pix(wcs_jacobian, world_center, pix_center,
                                       world_loc)

        # Note that the accuracy of the linear approximation isn't great.
        @test_approx_eq(pix_loc_test1, pix_loc)
        @test_approx_eq_eps(pix_loc_test2, pix_loc, 1e-2)
    end

    @test Model.pixel_world_jacobian(SampleData.wcs_id, pix_center) == [1.0 0.0; 0.0 1.0];

    test_jacobian(wcs, pix_center, world_center)
end

test_id_wcs()
test_linear_world_to_pix()


text:
#=
    sw_core_mod

    The kerenl c_sw calculates the horizontal transport for the set
        of variables residing on FV3's C-grid.

=#

module sw_core_mod

# Module to assign NetCDF variables, attributes, and dimensions into a Julia "Struct"
module InitializeState

  using OffsetArrays, NCDatasets, Parameters

  export State, current_state, assign_variables

  # Define the contents of the State Struct
  # The @with_kw "macro" is used to "unpack" the variables when the struct is passed as an argument to functions
  @with_kw struct State{T}

    isd        ::Int32
    ied        ::Int32
    jsd        ::Int32
    jed        ::Int32
    is         ::Int32
    ie         ::Int32
    js         ::Int32
    je         ::Int32
    nord       ::Int32
    npx        ::Int32
    npy        ::Int32
    npz        ::Int32
    dt2        ::Float64
    sw_corner  ::Int32
    se_corner  ::Int32
    ne_corner  ::Int32
    nw_corner  ::Int32

    ni         ::Int64
    nj         ::Int64
    nip1       ::Int64
    njp1       ::Int64
    nk         ::Int64
    nk9        ::Int64

    rarea      ::T
    rarea_c    ::T
    sin_sg     ::T
    cos_sg     ::T
    sina_v     ::T
    cosa_v     ::T
    sina_u     ::T
    cosa_u     ::T
    fC         ::T
    rdxc       ::T
    rdyc       ::T
    dx         ::T
    dy         ::T
    dxc        ::T
    dyc        ::T
    cosa_s     ::T
    rsin_u     ::T
    rsin_v     ::T
    rsin2      ::T
    dxa        ::T
    dya        ::T
    delpc      ::T
    delp       ::T
    ptc        ::T
    pt         ::T
    u          ::T
    v          ::T
    w          ::T
    uc         ::T
    vc         ::T
    ua         ::T
    va         ::T
    wc         ::T
    ut         ::T
    vt         ::T
    divg_d     ::T

  end

  # Assign NetCDF variables,dimensions and attributes to the struct
  function assign_variables(ds::NCDataset)

    isd =                 ds.attrib["isd"]
    ied =                 ds.attrib["ied"]
    jsd =                 ds.attrib["jsd"]
    jed =                 ds.attrib["jed"]
    is =                  ds.attrib["is"]
    ie =                  ds.attrib["ie"]
    js =                  ds.attrib["js"]
    je =                  ds.attrib["je"]
    nord =                ds.attrib["nord"]
    npx =                 ds.attrib["npx"]
    npy =                 ds.attrib["npy"]
    npz =                 ds.attrib["npz"]
    dt2 =                 ds.attrib["dt2"]
    sw_corner =           ds.attrib["sw_corner"]
    se_corner =           ds.attrib["se_corner"]
    nw_corner =           ds.attrib["nw_corner"]
    ne_corner =           ds.attrib["ne_corner"]
    ni =                  ds.dim["ni"]
    nj =                  ds.dim["nj"]
    nip1 =                ds.dim["nip1"]
    njp1 =                ds.dim["njp1"]
    nk =                  ds.dim["nk"]
    nk9 =                 ds.dim["nk9"]

    # Assign all of the NetCDF variables to a constant 'State' Struct
    global current_state = State{OffsetArray}(
      isd,
      ied,
      jsd,
      jed,
      is,
      ie,
      js,
      je,
      nord,
      npx,
      npy,
      npz,
      dt2,
      sw_corner,
      se_corner,
      nw_corner,
      ne_corner,
      ni,
      nj,
      nip1,
      njp1,
      nk,
      nk9,

      # Assign OffsetArrays
      # The first argument to OffsetArray is the NetCDF variable loaded in memory as a Julia array
      # The second, third (and fourth) args are the non-standardard indexed dimensions
      OffsetArray(Array(ds["rarea"]),        isd:ied,      jsd:jed     ),
      OffsetArray(Array(ds["rarea_c"]),      isd:ied + 1,  jsd:jed + 1 ),
      OffsetArray(Array(ds["sin_sg"]),       isd:ied,      jsd:jed, 1:9),
      OffsetArray(Array(ds["cos_sg"]),       isd:ied,      jsd:jed, 1:9),
      OffsetArray(Array(ds["sina_v"]),       isd:ied,      jsd:jed + 1 ),
      OffsetArray(Array(ds["cosa_v"]),       isd:ied,      jsd:jed + 1 ),
      OffsetArray(Array(ds["sina_u"]),       isd:ied + 1,  jsd:jed     ),
      OffsetArray(Array(ds["cosa_u"]),       isd:ied + 1,  jsd:jed     ),
      OffsetArray(Array(ds["fC"]),           isd:ied + 1,  jsd:jed + 1 ),
      OffsetArray(Array(ds["rdxc"]),         isd:ied + 1,  jsd:jed     ),
      OffsetArray(Array(ds["rdyc"]),         isd:ied,      jsd:jed + 1 ),
      OffsetArray(Array(ds["dx"]),           isd:ied,      jsd:jed + 1 ),
      OffsetArray(Array(ds["dy"]),           isd:ied + 1,  jsd:jed     ),
      OffsetArray(Array(ds["dxc"]),          isd:ied + 1,  jsd:jed     ),
      OffsetArray(Array(ds["dyc"]),          isd:ied,      jsd:jed + 1 ),
      OffsetArray(Array(ds["cosa_s"]),       isd:ied,      jsd:jed     ),
      OffsetArray(Array(ds["rsin_u"]),       isd:ied + 1,  jsd:jed     ),
      OffsetArray(Array(ds["rsin_v"]),       isd:ied,      jsd:jed + 1 ),
      OffsetArray(Array(ds["rsin2"]),        isd:ied,      jsd:jed     ),
      OffsetArray(Array(ds["dxa"]),          isd:ied,      jsd:jed     ),
      OffsetArray(Array(ds["dya"]),          isd:ied,      jsd:jed     ),
      OffsetArray(Array(ds["delpc"]),        isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["delp"]),         isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["ptc"]),          isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["pt"]),           isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["u"]),            isd:ied,      jsd:jed + 1, 1:npz),
      OffsetArray(Array(ds["v"]),            isd:ied + 1,  jsd:jed,     1:npz),
      OffsetArray(Array(ds["w"]),            isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["uc"]),           isd:ied + 1,  jsd:jed,     1:npz),
      OffsetArray(Array(ds["vc"]),           isd:ied,      jsd:jed + 1, 1:npz),
      OffsetArray(Array(ds["ua"]),           isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["va"]),           isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["wc"]),           isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["ut"]),           isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["vt"]),           isd:ied,      jsd:jed,     1:npz),
      OffsetArray(Array(ds["divg_d"]),       isd:ied + 1,  jsd:jed + 1, 1:npz)

    )

  end # function assign_variables

end # module InitializeState

# Continuation of the sw_core_mod

# include Julia pkgs and other modules
using Printf
using NCDatasets, Revise, OffsetArrays, Dates, Parameters, DataStructures
using .InitializeState

# export the functions in this module to use in sw_driver.jl
export current_state, assign_variables, print_state, write_state, c_sw!

function print_state(message::String, data::State, io::IOStream)

    println(io, "TEST")
    println(io, "TEST ", "="^115)
    @printf(io, "%-5s %19s\n", "TEST ", message)
    println(io, "TEST ", "="^115)
    @printf(io, "%-5s %14s %19s %19s %19s %19s %19s\n","TEST", "Variable", "Min", "Max", "First", "Last", "RMS")
    println(io, "TEST ", "-"^115)

    # print the variables
    print_2d_variable(io, "rarea", data.rarea)
    print_2d_variable(io, "rarea_c", data.rarea_c)
    print_3d_variable(io, "sin_sg", data.sin_sg)
    print_3d_variable(io, "cos_sg", data.cos_sg)
    print_2d_variable(io, "sina_v", data.sina_v)
    print_2d_variable(io, "cosa_v", data.cosa_v)
    print_2d_variable(io, "sina_u", data.sina_u)
    print_2d_variable(io, "cosa_u", data.cosa_u)
    print_2d_variable(io, "fC", data.fC)
    print_2d_variable(io, "rdxc", data.rdxc)
    print_2d_variable(io, "rdyc", data.rdyc)
    print_2d_variable(io, "dx", data.dx)
    print_2d_variable(io, "dy", data.dy)
    print_2d_variable(io, "dxc", data.dxc)
    print_2d_variable(io, "dyc", data.dyc)
    print_2d_variable(io, "cosa_s", data.cosa_s)
    print_2d_variable(io, "rsin_u", data.rsin_u)
    print_2d_variable(io, "rsin_v", data.rsin_v)
    print_2d_variable(io, "rsin2", data.rsin2)
    print_2d_variable(io, "dxa", data.dxa)
    print_2d_variable(io, "dya", data.dya)

    print_3d_variable(io,"delp", data.delp)
    print_3d_variable(io,"delpc", data.delpc)
    print_3d_variable(io,"pt", data.pt)
    print_3d_variable(io,"ptc", data.ptc)
    print_3d_variable(io,"u", data.u)
    print_3d_variable(io,"v", data.v)
    print_3d_variable(io,"w", data.w)
    print_3d_variable(io,"uc", data.uc)
    print_3d_variable(io,"vc", data.vc)
    print_3d_variable(io,"ua", data.ua)
    print_3d_variable(io,"va", data.va)
    print_3d_variable(io,"wc", data.wc)
    print_3d_variable(io,"ut", data.ut)
    print_3d_variable(io,"vt", data.vt)
    print_3d_variable(io,"divg_d", data.divg_d)

    println(io, "TEST ", "-"^115)
    println(io, "TEST")

    return nothing

end # function print_state

function print_2d_variable(io::IOStream,name::String, data::OffsetArray)

    min =     minimum(data)
    max =     maximum(data)
    first =   data[begin, begin]
    last =    data[end, end]
    rms =     sqrt(sum(data.^2.) / length(data))

    @printf(io, "%-5s %14s %19.10E %19.10E %19.10E %19.10E %19.10E\n","TEST", name, min, max , first, last, rms)

    return nothing

end # function print_2d_variable

function print_3d_variable(io::IOStream, name::String, data::OffsetArray)

    min =    minimum(data)
    max =    maximum(data)
    first =  data[begin, begin, begin]
    last =   data[end, end, end]
    rms =    sqrt(sum(data.^2.) / length(data))

    @printf(io, "%-5s %14s %19.10E %19.10E %19.10E %19.10E %19.10E\n","TEST", name, min, max , first, last, rms)

    return nothing

end # function print_3d_variable

# create a NetCDF file from the data in netCDF-4 format
function write_state(outputfilename::String, data::State)

    date = Dates.now()
    dateString = Dates.format(date, "yyyy-mm-dd HH:MM:SS")

    ds = NCDataset(outputfilename, "c" ; format=:netcdf4, attrib=OrderedDict(
    "creation_date"             => dateString,
    "kernel_name"               => "c_sw",
    "isd"                       => data.isd,
    "ied"                       => data.ied,
    "jsd"                       => data.jsd,
    "jed"                       => data.jed,
    "is"                        => data.is,
    "ie"                        => data.ie,
    "js"                        => data.js,
    "je"                        => data.je,
    "nord"                      => data.nord,
    "npx"                       => data.npx,
    "npy"                       => data.npy,
    "npz"                       => data.npz,
    "dt2"                       => data.dt2,
    "sw_corner"                 => data.sw_corner,
    "se_corner"                 => data.se_corner,
    "nw_corner"                 => data.nw_corner,
    "ne_corner"                 => data.ne_corner,
  ))

  # define Dimensions
    ds.dim["ni"] =    data.ni
    ds.dim["nj"] =    data.nj
    ds.dim["nip1"] =  data.nip1
    ds.dim["njp1"] =  data.njp1
    ds.dim["nk"] =    data.nk
    ds.dim["nk9"] =   data.nk9

  # Declare variables
    ncrarea =           defVar(ds, "rarea", Float64, ("ni", "nj"))
    ncrarea_c =         defVar(ds, "rarea_c", Float64, ("nip1", "njp1"))
    ncsin_sg =          defVar(ds, "sin_sg", Float64, ("ni", "nj", "nk9"))
    nccos_sg =          defVar(ds, "cos_sg", Float64, ("ni", "nj", "nk9"))
    ncsina_v =          defVar(ds, "sina_v", Float64, ("ni", "njp1"))
    nccosa_v =          defVar(ds, "cosa_v", Float64, ("ni", "njp1"))
    ncsina_u =          defVar(ds, "sina_u", Float64, ("nip1", "nj"))
    nccosa_u =          defVar(ds, "cosa_u", Float64, ("nip1", "nj"))
    ncfC =              defVar(ds, "fC", Float64, ("nip1", "njp1"))
    ncrdxc =            defVar(ds, "rdxc", Float64, ("nip1", "nj"))
    ncrdyc =            defVar(ds, "rdyc", Float64, ("ni", "njp1"))
    ncdx =              defVar(ds, "dx", Float64, ("ni", "njp1"))
    ncdy =              defVar(ds, "dy", Float64, ("nip1", "nj"))
    ncdxc =             defVar(ds, "dxc", Float64, ("nip1", "nj"))
    ncdyc =             defVar(ds, "dyc", Float64, ("ni", "njp1"))
    nccosa_s =          defVar(ds, "cosa_s", Float64, ("ni", "nj"))
    ncrsin_u =          defVar(ds, "rsin_u", Float64, ("nip1", "nj"))
    ncrsin_v =          defVar(ds, "rsin_v", Float64, ("ni", "njp1"))
    ncrsin2 =           defVar(ds, "rsin2", Float64, ("ni", "nj"))
    ncdxa =             defVar(ds, "dxa", Float64, ("ni", "nj"))
    ncdya =             defVar(ds, "dya", Float64, ("ni", "nj"))
    ncdelp =            defVar(ds, "delp", Float64, ("ni", "nj", "nk"))
    ncdelpc =           defVar(ds, "delpc", Float64, ("ni", "nj", "nk"))
    ncpt =              defVar(ds, "pt", Float64, ("ni", "nj", "nk"))
    ncptc =             defVar(ds, "ptc", Float64, ("ni", "nj", "nk"))
    ncu =               defVar(ds, "u", Float64, ("ni", "njp1", "nk"))
    ncv =               defVar(ds, "v", Float64, ("nip1", "nj", "nk"))
    ncw =               defVar(ds, "w", Float64, ("ni", "nj", "nk"))
    ncuc =              defVar(ds, "uc", Float64, ("nip1", "nj", "nk"))
    ncvc =              defVar(ds, "vc", Float64, ("ni", "njp1", "nk"))
    ncua =              defVar(ds, "ua", Float64, ("ni", "nj", "nk"))
    ncva =              defVar(ds, "va", Float64, ("ni", "nj", "nk"))
    ncwc =              defVar(ds, "wc", Float64, ("ni", "nj", "nk"))
    ncut =              defVar(ds, "ut", Float64, ("ni", "nj", "nk"))
    ncvt =              defVar(ds, "vt", Float64, ("ni", "nj", "nk"))
    ncdivg_d =          defVar(ds, "divg_d", Float64, ("nip1", "njp1", "nk"))

  # Define variables
    ncrarea[:] =   data.rarea
    ncrarea_c[:] = data.rarea_c
    ncsin_sg[:] =  data.sin_sg
    nccos_sg[:] =  data.cos_sg
    ncsina_v[:] =  data.sina_v
    nccosa_v[:] =  data.cosa_v
    ncsina_u[:] =  data.sina_u
    nccosa_u[:] =  data.cosa_u
    ncfC[:] =      data.fC
    ncrdxc[:] =    data.rdxc
    ncrdyc[:] =    data.rdyc
    ncdx[:] =      data.dx
    ncdy[:] =      data.dy
    ncdxc[:] =     data.dxc
    ncdyc[:] =     data.dyc
    nccosa_s[:] =  data.cosa_s
    ncrsin_u[:] =  data.rsin_u
    ncrsin_v[:] =  data.rsin_v
    ncrsin2[:] =   data.rsin2
    ncdxa[:] =     data.dxa
    ncdya[:] =     data.dya
    ncdelp[:] =    data.delp
    ncdelpc[:] =   data.delpc
    ncpt[:] =      data.pt
    ncptc[:] =     data.ptc
    ncu[:] =       data.u
    ncv[:] =       data.v
    ncw[:] =       data.w
    ncuc[:] =      data.uc
    ncvc[:] =      data.vc
    ncua[:] =      data.ua
    ncva[:] =      data.va
    ncwc[:] =      data.wc
    ncut[:] =      data.ut
    ncvt[:] =      data.vt
    ncdivg_d[:] =  data.divg_d

    close(ds)

end # function write_state

# Julia version of c_sw subroutine
function c_sw!(data::State, k::Int64)

  # Unpack arguments
  @unpack sw_corner, se_corner, nw_corner, ne_corner,
  rarea, rarea_c, sin_sg, cos_sg, sina_v, cosa_v,
  sina_u, cosa_u, fC, rdxc, rdyc, dx, dy, dxc, dyc,
  cosa_s, rsin_u, rsin_v, rsin2, dxa, dya,
  delpc, delp, ptc, pt, u, v, w, uc, vc, ua, va, wc, ut,
  vt, divg_d, dt2  = data

  # Unpack local variables
  @unpack isd, ied, jsd, jed, is, ie, js, je,
          nord, npx, npy, npz, dt2, ni, nj, nip1, njp1, nk, nk9  = data

  # re-dimension variables to the first two dimensions
  delpc =        @view delpc[:,:,k]
  delp =         @view delp[:,:,k]
  ptc =          @view ptc[:,:,k]
  pt =           @view pt[:,:,k]
  u =            @view u[:,:,k]
  v =            @view v[:,:,k]
  w =            @view w[:,:,k]
  uc =           @view uc[:,:,k]
  vc =           @view vc[:,:,k]
  ua =           @view ua[:,:,k]
  va =           @view va[:,:,k]
  wc =           @view wc[:,:,k]
  ut =           @view ut[:,:,k]
  vt =           @view vt[:,:,k]
  divg_d =       @view divg_d[:,:,k]

  #helper function to index dimensions
  indexdim(a, b) = b - a + 1

  # define local variables arrays and index them
  vort = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+1) , indexdim(js-1, je+1)), is-1:ie+1, js-1:je+1)
  ke   = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+1) , indexdim(js-1, je+1)), is-1:ie+1, js-1:je+1)
  fx   = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+2) , indexdim(js-1, je+1)), is-1:ie+2, js-1:je+1)
  fx1  = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+2) , indexdim(js-1, je+1)), is-1:ie+2, js-1:je+1)
  fx2  = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+2) , indexdim(js-1, je+1)), is-1:ie+2, js-1:je+1)
  fy   = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+1) , indexdim(js-1, je+2)), is-1:ie+1, js-1:je+2)
  fy1  = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+1) , indexdim(js-1, je+2)), is-1:ie+1, js-1:je+2)
  fy2  = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+1) , indexdim(js-1, je+2)), is-1:ie+1, js-1:je+2)

  # Convert corners from Int32 to Bool for Julia conditional statements
  sw_corner = !iszero(sw_corner)
  se_corner = !iszero(se_corner)
  ne_corner = !iszero(ne_corner)
  nw_corner = !iszero(nw_corner)

  iep1 = ie + 1
  jep1 = je + 1

  d2a2c_vect!(data, k)

  if nord > 0
      divergence_corner!(data, k)
  end

  for j = js-1:jep1
      for i = is-1:iep1+1
          if ut[i, j] > 0.0
              ut[i, j] = dt2 * ut[i, j] * dy[i, j] * sin_sg[i-1, j, 3]
          else
              ut[i, j] = dt2 * ut[i, j] * dy[i, j] * sin_sg[i, j, 1]
          end
      end
  end
  for j = js-1:je+2
      for i = is-1:iep1
          if vt[i, j] > 0.0
              vt[i, j] = dt2 * vt[i, j] * dx[i, j] * sin_sg[i, j-1, 4]
          else
              vt[i, j] = dt2 * vt[i, j] * dx[i, j] * sin_sg[i, j, 2]
          end
      end
  end

  # # ----------------
  # #  Transport delp:
  # # ----------------
  # # Xdir:
  fill2_4corners!(delp, pt, 1, sw_corner, se_corner, ne_corner, nw_corner, npx, npy)
  fill_4corners!(w, 1, sw_corner, se_corner, ne_corner, nw_corner, npx, npy)

  for j = js-1:je+1
      for i = is-1:ie+2
          if ut[i, j] > 0.
              fx1[i, j] = delp[i-1, j]
              fx[i, j]  = pt[i-1, j]
              fx2[i, j] = w[i-1, j]
          else
              fx1[i, j] = delp[i, j]
              fx[i, j]  = pt[i, j]
              fx2[i, j] = w[i, j]
          end
          fx1[i, j] = ut[i, j]  * fx1[i, j]
          fx[i, j]  = fx1[i, j] * fx[i, j]
          fx2[i, j] = fx1[i, j] * fx2[i, j]
      end
  end

  # # Ydir:
  fill2_4corners!(delp, pt, 2, sw_corner, se_corner, ne_corner, nw_corner, npx, npy)
  fill_4corners!(w, 2, sw_corner, se_corner, ne_corner, nw_corner, npx, npy)

  for j = js-1:je+2
      for i = is-1:ie+1
          if vt[i, j] > 0.
              fy1[i, j] = delp[i, j-1]
              fy[i, j]  = pt[i, j-1]
              fy2[i, j] = w[i, j-1]
          else
              fy1[i, j] = delp[i, j]
              fy[i, j]  = pt[i, j]
              fy2[i, j] = w[i, j]
          end
          fy1[i, j] = vt[i, j]  * fy1[i, j]
          fy[i, j]  = fy1[i, j] * fy[i, j]
          fy2[i, j] = fy1[i, j] * fy2[i, j]
      end
  end
  for j = js-1:je+1
      for i = is-1:ie+1
          delpc[i, j] = delp[i, j] + (fx1[i, j] - fx1[i+1, j] + fy1[i, j] - fy1[i, j+1]) *
              rarea[i, j]
          ptc[i, j] = (pt[i, j] * delp[i, j] + (fx[i, j] - fx[i+1, j] +
                       fy[i, j] - fy[i, j+1]) * rarea[i, j]) / delpc[i, j]
          wc[i, j] = (w[i, j] * delp[i, j] + (fx2[i, j] - fx2[i+1, j] +
                       fy2[i, j] - fy2[i, j+1]) * rarea[i, j]) / delpc[i, j]
      end
  end

  #   !------------
  #   ! Compute KE:
  #   !------------

  #   ! Since uc = u*, i.e. the covariant wind perpendicular to the face edge, if
  #   ! we want to compute kinetic energy we will need the true coordinate-parallel
  #   ! covariant wind, computed through u = uc*sina + v*cosa.
  #   !
  #   ! Use the alpha for the cell KE is being computed in.
  for j = js-1:jep1
      for i = is-1:iep1
          if ua[i, j] > 0.
              if i == 1
                  ke[1, j] = uc[1, j] * sin_sg[1, j, 1] + v[1, j] * cos_sg[1, j, 1]
              elseif i == npx
                  ke[i, j] = uc[npx, j] * sin_sg[npx, j, 1] + v[npx, j] * cos_sg[npx, j, 1]
              else
                  ke[i, j] = uc[i, j]
              end
          else
              if i == 0
                  ke[0, j] = uc[1, j] * sin_sg[0, j, 3] + v[1, j] * cos_sg[0, j, 3]
              elseif i == npx - 1
                  ke[i, j] = uc[npx, j] * sin_sg[npx-1, j, 3] + v[npx, j] * cos_sg[npx-1, j, 3]
              else
                  ke[i, j] = uc[i+1, j]
              end
          end
      end
  end
  for j = js-1:jep1
      for i = is-1:iep1
          if va[i, j] > 0.
              if j == 1
                  vort[i, 1] = vc[i, 1] * sin_sg[i, 1, 2] + u[i, 1] * cos_sg[i, 1, 2]
              elseif j == npy
                  vort[i, j] = vc[i, npy] * sin_sg[i, npy, 2] + u[i, npy] * cos_sg[i, npy, 2]
              else
                  vort[i, j] = vc[i, j]
              end
          else
              if j == 0
                  vort[i, 0] = vc[i, 1] * sin_sg[i, 0, 4] + u[i, 1] * cos_sg[i, 0, 4]
              elseif j == npy - 1
                  vort[i, j] = vc[i, npy] * sin_sg[i, npy-1, 4] + u[i, npy] * cos_sg[i, npy-1, 4]
              else
                  vort[i, j] = vc[i, j+1]
              end
          end
      end
  end

  dt4 = 0.5 * dt2
  for j = js-1:jep1
      for i = is-1:iep1
          ke[i, j] = dt4 * (ua[i, j] * ke[i, j] + va[i, j] * vort[i, j])
      end
  end

  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # Compute circulation on C grid
  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # To consider using true co - variant winds at face edges?
  for j = js-1:je+1
      for i = is:ie+1
          fx[i, j] = uc[i, j] * dxc[i, j]
      end
  end

  for j = js:je+1
      for i = is-1:ie+1
          fy[i, j] = vc[i, j] * dyc[i, j]
      end
  end

  for j = js:je+1
      for i = is:ie+1
          vort[i, j] = fx[i, j-1] - fx[i, j] - fy[i-1, j] + fy[i, j]
      end
  end

  # Remove the extra term at the corners:
  if sw_corner
      vort[1, 1] = vort[1, 1] + fy[0, 1]
  end
  if se_corner
      vort[npx, 1] = vort[npx, 1] - fy[npx, 1]
  end
  if ne_corner
      vort[npx, npy] = vort[npx, npy] - fy[npx, npy]
  end
  if nw_corner
      vort[1, npy] = vort[1, npy] + fy[0, npy]
  end

  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # Compute absolute vorticity
  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  for j = js:je+1
      for i = is:ie+1
          vort[i, j] = fC[i, j] + rarea_c[i, j] * vort[i, j]
      end
  end
  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # Transport absolute vorticity:
  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # To go from v to contravariant v at the edges, we divide by sin_sg;
  # but we  must multiply by sin_sg to get the proper flux.
  # These cancel, leaving us with fy1 = dt2 * v at the edges.
  # [For the same reason we only divide by sin instead of sin^2 in the interior]
  for j = js:je
      #DEC$ VECTOR ALWAYS
      for i = is:iep1
          if i == 1 || i == npx
              fy1[i, j] = dt2 * v[i, j]
          else
              fy1[i, j] = dt2 * (v[i, j] - uc[i, j] * cosa_u[i, j]) / sina_u[i, j]
          end
          if fy1[i, j] > 0.0
              fy[i, j] = vort[i, j]
          else
              fy[i, j] = vort[i, j+1]
          end
      end
  end
  for j = js:jep1
      if j == 1 || j == npy
          #DEC$ VECTOR ALWAYS
          for i = is:ie
              fx1[i, j] = dt2 * u[i, j]
              if fx1[i, j] > 0.0
                  fx[i, j] = vort[i, j]
              else
                  fx[i, j] = vort[i+1, j]
              end
          end
      else
          #DEC$ VECTOR ALWAYS
          for i = is:ie
              fx1[i, j] = dt2 * (u[i, j] - vc[i, j] * cosa_v[i, j]) / sina_v[i, j]
              if fx1[i, j] > 0.0
                  fx[i, j] = vort[i, j]
              else
                  fx[i, j] = vort[i+1, j]
              end
          end
      end
  end

  # Update time - centered winds on the C - Grid
  for j = js:je
      for i = is:iep1
          uc[i, j] = uc[i, j] + fy1[i, j] * fy[i, j] + rdxc[i, j] * (ke[i-1, j] - ke[i, j])
      end
  end
  for j = js:jep1
      for i = is:ie
          vc[i, j] = vc[i, j] - fx1[i, j] * fx[i, j] + rdyc[i, j] * (ke[i, j-1] - ke[i, j])
      end
  end

end # function c_sw

# divergence_corner
function divergence_corner!(data::State, k::Int64)

  # unpack arguments
  @unpack sw_corner, se_corner, ne_corner, nw_corner,
  rarea_c, sin_sg, cos_sg, dxc, dyc, u, v, ua, va, divg_d,
  is, ie, js, je, npx, npy = data

  # re-dimension variables to the first two dimensions
  u =         @view u[:,:,k]
  v =         @view v[:,:,k]
  ua =        @view ua[:,:,k]
  va =        @view va[:,:,k]
  divg_d =    @view divg_d[:,:,k]

  # Convert corners from Int32 to Bool for Julia conditional statements
  sw_corner = !iszero(sw_corner)
  se_corner = !iszero(se_corner)
  ne_corner = !iszero(ne_corner)
  nw_corner = !iszero(nw_corner)

  indexdim(a, b) = b - a + 1

  # local variables
  uf = OffsetArray(Array{Float64, 2}(undef, indexdim(is-2,ie+2) , indexdim(js-1, je+2)), is-2:ie+2, js-1:je+2)
  vf = OffsetArray(Array{Float64, 2}(undef, indexdim(is-1,ie+2) , indexdim(js-2, je+2)), is-1:ie+2, js-2:je+2)

  is2 = max(2, is)
  ie1 = min(npx - 1, ie + 1)

  #    9---4---8
  #    |       |
  #    1   5   3
  #    |       |
  #    6---2---7
  for j = js:je+1
    if j == 1 || j == npy
      for i = is-1:ie+1
        uf[i, j] = u[i, j] * dyc[i, j] * 0.5 * (sin_sg[i, j-1, 4] + sin_sg[i, j, 2])
      end
    else
      for i = is-1:ie+1
        uf[i, j] = (u[i, j] - 0.25 * (va[i, j-1] + va[i, j]) *
        (cos_sg[i, j-1, 4] + cos_sg[i, j, 2])) *
        dyc[i, j] * 0.5 * (sin_sg[i, j-1, 4] + sin_sg[i, j, 2])
      end
    end
  end

  for j = js-1:je+1
    for i = is2:ie1
      vf[i, j] = (v[i, j] - 0.25 * (ua[i-1, j] + ua[i, j]) *
      (cos_sg[i-1, j, 3] + cos_sg[i, j, 1])) *
      dxc[i, j] * 0.5 * (sin_sg[i-1, j, 3] + sin_sg[i, j, 1])
    end
    if is == 1
      vf[1, j] = v[1, j] * dxc[1, j] * 0.5 * (sin_sg[0, j, 3] + sin_sg[1, j, 1])
    end
    if ie + 1 == npx
      vf[npx, j] = v[npx, j] * dxc[npx, j] * 0.5 * (sin_sg[npx-1, j, 3] + sin_sg[npx, j, 1])
    end
  end

  for j = js:je+1
    for i = is:ie+1
      divg_d[i, j] = vf[i, j-1] - vf[i, j] + uf[i-1, j] - uf[i, j]
    end
  end

  # Remove the extra term at the corners:
  if sw_corner
    divg_d[1, 1] = divg_d[1, 1] - vf[1, 0]
  end
  if se_corner
    divg_d[npx, 1] = divg_d[npx, 1] - vf[npx, 0]
  end
  if ne_corner
    divg_d[npx, npy] = divg_d[npx, npy] + vf[npx, npy]
  end
  if nw_corner
    divg_d[1, npy] = divg_d[1, npy] + vf[1, npy]
  end

  for j = js:je+1
    for i = is:ie+1
      divg_d[i, j] = rarea_c[i, j] * divg_d[i, j]
    end
  end
  return nothing
end # function divergence_corner

#=
    d2a2c_vect

      There is a limit to how far this routine can fill uc and vc in the
      halo, and so either mpp_update_domains or some sort of boundary
      routine [extrapolation, outflow, interpolation from a nested grid]
      is needed after c_sw is completed if these variables are needed
      in the halo
=#
function d2a2c_vect!(data::State, k::Int64)

  @unpack sw_corner, se_corner, ne_corner, nw_corner,
  sin_sg, cosa_u, cosa_v, cosa_s, rsin_u, rsin_v,
  rsin2, dxa, dya, u, v, ua, va, uc, vc, ut, vt,
  is, js, ie, je, isd, ied, jsd, jed, npx, npy = data

  # "slice the 3d arrays to 2d arrays"
  u =       @view u[:,:,k]
  v =       @view v[:,:,k]
  uc =      @view uc[:,:,k]
  vc =      @view vc[:,:,k]
  ua =      @view ua[:,:,k]
  va =      @view va[:,:,k]
  ut =      @view ut[:,:,k]
  vt =      @view vt[:,:,k]

  # Convert corners from Int32 to Bool for Julia conditional statements
  sw_corner = !iszero(sw_corner)
  se_corner = !iszero(se_corner)
  ne_corner = !iszero(ne_corner)
  nw_corner = !iszero(nw_corner)

  # 4-pt Lagrange interpolation
  a1 =  0.5625
  a2 = -0.0625
  # volume-conserving cubic with 2nd drv=0 at end point:
  c1 = -2 / 14
  c2 = 11 / 14
  c3 =  5 / 14

  indexdim(a, b) = b - a + 1

  # Local Variables
  utmp = OffsetArray(fill(1.e8, (indexdim(isd,ied),indexdim(jsd,jed))), isd:ied, jsd:jed)
  vtmp = OffsetArray(fill(1.e8, (indexdim(isd,ied),indexdim(jsd,jed))), isd:ied, jsd:jed)

  # -  -  -  -  -  -  -  -  -  -
  # Interior:
  # -  -  -  -  -  -  -  -  -  -
  for j = max(4, js - 1) : min(npy-4, je+1)
    for i = max(4, isd) : min(npx-4, ied)
      utmp[i, j] = a2 * (u[i, j-1] + u[i, j+2]) + a1 * (u[i, j] + u[i, j+1])
    end
  end
  for j = max(4, jsd) : min(npy-4, jed)
    for i = max(4, is - 1) : min(npx-4, ie+1)
      vtmp[i, j] = a2 * (v[i-1, j] + v[i+2, j]) + a1 * (v[i, j] + v[i+1, j])
    end
  end

  # -  -  -  -  -  -  -  -  -  -
  # edges:
  # -  -  -  -  -  -  -  -  -  -
  if js == 1 || jsd < 4
    for j = jsd:3
      for i = isd:ied
        utmp[i, j] = 0.5 * (u[i, j] + u[i, j+1])
        vtmp[i, j] = 0.5 * (v[i, j] + v[i+1, j])
      end
    end
  end

  if je + 1 == npy || jed >= npy - 4
    for j = npy-3:jed
      for i = isd:ied
        utmp[i, j] = 0.5 * (u[i, j] + u[i, j+1])
        vtmp[i, j] = 0.5 * (v[i, j] + v[i+1, j])
      end
    end
  end

  if is == 1 || isd < 4
    for j = max(4, jsd): min(npy-4, jed)
      for i = isd:3
        utmp[i, j] = 0.5 * (u[i, j] + u[i, j+1])
        vtmp[i, j] = 0.5 * (v[i, j] + v[i+1, j])
      end
    end
  end

  if ie + 1 == npx || ied >= npx - 4
    for j = max(4, jsd): min(npy-4, jed)
      for i = npx-3:ied
        utmp[i, j] = 0.5 * (u[i, j] + u[i, j+1])
        vtmp[i, j] = 0.5 * (v[i, j] + v[i+1, j])
      end
    end
  end

  for j = js-2:je+2
    for i = is-2:ie+2
      ua[i, j] = (utmp[i, j] - vtmp[i, j] * cosa_s[i, j]) * rsin2[i, j]
      va[i, j] = (vtmp[i, j] - utmp[i, j] * cosa_s[i, j]) * rsin2[i, j]
    end
  end

  # A - > C
  # -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # Fix the edges
  # -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # Xdir:
  if sw_corner
    for i = -2:0
      utmp[i, 0] = -vtmp[0, 1-i]
    end
  end
  if se_corner
    for i = 0:2
      utmp[npx+i, 0] = vtmp[npx, i+1]
    end
  end
  if ne_corner
    for i = 0:2
      utmp[npx+i, npy] = -vtmp[npx, je-i]
    end
  end
  if nw_corner
    for i = -2:0
      utmp[i, npy] = vtmp[0, je+i]
    end
  end

  ifirst = max(3,     is-1)
  ilast  = min(npx-2, ie+2)

  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  # 4th order interpolation for interior points:
  # -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -  -
  for j = js-1:je+1
    for i = ifirst:ilast
      uc[i, j] = a1 * (utmp[i-1, j] + utmp[i, j]) + a2 * (utmp[i-2, j] + utmp[i+1, j])
      ut[i, j] = (uc[i, j] - v[i, j] * cosa_u[i, j]) * rsin_u[i, j]
    end
  end

  # Xdir:
  if sw_corner
    ua[-1, 0] = -va[0, 2]
    ua[ 0, 0] = -va[0, 1]
  end
  if se_corner
    ua[npx,   0] = va[npx, 1]
    ua[npx+1, 0] = va[npx, 2]
  end
  if ne_corner
    ua[npx,   npy] = -va[npx, npy-1]
    ua[npx+1, npy] = -va[npx, npy-2]
  end
  if nw_corner
    ua[-1, npy] = va[0, npy-2]
    ua[ 0, npy] = va[0, npy-1]
  end

  if is == 1
    for j = js-1:je+1
      uc[0, j] = c1 * utmp[-2, j] + c2 * utmp[-1, j] + c3 * utmp[0, j]
      ut[1, j] = edge_interpolate4!(ua[-1:2, j], dxa[-1:2, j])
      #Want to  the UPSTREAM value
      if ut[1, j] < 0.0
        uc[1, j] = ut[1, j] * sin_sg[0, j, 3]
      else
        uc[1, j] = ut[1, j] * sin_sg[1, j, 1]
      end
      uc[2, j] = c1 * utmp[3, j] + c2 * utmp[2, j] + c3 * utmp[1, j]
      ut[0, j] = (uc[0, j] - v[0, j] * cosa_u[0, j]) * rsin_u[0, j]
      ut[2, j] = (uc[2, j] - v[2, j] * cosa_u[2, j]) * rsin_u[2, j]
    end

    if ie + 1 == npx
      for j = js-1:je+1
        uc[npx-1, j] =
        c1 * utmp[npx-3, j] + c2 * utmp[npx-2, j] + c3 * utmp[npx-1, j]
        i = npx
        ut[i, j] = 0.25 * (-ua[i-2, j] + 3.0 * (ua[i-1, j] + ua[i, j]) - ua[i+1, j])
        ut[i, j] = edge_interpolate4!(ua[i-2:i+1, j], dxa[i-2:i+1, j])
        if ut[i, j] < 0.0
          uc[i, j] = ut[i, j] * sin_sg[i-1, j, 3]
        else
          uc[i, j] = ut[i, j] * sin_sg[i, j, 1]
        end
        uc[npx+1, j] = c3 * utmp[npx, j] + c2 * utmp[npx+1, j] + c1 * utmp[npx+2, j]
        ut[npx-1, j] = (uc[npx-1, j] - v[npx-1, j] * cosa_u[npx-1, j]) * rsin_u[npx-1, j]
        ut[npx+1, j] = (uc[npx+1, j] - v[npx+1, j] * cosa_u[npx+1, j]) * rsin_u[npx+1, j]
      end
    end

  end

  # -  -  -  -  -  -
  # Ydir:
  # -  -  -  -  -  -
  if sw_corner
    for j = -2:0
      vtmp[0, j] = -utmp[1-j, 0]
    end
  end
  if nw_corner
    for j = 0:2
      vtmp[0, npy+j] = utmp[j+1, npy]
    end
  end
  if se_corner
    for j = -2:0
      vtmp[npx, j] = utmp[ie+j, 0]
    end
  end
  if ne_corner
    for j = 0:2
      vtmp[npx, npy+j] = -utmp[ie-j, npy]
    end
  end
  if sw_corner
    va[0, -1] = -ua[2, 0]
    va[0,  0] = -ua[1, 0]
  end
  if se_corner
    va[npx,  0] = ua[npx-1, 0]
    va[npx, -1] = ua[npx-2, 0]
  end
  if ne_corner
    va[npx,   npy] = -ua[npx-1, npy]
    va[npx, npy+1] = -ua[npx-2, npy]
  end
  if nw_corner
    va[0,   npy] = ua[1, npy]
    va[0, npy+1] = ua[2, npy]
  end

  for j = js-1:je+2
    if j == 1
      for i = is-1:ie+1
        vt[i, j] = edge_interpolate4!(va[i, -1:2], dya[i, -1:2])
        if vt[i, j] < 0.0
          vc[i, j] = vt[i, j] * sin_sg[i, j-1, 4]
        else
          vc[i, j] = vt[i, j] * sin_sg[i, j, 2]
        end
      end
    elseif j == 0 || j == npy - 1
      for i = is-1:ie+1
        vc[i, j] = c1 * vtmp[i, j-2] + c2 * vtmp[i, j-1] + c3 * vtmp[i, j]
        vt[i, j] = (vc[i, j] - u[i, j] * cosa_v[i, j]) * rsin_v[i, j]
      end
    elseif j == 2 || j == npy + 1
      for i = is-1:ie+1
        vc[i, j] = c1 * vtmp[i, j+1] + c2 * vtmp[i, j] + c3 * vtmp[i, j-1]
        vt[i, j] = (vc[i, j] - u[i, j] * cosa_v[i, j]) * rsin_v[i, j]
      end
    elseif j == npy
      for i = is-1:ie+1
        vt[i, j] = 0.25 * (-va[i, j-2] + 3.0 * (va[i, j-1] + va[i, j]) - va[i, j+1])
        vt[i, j] = edge_interpolate4!(va[i, j-2:j+1], dya[i, j-2:j+1])
        if vt[i, j] < 0.0
          vc[i, j] = vt[i, j] * sin_sg[i, j-1, 4]
        else
          vc[i, j] = vt[i, j] * sin_sg[i, j, 2]
        end
      end
    else
      # 4th order interpolation for interior points:
      for i = is-1:ie+1
        vc[i, j] = a2 * (vtmp[i, j-2] + vtmp[i, j+1]) + a1 * (vtmp[i, j-1] + vtmp[i, j])
        vt[i, j] = (vc[i, j] - u[i, j] * cosa_v[i, j]) * rsin_v[i, j]
      end
    end
  end
  return nothing
end # function d2a2c_vect

# edge_interpolate4
function edge_interpolate4!(ua, dxa)

  u0L = 0.5 * ((2.0 * dxa[2] + dxa[1]) * ua[2] - dxa[2] * ua[1]) / (dxa[1] + dxa[2])
  u0R = 0.5 * ((2.0 * dxa[3] + dxa[4]) * ua[3] - dxa[3] * ua[4]) / (dxa[3] + dxa[4])

  return u0L + u0R

  # This is the original edge - interpolation code, which makes
  # a relatively small increase in the error in unstretched case 2.
  #
  # edge_interpolate4 = 0.25 * [3 * [ua[2] + ua[3]] - [ua[1] + ua[4]]]

end #function edge_interpolate4

#=
 fill2_4corners

    This function fills the 4 corners of the scalar fileds only as needed by c_core
=#
function fill2_4corners!(q1, q2, dir, sw_corner, se_corner, ne_corner, nw_corner, npx, npy)

  if dir == 1
    if sw_corner
      q1[-1, 0] = q1[0, 2]
      q1[ 0, 0] = q1[0, 1]
      q2[-1, 0] = q2[0, 2]
      q2[ 0, 0] = q2[0, 1]
    end
    if se_corner
      q1[npx+1, 0] = q1[npx, 2]
      q1[npx,   0] = q1[npx, 1]
      q2[npx+1, 0] = q2[npx, 2]
      q2[npx,   0] = q2[npx, 1]
    end
    if nw_corner
      q1[ 0, npy] = q1[0, npy-1]
      q1[-1, npy] = q1[0, npy-2]
      q2[ 0, npy] = q2[0, npy-1]
      q2[-1, npy] = q2[0, npy-2]
    end
    if ne_corner
      q1[npx,   npy] = q1[npx, npy-1]
      q1[npx+1, npy] = q1[npx, npy-2]
      q2[npx,   npy] = q2[npx, npy-1]
      q2[npx+1, npy] = q2[npx, npy-2]
    end

  end

  if dir == 2
    if sw_corner
      q1[0,  0] = q1[1, 0]
      q1[0, -1] = q1[2, 0]
      q2[0,  0] = q2[1, 0]
      q2[0, -1] = q2[2, 0]
    end
    if se_corner
      q1[npx,  0] = q1[npx-1, 0]
      q1[npx, -1] = q1[npx-2, 0]
      q2[npx,  0] = q2[npx-1, 0]
      q2[npx, -1] = q2[npx-2, 0]
    end
    if nw_corner
      q1[0, npy]   = q1[1, npy]
      q1[0, npy+1] = q1[2, npy]
      q2[0, npy]   = q2[1, npy]
      q2[0, npy+1] = q2[2, npy]
    end
    if ne_corner
      q1[npx, npy]   = q1[npx-1, npy]
      q1[npx, npy+1] = q1[npx-2, npy]
      q2[npx, npy]   = q2[npx-1, npy]
      q2[npx, npy+1] = q2[npx-2, npy]
    end

  end

    return nothing
end # function fill2_4corners

#=
 fill_4corners

    This routine fill the 4 corners of the scalar fileds only as needed by c_core
=#
function fill_4corners!(q, dir, sw_corner, se_corner, ne_corner, nw_corner, npx, npy)

  if dir == 1
    if sw_corner
      q[-1, 0] = q[0, 2]
      q[ 0, 0] = q[0, 1]
    end
    if se_corner
      q[npx+1, 0] = q[npx, 2]
      q[npx,   0] = q[npx, 1]
    end
    if nw_corner
      q[ 0, npy] = q[0, npy-1]
      q[-1, npy] = q[0, npy-2]
    end
    if ne_corner
      q[npx,   npy] = q[npx, npy-1]
      q[npx+1, npy] = q[npx, npy-2]
    end
  end
  if dir == 2
    if sw_corner
      q[0,  0] = q[1, 0]
      q[0, -1] = q[2, 0]
    end
    if se_corner
      q[npx,  0] = q[npx-1, 0]
      q[npx, -1] = q[npx-2, 0]
    end
    if nw_corner
      q[0, npy  ] = q[1, npy]
      q[0, npy+1] = q[2, npy]
    end
    if ne_corner
      q[npx, npy  ] = q[npx-1, npy]
      q[npx, npy+1] = q[npx-2, npy]
    end
  end
  return nothing
end #function fill_4corners

end # module sw_core_mod


text: descendingorder(n) = foldr( (x,accu) -> muladd(10,accu,x), sort!(digits(n)) )

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
using Test
using BigWig
import GenomicFeatures
using BioCore

using FormatSpecimens

@testset "BigWig" begin
    @testset "empty" begin
        buffer = IOBuffer()
        data = buffer.data
        writer = BigWig.Writer(buffer, [("chr1", 1000)])
        close(writer)
        reader = BigWig.Reader(IOBuffer(data))
        @test length(collect(reader)) == 0
    end

    @testset "small" begin
        buffer = IOBuffer()
        data = buffer.data
        writer = BigWig.Writer(buffer, [("chr1", 1000)])
        write(writer, ("chr1", 50, 100, 3.14))
        close(writer)
        reader = BigWig.Reader(IOBuffer(data))
        records = collect(reader)
        @test length(records) == 1
        @test BigWig.haschrom(records[1]) === BioCore.hasseqname(records[1]) === true
        @test BigWig.chrom(records[1]) == BioCore.seqname(records[1]) == "chr1"
        @test BigWig.haschromstart(records[1]) === BioCore.hasleftposition(records[1]) === true
        @test BigWig.chromstart(records[1]) === BioCore.leftposition(records[1]) === 50
        @test BigWig.haschromend(records[1]) === BioCore.hasrightposition(records[1]) === true
        @test BigWig.chromend(records[1]) === BioCore.rightposition(records[1]) === 100
        @test BigWig.hasvalue(records[1])
        @test BigWig.value(records[1]) === 3.14f0
        @test startswith(repr(records[1]), "BigWig.Record:\n")
        interval = convert(GenomicFeatures.Interval, records[1])
        @test GenomicFeatures.seqname(interval) == "chr1"
        @test GenomicFeatures.leftposition(interval) === 50
        @test GenomicFeatures.rightposition(interval) === 100
        @test GenomicFeatures.metadata(interval) === 3.14f0
        @test all(isnan.(BigWig.values(reader, "chr1", 1:49)))
        @test BigWig.values(reader, "chr1", 50:51) == [3.14f0, 3.14f0]
        @test BigWig.values(reader, "chr1", 99:100) == [3.14f0, 3.14f0]
        @test all(isnan.(BigWig.values(reader, "chr1", 101:200)))
        @test BigWig.values(reader, GenomicFeatures.Interval("chr1", 55, 56)) == [3.14f0, 3.14f0]

        # bedgraph (default)
        buffer = IOBuffer()
        data = buffer.data
        writer = BigWig.Writer(buffer, [("chr1", 1000)]; datatype=:bedgraph)
        write(writer, ("chr1",  1, 10, 0.0))
        write(writer, ("chr1", 15, 15, 1.0))
        write(writer, ("chr1", 90, 99, 2.0))
        close(writer)
        reader = BigWig.Reader(IOBuffer(data))
        records = collect(reader)
        @test length(records) == 3
        @test BigWig.chrom.(records) == ["chr1", "chr1", "chr1"]
        @test BigWig.chromstart.(records) == [1, 15, 90]
        @test BigWig.chromend.(records) == [10, 15, 99]
        @test BigWig.value.(records) == [0.0, 1.0, 2.0]

        # varstep
        buffer = IOBuffer()
        data = buffer.data
        writer = BigWig.Writer(buffer, [("chr1", 1000)]; datatype=:varstep)
        write(writer, ("chr1",  1, 10, 0.0))
        write(writer, ("chr1", 15, 24, 1.0))
        write(writer, ("chr1", 90, 99, 2.0))
        close(writer)
        reader = BigWig.Reader(IOBuffer(data))
        records = collect(reader)
        @test length(records) == 3
        @test BigWig.chrom.(records) == ["chr1", "chr1", "chr1"]
        @test BigWig.chromstart.(records) == [1, 15, 90]
        @test BigWig.chromend.(records) == [10, 24, 99]
        @test BigWig.value.(records) == [0.0, 1.0, 2.0]

        # fixedstep
        buffer = IOBuffer()
        data = buffer.data
        writer = BigWig.Writer(buffer, [("chr1", 1000)]; datatype=:fixedstep)
        write(writer, ("chr1",  1,  5, 0.0))
        write(writer, ("chr1", 11, 15, 1.0))
        write(writer, ("chr1", 21, 25, 2.0))
        close(writer)
        reader = BigWig.Reader(IOBuffer(data))
        records = collect(reader)
        @test length(records) == 3
        @test BigWig.chrom.(records) == ["chr1", "chr1", "chr1"]
        @test BigWig.chromstart.(records) == [1, 11, 21]
        @test BigWig.chromend.(records) == [5, 15, 25]
        @test BigWig.value.(records) == [0.0, 1.0, 2.0]
    end

    @testset "large" begin
        buffer = IOBuffer()
        data = buffer.data
        binsize = 32
        writer = BigWig.Writer(buffer, [("chr1", 100_000), ("chr2", 90_000)], binsize=binsize)
        for i in 1:10_000
            write(writer, ("chr1", (i-1)*10+1, i*10, log(i)))
        end
        n = 0
        p = 1
        while p ≤ 90_000
            sz = min(rand(1:100), 90_000 - p)
            write(writer, ("chr2", p, p+sz, log(p)))
            n += 1
            p += sz + 1
        end
        close(writer)
        reader = BigWig.Reader(IOBuffer(data))
        records = collect(reader)
        @test length(records) == 10_000 + n
        records = collect(GenomicFeatures.eachoverlap(reader, GenomicFeatures.Interval("chr1", 50_001, 50_165)))
        @test length(records) == 17
        @testset for bin in [1, 5, 10, 51, 300]
            for scale in 1:2
                binsize_scaled = binsize * BigWig.ZOOM_SCALE_FACTOR^(scale-1)
                chromstart = (bin - 1) * binsize_scaled + 1
                chromend = bin * binsize_scaled
                @test BigWig.coverage(reader, "chr1", chromstart, chromend; usezoom=false) == BigWig.coverage(reader, "chr1", chromstart, chromend; usezoom=true)
                @test BigWig.mean(reader, "chr1", chromstart, chromend; usezoom=false)     ≈  BigWig.mean(reader, "chr1", chromstart, chromend; usezoom=true)
                @test BigWig.minimum(reader, "chr1", chromstart, chromend; usezoom=false)  ≈  BigWig.minimum(reader, "chr1", chromstart, chromend; usezoom=true)
                @test BigWig.maximum(reader, "chr1", chromstart, chromend; usezoom=false)  ≈  BigWig.maximum(reader, "chr1", chromstart, chromend; usezoom=true)
                # TODO: use more stable algorithm?
                #@test_approx_eq BigWig.std(reader, "chr1", chromstart, chromend; usezoom=false)     BigWig.std(reader, "chr1", chromstart, chromend; usezoom=true)
            end
        end
    end

    @testset "round trip" begin
        function test_round_trip(filepath)
            reader = open(BigWig.Reader, filepath)
            buffer = IOBuffer()
            data = buffer.data
            writer = BigWig.Writer(buffer, BigWig.chromlist(reader))
            original = []
            for record in reader
                t = (BigWig.chrom(record), BigWig.chromstart(record), BigWig.chromend(record), BigWig.value(record))
                write(writer, t)
                push!(original, t)
            end
            close(writer)
            close(reader)

            reader = BigWig.Reader(IOBuffer(data))
            copy = []
            for record in reader
                t = (BigWig.chrom(record), BigWig.chromstart(record), BigWig.chromend(record), BigWig.value(record))
                push!(copy, t)
            end
            close(reader)

            @test original == copy
        end

        dir = path_of_format("BBI")
        for specimen in list_valid_specimens("BBI")
            if hastag(specimen, "bigwig")
                test_round_trip(joinpath(dir, specimen["filename"]))
            end
        end

    end
end


text:
T = Float64
# vis = Visualizer()
# open(vis)
include(joinpath(module_dir(), "src", "dynamics", "quadruped", "visuals.jl"))

# s = get_simulation("quadruped", "flat_2D_lc", "flat")

ref_traj = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
    joinpath(module_dir(), "src/dynamics/quadruped/gaits/gait2.jld2"),
    load_type = :split_traj_alt))

im_traj = ImplicitTraj(ref_traj, s, ip_type=:mehrotra)

H = 10
H = ref_traj.H
h = ref_traj.h
q0_sim = SVector{s.model.dim.q}(deepcopy(ref_traj.q[1]))
q1_sim = SVector{s.model.dim.q}(deepcopy(ref_traj.q[2]))

p = open_loop_policy(ref_traj.u; N_sample = 1)

sim_b = simulator(s, q0_sim, q1_sim, h, H, p=p, ip_type=:interior_point,
    sim_opts = SimulatorOptions(warmstart=true))
sim_m = simulator(s, q0_sim, q1_sim, h, H, p=p, ip_type=:mehrotra,
    ip_opts = InteriorPointOptions(max_iter=100),
    sim_opts = SimulatorOptions(warmstart=false))


text:
function rel_energy_err_sde(sol)

    NQ = ndims(sol.q)

    if NQ==2

        en_ref  = 0.5 * ( sol.q.d[1,1]^2 + sol.q.d[2,1]^2 )
        en_last = 0.5 * ( sol.q.d[1,end]^2 + sol.q.d[2,end]^2 )

        return abs((en_last-en_ref)/en_ref)

    elseif NQ==3

        en_ref  = 0.5 * ( sol.q.d[1,1,:].^2 + sol.q.d[2,1,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:].^2 + sol.q.d[2,end,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    elseif NQ==4

        en_ref  = 0.5 * ( sol.q.d[1,1,:,:].^2 + sol.q.d[2,1,:,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:,:].^2 + sol.q.d[2,end,:,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    end
end

function rel_energy_err_psde(sol)

    NQ = ndims(sol.q)

    if NQ==2

        en_ref  = 0.5 * ( sol.q.d[1,1]^2 + sol.p.d[1,1]^2 )
        en_last = 0.5 * ( sol.q.d[1,end]^2 + sol.p.d[1,end]^2 )

        return abs((en_last-en_ref)/en_ref)

    elseif NQ==3

        en_ref  = 0.5 * ( sol.q.d[1,1,:].^2 + sol.p.d[1,1,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:].^2 + sol.p.d[1,end,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    elseif NQ==4

        en_ref  = 0.5 * ( sol.q.d[1,1,:,:].^2 + sol.p.d[1,1,:,:].^2 )
        en_last = 0.5 * ( sol.q.d[1,end,:,:].^2 + sol.p.d[1,end,:,:].^2 )

        return maximum(abs.( (en_last .- en_ref) ./ en_ref ))

    end
end


text:
include("take_step.jl")
include("IPM_tools.jl")
include("terminate.jl")
include("display_progress.jl")
include("one_phase.jl")
include("delta_strategy.jl")


text:
#unum-sqrs.jl

#square and square root

#computes the square of a unum x.  For unums there is no difference from multiplication.
function sqr{ESS,FSS}(x::Unum{ESS,FSS})
  return x * x
end

#overload the power_by_squaring function to allow for the defalt integer power
#thing to work correctly.

function _rawpow(x, p::Integer)
  if p == 1
    return copy(x)
  elseif p == 0
    return one(U)
  elseif p == 2
    return x * x
  elseif iseven(p)
    t = _rawpow(x, p ÷ 2)
    return t * t
  else
    return x * _rawpow(x, p - 1)
  end
end

function pow(x, p::Integer)
  if p == 1
    return copy(x)
  elseif p == 0
    return one(x)
  elseif p == 2
    return sqr(x)
  elseif p < 0
    throw(DomainError())
  end
  #set up the number of iterations we do the multiply.
  iseven(p) ? sqr(_rawpow(x, p ÷ 2)) : x * _rawpow(x, p - 1)
end

import Base.^
@universal ^(x::Unum, p::Integer) = pow(x, p)
@universal ^(x::Ubound, p::Integer) = pow(x, p)

export sqr

max_stars_repo_name: REX-Computing/unumjl

max_issues_repo_name: REX-Computing/unumjl

max_forks_repo_name: REX-Computing/unumjl


text:
#=

  Gaussian mixture model.

  cf ~/blog/gaussian_mixture_model.blog
     ~/psi/gaussian_mixture_model.psi
     ~/webppl/gaussian_mixture_model.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function gaussian_mixture_model()
    x = [0.2,1.0,0.5,0.6]
    n = length(x)

    p ~ Beta(0.5, 1)
    a ~ Uniform(-1, 1)
    b ~ Uniform(-1, 1)

    z = tzeros(n)
    x = tzeros(n)
    for i in 1:n
        z[i] ~ Bernoulli(p)
        if z[i] == 1
            x[i] ~ Normal(a, 1.0)
        else
            x[i] ~ Normal(b, 1.0)
        end
    end

    return a > b

end

model = gaussian_mixture_model()

num_chains = 4

# chains = sample(model, Prior(), 10_000)
# chains = sample(model, MH(), MCMCThreads(), 100_000, num_chains)
# chains = sample(model, MH(), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, MH(), 10_000)

# chains = sample(model, PG(15), MCMCThreads(), 1_000, num_chains)

chains = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, SMC(1000), 10_000)
# chains = sample(model, IS(), 10_000)
#
#chains = sample(model, Gibbs(HMC(0.1,5,:a,:b),PG(15,:p)), 10_000)
# chains = sample(model, Gibbs(NUTS(1000,0.65,:a,:b),PG(15,:p)), 10_000)
# chains = sample(model, Gibbs(HMC(0.1,5,:a,:b),SMC(1000,:p)), 10_000) # Nope

display(chains)
# display(plot(chains))

# show_var_dist_pct(chains,:a,20)
# show_var_dist_pct(chains,:b,20)
# show_var_dist_pct(chains,:p,20)

genq = generated_quantities(model, chains)
show_var_dist_pct(genq,20)

max_stars_repo_name: Wikunia/hakank

max_issues_repo_name: Wikunia/hakank

max_forks_repo_name: Wikunia/hakank


text:
using JETTest, Test

@testset "JETTest.jl" begin
    @testset "dispatch.jl" begin
        include("test_dispatch.jl")
    end

    @testset "self-testing" begin
        frame_filter = module_filter(JETTest)
        function function_filter(@nospecialize(ft))
            return ft !== typeof(Core.Compiler.widenconst) # `widenconst` is very untyped, ignore
        end

        @test_nodispatch frame_filter=frame_filter function_filter=function_filter skip_nonconcrete_calls=false analyze_dispatch(sin, (Int,))
    end
end