text:
mutable struct gpuLDA <: GPUTopicModel
	K::Int
	M::Int
	V::Int
	N::Vector{Int}
	C::Vector{Int}
	B::Int
	corp::Corpus
	batches::Vector{UnitRange{Int}}
	topics::VectorList{Int}
	alpha::Vector{Float32}
	beta::Matrix{Float32}
	gamma::VectorList{Float32}
	phi::MatrixList{Float32}
	Elogtheta::VectorList{Float32}
	Elogthetasum::Vector{Float32}
	newbeta::Void
	Npsums::VectorList{Int}
	Jpsums::VectorList{Int}
	terms::VectorList{Int}
	counts::VectorList{Int}
	words::VectorList{Int}
	device::cl.Device
	context::cl.Context
	queue::cl.CmdQueue
	betakern::cl.Kernel
	betanormkern::cl.Kernel
	newbetakern::cl.Kernel
	gammakern::cl.Kernel
	phikern::cl.Kernel
	phinormkern::cl.Kernel
	Elogthetakern::cl.Kernel
	Elogthetasumkern::cl.Kernel
	Npsumsbuf::cl.Buffer{Int}
	Jpsumsbuf::cl.Buffer{Int}
	termsbuf::cl.Buffer{Int}
	countsbuf::cl.Buffer{Int}
	wordsbuf::cl.Buffer{Int}
	alphabuf::cl.Buffer{Float32}
	betabuf::cl.Buffer{Float32}
	newbetabuf::cl.Buffer{Float32}
	gammabuf::cl.Buffer{Float32}
	phibuf::cl.Buffer{Float32}
	Elogthetabuf::cl.Buffer{Float32}
	Elogthetasumbuf::cl.Buffer{Float32}
	elbo::Float32
	newelbo::Float32

	function gpuLDA(corp::Corpus, K::Integer, batchsize::Integer=length(corp))
		@assert !isempty(corp)
		@assert all(ispositive.([K, batchsize]))
		checkcorp(corp)

		M, V, U = size(corp)
		N = [length(doc) for doc in corp]
		C = [size(doc) for doc in corp]

		batches = partition(1:M, batchsize)
		B = length(batches)

		topics = [collect(1:V) for _ in 1:K]

		alpha = ones(K)
		beta = rand(Dirichlet(V, 1.0), K)'
		newbeta = Array{Float32}(0, 0)
		gamma = [ones(K) for _ in 1:M]
		phi = [ones(K, N[d]) / K for d in batches[1]]

		model = new(K, M, V, N, C, B, copy(corp), batches, topics, alpha, beta, gamma, phi)
		fixmodel!(model, check=false)

		for (b, batch) in enumerate(batches)
			model.phi = [ones(K, N[d]) / K for d in batch]
			model.Elogtheta = [digamma.(ones(K)) - digamma(K) for _ in batch]
			updateNewELBO!(model, b)
		end
		model.phi = [ones(K, N[d]) / K for d in batches[1]]
		model.Elogtheta = [digamma.(ones(K)) - digamma(K) for _ in batches[1]]
		updateELBO!(model)
		return model
	end
end

function Elogptheta(model::gpuLDA, d::Int, m::Int)
	x = lgamma(sum(model.alpha)) - sum(lgamma.(model.alpha)) + dot(model.alpha - 1, model.Elogtheta[m])
	return x
end

function Elogpz(model::gpuLDA, d::Int, m::Int)
	counts = model.corp[d].counts
	x = dot(model.phi[m] * counts, model.Elogtheta[m])
	return x
end

function Elogpw(model::gpuLDA, d::Int, m::Int)
	terms, counts = model.corp[d].terms, model.corp[d].counts
	x = sum(model.phi[m] .* log.(@boink model.beta[:,terms]) * counts)
	return x
end

function Elogqtheta(model::gpuLDA, d::Int)
	x = -entropy(Dirichlet(model.gamma[d]))
	return x
end

function Elogqz(model::gpuLDA, d::Int, m::Int)
	counts = model.corp[d].counts
	x = -sum([c * entropy(Categorical(model.phi[m][:,n])) for (n, c) in enumerate(counts)])
	return x
end

function updateELBO!(model::gpuLDA)
	model.elbo = model.newelbo
	model.newelbo = 0
	return model.elbo
end

function updateNewELBO!(model::gpuLDA, b::Int)
	batch = model.batches[b]
	for (m, d) in enumerate(batch)
		model.newelbo += (Elogptheta(model, d, m)
						+ Elogpz(model, d, m)
						+ Elogpw(model, d, m)
						- Elogqtheta(model, d)
						- Elogqz(model, d, m))
	end
end

function updateAlpha!(model::gpuLDA, niter::Integer, ntol::Real)
	"Interior-point Newton method with log-barrier and back-tracking line search."

	@host model.alphabuf
	@host model.Elogthetasumbuf

	nu = model.K
	for _ in 1:niter
		rho = 1.0
		alphaGrad = Float32[nu / model.alpha[i] + model.M * (digamma(sum(model.alpha)) - digamma(model.alpha[i])) for i in 1:model.K] + model.Elogthetasum
		alphaHessDiag = -(model.M * trigamma.(model.alpha) + (nu ./ model.alpha.^2))
		p = (alphaGrad - sum(alphaGrad ./ alphaHessDiag) / (1 / (model.M * trigamma(sum(model.alpha))) + sum(1 ./ alphaHessDiag))) ./ alphaHessDiag

		while minimum(model.alpha - rho * p) < 0
			rho *= 0.5
		end
		model.alpha -= rho * p

		if (norm(alphaGrad) < ntol) & ((nu / model.K) < ntol)
			break
		end
		nu *= 0.5
	end
	@bumper model.alpha

	model.Elogthetasum = zeros(model.K)
	@buf model.alpha
	@buf model.Elogthetasum
end

const LDA_BETA_cpp =
"""
kernel void
updateBeta(long K,
			global float *newbeta,
			global float *beta)

			{
			long i = get_global_id(0);
			long j = get_global_id(1);

			beta[K * j + i] = newbeta[K * j + i];
			newbeta[K * j + i] = 0.0f;
			}
			"""

const LDA_BETA_NORM_cpp =
"""
kernel void
normalizeBeta(long K,
				long V,
				global float *beta)

				{
				long i = get_global_id(0);

				float normalizer = 0.0f;

				for (long j=0; j<V; j++)
					normalizer += beta[K * j + i];

				for (long j=0; j<V; j++)
					beta[K * j + i] /= normalizer;
				}
				"""

function updateBeta!(model::gpuLDA)
	model.queue(model.betakern, (model.K, model.V), nothing, model.K, model.newbetabuf, model.betabuf)
	model.queue(model.betanormkern, model.K, nothing, model.K, model.V, model.betabuf)
end

const LDA_NEWBETA_cpp =
"""
kernel void
updateNewbeta(long K,
				const global long *Jpsums,
				const global long *counts,
				const global long *words,
				const global float *phi,
				global float *newbeta)

				{
				long i = get_global_id(0);
				long j = get_global_id(1);

				float acc = 0.0f;

				for (long w=Jpsums[j]; w<Jpsums[j+1]; w++)
					acc += counts[words[w]] * phi[K * words[w] + i];

				newbeta[K * j + i] += acc;
				}
				"""

function updateNewBeta!(model::gpuLDA)
	model.queue(model.newbetakern, (model.K, model.V), nothing, model.K, model.Jpsumsbuf, model.countsbuf, model.wordsbuf, model.phibuf, model.newbetabuf)
end

const LDA_GAMMA_cpp =
"""
kernel void
updateGamma(long F,
			long K,
			const global long *Npsums,
			const global long *counts,
			const global float *alpha,
			const global float *phi,
			global float *gamma)

			{
			long i = get_global_id(0);
			long d = get_global_id(1);

			float acc = 0.0f;

			for (long n=Npsums[d]; n<Npsums[d+1]; n++)
				acc += phi[K * n + i] * counts[n];

			gamma[K * (F + d) + i] = alpha[i] + acc + $(EPSILON32);
			}
			"""

function updateGamma!(model::gpuLDA, b::Int)
	batch = model.batches[b]
	model.queue(model.gammakern, (model.K, length(batch)), nothing, batch[1] - 1, model.K, model.Npsumsbuf, model.countsbuf, model.alphabuf, model.phibuf, model.gammabuf)
end

const LDA_PHI_cpp =
"""
kernel void
updatePhi(long K,
			const global long *Npsums,
			const global long *terms,
			const global float *beta,
			const global float *Elogtheta,
			global float *phi)

			{
			long i = get_global_id(0);
			long d = get_global_id(1);

			for (long n=Npsums[d]; n<Npsums[d+1]; n++)
				phi[K * n + i] = beta[K * terms[n] + i] * exp(Elogtheta[K * d + i]);
			}
			"""

const LDA_PHI_NORM_cpp =
"""
kernel void
normalizePhi(long K,
				global float *phi)

				{
				long dn = get_global_id(0);

				float normalizer = 0.0f;

				for (long i=0; i<K; i++)
					normalizer += phi[K * dn + i];

				for (long i=0; i<K; i++)
					phi[K * dn + i] /= normalizer;
				}
				"""

function updatePhi!(model::gpuLDA, b::Int)
	batch = model.batches[b]
	model.queue(model.phikern, (model.K, length(batch)), nothing, model.K, model.Npsumsbuf, model.termsbuf, model.betabuf, model.Elogthetabuf, model.phibuf)
	model.queue(model.phinormkern, sum(model.N[batch]), nothing, model.K, model.phibuf)
end

const LDA_ELOGTHETA_cpp =
"""
$(DIGAMMA_cpp)

kernel void
updateElogtheta(long F,
				long K,
				const global float *gamma,
				global float *Elogtheta)

				{
				long d = get_global_id(0);

				float acc = 0.0f;

				for (long i=0; i<K; i++)
					acc += gamma[K * (F + d) + i];

				acc = digamma(acc);

				for (long i=0; i<K; i++)
					Elogtheta[K * d + i] = digamma(gamma[K * (F + d) + i]) - acc;
				}
				"""

function updateElogtheta!(model::gpuLDA, b::Int)
	batch = model.batches[b]
	model.queue(model.Elogthetakern, length(batch), nothing, batch[1] - 1, model.K, model.gammabuf, model.Elogthetabuf)
end

const LDA_ELOGTHETASUM_cpp =
"""
kernel void
updateElogthetasum(long K,
					long D,
					const global float *Elogtheta,
					global float *Elogthetasum)

					{
					long i = get_global_id(0);

					float acc = 0.0f;

					for (long d=0; d<D; d++)
						acc += Elogtheta[K * d + i];

					Elogthetasum[i] += acc;
					}
					"""

function updateElogthetasum!(model::gpuLDA, b::Int)
	batch = model.batches[b]
	model.queue(model.Elogthetasumkern, model.K, nothing, model.K, length(batch), model.Elogthetabuf, model.Elogthetasumbuf)
end

function train!(model::gpuLDA; iter::Integer=150, tol::Real=1.0, niter::Integer=1000, ntol::Real=1/model.K^2, viter::Integer=10, vtol::Real=1/model.K^2, chkelbo::Int=1)
	@assert all(.!isnegative.([tol, ntol, vtol]))
	@assert all(ispositive.([iter, niter, viter, chkelbo]))
	lowVRAM = model.B > 1

	for k in 1:iter
		chk = (k % chkelbo == 0)
		for b in 1:model.B
			for _ in 1:viter
				oldElogtheta = @host b model.Elogthetabuf
				updateElogtheta!(model, b)
				updatePhi!(model, b)
				updateGamma!(model, b)
				Elogtheta = @host b model.Elogthetabuf
				if sum([norm(diff) for diff in oldElogtheta - Elogtheta]) < length(model.batches[b]) * vtol
					break
				end
			end
			updateElogtheta!(model, b)
			updateElogthetasum!(model, b)
			updateNewBeta!(model)
			if chk
				updateHost!(model, b)
				updateNewELBO!(model, b)
			end
			lowVRAM && updateBuf!(model, b)
		end
		updateAlpha!(model, niter, ntol)
		updateBeta!(model)
		if checkELBO!(model, k, chk, tol)
			break
		end
	end
	updateHost!(model, 1)
	model.topics = [reverse(sortperm(vec(model.beta[i,:]))) for i in 1:model.K]
	nothing
end


text:
"""
    rungekuttasistem(f,a,b,y0,M)
berisi program untuk mencari solusi sistem persamaan differensial `x'=f(t,z)`
dan `y'=g(t,z)` dengan masalah
nilai awal `x(a)=x0` dan `y(a)=y0` pada interval `[a,b]`. Program ini secara default
berisi 5 masukan, yaitu fungsi yang berisi `f(t,z)` dan `g(t,z)`, titik ujung interval `[a,b]`,
nilai awal `z0` yang berisi `[x0,y0]`, dan jumlah sub-interval `M`. Program dapat digunakan juga
untuk sistem persamaan diferensial lebih dari 2 persamaan.

# Example
```jl
julia> f(t,z) = [ z[1]+2*z[2] , 3*z[1]+2*z[2] ];

julia> sol = rungekuttasistem(f,0,0.2,[6,4],10)
11×3 Array{Float64,2}:
 0.0    6.0       4.0
 0.02   6.29355   4.53932
 0.04   6.61562   5.11949
 0.06   6.96853   5.74397
 ⋮
 0.14   8.74141   8.76532
 0.16   9.29021   9.6746
 0.18   9.88827  10.6561
 0.2   10.5396   11.7158
```
return solusi masalah nilai awal `sol`.
"""
function rungekuttasistem(f,a,b,y0,M)
  M = Int(M)
  h = (b-a)/M;
  T = a:h:b;
  Y = Array{Float64}(undef,M+1,length(y0))
  Y[1,:] = y0;
  for k = 1:M
    f1 = f(T[k]     ,Y[k,:]        );
    f2 = f(T[k]+h/2 ,Y[k,:]+f1*h/2 );
    f3 = f(T[k]+h/2 ,Y[k,:]+f2*h/2 );
    f4 = f(T[k]+h   ,Y[k,:]+f3*h   );
    Y[k+1,:] = Y[k,:] + h/6*(f1+2*f2+2*f3+f4);
  end
  sol = [T Y];
end

ext: jl


text:
using CmdStan, Distributions, DiffEqBayes

println("Starting the test")
@test stan_string(Bernoulli(1)) == "bernoulli(1.0)"
@test stan_string(Bernoulli) == "bernoulli"

@test stan_string(Binomial(5,0.3)) == "binomial(5, 0.3)"
@test stan_string(Binomial) == "binomial"

@test stan_string(BetaBinomial(5,1,1)) == "beta_binomial(5, 1.0, 1.0)"
@test stan_string(BetaBinomial) == "beta_binomial"

@test stan_string(Hypergeometric(5,5,3)) == "hypergeometric(5, 5, 3)"
@test stan_string(Hypergeometric) == "hypergeometric"

@test stan_string(NegativeBinomial(5,0.3)) == "neg_binomial(5.0, 0.3)"
@test stan_string(NegativeBinomial) == "neg_binomial"

@test stan_string(Poisson(5)) == "poisson(5.0)"
@test stan_string(Poisson) == "poisson"

@test stan_string(Normal(0,1)) == "normal(0.0, 1.0)"
@test stan_string(Normal) == "normal"

@test stan_string(TDist(5)) == "student_t(5.0,0,1)"
@test stan_string(TDist) == "student_t"

@test stan_string(Cauchy(0,1)) == "cauchy(0.0, 1.0)"
@test stan_string(Cauchy) == "cauchy"

@test stan_string(Laplace(0,1)) == "double_exponential(0.0, 1.0)"
@test stan_string(Laplace) == "double_exponential"

@test stan_string(Distributions.Logistic(0,1)) == "logistic(0.0, 1.0)"
@test stan_string(Distributions.Logistic) == "logistic"

@test stan_string(Gumbel(0,1)) == "gumbel(0.0, 1.0)"
@test stan_string(Gumbel) == "gumbel"

@test stan_string(LogNormal(0,1)) == "lognormal(0.0, 1.0)"
@test stan_string(LogNormal) == "lognormal"

@test stan_string(Chisq(5)) == "chi_square(5.0)"
@test stan_string(Chisq) == "chi_square"

@test stan_string(Exponential(5)) == "exponential(5.0)"
@test stan_string(Exponential) == "exponential"

@test stan_string(Gamma(2,3)) == "gamma(2.0, 3.0)"
@test stan_string(Gamma) == "gamma"

@test stan_string(InverseGamma(2,3)) == "inv_gamma(2.0, 3.0)"
@test stan_string(InverseGamma) == "inv_gamma"

@test stan_string(Weibull(1,1)) == "weibull(1.0, 1.0)"
@test stan_string(Weibull) == "weibull"

@test stan_string(Frechet(1,1)) =="frechet(1.0, 1.0)"
@test stan_string(Frechet) == "frechet"

@test stan_string(Rayleigh(5)) == "rayleigh(5.0)"
@test stan_string(Rayleigh) == "rayleigh"

@test stan_string(Pareto(2,3)) == "pareto(2.0, 3.0)"
@test stan_string(Pareto) == "pareto"

@test stan_string(GeneralizedPareto(0,1,2)) == "pareto_type_2(0.0, 1.0, 2.0)"
@test stan_string(GeneralizedPareto) == "pareto_type_2"

@test stan_string(Beta(3,3)) == "beta(3.0, 3.0)"
@test stan_string(Beta) == "beta"

@test stan_string(Uniform(1,4)) =="uniform(1.0, 4.0)"
@test stan_string(Uniform) == "uniform"

@test stan_string(VonMises(0,2)) == "von_mises(0.0, 2.0)"
@test stan_string(VonMises) == "von_mises"


text:
"Abstractly interpreted PDDL domain."
struct AbstractedDomain{D <: Domain} <: Domain
    domain::D
    interpreter::AbstractInterpreter
end

AbstractedDomain(domain::Domain; options...) =
    AbstractedDomain(domain, AbstractInterpreter(; options...))

Base.copy(domain::AbstractedDomain) = deepcopy(domain)

get_name(domain::AbstractedDomain) = get_name(domain.domain)

get_requirements(domain::AbstractedDomain) = get_requirements(domain.domain)

get_typetree(domain::AbstractedDomain) = get_typetree(domain.domain)

get_constants(domain::AbstractedDomain) = get_constants(domain.domain)

get_constypes(domain::AbstractedDomain) = get_constypes(domain.domain)

get_predicates(domain::AbstractedDomain) = get_predicates(domain.domain)

get_functions(domain::AbstractedDomain) = get_functions(domain.domain)

get_funcdefs(domain::AbstractedDomain) = get_funcdefs(domain.domain)

get_fluents(domain::AbstractedDomain) = get_fluents(domain.domain)

get_axioms(domain::AbstractedDomain) = get_axioms(domain.domain)

get_actions(domain::AbstractedDomain) = get_actions(domain.domain)


text:
# this file downloads files which are used during testing the package
# if they are already present and their checksum matches, they are not downloaded again

using Pkg.PlatformEngines

probe_platform_engines!()  # for download

const testdatadir = @__DIR__

REPO_URL = "https://github.com/yeesian/ArchGDALDatasets/blob/master/"

# remote files with SHA-2 256 hash
"""
To add more files, follow the below steps to generate the SHA
```
julia> using SHA
julia> open(filepath/filename) do f
           bytes2hex(sha256(f))
       end
```
"""
remotefiles = [
    ("data/multi_geom.csv", "00520017658b66ff21e40cbf553672fa8e280cddae6e7a5d1f8bd36bcd521770"),
    ("data/missing_testcase.csv", "d49ba446aae9ef334350b64c876b4de652f28595fdecf78bea4e16af4033f7c6"),
    ("data/point.geojson", "8744593479054a67c784322e0c198bfa880c9388b39a2ddd4c56726944711bd9"),
    ("data/utmsmall.tif", "f40dae6e8b5e18f3648e9f095e22a0d7027014bb463418d32f732c3756d8c54f"),
    ("gdalworkshop/world.tif", "b376dc8af62f9894b5050a6a9273ac0763ae2990b556910d35d4a8f4753278bb"),
    ("ospy/data1/sites.dbf", "7df95edea06c46418287ae3430887f44f9116b29715783f7d1a11b2b931d6e7d"),
    ("ospy/data1/sites.prj", "81fb1a246728609a446b25b0df9ede41c3e7b6a133ce78f10edbd2647fc38ce1"),
    ("ospy/data1/sites.sbn", "198d9d695f3e7a0a0ac0ebfd6afbe044b78db3e685fffd241a32396e8b341ed3"),
    ("ospy/data1/sites.sbx", "49bbe1942b899d52cf1d1b01ea10bd481ec40bdc4c94ff866aece5e81f2261f6"),
    ("ospy/data1/sites.shp", "69af5a6184053f0b71f266dc54c944f1ec02013fb66dbb33412d8b1976d5ea2b"),
    ("ospy/data1/sites.shx", "1f3da459ccb151958743171e41e6a01810b2a007305d55666e01d680da7bbf08"),
    ("ospy/data2/ut_counties.txt", "06585b736091f5bbc62eb040918b1693b2716f550ab306026732e1dfa6cd49a7"),
    ("ospy/data3/cache_towns.dbf", "2344b5195e1a7cbc141f38d6f3214f04c0d43058309b162e877fca755cd1d9fa"),
    ("ospy/data3/cache_towns.sbn", "217e938eb0bec1cdccf26d87e5127d395d68b5d660bc1ecc1d7ec7b3f052f4e3"),
    ("ospy/data3/cache_towns.sbx", "e027b3f67bbb60fc9cf67ab6f430b286fd8a1eaa6c344edaa7da4327485ee9f2"),
    ("ospy/data3/cache_towns.shp", "635998f789d349d80368cb105e7e0d61f95cc6eecd36b34bf005d8c7e966fedb"),
    ("ospy/data3/cache_towns.shx", "0cafc504b829a3da2c0363074f775266f9e1f6aaaf1e066b8a613d5862f313b7"),
    ("ospy/data4/aster.img", "2423205bdf820b1c2a3f03862664d84ea4b5b899c57ed33afd8962664e80a298"),
    ("ospy/data4/aster.rrd", "18e038aabe8fd92b0d12cd4f324bb2e0368343e20cc41e5411a6d038108a25cf"),
    ("ospy/data5/doq1.img", "70b8e641c52367107654962e81977be65402aa3c46736a07cb512ce960203bb7"),
    ("ospy/data5/doq1.rrd", "f9f2fe57d789977090ec0c31e465052161886e79a4c4e10805b5e7ab28c06177"),
    ("ospy/data5/doq2.img", "1e1d744f17e6a3b97dd9b7d8705133c72ff162613bae43ad94417c54e6aced5d"),
    ("ospy/data5/doq2.rrd", "8274dad00b27e008e5ada62afb1025b0e6e2ef2d2ff2642487ecaee64befd914"),
]

for (f, sha) in remotefiles
    # create the directories if they don't exist
    currdir = dirname(f)
    isdir(currdir) || mkpath(currdir)
    # download the file if it is not there or if it has a different checksum
    currfile = normpath(joinpath(testdatadir, f))
    url = REPO_URL * f * "?raw=true"
    PlatformEngines.download_verify(url, sha, currfile; force=true)
end


text:
# Functions to create standard graph matrices
#import Base.eigvals
export adjacency, char_poly, laplace, incidence, dist_matrix, eigvals
using SparseArrays
# Adjaceny Matrix

"""
`adjacency(G)` returns the adjacency matrix of `G`.

Note: If the vertices can be sorted by `sort`, then the first row of
the adjacency matrix corresponds to the first vertex (in order) in `G`
and so forth. However, if the vertices are not sortable in this way,
the mapping between vertices and rows/columns of the matrix is
unpredictable.
"""
function adjacency(G::SimpleGraph)
    n = NV(G)
    A = zeros(Int,(n,n))

    # create a table from V to 1:n
    d = vertex2idx(G)

    for e in G.E
        i = d[e[1]]
        j = d[e[2]]
        A[i,j]=1
        A[j,i]=1
    end

    return A
end

# Laplace matrix

"""
`laplace(G)` returns the Laplacian matrix of `G`. This is the
adjacency matrix minus the (diagonal) degree matrix. See `adjacency`
to understand how vertices correspond to rows/columns of the resulting
matrix.
"""
function laplace(G::SimpleGraph)
    A = adjacency(G)
    d = collect(sum(A,dims=1))[:]
    # D = diagm(d) # deprecated version
    D = Matrix(Diagonal(d))
    L = D-A
    return L
end

# incidence matrix
"""

`incidence(G)` returns the vertex-edge incidence matrix of `G`.

Notes:

* The result is a sparse matrix. Wrap in `full` to convert to nonsparse.

* Each column of the matrix has exactly one `+1` and one `-1`. If `G`
is undirected and an unsigned incidence matrix is desired, use
`incidence(G,false)`.
"""
function incidence(G::SimpleGraph, signed::Bool = true)
    n = NV(G)
    m = NE(G)
    M = spzeros(Int,n,m)
    d = vertex2idx(G)
    E = elist(G)
    a = 1
    b = signed ? -1 : 1

    idx = 0
    for e in E
        i = d[e[1]]
        j = d[e[2]]
        idx += 1
        M[i,idx] = a
        M[j,idx] = b
    end

    return M
end

# Create the n-by-n distance matrix
"""
`dist_matrix(G)` returns a matrix whose `i,j`-entry is the distance
from the `i`th vertex to the `j`th vertex. If there is no `i,j`-path,
that entry is `-1`.
"""
function dist_matrix(G::AbstractSimpleGraph)
    if cache_check(G,:dist_matrix)
      return cache_recall(G,:dist_matrix)
    end
    vtcs = vlist(G)
    n = length(vtcs)
    dd = dist(G)

    A = zeros(Int,n,n)

    for i = 1:n
        u = vtcs[i]
        for j = 1:n
            v = vtcs[j]
            A[i,j] = dd[(u,v)]
        end
    end
    cache_save(G,:dist_matrix,A)
    return A
end

"""
`char_poly(G)` returns the characteristic polynomial of
`adjacency(G)`. Use `char_poly(G,function)` for other
possible integer matrix functions such as `laplace`.
"""
function char_poly(G::AbstractSimpleGraph, func::Function=adjacency)
    if cache_check(G,:char_poly) && func==adjacency
      return cache_recall(G,:char_poly)
    end
    evs = eigvals(G,func)
    P = poly(evs)
    cs = round.(Int,real(coeffs(P)))
    P =  Poly(cs)
    if func==adjacency
        cache_save(G,:char_poly,P)
    end
    return P
end

"""
`eigvals(G)` for a `SimpleGraph` returns the eigenvalues of `G`'s
adjacency matrix. More generally, `eigvals(G,mat)` returns the eigenvalues
of `mat(G)` where `mat` is a matrix-valued function of `G`. In particular,
one can use `mat(G,laplace)` to find the eigenvalues of `G`'s Laplacian.
"""
function LinearAlgebra.eigvals(G::SimpleGraph, mat::Function = adjacency)
    return eigvals(mat(G))
end


text:
module AutomatonEdification

include("datasets/mnist.jl")

include("clusters/clusters.jl")

include("optimization/loss_functions.jl")

include("linear_classifiers/linear_classifiers.jl")
include("linear_classifiers/centroid_linear_classifier_impl.jl")
#=
Executables:
include("src/linear_classifiers/centroid_linear_classifier_executable.jl")
=#

include("neural_networks/activation_functions.jl")
include("neural_networks/neural_networks.jl")
#=
Executables:
include("src/neural_networks/xor_neural_network_executable.jl")
include("src/neural_networks/mnist_neural_network_executable.jl")
=#

end # module

max_stars_repo_name: brianxie/automaton-edification

max_issues_repo_name: brianxie/automaton-edification

max_forks_repo_name: brianxie/automaton-edification


text:
#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
    vecscale(svec::Vector{F})::Matrix{F} where {F<:AbstractFloat}

Transformation matrix for a scaling.

Returns a (4x4) Transformation matrix for scaling.

The matrix is:

    [ sx   0   0   0]
    [  0  sy   0   0]
    [  0   0  sz   0]
    [  0   0   0   1]

# Arguments:
- `svec`: A vectors defining the scaling along the x,y and z axes.
  i.e. `[sx; sy; sy]`

# Output:
- `st`: Scaling Transformation Matrix

# Examples:
```julia
julia> ss = vecscale(svec)
```

Scale up the NURBS line `[0.0;0.0;0.0]` - `[1.0;1.0;1.0]` by 3 along the
x-axis, 2 along the y-axis and 4 along the z-axis.

```julia
julia> line = nrbline(vec([0.0 0.0 0.0]),vec([1.0 1.0 1.0]));
julia> scale = vecscale(vec([3.0 2.0 4.0]));
julia> sline = nrbtform(line, scale)
```

# Reference
- see also: [`nrbtform`](@ref)
"""
function vecscale(svec::Vector{F})::Matrix{F} where {F<:AbstractFloat}
if length(svec) < 1 || length(svec) > 3
  throw(ArgumentError("Translation vector required"));
end
s = [svec;1;1];
ss = [s[1]    0    0   0;
      0    s[2]    0   0;
      0       0 s[3]   0;
      0       0    0   1];
return ss;
end # vecscale


text:
using Test
using BartIO
using PyCall

include("testIO.jl");


text:
function getsum(a::Int, b::Int)::Int
    start,stop = minmax(a,b)
    N = stop-start+1
    (N*(N-1))÷2 + N*start
end

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
#! /usr/bin/julia

# Rosetta Code, Brownian tree

using Color, Images, FixedPointNumbers

const W = 512
const W0 = W>>1
const H = 512
const H0 = H>>1

type Point{T<:Integer}
    x::T
    y::T
end

type Line{T<:Point}
    a::T
    b::T
    stride = Array{Int,1}
end

function start(l::Line)
    n = l.a*l.sa + l.b*l.sa
    return Segment(l.a, l.b, 1, n)
end

function next(l::Line, s::Segment)
    a = s.a + 1

function Line(u::Point, v::Point, img::Image)
    if abs(u.x-v.x) > abs(u.y-v.y)
        if v.x > u.x
            x = [u.x, v.x]
            y = [u.y, v.y]
        else
            x = [v.x, u.x]
            y = [v.y, u.y]
        end
    else
        if v.y > u.y
            x = [u.y, v.y]
            y = [u.x, v.x]
        else
            x = [v.y, u.y]
            y = [v.x, v.y]
        end
    end
    a = x[1]
    b = y[1]
    db = y[2] - y[1]

function motecolor(x::Int, y::Int)
    h = clamp(180*(atan2(y-H0, x-W0)/pi + 1.0), 0.0, 360.0)
    return HSV(h, 0.5, 0.5)
end

img = Image(zeros(RGB{Ufixed8}, H, W))
img["x", W0, "y", H0] = RGB(1, 1, 1)
isfree = trues(W, H)
isfree[W0, H0] = false
for i in 1:N
    x = rand(1:W)
    y = rand(1:H)
    isfree[x, y] || continue
    mc = motecolor(x, y)
    for j in 1:10^3
        xp = x + rand(-1:1)
        yp = y + rand(-1:1)
        iscontained = 0 < xp <= W && 0 < yp <= H
        if  iscontained && isfree[xp, yp]
            x = xp
            y = yp
            continue
        else
            if SIDESTICK || (iscontained && !isfree[xp, yp])
                isfree[x, y] = false
                img["x", x, "y", y] = mc
            end
            break
        end
    end
end

imwrite(img, "brownian_tree.png")

max_stars_repo_name: MichaeLeroy/rosetta-code

max_issues_repo_name: MichaeLeroy/rosetta-code

max_forks_repo_name: MichaeLeroy/rosetta-code


text:
using Turing, Distributions
using DataFrames
using MCMCChains, Plots, StatsPlots
using Distances

using Random
Random.seed!(0)

data = DataFrame(R = Float64[0, 0.04, 0.07, 0.11, 0.165, 0.225, 0.27, 0.31, 0.315],
                 S = Float64[0,   30,   70,  140,   300,   600, 1200, 2600,  3800]);

michaelis_menten(b1, b2, s) = (b1 * s) / (b2 + s)

# Bayesian non-linear regression, Michaelis-Menten kinetics.
@model function nonlinear_regression_MM(R, S)
    # Set the priors on our coefficients.
    b1 ~ truncated(Gamma(1.5, 0.8),   0.001, Inf)
    #b2 ~ truncated(Gamma(1.5, 0.002), 0.001, Inf)
    b2 ~ Normal(500, 20)

    # Set variance prior. Default in brms is a Normal(0, 1)??
    σ ~ truncated(Normal(0, 1), 0, Inf)

    # Calculate all the mu terms.
    for i in eachindex(R)
        r = michaelis_menten(b1, b2, S[i]) # (b1 * S[i]) / (b2 + S[i]) # Michaelis-Menten kinetics model
        R[i] ~ Normal(r, sqrt(σ))
    end
end

model = nonlinear_regression_MM(data.R, data.S)
# All of these fail:
chain = sample(model, NUTS(0.65), 4_000)
chain = sample(model, NUTS(20, 0.95), 4_000)
# Doesn't help:
#using DynamicHMC
#chain = sample(model, DynamicNUTS(), 4_000)
#chain = sample(model, IS(), 4_000)
chain = sample(model, HMCDA(200, 0.65, 0.3), 4_000)
# Doesn't work:
chain = sample(model, PG(), 4_000)

# Make predictions given the parameter values of the chain. Skip the initial half.
function prediction_mm(chain, S, numskip = div(length(chain), 2))
    b1vals = get(chain, :b1).b1[(numskip+1):end]
    b2vals = get(chain, :b2).b2[(numskip+1):end]
    medvals, qlows, qhighs = Float64[], Float64[], Float64[]
    for s in S
        preds = michaelis_menten.(b1vals, b2vals, s)
        qs = quantile(preds, [0.025, 0.50, 0.975])
        push!(qlows, qs[1])
        push!(medvals, qs[2])
        push!(qhighs, qs[3])
    end
    return medvals, qlows, qhighs
end

preds, qlows, qhighs = prediction_mm(chain, data.S)
msd(data.R, preds)

# Let's use BBO to find a good mode and then use that to set the prior on b2. We use OnlineStats
# to save the 5% quantile (since minimizing) and save all values found during optimization that
# are at least as good as it.
using BlackBoxOptim
const S = data.S
const R = data.R

# Minimize the mean squared distance, msd, while saving all param values that lead to top 5%
# fitness:
function fitness(x::Vector{Float64})
    preds = michaelis_menten.(x..., S)
    msd(preds, R)
end

# We assume b1 and b2 are positive and < 10k:
res = bboptimize(fitness; Method = :generating_set_search, SearchRange = (1e-10, 1000.0), NumDimensions = 2, MaxTime = 2.0);
modeb1, modeb2 = best_candidate(res) # 0.34, 298.3

# Now we can set a prior which is quite wide around the mode found:
@model function nonlinear_regression_MM2(R, S)
    # Set the priors on our coefficients.
    b1 ~ truncated(Normal(0.34, 0.20), 0.001, Inf)
    b2 ~ truncated(Normal(298.3, 50.0), 0.001, Inf)

    # Set variance prior. Default in brms is a Normal(0, 1)??
    σ ~ truncated(Normal(0, 1), 0, Inf)

    # Calculate all the mu terms.
    for i in eachindex(R)
        r = michaelis_menten(b1, b2, S[i]) # (b1 * S[i]) / (b2 + S[i]) # Michaelis-Menten kinetics model
        R[i] ~ Normal(r, sqrt(σ))
    end
end

model2 = nonlinear_regression_MM2(data.R, data.S)
chain = sample(model2, NUTS(0.65), 4_000)

# Another idea is to use BBO also to find the largest ranges of values that give distances
# within some fixed factor of the mode.
const BestFitness = msd(michaelis_menten.(modeb1, modeb2, S), R)
const FactorWorse = 10
function fitnessrange(x::Vector{Float64})
    b1a, b1b = x[1:2]
    b2a, b2b = x[3:4]
    predsa = michaelis_menten.(b1a, b2a, S)
    da = msd(predsa, R)
    predsb = michaelis_menten.(b1b, b2b, S)
    db = msd(predsb, R)

    # Penalty if msd distance is more than FactorWorse worse than best
    penaltya = (da > FactorWorse*BestFitness) ? 1000.0 + abs(da - FactorWorse*BestFitness) : 0.0
    penaltyb = (db > FactorWorse*BestFitness) ? 1000.0 + abs(db - FactorWorse*BestFitness) : 0.0

    # We want to maximize the size of the ranges so fitness is the negative of that (since minimizing overall)
    rangesize = abs(b1a - b1b) + abs(b2a - b2b)

    penaltya + penaltyb - 0.0001 * rangesize
end

res = bboptimize(fitnessrange; Method = :generating_set_search, SearchRange = (1e-10, 1000.0), NumDimensions = 4, MaxTime = 2.0);
# When I run this I get the best candidate something like: 0.370357, 0.314395, 440.648, 200.681
# So let's set Uniform(0.31, 0.38) and Uniform(200, 440) priors.

@model function nonlinear_regression_MM3(R, S)
    # Set the priors on our coefficients.
    #b1 ~ Uniform(0.31, 0.38)
    #b2 ~ Uniform(200.0, 440.0)
    b1 ~ Uniform(0.01, 1.0)
    b2 ~ Uniform(10.0, 500.0)

    # Set variance prior. Default in brms is a Normal(0, 1)??
    σ ~ truncated(Normal(0, 1), 0, Inf)

    # Calculate all the mu terms.
    for i in eachindex(R)
        r = michaelis_menten(b1, b2, S[i]) # (b1 * S[i]) / (b2 + S[i]) # Michaelis-Menten kinetics model
        R[i] ~ Normal(r, sqrt(σ))
    end
end

model3 = nonlinear_regression_MM3(data.R, data.S)
chain = sample(model3, NUTS(0.65), 4_000)

max_stars_repo_name: robertfeldt/TuringTutorials

max_issues_repo_name: robertfeldt/TuringTutorials

max_forks_repo_name: robertfeldt/TuringTutorials


text:
module CLIWrappers

using Random
using Dates
using DataFrames
using CSV

using ..DecisionTreeLearners
using ..Pipelines
using ..AbsTypes
using ..Utils

using ..ValDateFilters
using ..Statifiers
using ..Monotonicers

export tsmlrun

const COMMONARG = Dict(
                    :dateformat=>"dd/mm/yyyy HH:MM",
                    :dateinterval=>Dates.Hour(1),
                    :processmissing=>true,
                    :strict => false
                  )

function tsmlrun(inputname::AbstractString,outputname::AbstractString="",datefmt::AbstractString="dd/mm/yyyy HH:MM",otype::AbstractString="table")
   if otype == "table"
      imputedoutput(inputname,outputname,datefmt)
   elseif otype == "stat"
      imputedstat(inputname,outputname,datefmt)
   else
      throw(ArgumentError("wrong output type"))
   end
end

function rawstat(inputname::AbstractString,outputname::AbstractString="",datefmt::AbstractString="dd/mm/yyyy HH:MM")
   isfile(inputname) || throw(ArgumentError("input file name does not exist"))
   csvreader = CSVDateValReader(Dict(:filename=>inputname,:dateformat=>datefmt))
   stfier = Statifier()
   lpipe = @pipeline csvreader |> stfier
   res=fit_transform!(lpipe)
   if outputname != ""
      res |> CSV.write(outputname)
   end
   return res
end

function aggregatedstat(inputname::AbstractString,outputname::AbstractString="",datefmt::AbstractString="dd/mm/yyyy HH:MM")
   isfile(inputname) || throw(ArgumentError("input file name does not exist"))
   csvreader = CSVDateValReader(Dict(:filename=>inputname,:dateformat=>datefmt))
   valgator = DateValgator(COMMONARG)
   stfier = Statifier()
   lpipe = @pipeline csvreader |> valgator |> stfier
   res=fit_transform!(lpipe)
   if outputname != ""
      res |> CSV.write(outputname)
   end
   return res
end

function aggregatedoutput(inputname::AbstractString,outputname::AbstractString="",datefmt::AbstractString="dd/mm/yyyy HH:MM")
   isfile(inputname) || throw(ArgumentError("input file name does not exist"))
   csvreader = CSVDateValReader(Dict(:filename=>inputname,:dateformat=>datefmt))
   valgator = DateValgator(COMMONARG)
   lpipe = @pipeline csvreader |> valgator
   res=fit_transform!(lpipe)
   if outputname != ""
      res |> CSV.write(outputname)
   end
   return res
end

function imputedstat(inputname::AbstractString,outputname::AbstractString="",datefmt::AbstractString="dd/mm/yyyy HH:MM")
   isfile(inputname) || throw(ArgumentError("input file name does not exist"))
   csvreader = CSVDateValReader(Dict(:filename=>inputname,:dateformat=>datefmt))
   valgator = DateValgator(COMMONARG)
   valnner = DateValNNer(COMMONARG)
   mononicer = Monotonicer()
   stfier = Statifier()
   lpipe = @pipeline csvreader |> valgator |> valnner |> mononicer |> stfier
   res=fit_transform!(lpipe)
   if outputname != ""
      res |> CSV.write(outputname)
   end
   return res
end

function imputedoutput(inputname::AbstractString,outputname::AbstractString="",datefmt::AbstractString="dd/mm/yyyy HH:MM")
   isfile(inputname) || throw(ArgumentError("input file name does not exist"))
   csvreader = CSVDateValReader(Dict(:filename=>inputname,:dateformat=>datefmt))
   valgator = DateValgator(COMMONARG)
   valnner = DateValNNer(COMMONARG)
   mononicer = Monotonicer()
   stfier = Statifier()
   lpipe = @pipeline csvreader |> valgator |> valnner |> mononicer
   res=fit_transform!(lpipe)
   if outputname != ""
      res |> CSV.write(outputname)
   end
   return res
end

end


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

@testset ExtendedTestSet "ImagePreprocessing" begin
    step = ImagePreprocessing((0, 0, 0), (.5, .5, .5))
    image = rand(RGB, 100, 100)
    x = FastAI.run(step, Training(), image)
    @test size(x) == (100, 100, 3)
    @test eltype(x) == Float32

    image2 = rand(RGB, 100, 100)
    buf = copy(x)
    @test_nowarn FastAI.run!(buf, step, Training(), image2)
    @test !(buf ≈ x)
end


text:
using Revise
using BridgeLandmarks
const BL = BridgeLandmarks
using Distributions
using LinearAlgebra
using Random
using StaticArrays

# A small example useful for testing purposes

n = 27
xobs0 = [PointF(2.0cos(t), sin(t))/4.0  for t in collect(0:(2pi/n):2pi)[2:end]]
θ, ψ =  π/6, 0.1
rot =  SMatrix{2,2}(cos(θ), sin(θ), -sin(θ), cos(θ))
stretch = SMatrix{2,2}(1.0 + ψ, 0.0, 0.0, 1.0 - ψ)
shift = [0.5, -1.0]
xobsT = [rot * stretch * xobs0[i] + shift  for i in 1:n]

pars = Pars_ms(γinit=1.0/√n, aprior=Pareto(1.0, 0.1), η =  n -> 0.0, dt = 0.01, ρlowerbound=0.9)
pars = Pars_ahs(db=[3.0, 2.0],stdev=.5,γinit=.1, aprior=Pareto(1.0, 0.1), η =  n -> 0.0, dt = 0.01, ρlowerbound=0.9)

Σobs = fill([pars.σobs^2 * one(UncF) for i in 1:n],2) # noise on observations

T = 1.0; t = 0.0:pars.dt:T; tt = tc(t,T)
obs_atzero = true
fixinitmomentato0 = false
ITER = 100
subsamples = 0:pars.skip_saveITER:ITER
obsinfo = BL.set_obsinfo(xobs0,[xobsT],Σobs, obs_atzero,fixinitmomentato0)

ainit = 0.5*mean(norm.(diff(xobs0)))
cinit = pars.cinit
γinit = pars.γinit
if pars.model == :ms
    P = MarslandShardlow(ainit, cinit, γinit, 0.0, n)
elseif pars.model == :ahs
    nfsinit = construct_nfs(pars.db, pars.stdev, γinit)
    P = Landmarks(ainit, cinit, n, pars.db , pars.stdev, nfsinit)
end

xinit = State(xobs0, zeros(PointF,P.n))
mT = zeros(PointF, n)

nshapes = 1
gramT_container = [BL.gram_matrix([xobsT][k], P) for k ∈ 1:nshapes]  # gram matrices at observations at time T

Paux = [BL.auxiliary(P, State(obsinfo.xobsT[k],mT), gramT_container[k]) for k ∈ 1:nshapes]
guidrec = [BL.GuidRecursions(t,obsinfo)  for _ ∈ 1:nshapes]  # initialise guiding terms
Paux = [BL.auxiliary(P, State(obsinfo.xobsT[k],mT)) for k ∈ 1:nshapes] # auxiliary process for each shape
Q = BL.GuidedProposal(P,Paux,t,obsinfo.xobs0,obsinfo.xobsT,guidrec,nshapes,[mT for _ ∈ 1:nshapes])
Q = BL.update_guidrec!(Q, obsinfo)   # compute backwards recursion

X = [BL.initSamplePath(t, xinit) for _ ∈ 1:nshapes]
W = [BL.initSamplePath(t,  zeros(PointF, BL.dimwiener(P))) for _ ∈ 1:nshapes]
for k ∈ 1:nshapes   sample!(W[k], BL.Wiener{Vector{PointF}}())  end

ll, X = BL.gp!(BL.LeftRule(), X, xinit, W, Q; skip=sk)
#only one shape
Xpath = X[1]

a = [extract_initial_and_endstate(5, Xpath)]
push!(a, extract_initial_and_endstate(555, Xpath))
A =vcat(a...)
B = DataFrame(A)
rename!(B, ["iterate", "time0", "time1"])
N = div(nrow(B),2d)
B[!,:type] = repeat(["pos1", "pos2", "mom1", "mom2"], N)

function testf!(X, xinit, W, Q, R; skip=sk)
    for i in 1:R
        ll, X = BL.gp!(BL.LeftRule(), X, xinit, W, Q; skip=sk)
    end
    ll, X
end
@btime testf!(X, xinit, W, Q, 1; skip=sk)
@time testf!(X, xinit, W, Q, 10; skip=sk)

function testf2!(X, xinit, W, Q, R; skip=sk)
    for i in 1:R
        BL.gp!(BL.LeftRule(), X, xinit, W, Q; skip=sk)
    end
    nothing
end
@btime testf2!(X, xinit, W, Q, 1; skip=sk)

# check vecofpoints2state

xv = rand(PointF, 10)

function opt1(xv)
    BL.deepvec2state(BL.deepvec(xv))
end

@btime  opt1(xv)
@btime  BL.vecofpoints2state(xv)

nshapes = 1
T = 1.0; t = 0.0:pars.dt:T; tt = tc(t,T)
#W = [BL.initSamplePath(t,  zeros(PointF, BL.dimwiener(P))) for _ ∈ 1:nshapes]

W = BL.initSamplePath(t,  zeros(PointF, BL.dimwiener(P)))
Wnew = BL.initSamplePath(t,  zeros(PointF, BL.dimwiener(P)))
Wᵒ = BL.initSamplePath(t,  zeros(PointF, BL.dimwiener(P)))

function update_path2!(W, Wnew,Wᵒ)
    ρ = 0.9
    ρ_ = sqrt(1.0-ρ^2)
    # From current state (x,W) with loglikelihood ll, update to (x, Wᵒ)
#    sample!(Wnew, Bridge.Wiener{Vector{PointF}}())
    for i ∈ eachindex(W.yy)
        for j ∈ eachindex(W.yy[1])
            Wᵒ.yy[i][j] = ρ * W.yy[i][j] + ρ_ * Wnew.yy[i][j]
        end
    end
end

function ff2(W, Wnew, Wᵒ, R)
    for k in 1:R
        update_path2!(W, Wnew, Wᵒ)
    end
end
@btime ff2(W, Wnew, Wᵒ, 1)
@btime ff2(W, Wnew, Wᵒ, 2)
@btime ff2(W, Wnew, Wᵒ, 3)

function update_path3!(W, Wnew,Wᵒ)
    ρ = 0.9
    ρ_ = sqrt(1.0-ρ^2)
    # From current state (x,W) with loglikelihood ll, update to (x, Wᵒ)
#    sample!(Wnew, Bridge.Wiener{Vector{PointF}}())
    Wᵒ.yy .= ρ * W[1].yy .+ ρ_ * Wnew.yy
end

function ff3(W, Wnew, Wᵒ, R)
    for k in 1:R
        update_path2!(W, Wnew, Wᵒ)
    end
end
@btime ff3(W, Wnew, Wᵒ, 1)
@btime ff3(W, Wnew, Wᵒ, 2)

@btime landmarksmatching(xobs0,xobsT; ITER=10, pars=p_ms, updatescheme=ups, printskip=printskip, outdir=outdir_ms)
#  259.032 ms (2205591 allocations: 231.72 MiB)
# 273.766 ms (2205601 allocations: 231.72 MiB)
#  1.023 s (21896571 allocations: 1.04 GiB)  (also calling gp!)
#  201.703 ms (236768 allocations: 151.89 MiB)
@btime landmarksmatching(xobs0,xobsT; ITER=20, pars=p_ms, updatescheme=ups, printskip=printskip, outdir=outdir_ms)
#  268.554 ms (2229751 allocations: 232.47 MiB)
#   264.369 ms (2229771 allocations: 232.47 MiB)
#  1.863 s (41611713 allocations: 1.84 GiB)  (also calling gp!)
# 220.015 ms (261200 allocations: 155.45 MiB)

struct MatArr{T}
    x::Vector{Matrix{T}}
end

u1 = MatArr(fill(zeros(Float64,2,3),2))
u1.x[1] = reshape(rand(4),2,2)


text:
using SciPy
using Documenter
using PyPlot # for figure generation
using Distributions

makedocs(;
    modules=[SciPy],
    authors="Atsushi Sakai <asakai.amsl+github@gmail.com> and contributors",
    repo="https://github.com/AtsushiSakai/SciPy.jl/blob/{commit}{path}#L{line}",
    sitename="SciPy.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://AtsushiSakai.github.io/SciPy.jl",
        assets=String["assets/favicon.ico"],
    ),
    pages=[
        "Home" => "index.md",
        "Tips" => "tips.md",
        "Examples" => Any[
            "cluster" => "examples/cluster.md",
            "constants" => "examples/constants.md",
            "fft" => "examples/fft.md",
            "interpolate" => "examples/interpolate.md",
            "io" => "examples/io.md",
            "linalg" => "examples/linalg.md",
            "ndimage" => "examples/ndimage.md",
            "odr" => "examples/odr.md",
            "optimize" => "examples/optimize.md",
            "signal" => "examples/signal.md",
            "sparse" => "examples/sparse.md",
            "spatial" => "examples/spatial.md",
            "special" => "examples/special.md",
            "stats" => "examples/stats.md",
        ],
        "API References" => "api_references.md",
    ],
    doctest = false, # workaround for https://github.com/JuliaDocs/Documenter.jl/issues/1394
)

deploydocs(;
    repo="github.com/AtsushiSakai/SciPy.jl",
)


text:
module BasicAuthRequest

using ..Base64
using URIs
import ..Messages: setheader, hasheader
import ..@debug, ..DEBUG_LEVEL

export basicauthlayer
"""
    basicauthlayer(req) -> HTTP.Response

Add `Authorization: Basic` header using credentials from url userinfo.
"""
function basicauthlayer(handler)
    return function(req; basicauth::Bool=true, kw...)
        if basicauth
            userinfo = unescapeuri(req.url.userinfo)
            if !isempty(userinfo) && !hasheader(req.headers, "Authorization")
                @debug 1 "Adding Authorization: Basic header."
                setheader(req.headers, "Authorization" => "Basic $(base64encode(userinfo))")
            end
        end
        return handler(req; kw...)
    end
end

end # module BasicAuthRequest


text:
module SuperModels

using MPI

using ClimateMachine

using ClimateMachine: Settings
using ...BalanceLaws: parameter_set

using ...DGMethods.NumericalFluxes

using ..HydrostaticBoussinesq:
    HydrostaticBoussinesqModel, NonLinearAdvectionTerm, Forcing

using ..OceanProblems: InitialValueProblem, InitialConditions
using ..Ocean: FreeSlip, Impenetrable, Insulating, OceanBC, Penetrable
using ..CartesianFields: SpectralElementField

using ...Mesh.Filters: CutoffFilter, ExponentialFilter
using ...Mesh.Grids: polynomialorders, DiscontinuousSpectralElementGrid

using ClimateMachine:
    LS3NRK33Heuns,
    OceanBoxGCMConfigType,
    OceanBoxGCMSpecificInfo,
    DriverConfiguration

import ClimateMachine: SolverConfiguration

#####
##### It's super good
#####

struct HydrostaticBoussinesqSuperModel{D, G, E, S, F, N, T, C}
    domain::D
    grid::G
    equations::E
    state::S
    fields::F
    numerical_fluxes::N
    timestepper::T
    solver_configuration::C
end

"""
    HydrostaticBoussinesqSuperModel(; domain, time_step, parameters,
         initial_conditions = InitialConditions(),
                  advection = (momentum = NonLinearAdvectionTerm(), tracers = NonLinearAdvectionTerm()),
         turbulence_closure = (νʰ=0, νᶻ=0, κʰ=0, κᶻ=0),
                   coriolis = (f₀=0, β=0),
        rusanov_wave_speeds = (cʰ=0, cᶻ=0),
                   buoyancy = (αᵀ=0,)
           numerical_fluxes = ( first_order = RusanovNumericalFlux(),
                               second_order = CentralNumericalFluxSecondOrder(),
                                   gradient = CentralNumericalFluxGradient())
                timestepper = ClimateMachine.ExplicitSolverType(solver_method=LS3NRK33Heuns),
    )

Builds a `SuperModel` that solves the Hydrostatic Boussinesq equations.
"""
function HydrostaticBoussinesqSuperModel(;
    domain,
    parameters,
    time_step, # We don't want to have to provide this here, but neverthless it's required.
    initial_conditions = InitialConditions(),
    advection = (
        momentum = NonLinearAdvectionTerm(),
        tracers = NonLinearAdvectionTerm(),
    ),
    turbulence_closure = (νʰ = 0, νᶻ = 0, κʰ = 0, κᶻ = 0),
    coriolis = (f₀ = 0, β = 0),
    rusanov_wave_speeds = (cʰ = 0, cᶻ = 0),
    buoyancy = (αᵀ = 0,),
    forcing = Forcing(),
    numerical_fluxes = (
        first_order = RusanovNumericalFlux(),
        second_order = CentralNumericalFluxSecondOrder(),
        gradient = CentralNumericalFluxGradient(),
    ),
    timestepper = ClimateMachine.ExplicitSolverType(
        solver_method = LS3NRK33Heuns,
    ),
    filters = nothing,
    modeldata = NamedTuple(),
    array_type = Settings.array_type,
    mpicomm = MPI.COMM_WORLD,
    init_on_cpu = true,
    boundary_tags = ((0, 0), (0, 0), (1, 2)),
    boundary_conditions = (
        OceanBC(Impenetrable(FreeSlip()), Insulating()),
        OceanBC(Penetrable(FreeSlip()), Insulating()),
    ),
)

    #####
    ##### Build the grid
    #####

    # Change global setting if its set here
    Settings.array_type = array_type

    grid = DiscontinuousSpectralElementGrid(
        domain;
        boundary_tags = boundary_tags,
        mpicomm = mpicomm,
        array_type = array_type,
    )

    FT = eltype(domain)

    #####
    ##### Construct generic problem type InitialValueProblem
    #####

    problem = InitialValueProblem{FT}(
        dimensions = (domain.L.x, domain.L.y, domain.L.z),
        initial_conditions = initial_conditions,
        boundary_conditions = boundary_conditions,
    )

    #####
    ##### Build HydrostaticBoussinesqModel/Equations
    #####

    equations = HydrostaticBoussinesqModel{eltype(domain)}(
        parameters,
        problem,
        momentum_advection = advection.momentum,
        tracer_advection = advection.tracers,
        forcing = forcing,
        cʰ = convert(FT, rusanov_wave_speeds.cʰ),
        cᶻ = convert(FT, rusanov_wave_speeds.cᶻ),
        αᵀ = convert(FT, buoyancy.αᵀ),
        νʰ = convert(FT, turbulence_closure.νʰ),
        νᶻ = convert(FT, turbulence_closure.νᶻ),
        κʰ = convert(FT, turbulence_closure.κʰ),
        κᶻ = convert(FT, turbulence_closure.κᶻ),
        fₒ = convert(FT, coriolis.f₀),
        β = FT(coriolis.β),
    )

    ####
    #### "modeldata"
    ####
    #### OceanModels require filters (?). If one was not provided, we build a default.
    ####

    # Default vertical filter and horizontal exponential filter:
    if isnothing(filters)
        filters = (
            vert_filter = CutoffFilter(grid, polynomialorders(grid)),
            exp_filter = ExponentialFilter(grid, 1, 8),
        )
    end

    modeldata = merge(modeldata, filters)

    ####
    #### We build a DriverConfiguration here for the purposes of building
    #### a SolverConfiguration. Then we throw it away.
    ####

    driver_configuration = DriverConfiguration(
        OceanBoxGCMConfigType(),
        "",
        (domain.Np, domain.Np),
        eltype(domain),
        array_type,
        parameter_set(equations),
        equations,
        MPI.COMM_WORLD,
        grid,
        numerical_fluxes.first_order,
        numerical_fluxes.second_order,
        numerical_fluxes.gradient,
        nothing,
        nothing, # filter
        OceanBoxGCMSpecificInfo(),
    )

    ####
    #### Pass through the SolverConfiguration interface so that we use
    #### the checkpointing infrastructure
    ####

    solver_configuration = ClimateMachine.SolverConfiguration(
        zero(FT),
        convert(FT, time_step),
        driver_configuration,
        init_on_cpu = init_on_cpu,
        ode_dt = convert(FT, time_step),
        ode_solver_type = timestepper,
        Courant_number = 0.4,
        modeldata = modeldata,
    )

    state = solver_configuration.Q

    u = SpectralElementField(domain, grid, state, 1)
    v = SpectralElementField(domain, grid, state, 2)
    η = SpectralElementField(domain, grid, state, 3)
    θ = SpectralElementField(domain, grid, state, 4)

    fields = (u = u, v = v, η = η, θ = θ)

    return HydrostaticBoussinesqSuperModel(
        domain,
        grid,
        equations,
        state,
        fields,
        numerical_fluxes,
        timestepper,
        solver_configuration,
    )
end

current_time(model::HydrostaticBoussinesqSuperModel) =
    model.solver_configuration.solver.t
Δt(model::HydrostaticBoussinesqSuperModel) =
    model.solver_configuration.solver.dt
current_step(model::HydrostaticBoussinesqSuperModel) =
    model.solver_configuration.solver.steps

end # module


text:
# ------------------------------------------------------------------
_load(file) = endswith(file, ".jls") ? deserialize(file) : FileIO.load(file)

# ------------------------------------------------------------------
_save(file, dat) = endswith(file, ".jls") ? serialize(file, dat) : FileIO.save(file, dat)


text:
module Godunov
using LinearAlgebra
#using StaticArrays
"""
    For now, we assume a constant (fixed) viscosity "ν".
"""

ν = 3.01 / 3
s = 3.01 # relaxation time
q = 3*ν/s

"""
A planetary particulate disk is called a "ring", but this terminology is confusing.
"""
mutable struct ring
    N::Integer              # number of cells
    β::Float64              # power-law index for viscosity
    Δt::Float64             # time step
    X_::Array{Float64,1}    # positions (nodes)
    ΔX_::Array{Float64,1}   # cell sizes (cell)
    U_::Array{Float64,1}    # velocity (on nodes)
    Σ_::Array{Float64,1}    # surface density
    J_::Array{Float64,1}    # mass flux
    r₊::Array{Float64,1}    # 2-component eigenvector, right-travelling wave (relative to fluid) [NB: independent of position]
    r₋::Array{Float64,1}    # 2-component eigenvector, left-travelling wave (relative to fluid) [NB: independent of position]
    λ₊::Array{Float64,1}    # eigenvalues, right-travelling (relative to fluid) wave
    λ₋::Array{Float64,1}    # eigenvalues, left-travelling (relative to fluid) wave
    Λ₊R::Any                # diagonal matrix of eigenvalues, Right-travelling (absolute) right-travelling (relative)
    Λ₊L::Any                # diagonal matrix of eigenvalues, Left-travelling (absolute) right-travelling (relative)
    Λ₋R::Any                #
    Λ₋L::Any                #
    α⁺::Array{Float64,1}    #
    α⁻::Array{Float64,1}    #
    t::Float64              # time
end
function ring(N=5;β=0,xf=3.0,xi=0.01,vfield="standard")#10*xf/N)
    function setEs(U_::Array{Float64,1})
        N = length(U_)
        r₊ = [1.0, √q]
        r₋ = [1.0,-√q]
        λ₊ = zeros(N)
        λ₋ = zeros(N)
        for i in 1:N
            λ₊[i] = U_[i] + √q
            λ₋[i] = U_[i] - √q
        end
        Λ₊R =  Diagonal(λ₊ .* (λ₊ .> 0))
        Λ₊L = -Diagonal(λ₊ .* (λ₊ .< 0))
        Λ₋R =  Diagonal(λ₋ .* (λ₋ .> 0))
        Λ₋L = -Diagonal(λ₋ .* (λ₋ .< 0))
        return (r₊,r₋,λ₊,λ₋,Λ₊R,Λ₊L,Λ₋R,Λ₋L)
    end
    """ in progress.... """
    function setU(U_,X_,vfield)
        if vfield=="standard"
            U_ = 1.0 ./ X_.^4
        elseif vfield=="still"
            U_ *= 0.0
#        elseif vfield =
        end
    end
    xtemp_ = collect(range(xi,xf,length=N+1))
    X_ = xtemp_[2:end]
    ΔX_ = diff(xtemp_)
    U_ = 1.0 ./ X_.^4
    #U_ *= 0.0;
    U_ .= -0.1
    Σ_ = ones(N); Σ_[1] = 0.0;
    Σ_[1:(2*N÷5)] .= 0.0; Σ_[(3*N÷5):end] .= 0.0
    J_ = zeros(N)
    (r₊,r₋,λ₊,λ₋,Λ₊R,Λ₊L,Λ₋R,Λ₋L) = setEs(U_)
    FUDGE = 0.1
    Δt = 0.5 / max( maximum(abs.(λ₊./ΔX_)), maximum(abs.(λ₋./ΔX_))) * FUDGE
    α⁺=zeros(N)
    α⁻=zeros(N)
    return ring(N,β,Δt,X_,ΔX_,U_,Σ_,J_,r₊,r₋,λ₊,λ₋,Λ₊R,Λ₊L,Λ₋R,Λ₋L,α⁺,α⁻,0.0)
end

"""
getR(): get right eigenvectors (two 2-component vectors), and eigenvalues (2*N).
"""

function getAlphas!(r::ring)
    N = r.N; Σ_ = r.Σ_; J_ = r.J_; α⁺ = r.α⁺; α⁻ = r.α⁻
    for i in 1:N
        if i < N
            ΔΣ = Σ_[i+1] - Σ_[i]
            ΔJ = J_[i+1] - J_[i]
        else
            ΔΣ = Σ_[1] - Σ_[i] # 1.0 - Σ_[i]
            ΔJ = J_[1] - J_[i] #0.0 - J_[i]
        end
        #print(0.5 * [1  1/√q] ⋅ [ΔΣ, ΔJ])
        α⁺[i] = 0.5 * [1  1/√q] ⋅ [ΔΣ, ΔJ]
        α⁻[i] = 0.5 * [1 -1/√q] ⋅ [ΔΣ, ΔJ]
    #=    α⁺[i] = [ √q , 1] ⋅ [ΔΣ, ΔJ] ./ √(1+q)
        α⁻[i] = [ √q ,-1] ⋅ [ΔΣ, ΔJ] ./ √(1+q) =#
    end
end#function

function step!(r::ring)
    getAlphas!(r)
    Δt=r.Δt; Σ_=r.Σ_; J_=r.J_; α⁺=r.α⁺; α⁻=r.α⁻; r₋=r.r₋; r₊=r.r₊
    Λ₊L=r.Λ₊L; Λ₊R=r.Λ₊R; Λ₋L=r.Λ₋L; Λ₋R=r.Λ₋R
    N = length(Σ_)
    """ LHS: """
    Σ_[1:end  ] += Δt * ((Λ₋L * α⁻ * r₋[1])                     )[1:end  ]
    Σ_[1:end  ] += Δt * (                     (Λ₊L * α⁺ * r₊[1]))[1:end  ]
    Σ_[2:  end] += Δt * (-(Λ₋R * α⁻ * r₋[1])                     )[1:end-1]
    Σ_[2:  end] += Δt * (                     -(Λ₊R * α⁺ * r₊[1]))[1:end-1]
    #Σ_ = Σ_.* (Σ_ .> 0)  # sanity - kills instability?
    J_[1:end  ] += Δt * ((Λ₋L * α⁻ * r₋[2])                     )[1:end  ]
    J_[1:end  ] += Δt * (                      (Λ₊L * α⁺ * r₊[2]))[1:end  ]
    J_[2:  end] += Δt * (-(Λ₋R * α⁻ * r₋[2])                     )[1:end-1]
    J_[2:  end] += Δt * (                     -(Λ₊R * α⁺ * r₊[2]))[1:end-1]
    #=
    #Threads.@threads for i in 1:N
    @simd for i in 1:N
        if i<N
            Σ_[i] += Δt * ((Λ₋L[i] * α⁻[i] * r₋[1])                     )
            Σ_[i] += Δt * (                     (Λ₊L[i] * α⁺[i] * r₊[1]))
        end
        if i>1
            Σ_[i] += Δt * (-(Λ₋R[i-1] * α⁻[i-1] * r₋[1])                     )
            Σ_[i] += Δt * (                     -(Λ₊R[i-1] * α⁺[i-1] * r₊[1]))
        end
        if i<N
            J_[i] += Δt * ((Λ₋L[i] * α⁻[i] * r₋[2])                     )
            J_[i] += Δt * (                     (Λ₊L[i] * α⁺[i] * r₊[2]))
        end
        if i>1
            J_[i] += Δt * (-(Λ₋R[i-1] * α⁻[i-1] * r₋[2])                     )
            J_[i] += Δt * (                     -(Λ₊R[i-1] * α⁺[i-1] * r₊[2]))
        end
    r.Σ_[i] = Σ_[i]
    r.J_[i] = J_[i]
    end
    =#
    """ RHS: """

    J_[1:end] -= Δt * J_[1:end] / s # forwards Euler

    r.t += Δt
end

end#module

max_stars_repo_name: diakosmos/razor

max_issues_repo_name: diakosmos/razor

max_forks_repo_name: diakosmos/razor


text:
abstract type BATFitBackend <: FitBackend end
backend_type(::Val{:BAT}) = BATFitBackend
fit(::Type{BATFitBackend}, args...; kwargs...) = bat_fit(args...; kwargs...)

function bat_fit(UvD::Type{<:UvSpectrumDensity}, h::Histogram{<:Any, 1};
            bounds = initial_parameter_guess(UvD, h)[end],
            rng_seed = Philox4x((321, 456)),
            prior = missing,
            nsamples::Int = 10^5,
            nchains::Int = 4,
            nsteps_init::Int = 2000,
            nsteps_per_cycle::Int = 2000,
            strict = true,
            max_ncycles::Int = 1,
            threshold::Real = 1.1)
    hp = HistLLHPrecalulations(h)
    sd = SpectrumDensity{typeof(hp), UvD}(hp)
    posterior = BAT.PosteriorDensity(sd, ismissing(prior) ? BAT.NamedTupleDist(bounds) : prior )
    posterior_transform = BAT.bat_transform(BAT.PriorToGaussian(), posterior, BAT.PriorSubstitution()).result
    trafo = BAT.trafoof(posterior_transform.likelihood)

    mcmcalgo = BAT.HamiltonianMC()

    convergence = BAT.BrooksGelmanConvergence(
        threshold = threshold,
        corrected = false
    )

    init = BAT.MCMCChainPoolInit(
        init_tries_per_chain = 8..128,
        nsteps_init = nsteps_init,
    )

    burnin = BAT.MCMCMultiCycleBurnin(
        nsteps_per_cycle = nsteps_per_cycle,
        max_ncycles = max_ncycles
    )

    bat_samples = inv(trafo).(
            BAT.bat_sample(
                rng_seed, posterior_transform,
                BAT.MCMCSampling(
                    mcalg = mcmcalgo,
                    nchains = nchains,
                    nsteps = nsamples,
                    init = init,
                    burnin = burnin,
                    strict = strict,
                    convergence = convergence,
                )
            ).result
        )
    return UvD(mode(bat_samples)), bat_samples
end


text:
let
    pomdp = BabyPOMDP()
    up = updater(pomdp)

    bmdp = GenerativeBeliefMDP(pomdp, up)
    b = initialstate(bmdp, Random.GLOBAL_RNG)
    @inferred generate_sr(bmdp, b, true, MersenneTwister(4))
end


text:
export laplacian_op, poisson_op, sat_op, upwlap_op, upwps_op, fwi_op, fwi_obs_op

OPS_DIR = @__DIR__

@doc raw"""
    fwi_op(lambda::Union{PyObject, Array{Float64}},mu::Union{PyObject, Array{Float64}},
        den::Union{PyObject, Array{Float64}},stf::Union{PyObject, Array{Float64}},
        gpu_id::Union{PyObject, Integer},shot_ids::Union{PyObject, Array{T}},para_fname::String) where T<:Integer

Computes the FWI loss function.

- `lambda` : Lame's first parameter (unit: MPa)
- `mu` : Lame's second parameter (shear modulus, unit: MPa)
- `den` : Density
- `stf` : Source time functions
- `gpu_id` : The ID of GPU to run this FWI operator
- `shot_ids` : The source function IDs (determining the location of sources)
- `para_fname` : Parameter file location
"""
function fwi_op(lambda::Union{PyObject, Array{Float64}},mu::Union{PyObject, Array{Float64}},
        den::Union{PyObject, Array{Float64}},stf::Union{PyObject, Array{Float64}},
        gpu_id::Union{PyObject, Integer},shot_ids::Union{PyObject, Array{T}},para_fname::String) where T<:Integer
    lambda = convert_to_tensor(lambda, dtype=Float64)
    mu = convert_to_tensor(mu, dtype=Float64)
    den = convert_to_tensor(den, dtype=Float64)
    stf = convert_to_tensor(stf, dtype=Float64)
    gpu_id = convert_to_tensor(gpu_id, dtype=Int32)
    shot_ids = convert_to_tensor(shot_ids, dtype=Int32)
    fwi_op = load_op_and_grad("$OPS_DIR/FWI/build/libFwiOp", "fwi_op")
    fwi_op(lambda,mu,den,stf,gpu_id,shot_ids,para_fname)
end

@doc raw"""
    fwi_obs_op(lambda::Union{PyObject, Array{Float64}},mu::Union{PyObject, Array{Float64}},
        den::Union{PyObject, Array{Float64}},stf::Union{PyObject, Array{Float64}},
        gpu_id::Union{PyObject, Integer},shot_ids::Union{PyObject, Array{T}},para_fname::String) where T<:Integer

Generates the observation data and store them as files which will be used by [`fwi_op`](@ref)
For the meaning of parameters, see [`fwi_op`](@ref).
"""
function fwi_obs_op(lambda::Union{PyObject, Array{Float64}},mu::Union{PyObject, Array{Float64}},
    den::Union{PyObject, Array{Float64}},stf::Union{PyObject, Array{Float64}},
    gpu_id::Union{PyObject, Integer},shot_ids::Union{PyObject, Array{T}},para_fname::String) where T<:Integer
    lambda = convert_to_tensor(lambda, dtype=Float64)
    mu = convert_to_tensor(mu, dtype=Float64)
    den = convert_to_tensor(den, dtype=Float64)
    stf = convert_to_tensor(stf, dtype=Float64)
    gpu_id = convert_to_tensor(gpu_id, dtype=Int32)
    shot_ids = convert_to_tensor(shot_ids, dtype=Int32)
    fwi_obs_op = load_op("$OPS_DIR/FWI/build/libFwiOp", "fwi_obs_op")
    fwi_obs_op(lambda, mu, den, stf, gpu_id, shot_ids, para_fname)
end

@doc raw"""
    laplacian_op(coef::Union{PyObject, Array{Float64}}, f::Union{PyObject, Array{Float64}},
            h::Union{PyObject, Float64}, ρ::Union{PyObject, Float64})

Computes the Laplacian of function $f(\mathbf{x})$; here ($\mathbf{x}=[z\quad x]^T$)
```math
-\nabla\cdot\left(c(\mathbf{x}) \nabla \left(u(\mathbf{x}) -\rho \begin{bmatrix}z \\ 0\end{bmatrix}  \right)\right)
```
"""
function laplacian_op(coef::Union{PyObject, Array{Float64}}, f::Union{PyObject, Array{Float64}},
    h::Union{PyObject, Float64}, ρ::Union{PyObject, Float64})
    coef = convert_to_tensor(coef, dtype=Float64)
    f = convert_to_tensor(f, dtype=Float64)
    h = convert_to_tensor(h, dtype=Float64)
    ρ = convert_to_tensor(ρ, dtype=Float64)
    laplacian = load_op_and_grad("$OPS_DIR/Laplacian/build/libLaplacian", "laplacian")
    laplacian_op(coef, f, h, ρ)
end

@doc raw"""
    poisson_op(c::Union{PyObject, Float64}, g::Union{PyObject, Float64},
        h::Union{PyObject, Float64}, ρ::Union{PyObject, Float64}, index::Union{Integer, PyObject}=0)

Solves the Poisson equation ($\mathbf{x}=[z\quad x]^T$)

$\begin{aligned}
-\nabla\cdot\left(c(\mathbf{x}) \nabla \left(u(\mathbf{x}) -\rho  \begin{bmatrix}z \\ 0\end{bmatrix}   \right)\right) &=  g(\mathbf{x}) & \mathbf{x}\in \Omega\\
\frac{\partial u(x)}{\partial n} &=  0 & \mathbf{x}\in \Omega\\
\end{aligned}$

Here $\Omega=[0,n_zh]\times [0, n_xh]$. The equation is solved using finite difference method, where the step size in each direction is $h$. Mathematically, the solution to the PDE is determined up to a constant. Numerically, we discretize the equation with the scheme
```math
(A+E_{11})\mathbf{u} = \mathbf{f}
```
where $A$ is the finite difference coefficient matrix,
```math
(E_{11})_{ij} = \left\{ \begin{matrix}1 & i=j=1 \\ 0 & \mbox{ otherwise }\end{matrix}\right.
```

- `index` : `Int32`, when `index=1`, `SparseLU` is used to solve the linear system; otherwise the function invokes algebraic multigrid method from `amgcl`.
"""
function poisson_op(c::Union{PyObject, Array{Float64}}, g::Union{PyObject, Array{Float64}},
    h::Union{PyObject, Float64},
    ρ::Union{PyObject, Float64}, index::Union{PyObject, Integer}=0)
    c = convert_to_tensor(c, dtype=Float64)
    g = convert_to_tensor(g, dtype=Float64)
    h = convert_to_tensor(h, dtype=Float64)
    ρ = convert_to_tensor(ρ, dtype=Float64)
    index = convert_to_tensor(index, dtype=Int64)
    poisson_op = load_op_and_grad("$OPS_DIR/Poisson/build/libPoissonOp", "poisson_op")
    poisson_op(c, g, h, ρ, index)
end

@doc raw"""
    sat_op(s0::Union{PyObject, Array{Float64}},pt::Union{PyObject, Array{Float64}},
        permi::Union{PyObject, Array{Float64}},poro::Union{PyObject, Array{Float64}},
        qw::Union{PyObject, Array{Float64}},qo::Union{PyObject, Array{Float64}},
        muw::Union{PyObject, Float64},muo::Union{PyObject, Float64},
        sref::Union{PyObject, Array{Float64}},dt::Union{PyObject, Float64},h::Union{PyObject, Float64})

Solves the following discretized equation
```math
\phi (S_2^{n + 1} - S_2^n) - \nabla  \cdot \left( {{m_2}(S_2^{n + 1})K\nabla \Psi _2^n} \right) \Delta t= \left(q_2^n + q_1^n \frac{m_2(S^{n+1}_2)}{m_1(S^{n+1}_2)}\right) \Delta t
```
where
```math
m_2(s) = \frac{s^2}{\mu_w}\qquad m_1(s) = \frac{(1-s)^2}{\mu_o}
```
This is a nonlinear equation and is solved with the Newton-Raphson method.

- `s0` : $n_z\times n_x$, saturation of fluid, i.e., $S_2^n$
- `pt` : $n_z\times n_x$, potential of fluid, i.e., $\Psi_2^n$
- `permi` : $n_z\times n_x$, permeability, i.e., $K$
- `poro` : $n_z\times n_x$, porosity, i.e., $\phi$
- `qw` : $n_z\times n_x$, injection or production rate of ﬂuid 1, $q_2^n$
- `qo` : $n_z\times n_x$, injection or production rate of ﬂuid 2, $q_1^n$
- `muw` : viscosity of fluid 1, i.e., $\mu_w$
- `muo` : viscosity of fluid 2, i.e., $\mu_o$
- `sref` : $n_z\times n_x$, initial guess for $S_2^{n+1}$
- `dt` : Time step size
- `h` : Spatial step size
"""
function sat_op(s0::Union{PyObject, Array{Float64}},pt::Union{PyObject, Array{Float64}},
    permi::Union{PyObject, Array{Float64}},poro::Union{PyObject, Array{Float64}},
    qw::Union{PyObject, Array{Float64}},qo::Union{PyObject, Array{Float64}},
    muw::Union{PyObject, Float64},muo::Union{PyObject, Float64},
    sref::Union{PyObject, Array{Float64}},dt::Union{PyObject, Float64},h::Union{PyObject, Float64})
    s0 = convert_to_tensor(s0, dtype=Float64)
    pt = convert_to_tensor(pt, dtype=Float64)
    permi = convert_to_tensor(permi, dtype=Float64)
    poro = convert_to_tensor(poro, dtype=Float64)
    qw = convert_to_tensor(qw, dtype=Float64)
    qo = convert_to_tensor(qo, dtype=Float64)
    muw = convert_to_tensor(muw, dtype=Float64)
    muo = convert_to_tensor(muo, dtype=Float64)
    sref = convert_to_tensor(sref, dtype=Float64)
    dt = convert_to_tensor(dt, dtype=Float64)
    h = convert_to_tensor(h, dtype=Float64)
    sat_op = load_op_and_grad("$OPS_DIR/Saturation/build/libSatOp", "sat_op")
    sat_op(s0,pt,permi,poro,qw,qo,muw,muo,sref,dt,h)
end

@doc raw"""
    upwlap_op(perm::Union{PyObject, Array{Float64}},
        mobi::Union{PyObject, Array{Float64}},
        func::Union{PyObject, Array{Float64}},
        h::Union{PyObject, Float64},
        rhograv::Union{PyObject, Float64})

Computes the Laplacian of function $f(\mathbf{x})$; here $\mathbf{x}=[z\quad x]^T$.
```math
\nabla\cdot\left(m(\mathbf{x})K(\mathbf{x}) \nabla \left(f(\mathbf{x}) -\rho \begin{bmatrix}z \\ 0\end{bmatrix}  \right)\right)
```
The permeability on the computational grid is computed with Harmonic mean;
the mobility is computed with upwind scheme.

- `perm` : $n_z\times n_x$, permeability of fluid, i.e., $K$
- `mobi` : $n_z\times n_x$, mobility of fluid, i.e., $m$
- `func` : $n_z\times n_x$, potential of fluid, i.e., $f$
- `h` : `Float64`, spatial step size
- `rhograv` : `Float64`, i.e., $\rho$
"""
function upwlap_op(perm::Union{PyObject, Array{Float64}},
    mobi::Union{PyObject, Array{Float64}},
    func::Union{PyObject, Array{Float64}},
    h::Union{PyObject, Float64},rhograv::Union{PyObject, Float64})
    perm = convert_to_tensor(perm, dtype=Float64)
    mobi = convert_to_tensor(mobi, dtype=Float64)
    func = convert_to_tensor(func, dtype=Float64)
    h = convert_to_tensor(h, dtype=Float64)
    rhograv = convert_to_tensor(rhograv, dtype=Float64)
    upwlap_op = load_op_and_grad("$OPS_DIR/Upwlap/build/libUpwlapOp", "upwlap_op")
    upwlap_op(perm,mobi,func,h,rhograv)
end

@doc raw"""

Solves the Poisson equation
```math
-\nabla\cdot\left(m(\mathbf{x})K(\mathbf{x}) \nabla \left(u(\mathbf{x}) -\rho  \begin{bmatrix}z \\ 0\end{bmatrix}   \right)\right) =  g(\mathbf{x})
```
See [`upwps_op`](@ref) for detailed description.

- `perm` : $n_z\times n_x$, permeability of fluid, i.e., $K$
- `mobi` : $n_z\times n_x$, mobility of fluid, i.e., $m$
- `src` : $n_z\times n_x$, source function, i.e., $g(\mathbf{x})$
- `funcref` : $n_z\times n_x$, currently it is not not used
- `h` : `Float64`, spatial step size
- `rhograv` : `Float64`, i.e., $\rho$
- `index` : `Int32`, when `index=1`, `SparseLU` is used to solve the linear system; otherwise the function invokes algebraic multigrid method from `amgcl`.
"""
function upwps_op(perm::Union{PyObject, Array{Float64}},mobi::Union{PyObject, Array{Float64}},
    src,funcref,h::Union{PyObject, Float64},rhograv::Union{PyObject, Float64},index::Union{PyObject, Integer})
    perm = convert_to_tensor(perm, dtype=Float64)
    mobi = convert_to_tensor(mobi, dtype=Float64)
    funcref = convert_to_tensor(funcref, dtype=Float64)
    h = convert_to_tensor(h, dtype=Float64)
    rhograv = convert_to_tensor(rhograv, dtype=Float64)
    index = convert_to_tensor(index, dtype=Int64)
    upwps_op = load_op_and_grad("$OPS_DIR/Upwps/build/libUpwpsOp", "upwps_op")
    upwps_op(perm,mobi,src,funcref,h,rhograv,index)
end


text:
using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack
using StaticArrays

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DQuad
using UniformQuadMesh

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
include("../EntropyStableEuler.jl/src/logmean.jl")
using EntropyStableEuler
using EntropyStableEuler.Fluxes2D

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,rhou,E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end

function pfun_nd(rho, rhou, rhov, E)
    rhoUnorm2 = (rhou^2+rhov^2)/rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function pfun_nd(rho, rhou, E)
    rhoUnorm2 = (rhou^2)/rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function primitive_to_conservative_hardcode(rho,u,v,p)
    Unorm = u^2+v^2
    E = @. p/(γ-1) + .5*rho*Unorm
    return (rho,rho*u,rho*v,E)
end

function euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
    rhoR,uR,vR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))
    vavg   = (@. .5*(vL+vR))

    unorm = (@. uL*uR + vL*vR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg)
    FxS4 = (@. f4aux*uavg)

    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FxS3)
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. f4aux*vavg)
    return (FxS1,FxS2,FxS3,FxS4),(FyS1,FyS2,FyS3,FyS4)
end

function euler_fluxes(rhoL,uL,vL,betaL,rhoR,uR,vR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_2D(rhoL,uL,vL,betaL,rhologL,betalogL,
                           rhoR,uR,vR,betaR,rhologR,betalogR)
end

const TOL = 1e-14
const Nc = 4 # number of components

"Approximation parameters"
N = 2
K1D = 20
T = 0.6

# Initial condition 2D shocktube
const γ = 1.4
const rhoL = 120.0
const uL   = 0.0
const vL   = 0.0
const pL   = rhoL/γ
const rhoR = 1.2
const uR   = 0.0
const vR   = 0.0
const pR   = rhoR/γ

const Re = 10^4
const mu = 1/Re
const lambda = 2/3*mu#0.0
const Pr = 0.73
const cp = γ/(γ-1)
const cv = 1/(γ-1)
const kappa = mu*cp/Pr

"Mesh related variables"
VX, VY, EToV = uniform_quad_mesh(2*K1D,K1D)
# # uniform mesh
# @. VX = (VX+1)/2
# @. VY = (VY+1)/4
@. VX = sin(.5*pi*(VX+1)/2)
@. VY = ((VY+1)/2)^2/2

rd = init_reference_quad(N,gauss_lobatto_quad(0,0,N))
"Initialize reference element"
r,_ = gauss_lobatto_quad(0,0,N)   # Reference nodes
VDM = vandermonde_1D(N,r)         # modal to nodal
Dr = grad_vandermonde_1D(N,r)/VDM # nodal differentiation
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1]) # nodal linear interpolation

Nq = N
rq,wq = gauss_lobatto_quad(0,0,Nq)
M = diagm(wq)
Mlump = zeros(size(M))
Mlump_inv = zeros(size(M))
for i = 1:Nq+1
    Mlump[i,i] = sum(M[i,:])
    Mlump_inv[i,i] = 1.0/Mlump[i,i]
end

# 1D operators
M1D = Mlump
M1D_inv = Mlump_inv
Q1D = M*Dr
B1D = zeros(N+1,N+1)
B1D[1,1] = -1
B1D[end,end] = 1
L = Array(spdiagm(0=>-2*ones(N+1), 1=>ones(N), -1=>ones(N)))
L[1,1] = -1
L[end,end] = -1
psi = pinv(L)*-1/2*B1D*ones(N+1)
S0 = zeros(N+1,N+1)
for i = 1:N+1
    for j = 1:N+1
        if L[i,j] != 0
            S0[i,j] = psi[j] - psi[i]
        end
    end
end
Q01D = S0+1/2*B1D

# Drop zeros
Q1D = Matrix(droptol!(sparse(Q1D),TOL))
Q01D = Matrix(droptol!(sparse(Q01D),TOL))
B1D = Matrix(droptol!(sparse(B1D),TOL))
S01D = Matrix(droptol!(sparse(S0),TOL))
S1D = Matrix(droptol!(sparse((Q1D-Q1D')/2),TOL))

# Tensor product operators
M = droptol!(sparse(kron(M1D,M1D)),TOL)
Minv = spdiagm(0 => 1 ./ diag(M))
Qr = droptol!(sparse(kron(M1D,Q1D)),TOL)
Qs = droptol!(sparse(kron(Q1D,M1D)),TOL)
Qr0 = droptol!(sparse(kron(M1D,Q01D)),TOL)
Qs0 = droptol!(sparse(kron(Q01D,M1D)),TOL)
Sr = droptol!(sparse((Qr-Qr')/2),TOL)
Ss = droptol!(sparse((Qs-Qs')/2),TOL)
S0r = droptol!(sparse(kron(M1D,S01D)),TOL)
S0s = droptol!(sparse(kron(S01D,M1D)),TOL)
Br = droptol!(sparse(Qr+Qr'),TOL)
Bs = droptol!(sparse(Qs+Qs'),TOL)

@unpack Vf,Dr,Ds,LIFT = rd

md = init_mesh((VX,VY),EToV,rd)
@unpack rxJ,sxJ,ryJ,syJ,sJ,J,xf,yf,mapM,mapP,mapB,nxJ,nyJ,x,y = md
xb,yb = (x->x[mapB]).((xf,yf))

# 2D shocktube
topwall  = mapB[findall(@. abs(yb-1/2) < TOL)]
wall     = mapB[findall(@. abs(yb-1/2) > TOL)]
boundary = mapB
nx_top  = nxJ[topwall]./sJ[topwall]
ny_top  = nyJ[topwall]./sJ[topwall]
nx_wall = nxJ[wall]./sJ[wall]
ny_wall = nyJ[wall]./sJ[wall]
nx_b    = nxJ[boundary]./sJ[boundary]
ny_b    = nyJ[boundary]./sJ[boundary]

# Make domain periodic
@unpack Nfaces,Vf = rd
@unpack xf,yf,K,mapM,mapP,mapB = md
# LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
# mapPB = build_periodic_boundary_maps(xf,yf,LX,LY,Nfaces*K,mapM,mapP,mapB)
# mapP[mapB] = mapPB
# @pack! md = mapP

function impose_BCs_inviscid_Ub!(QP,Qf,boundary,nx_b,ny_b)
    # No-slip at walls
    u_1 = Qf[2][boundary]
    u_2 = Qf[3][boundary]
    n_1 = nx_b
    n_2 = ny_b

    Un = @. u_1*n_1 + u_2*n_2
    Ut = @. u_1*n_2 - u_2*n_1

    # ρ^+ = ρ, p^+ = p (beta^+ = beta)
    @. QP[1][boundary] = Qf[1][boundary]
    @. QP[4][boundary] = Qf[4][boundary]

    # u_n^+ = -u_n, u_t^+ = u_t
    @. QP[2][boundary] = 1/(-n_1^2-n_2^2)*(n_1*Un-n_2*Ut)
    @. QP[3][boundary] = 1/(-n_1^2-n_2^2)*(n_2*Un+n_1*Ut)
end

function impose_BCs_inviscid_U!(QP,Qf,UbP,boundary,nx_b,ny_b)
    rho  = UbP[1][boundary]
    u    = UbP[2][boundary]
    v    = UbP[3][boundary]
    beta = UbP[4][boundary]
    p    = rho./beta./2.0
    E    = @. p/(γ-1) + .5*rho*(u^2+v^2)

    QP[1][boundary] .= rho
    QP[2][boundary] .= rho.*u
    QP[3][boundary] .= rho.*v
    QP[4][boundary] .= E
end

function impose_BCs_flux!(flux_x_P,flux_y_P,UbP,boundary,nx_b,ny_b)
    rho  = UbP[1][boundary]
    u    = UbP[2][boundary]
    v    = UbP[3][boundary]
    beta = UbP[4][boundary]
    p    = rho./beta./2.0
    E    = @. p/(γ-1) + .5*rho*(u^2+v^2)

    @. flux_x_P[1][boundary] = rho*u
    @. flux_x_P[2][boundary] = rho*u^2+p
    @. flux_x_P[3][boundary] = rho*u*v
    @. flux_x_P[4][boundary] = u*(E+p)

    @. flux_y_P[1][boundary] = rho*v
    @. flux_y_P[2][boundary] = rho*u*v
    @. flux_y_P[3][boundary] = rho*v^2+p
    @. flux_y_P[4][boundary] = v*(E+p)
end

function impose_BCs_lam!(lamP,lam,topwall,wall)
    lamP[topwall] .= 0.0
    lam[topwall]  .= 0.0
    lamP[wall]    .= 0.0
    lam[wall]     .= 0.0
end

function impose_BCs_entropyvars!(VUP,VUf,topwall,nx_top,ny_top,wall,nx_wall,ny_wall)
    # adiabatic on wall
    @. VUP[2][wall] = -VUf[2][wall]
    @. VUP[3][wall] = -VUf[3][wall]
    @. VUP[4][wall] =  VUf[4][wall]

    # Reflective on top
    v_1 = VUf[2][topwall]
    v_2 = VUf[3][topwall]
    n_1 = nx_top
    n_2 = ny_top

    VUn = @. v_1*n_1 + v_2*n_2
    VUt = @. v_1*n_2 - v_2*n_1

    # v_4^+ = v_4
    @. VUP[4][topwall] = VUf[4][topwall]
    # v_n^+ = -v_n, v_t^+ = v_t
    @. VUP[2][topwall] = VUf[2][topwall] - 2*VUn*n_1
    @. VUP[3][topwall] = VUf[3][topwall] - 2*VUn*n_2
end

function impose_BCs_stress!(σxP,σyP,σxf,σyf,topwall,nx_top,ny_top,wall,nx_wall,ny_wall)
    # Adiabatic no-slip BC
    @. σxP[2][wall] = σxf[2][wall]
    @. σyP[2][wall] = σyf[2][wall]
    @. σxP[3][wall] = σxf[3][wall]
    @. σyP[3][wall] = σyf[3][wall]
    @. σxP[4][wall] = -σxf[4][wall]
    @. σyP[4][wall] = -σyf[4][wall]

    # Reflective on top
    sigma_x_1 = σxf[2][topwall]
    sigma_x_2 = σxf[3][topwall]
    sigma_y_1 = σyf[2][topwall]
    sigma_y_2 = σyf[3][topwall]
    n_1 = nx_top
    n_2 = ny_top

    σn_x = @. sigma_x_1*n_1 + sigma_x_2*n_2
    σn_y = @. sigma_y_1*n_1 + sigma_y_2*n_2

    @. σxP[2][topwall] = -σxf[2][topwall]+2*n_1*σn_x
    @. σyP[2][topwall] = -σyf[2][topwall]+2*n_1*σn_y
    @. σxP[3][topwall] = -σxf[3][topwall]+2*n_2*σn_x
    @. σyP[3][topwall] = -σyf[3][topwall]+2*n_2*σn_y
    @. σxP[4][topwall] = -σxf[4][topwall]
    @. σyP[4][topwall] = -σyf[4][topwall]
end

function init_visc_fxn(λ,μ,Pr)
    let λ=-λ,μ=μ,Pr=Pr
        function viscous_matrices!(Kxx,Kxy,Kyy,v)
            v1,v2,v3,v4 = v
            inv_v4_cubed = @. 1/(v4^3)
            λ2μ = (λ+2.0*μ)
            Kxx[2,2] = inv_v4_cubed*-λ2μ*v4^2
            Kxx[2,4] = inv_v4_cubed*λ2μ*v2*v4
            Kxx[3,3] = inv_v4_cubed*-μ*v4^2
            Kxx[3,4] = inv_v4_cubed*μ*v3*v4
            Kxx[4,2] = inv_v4_cubed*λ2μ*v2*v4
            Kxx[4,3] = inv_v4_cubed*μ*v3*v4
            Kxx[4,4] = inv_v4_cubed*-(λ2μ*v2^2 + μ*v3^2 - γ*μ*v4/Pr)

            Kxy[2,3] = inv_v4_cubed*-λ*v4^2
            Kxy[2,4] = inv_v4_cubed*λ*v3*v4
            Kxy[3,2] = inv_v4_cubed*-μ*v4^2
            Kxy[3,4] = inv_v4_cubed*μ*v2*v4
            Kxy[4,2] = inv_v4_cubed*μ*v3*v4
            Kxy[4,3] = inv_v4_cubed*λ*v2*v4
            Kxy[4,4] = inv_v4_cubed*(λ+μ)*(-v2*v3)

            Kyy[2,2] = inv_v4_cubed*-μ*v4^2
            Kyy[2,4] = inv_v4_cubed*μ*v2*v4
            Kyy[3,3] = inv_v4_cubed*-λ2μ*v4^2
            Kyy[3,4] = inv_v4_cubed*λ2μ*v3*v4
            Kyy[4,2] = inv_v4_cubed*μ*v2*v4
            Kyy[4,3] = inv_v4_cubed*λ2μ*v3*v4
            Kyy[4,4] = inv_v4_cubed*-(λ2μ*v3^2 + μ*v2^2 - γ*μ*v4/Pr)
        end
        return viscous_matrices!
    end
end

viscous_matrices! = init_visc_fxn(lambda,mu,Pr)

# Initial condition 2D shocktube
rho_init(x) = (x < 0.5) ? rhoL : rhoR
p_init(x)   = (x < 0.5) ? pL   : pR

rho = @. rho_init(x)
u   = zeros(size(x))
v   = zeros(size(x))
p   = @. p_init(x)

U = primitive_to_conservative_hardcode.(rho,u,v,p)
rho = [x[1] for x in U]
rhou = [x[2] for x in U]
rhov = [x[3] for x in U]
E = [x[4] for x in U]
U = (rho,rhou,rhov,E)

Np = (N+1)*(N+1)
face_idx = [1:N+1; (N+1):(N+1):Np; Np:-1:Np-N; Np-N:-(N+1):1]
x_idx    = [(N+2):(2*N+2); (3*N+4):(4*N+4)]
y_idx    = [1:(N+1); (2*N+3):(3*N+3)]
S0r1 = sum(S0r,dims=2)
S0s1 = sum(S0s,dims=2)
K      = size(U[1],2)
Np     = (N+1)*(N+1)
Nfaces = 4
Nfp    = Nfaces*(N+1)

# Preallocation
VU     = zero.(U)
flux_x = zero.(U)
flux_y = zero.(U)
lam    = zeros(Nfp,K)
LFc    = zeros(Nfp,K)
Ub     = zero.(U)
sigma_x = zero.(U)
sigma_y = zero.(U)

function zhang_wavespd(rhoL,rhouL,rhovL,EL,sigma2L,sigma3L,sigma4L,pL)
    uL   = rhouL/rhoL
    eL   = (EL - .5*rhoL*uL^2)/rhoL
    tauL = sigma2L
    qL   = uL*tauL-sigma3L
    wavespdL = abs(uL)+1/(2*rhoL^2*eL)*(sqrt(rhoL^2*qL^2+2*rhoL^2*eL*abs(tauL-pL)^2)+rhoL*abs(qL))
    return wavespdL
end

function limiting_param(U_low, P_ij)
    l = 1.0
    # Limit density
    if U_low[1] + P_ij[1] < -TOL
        l = min(abs(U_low[1])/(abs(P_ij[1])+1e-14), 1.0)
    end

    # limiting internal energy (via quadratic function)
    a = P_ij[1]*P_ij[4]-1.0/2.0*(P_ij[2]^2+P_ij[3]^2)
    b = U_low[4]*P_ij[1]+U_low[1]*P_ij[4]-U_low[2]*P_ij[2]-U_low[3]*P_ij[3]
    c = U_low[4]*U_low[1]-1.0/2.0*(U_low[2]^2+U_low[3]^2)

    l_eps_ij = 1.0
    if b^2-4*a*c >= 0
        r1 = (-b+sqrt(b^2-4*a*c))/(2*a)
        r2 = (-b-sqrt(b^2-4*a*c))/(2*a)
        if r1 > TOL && r2 > TOL
            l_eps_ij = min(r1,r2)
        elseif r1 > TOL && r2 < -TOL
            l_eps_ij = r1
        elseif r2 > TOL && r1 < -TOL
            l_eps_ij = r2
        end
    end

    l = min(l,l_eps_ij)
    return l
end

function rhs_IDP_fixdt!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y,J,rxJ,syJ,sJ)
    # TODO: hardcoded!
    K      = size(U[1],2)
    Np     = (N+1)*(N+1)
    Nfaces = 4
    Nfp    = Nfaces*(N+1)

    p = pfun_nd.(U[1],U[2],U[3],U[4])
    #@. p = (γ-1)*(U[4]-.5*(U[2]^2+U[3]^2)/U[1])
    @. flux_x[1] = U[2]
    @. flux_x[2] = U[2]^2/U[1]+p
    @. flux_x[3] = U[2]*U[3]/U[1]
    @. flux_x[4] = U[4]*U[2]/U[1]+p*U[2]/U[1]
    @. flux_y[1] = U[3]
    @. flux_y[2] = U[2]*U[3]/U[1]
    @. flux_y[3] = U[3]^2/U[1]+p
    @. flux_y[4] = U[4]*U[3]/U[1]+p*U[3]/U[1]

    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[3]/U[1]
    @. Ub[4] = U[1]/(2*p)
    Ubf = (x->Vf*x).(Ub)
    UbP = (x->x[mapP]).(Ubf)
    impose_BCs_inviscid_Ub!(UbP,Ubf,boundary,nx_b,ny_b)

    Uf = (x->Vf*x).(U)
    UP = (x->x[mapP]).(Uf)
    impose_BCs_inviscid_U!(UP,Uf,UbP,boundary,nx_b,ny_b)

    # simple lax friedrichs dissipation
    (rhoM,rhouM,rhovM,EM) = Uf
    rhoUM_n = @. (rhouM*nxJ + rhovM*nyJ)/sJ
    @. lam  = abs(sqrt(abs(rhoUM_n/rhoM))+sqrt(γ*(γ-1)*(EM-.5*rhoUM_n^2/rhoM)/rhoM))
    lamP = lam[mapP]
    impose_BCs_lam!(lamP,lam,topwall,wall)
    @. LFc = max(lam,lamP)*sJ

    flux_x_f = (x->Vf*x).(flux_x)
    flux_x_P = (x->x[mapP]).(flux_x_f)
    flux_y_f = (x->Vf*x).(flux_y)
    flux_y_P = (x->x[mapP]).(flux_y_f)
    impose_BCs_flux!(flux_x_P,flux_y_P,UbP,boundary,nx_b,ny_b)

    U_low    = [zeros(Float64,Np) for i = 1:Nc]
    U_low_i  = zeros(Float64,Nc)
    F_low    = [zeros(Float64,Np,Np) for i = 1:Nc]
    F_high   = [zeros(Float64,Np,Np) for i = 1:Nc]
    F_P      = [zeros(Float64,Nfp) for i = 1:Nc]
    rhsU  = [zeros(Float64,Np,K) for i = 1:Nc]
    L     = zeros(Float64,Np,Np)
    P_ij  = zeros(Float64,Nc)

    # =======================
    # Calculate viscous part
    # =======================
    # VU  = v_ufun(U...)
    # VUf = (x->Vf*x).(VU)
    # VUP = (x->x[mapP]).(VUf)
    vector_norm(U) = sum((x->x.^2).(U))
    @. VU[3] = U[2]^2+U[3]^2 # rhoUnorm
    @. VU[4] = U[4]-.5*VU[3]/U[1] # rhoe
    @. VU[1] = log((γ-1)*VU[4]/(U[1]^γ)) # sU
    @. VU[1] = (-U[4]+VU[4]*(γ+1-VU[1]))/VU[4]
    @. VU[2] = U[2]/VU[4]
    @. VU[3] = U[3]/VU[4]
    @. VU[4] = -U[1]/VU[4]
    VUf = (x->Vf*x).(VU)
    VUP = (x->x[mapP]).(VUf)
    impose_BCs_entropyvars!(VUP,VUf,topwall,nx_top,ny_top,wall,nx_wall,ny_wall)

    surfx(uP,uf) = LIFT*(@. .5*(uP-uf)*nxJ)
    surfy(uP,uf) = LIFT*(@. .5*(uP-uf)*nyJ)
    volx(ur) = @. rxJ*ur
    voly(us) = @. syJ*us
    VUr = (x->Dr*x).(VU)
    VUs = (x->Ds*x).(VU)
    VUx = volx.(VUr) .+ surfx.(VUP,VUf)
    VUy = voly.(VUs) .+ surfy.(VUP,VUf)
    # VUx = (x->x./J).(VUx)
    # VUy = (x->x./J).(VUy)
    for c = 1:Nc
        @. VUx[c] = VUx[c]/J
        @. VUy[c] = VUy[c]/J
    end

    # initialize sigma_x,sigma_y = viscous rhs
    Kxx,Kxy,Kyy = ntuple(x->MMatrix{4,4}(zeros(Nc,Nc)),3)
    sigma_x_e = zero.(getindex.(VU,:,1))
    sigma_y_e = zero.(getindex.(VU,:,1))
    for e = 1:K
        fill!.(sigma_x_e,0.0)
        fill!.(sigma_y_e,0.0)

        # mult by matrices and perform local projections
        for i = 1:Np
            vxi = getindex.(VUx,i,e)
            vyi = getindex.(VUy,i,e)
            viscous_matrices!(Kxx,Kxy,Kyy,getindex.(VU,i,e))

            for col = 1:Nc
                vxi_col = vxi[col]
                vyi_col = vyi[col]
                for row = 1:Nc
                    sigma_x_e[row][i] += Kxx[row,col]*vxi_col + Kxy[row,col]*vyi_col
                    sigma_y_e[row][i] += Kxy[col,row]*vxi_col + Kyy[row,col]*vyi_col
                end
            end
        end
        setindex!.(sigma_x,sigma_x_e,:,e)
        setindex!.(sigma_y,sigma_y_e,:,e)
    end

    sxf = (x->Vf*x).(sigma_x)
    syf = (x->Vf*x).(sigma_y)
    sxP = (x->x[mapP]).(sxf)
    syP = (x->x[mapP]).(syf)
    impose_BCs_stress!(sxP,syP,sxf,syf,topwall,nx_top,ny_top,wall,nx_wall,ny_wall)

    # =====================
    # Loop through elements
    # =====================
    for k = 1:K
        J_k = J[1,k]
        rxJ_k = rxJ[1,k]
        syJ_k = syJ[1,k]
        for c = 1:Nc
            F_low[c]  .= 0.0
            F_high[c] .= 0.0
            F_P[c]    .= 0.0
            U_low[c]  .= 0.0
        end
        L .= 0.0

        # Calculate low order algebraic flux
        for i = 1:Np
            for j = 1:Np
                c_ij_norm = sqrt(rxJ_k^2*S0r[i,j]^2+syJ_k^2*S0s[i,j]^2)
                if abs(c_ij_norm) >= TOL
                    n_ij_x = rxJ_k*S0r[i,j]/c_ij_norm
                    n_ij_y = syJ_k*S0s[i,j]/c_ij_norm
                    wavespd_i = wavespeed_1D(U[1][i,k],n_ij_x*U[2][i,k]+n_ij_y*U[3][i,k],U[4][i,k])
                    wavespd_j = wavespeed_1D(U[1][j,k],n_ij_x*U[2][j,k]+n_ij_y*U[3][j,k],U[4][j,k])
                    wavespd = max(wavespd_i,wavespd_j)
                    d_ij = wavespd*c_ij_norm

                    for c = 1:Nc
                        F_low[c][i,j] = (rxJ_k*S0r[i,j]*(flux_x[c][i,k]+flux_x[c][j,k])
                                        +syJ_k*S0s[i,j]*(flux_y[c][i,k]+flux_y[c][j,k])
                                        -rxJ_k*S0r[i,j]*(sigma_x[c][i,k]+sigma_x[c][j,k])
                                        -syJ_k*S0s[i,j]*(sigma_y[c][i,k]+sigma_y[c][j,k])
                                        -d_ij*(U[c][j,k]-U[c][i,k]))
                    end
                end
            end
        end

        # Calculate high order algebraic flux
        for i = 1:Np-1
            for j = i+1:Np
                # TODO: can preallocate nonzero entries
                if Sr[i,j] != 0.0 || Ss[i,j] != 0.0
                    F1,F2 = euler_fluxes(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[4][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],Ub[4][j,k])
                    for c = 1:Nc
                        val = (2*rxJ_k*Sr[i,j]*F1[c]+2*syJ_k*Ss[i,j]*F2[c]
                                -rxJ_k*S0r[i,j]*(sigma_x[c][i,k]+sigma_x[c][j,k])
                                -syJ_k*S0s[i,j]*(sigma_y[c][i,k]+sigma_y[c][j,k]))
                        F_high[c][i,j] = val
                        F_high[c][j,i] = -val
                    end
                end
            end
        end

        # Calculate interface fluxes
        for i = 1:Nfp
            Jf = sJ[i,k]
            S0r_ij = -S0r1[face_idx[i]]
            S0s_ij = -S0s1[face_idx[i]]

            # flux in x direction
            if i in x_idx
                wavespd_M = wavespeed_1D(Uf[1][i,k],Uf[2][i,k],Uf[4][i,k])
                wavespd_P = wavespeed_1D(UP[1][i,k],UP[2][i,k],UP[4][i,k])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0r_ij)
                for c = 1:Nc
                    F_P[c][i] = (Jf*S0r_ij*(flux_x_f[c][i,k]+flux_x_P[c][i,k])
                                -Jf*S0r_ij*(sxf[c][i,k]+sxP[c][i,k])
                                -LFc[i,k]*abs(S0r_ij)*(UP[c][i,k]-Uf[c][i,k]))
                end
            end

            # flux in y direction
            if i in y_idx
                wavespd_M = wavespeed_1D(Uf[1][i,k],Uf[3][i,k],Uf[4][i,k])
                wavespd_P = wavespeed_1D(UP[1][i,k],UP[3][i,k],UP[4][i,k])
                wavespd = max(wavespd_M,wavespd_P)
                d_ij = wavespd*abs(S0s_ij)
                for c = 1:Nc
                    F_P[c][i] = (Jf*S0s_ij*(flux_y_f[c][i,k]+flux_y_P[c][i,k])
                                -Jf*S0s_ij*(syf[c][i,k]+syP[c][i,k])
                                -LFc[i,k]*abs(S0s_ij)*(UP[c][i,k]-Uf[c][i,k]))
                end
            end
        end

        # Calculate low order solution
        for i = 1:Np
            for j = 1:Np
                for c = 1:Nc
                    U_low[c][i] += F_low[c][i,j]
                end
            end
        end
        for i = 1:Nfp
            for c = 1:Nc
                U_low[c][face_idx[i]] += F_P[c][i]
            end
        end
        for c = 1:Nc
            for i = 1:Np
                U_low[c][i] = U[c][i,k]-dt*Minv[i,i]/J_k*U_low[c][i]
            end
        end

        # Calculate limiting parameters
        for i = 1:Np
            m_i = J_k/Minv[i,i]
            lambda_j = 1/(Np-1)
            for c = 1:Nc
                U_low_i[c] = U_low[c][i]
            end
            for j = 1:Np
                if i != j
                    for c = 1:Nc
                        P_ij[c] = dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                    end
                    L[i,j] = limiting_param(U_low_i, P_ij)
                end
            end
        end

        # Symmetrize limiting parameters
        for i = 1:Np
            for j = i+1:Np
                l = min(L[i,j],L[j,i])
                L[i,j] = l
                L[j,i] = l
            end
        end

        #L = zeros(size(L))

        # l_e = 1.0
        # for i = 1:Np
        #     for j = 1:Np
        #         if i != j
        #             l_e = min(L[i,j],l_e)
        #         end
        #     end
        # end
        # L .= max(l_e,0.0)

        for c = 1:Nc
            for i = 1:Np
                for j = 1:Np
                    rhsU[c][i,k] += (L[i,j]-1)*F_low[c][i,j]-L[i,j]*F_high[c][i,j]
                end
            end
            for i = 1:Nfp
                rhsU[c][face_idx[i],k] -= F_P[c][i]
            end
        end
    end

    # for c = 1:Nc
    #     rhsU[c] .= 1/J*Minv*rhsU[c]
    # end
    for c = 1:Nc
        rhsU[c] .= Minv*rhsU[c]
    end
    # for k = 1:K
    #     for c = 1:Nc
    #         rhsU[c][:,k] = 1/J[1,k]*Minv*rhsU[c][:,k]
    #     end
    # end
    for k = 1:K
        J_k = J[1,k]
        for i = 1:Np
            for c = 1:Nc
                rhsU[c][i,k] = 1/J_k*rhsU[c][i,k]
            end
        end
    end

    return rhsU
end

function rhs_IDP_vardt!(U,N,K1D,M,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y,J,rxJ,syJ,sJ)
    # TODO: hardcoded!
    K      = size(U[1],2)
    Np     = (N+1)*(N+1)
    Nfaces = 4
    Nfp    = Nfaces*(N+1)

    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[3]/U[1]
    @. Ub[4] = U[1]/(2*p)
    Ubf = (x->Vf*x).(Ub)
    UbP = (x->x[mapP]).(Ubf)
    impose_BCs_inviscid_Ub!(UbP,Ubf,boundary,nx_b,ny_b)

    Uf = (x->Vf*x).(U)
    UP = (x->x[mapP]).(Uf)
    impose_BCs_inviscid_U!(UP,Uf,UbP,boundary,nx_b,ny_b)

    # simple lax friedrichs dissipation
    (rhoM,rhouM,rhovM,EM) = Uf
    rhoUM_n = @. (rhouM*nxJ + rhovM*nyJ)/sJ
    @. lam  = abs(sqrt(abs(rhoUM_n/rhoM))+sqrt(γ*(γ-1)*(EM-.5*rhoUM_n^2/rhoM)/rhoM))
    lamP = lam[mapP]
    impose_BCs_lam!(lamP,lam,topwall,wall)
    @. LFc = max(lam,lamP)*sJ

    # =======================
    # Determine timestep size
    # =======================
    d_ii_arr = zeros(Np,K)
    for k = 1:K
        J_k = J[1,k]
        rxJ_k = rxJ[1,k]
        syJ_k = syJ[1,k]
        for i = 1:Np
            for j = 1:Np
                c_ij_norm = sqrt(rxJ_k^2*S0r[i,j]^2+syJ_k^2*S0s[i,j]^2)
                if abs(c_ij_norm) >= TOL
                    n_ij_x = rxJ_k*S0r[i,j]/c_ij_norm
                    n_ij_y = syJ_k*S0s[i,j]/c_ij_norm
                    wavespd_i = wavespeed_1D(U[1][i,k],n_ij_x*U[2][i,k]+n_ij_y*U[3][i,k],U[4][i,k])
                    wavespd_j = wavespeed_1D(U[1][j,k],n_ij_x*U[2][j,k]+n_ij_y*U[3][j,k],U[4][j,k])
                    wavespd = max(wavespd_i,wavespd_j)
                    d_ij = wavespd*c_ij_norm
                    d_ii_arr[i,k] -= d_ij
                end
            end
        end

        for i = 1:Nfp
            S0r_ij = -S0r1[face_idx[i]]
            S0s_ij = -S0s1[face_idx[i]]
            if i in x_idx
                d_ii_arr[face_idx[i]] -= LFc[i,k]*abs(S0r_ij)
            end
            if i in y_idx
                d_ii_arr[face_idx[i]] -= LFc[i,k]*abs(S0s_ij)
            end
        end
    end

    dt = Inf
    for k = 1:K
        for i = 1:Np
            J_k = J[i,k]
            dt = min(-.5*J_k*M[i,i]/d_ii_arr[i,k],dt)
        end
    end
    #dt = minimum(-J/2*M*(1 ./d_ii_arr))

    return rhs_IDP_fixdt!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y,J,rxJ,syJ,sJ), dt
end

# Time stepping
"Time integration"
t = 0.0
U = collect(U)
resU = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]
resW = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]
resZ = [zeros(size(x)),zeros(size(x)),zeros(size(x)),zeros(size(x))]

gr(aspect_ratio=:equal,legend=false,
   markerstrokewidth=0,markersize=2,axis=nothing)
const GIFINTERVAL = 100
#plot()

#plotting nodes
@unpack VDM = rd
rp,sp = equi_nodes_2D(10)
Vp = vandermonde_2D(N,rp,sp)/VDM
gr(aspect_ratio=:equal,legend=false,
   markerstrokewidth=0,markersize=2)

# dt = 1e-4
# rhs_IDP_fixdt!(resW,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
# @btime rhs_IDP_fixdt!(resW,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
# # @profiler rhs_IDP!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);

dt_hist = []
anim = Animation()
i = 1

while t < T
     # SSPRK(3,3)
    # dt = min(1e-4,T-t)
    #rhsU = rhs_IDP!(U,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y);
    rhsU,dt = rhs_IDP_vardt!(U,N,K1D,M,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y,J,rxJ,syJ,sJ);
    dt = min(500*dt,T-t)
    @. resW = U + dt*rhsU
    rhsU = rhs_IDP_fixdt!(resW,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y,J,rxJ,syJ,sJ);
    @. resZ = resW+dt*rhsU
    @. resW = 3/4*U+1/4*resZ
    rhsU = rhs_IDP_fixdt!(resW,N,K1D,Minv,Vf,Dr,Ds,nxJ,nyJ,Sr,Ss,S0r,S0s,S0r1,S0s1,LIFT,mapP,face_idx,x_idx,y_idx,boundary,nx_b,ny_b,topwall,nx_top,ny_top,wall,nx_wall,ny_wall,dt,VU,flux_x,flux_y,lam,LFc,Ub,sigma_x,sigma_y,J,rxJ,syJ,sJ);
    @. resZ = resW+dt*rhsU
    @. U = 1/3*U+2/3*resZ

    push!(dt_hist,dt)
    global t = t + dt
    println("Current time $t with time step size $dt, and final time $T, at step $i")
    global i = i + 1

    # if mod(i,100) == 1
    #     rho = U[1]
    #     rhou = U[2]
    #     rhov = U[3]
    #     E = U[4]
    #     rx = 4*K1D
    #     sy = 4*K1D
    #
    #     g = sqrt.((rx*Dr*rho).^2 .+ (sy*Ds*rho).^2)
    #     g_min = minimum(g)
    #     g_max = maximum(g)
    #
    #     xp = Vp*x
    #     yp = Vp*y
    #     vv = @. exp(-10*(g-g_min)/(g_max-g_min))
    #     vv = Vp*vv
    #     scatter(xp,yp,vv,zcolor=vv,camera=(0,90),colorbar=:right)
    #     frame(anim)
    # end
end

################
### Plotting ###
################

# gif(anim,"~/Desktop/tmp.gif",fps=15)

rho = U[1]
rhou = U[2]
rhov = U[3]
E = U[4]

rx = rxJ./J
sy = syJ./J

g = sqrt.((rx.*(Dr*rho)).^2 .+ (sy.*(Ds*rho)).^2)
g_min = minimum(g)
g_max = maximum(g)

# scatter(x,y,U[1],zcolor=U[1],camera=(0,90),colorbar=:right)
xp = x
yp = y
#vv = Vp*U[1]
vv = @. exp(-10*(g-g_min)/(g_max-g_min))
color_schlieren = cgrad(:Blues_9,20,categorical=true,rev=true)
scatter(xp,yp,vv,zcolor=vv,camera=(0,90),ms=1,colorbar=:right,color=color_schlieren)
savefig("~/Desktop/N=$N,K1D=$K1D,T=$T,2Dshocktube.png")
#scatter(xp,yp,exact_rho_p,zcolor=exact_rho_p,camera=(0,90),colorbar=:right)

max_stars_repo_name: yiminllin/ESDG-PosLimit

max_issues_repo_name: yiminllin/ESDG-PosLimit

max_forks_repo_name: yiminllin/ESDG-PosLimit


text:
module TestIndex
using Base.Test
using DataFrames, DataFrames.Index

i = Index()
push!(i, :A)
push!(i, :B)

inds = {1,
        1.0,
        :A,
        [true],
        trues(1),
        [1],
        [1.0],
        1:1,
        1.0:1.0,
        [:A],
        @data([true]),
        @data([1]),
        @data([1.0]),
        @data([:A]),
        DataArray([:A]),
        PooledDataArray([true]),
        @pdata([1]),
        @pdata([1.0]),
        @pdata([:A]),
        PooledDataArray([:A])}

for ind in inds
    if isequal(ind, :A) || ndims(ind) == 0
        @test isequal(i[ind], 1)
    else
        @test (i[ind] == [1])
    end
end

@test names(i) == [:A,:B]
@test names!(i, [:a,:b]) == Index([:a,:b])
@test rename(i, [:a=>:A, :b=>:B]) == Index([:A,:B])
@test rename(i, :a, :A) == Index([:A,:b])
@test rename(i, [:a], [:A]) == Index([:A,:b])
# @test rename(i, uppercase) == Index([:A,:B])
@test delete!(i, :a) == Index([:b])
push!(i, :C)
@test delete!(i, 1) == Index([:C])

end


text:
# Autogenerated wrapper script for SHTOOLS_jll for x86_64-w64-mingw32-libgfortran4
export libSHTOOLS, libSHTOOLS_mp

using CompilerSupportLibraries_jll
using FFTW_jll
using OpenBLAS32_jll
JLLWrappers.@generate_wrapper_header("SHTOOLS")
JLLWrappers.@declare_library_product(libSHTOOLS, "libSHTOOLS.dll")
JLLWrappers.@declare_library_product(libSHTOOLS_mp, "libSHTOOLS-mp.dll")
function __init__()
    JLLWrappers.@generate_init_header(CompilerSupportLibraries_jll, FFTW_jll, OpenBLAS32_jll)
    JLLWrappers.@init_library_product(
        libSHTOOLS,
        "bin\\libSHTOOLS.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libSHTOOLS_mp,
        "bin\\libSHTOOLS-mp.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using Documenter, DocumenterTools, StructuralDDP, StructuralDDPModels

makedocs(sitename="StructuralDDP.jl",
	modules = [StructuralDDP, StructuralDDPModels],
	authors = "Pascal Golec",
	doctest = false)

deploydocs(#deps = Deps.pip("mkdocs", "python-markdown-math"),
    repo = "github.com/pascalgolec/StructuralDDP.jl.git",
	# target = "build",
	)


text:
module rocFFT

using ..AMDGPU
using ..AMDGPU: librocfft
using ..HIP: hipStreamSynchronize

using CEnum

include("librocfft_common.jl")
include("error.jl")
version() = VersionNumber(rocfft_version_major, rocfft_version_minor, rocfft_version_patch)

include("librocfft.jl")
include("util.jl")
include("fft.jl")

rocfft_setup()
atexit(rocfft_cleanup)

end


text:
"""
Just some useful functions
"""

using BSON
using DelimitedFiles
using StaticArrays
using Setfield

#const σₓ = SMatrix{} # etc

"""
Compute the commutator between two matrices A and B
"""
function commutator(A, B)
    A * B - B * A
end

# mul!ing over a in-place commutator :D (I need to get out more)
function commutator!(A, B, store)
    mul!(store, A, B)

    A * B - B * A
end

"""
Function that will take an optimal pulse and a solution result and simulate the pulse again returning the output so you can check the gate is correct
"""
function test_pulse(prob, solres)
    @show "not implemented"

end

"""
Initialise an ensemble from an ensemble problem definition. Right now this creates an array of problems just now, not sure if this is the best idea though.
"""
function init_ensemble(ens)
    ensemble_problem_array = [deepcopy(ens.prob) for k = 1:ens.n_ens]
    for k = 1:ens.n_ens
        prob_to_update = ensemble_problem_array[k]
        prob_to_update = @set prob_to_update.A = ens.A_g(k)
        prob_to_update = @set prob_to_update.B = ens.B_g(k)
        prob_to_update = @set prob_to_update.Xi = ens.XiG(k)
        prob_to_update = @set prob_to_update.Xt = ens.XtG(k)
        ensemble_problem_array[k] = prob_to_update
    end
    ensemble_problem_array # vector{ens.problem}
end

"""
Save a SolutionResult to file
"""
function save(solres, file_path)
    # convert the solres file into a dict for saving

    solres_dict = Dict()
    solres_keys = fieldnames(typeof(solres))[2:end]

    for k in solres_keys
        push!(solres_dict, k => getfield(solres, k))
    end

    bson(file_path, solres_dict)
end

"""
Load a SolutionResult from file
"""
function load(file_path)
    solres_dict = BSON.load(file_path)
    sol = SolutionResult(
        nothing,
        solres_dict[:fidelity],
        solres_dict[:problem_info],
        solres_dict[:optimised_pulses],
    )
end

"""
Write an array called pulse to a file, could also write a time array if that's necessary.
"""
function pulse_to_file(pulse, file_path)
    # can maybe check the shape and reshape so that its always time going down the file
    open(file_path, "w") do io
        writedlm(io, pulse')
    end
end

function pulse_to_file(pulse, file_path, duration)
    # can maybe check the shape and reshape so that its always time going down the file
    N = length(pulse)
    time_array = collect(range(0, duration, length = N))
    open(file_path, "w") do io
        writedlm(io, [time_array pulse'])
    end
end