text:
# precomputations.jl: contains all functions doing the precomputations
# About
#   Author       - Peter Opsomer (Peter.Opsomer@cs.kuleuven.be)
#   History      - Created October 2014, last edit October 2015

# Compute the contour integral of the given function using the trapezoidal rule around the interval.
# Input
#   f   - the integrand
#   rho - parameter for the Bernstein ellipse
#   M   - number of terms, use successive doubling from [F. Bornemann 2011] if NaN]
# Output
#   cont- the contour integral
function trap_rule(f,rho,M=NaN)
if isnan(M)
	M = 8; tol = 1e-15;
	function thi(k,l) 2im*pi*(1:l:k)/k; end
	function bernst(k,l) rho/2*exp(thi(k,l)) + 1/rho/2*exp(-thi(k,l)); end
	s = f(bernst(M,1)).*(rho*1im/2*exp(thi(M,1)) -1im/rho/2*exp(-thi(M,1)) )/1im; cont = mean(s); err1 = NaN;
	while M < 1e6
		M = 2*M;
		s = [s; f(bernst(M,2)).*(rho*1im/2*exp(thi(M,2)) -1im/rho/2*exp(-thi(M,2)) )/1im];
		val = mean(s); kappa = mean(abs(s))/abs(val);
		err0 = abs(val-cont)/abs(val);
		if (err0^3/err1^2 <= kappa*tol) || (~isfinite(kappa)) break; end
		cont = val; err1 = err0;
	end
elseif true
	ths = linspace(0,2*pi,M+1); # = theta_k, k = 0 to M-1
	ths = ths[1:(end-1)]
	hh = ths[2]-ths[1]; # To be sure to have the right distance
	zs = rho/2*exp(1im*ths) + 1/rho/2*exp(-1im*ths);
	cont = hh*sum(f(zs).*(rho*1im/2*exp(1im*ths) -1im/rho/2*exp(-1im*ths) ) )/(2*pi*1im );
else
	# Below an ellipse, but could also need other parametrisations, for example when evaluating psi(x) for x further from the interval than the point where log(h) is not analytic
    r1 = 1.8;
    r2 = 0.8;
    zs = r1*cos(ths) +r2*1im*sin(ths);
    cont = hh*sum(f(zs).*(-r1*sin(ths) +r2*1im*cos(ths) ) )/(2*pi*1im );
end
if rho < 1 # Bernstein in negative direction but need positive direction
    cont = -cont;
end
return cont
end

# Compute the contour integrals c_n, d_n, D_\infty, psi(z) [, its derivative and its contour integral].
# Input
#   alpha, beta - Parts of the weight w(x) = (1-x).^alpha.*(1+x).^beta.*h(x)
#   h           - Function for analytic function h(x)
#   nrT         - Number of terms, where 1 gives leading order term
#   [rho        - Berstein ellipse parameter: small to not encircle where log(h) is analytic and large to encircle the point x where to evaluate psi(x). Choose it to optimise the conditioning number as in [F. Bornemann, 2011] & [Wang & Huybrechs 2015], so it is a vector with rho(n+2) optimal for log(h)/sqrt(zeta-1)/sqrt(zeta+1)/(z-zeta)^(n+1) with rho+1/rho going to twice the distance of where log h is not analytic to the interval, or it is a scalar which is acceptable for all n.]
#   [M          - Number of points in trapezoidal rules, with M(n+2) optimal for c&d(n+1) and increasing with n, or M a scalar or NaN for successive doubling]
# Output
#   c, d        - The c- and d-coefficients, size(c&d) == ceil(Integer, nrT/2)
#   Dinf        - Limit of the Szego function
#   psi         - Anonymous function for the phase function
#   [dpsi       - Derivative of the phase function]
#   [contpsi	- Contour integral in psi, contpsi(x,p) has (zeta-x)^p]
function contour_integrals(alpha,beta,h,nrT,rho=false,M=NaN)
c = (1.0+1im)*zeros(ceil(Integer, nrT/2),1);
d = (1.0+1im)*zeros(ceil(Integer, nrT/2),1);
ra = 2*rand()-1; mm = 2; cc = 7;
# Checking for known exact solutions by comparing value at one random z = ra
if h(ra) == 1
    print("Using exact result for Legendre weight in contour_integrals",'\n')
    Dinf = 2.0^(-alpha/2-beta/2);
    function psi(x) 1/2.*( (alpha + beta).*acos(x+0.0im)-alpha*pi); end
    function dpsi(x) 1/2.*(alpha+beta).*(-1./(sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x) ) ); end
    function contpsi(x,p) 0; end
elseif abs(h(ra) - exp(-cc*ra.^(2*mm)) ) <100*eps()
    print("Using exact result for Gibbs weight in contour_integrals",'\n')
    # binomial is only implemented for integers in Julia version 0.3.1
    function binom(x,n) prod(x:-1:(x-n+1) )/prod(1:n); end
    for n = 0:(2*mm-1)
	for j = 0:floor( (2*mm-n-1)/2)
	    # Actually the second binomial only has integers but seems to be converted to floats
	    c[n+1] = c[n+1] -cc*binom(j-1/2,j)*binom(2*mm-1-2*j,2*mm-n-1-2*j);
	end
	d[n+1] = (-1)^(n+1)*c[n+1];
    end
    Dinf = 2.0^(-alpha/2-beta/2)*exp(-cc/2*binom(mm-1/2,mm));
    degs = (0:(length(c)-1) );
    function psi(x) 1/2.*( (alpha + beta).*acos(x+0.0im)-alpha*pi) +sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x)/2.*sum( c.*(x-1.0+0.0im).^(degs) ) ; end
    function dpsi(x) 1/2.*(alpha+beta).*(-1./(sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x) ) ) -x./sqrt(1.0+0.0im-x)./sqrt(1.0+0.0im+x)/2.*sum(c.*(x-1.0+0.0im).^degs) +sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x)/2.*sum(degs[2:end].*c[2:end].*(x-1.0+0.0im).^(degs[2:end]-1) ); end
   function contpsi(x,p) 2im*pi*sum( (degs[p:end]).^(p-1).*c[p:end].*(0im+x-1).^(degs[p:end]-p+1) ); end
else # Trapezoidal rules
    if abs(h(ra) - 1./sqrt(ra+3)) < 100*eps()
	print("Using trapezoidal rules for Fourier weight in contour_integrals", '\n')
	rho = 4; M = NaN; # M = 80; # But determine M automatically later
    elseif rho == false
	rho = 1.25; # > 1 so positive direction
    end
    lr = length(rho);
    if isnan(M)
	for nr = 1:(ceil(Integer, nrT/2)-1) # nr = n+1 in c_n = c(nr)
		c[nr] = trap_rule(z -> log(h(0.0im+z ) +0.0im)./( sqrt(z -1.0+0.0im).*sqrt(z +1.0+0.0im).*(z -1.0+0.0im).^nr ), rho[min(nr+1,lr)],M);
		d[nr] = trap_rule(z -> log(h(0.0im+z ) +0.0im)./( sqrt(z -1.0+0.0im).*sqrt(z +1.0+0.0im).*(z +1.0+0.0im).^nr ), rho[min(nr+1,lr)],M);
	end
	Dinf = real(2.0^( -(alpha+beta)./2).*exp(trap_rule(z -> log(h(0.0im+z)+0.0im )./sqrt(z -1.0+0.0im)./sqrt(z +1.0+0.0im), rho[1],M)/2) );

	function psi(x) 1/2.*( (alpha + beta).*acos(x+0.0im)-alpha*pi) + sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x).*trap_rule(z -> log(h(0.0im+z)+0.0im )./sqrt(z-1.0+0.0im)./sqrt(z+1.0+0.0im)./(0.0im+z-x), rho[min(2,lr)], M)/2; end

	function dpsi(x) 1/2.*(alpha + beta).*(-1)./(sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im +x)) -x/(sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x)*2)*trap_rule(z -> log(h(0.0im+z) +0.0im)./sqrt(z-1.0+0.0im)./sqrt(z+1.0+0.0im)./(0.0im+z-x),rho[min(2,lr)],M) + sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x)/2*trap_rule(z -> log(h(0.0im+z)+0.0im )./sqrt(z-1.0+0.0im)./sqrt(z+1.0+0.0im)./(z-x+0.0im).^2,rho[min(3,lr)],M); end
	function contpsi(x,p) 2im*pi*trap_rule(z -> log(h(0.0im+z) )./sqrt(0.0im+z-1)./sqrt(0.0im+z+1)./(0.0im+z-x).^p, rho[min(2,lr)], M); end

   else
	lm = length(M);
	for nr = 1:(ceil(Integer, nrT/2)-1) # nr = n+1 in c_n = c(nr)
		c[nr] = trap_rule(z -> log(h(0.0im+z ) +0.0im)./( sqrt(z -1.0+0.0im).*sqrt(z +1.0+0.0im).*(z -1.0+0.0im).^nr ), rho[min(nr+1,lr)],M[min(nr+1, lm)]);
		d[nr] = trap_rule(z -> log(h(0.0im+z ) +0.0im)./( sqrt(z -1.0+0.0im).*sqrt(z +1.0+0.0im).*(z +1.0+0.0im).^nr ), rho[min(nr+1,lr)],M[min(nr+1, lm)]);
	end
	Dinf = real(2.0^( -(alpha+beta)./2).*exp(trap_rule(z -> log(h(0.0im+z)+0.0im )./sqrt(z -1.0+0.0im)./sqrt(z +1.0+0.0im), rho[1],M[1])/2) );

	function psi(x) 1/2.*( (alpha + beta).*acos(x+0.0im)-alpha*pi) + sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x).*trap_rule(z -> log(h(0.0im+z)+0.0im )./sqrt(z-1.0+0.0im)./sqrt(z+1.0+0.0im)./(0.0im+z-x), rho[min(2, lm)], M[min(2, lm)])/2; end

	function dpsi(x) 1/2.*(alpha + beta).*(-1)./(sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im +x)) -x/(sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x)*2)*trap_rule(z -> log(h(0.0im+z) +0.0im)./sqrt(z-1.0+0.0im)./sqrt(z+1.0+0.0im)./(0.0im+z-x),rho[min(2,lr)],M[min(2, lm)]) + sqrt(1.0+0.0im-x).*sqrt(1.0+0.0im+x)/2*trap_rule(z -> log(h(0.0im+z)+0.0im )./sqrt(z-1.0+0.0im)./sqrt(z+1.0+0.0im)./(z-x+0.0im).^2,rho[min(3,lr)],M[min(3, lm)]); end
	function contpsi(x,p) 2im*pi*trap_rule(z -> log(h(0.0im+z) )./sqrt(0.0im+z-1)./sqrt(0.0im+z+1)./(0.0im+z-x).^p, rho[min(2,lr)], M[min(3, lm)]); end
   end
end
return (c, d, Dinf, psi, dpsi,contpsi)
end

# Compute the W- or V-matrices to construct the asymptotic expansion of R
# using the procedure with the convolutions as explained in the paper.
# Input
#   q            - alpha (when computing Wright) or beta (when computing Wleft)
#   t            - beta or alpha
#   Dinf         - Limit of the Szego function
#   cd           - c-or d- Coefficients needed for the U-matrices
#   maxOrder     - The maximum order of the error
#   r            - 1 when computing Wright, -1 when computing Wleft
#   isW          - true if the result are the Ws-, false if the V-s
#   [mos         - Maximum orders for each k, default ceil(Integer, (maxOrder-2)/2) for all k]
# Output
#   WVc          - Coefficient matrices for (z \pm 1)^m of s_k(z) or Delta_k(z)
function WV(q,t,Dinf,cd, maxOrder,r,isW, mos=ceil(Integer,  (maxOrder-2)/2)*ones(maxOrder-1) )
if maxOrder <= 1
	return (1+1im)*zeros(2,2,0,0);
end
mosDefault = norm(mos-ceil(Integer,  (maxOrder-2)/2)*ones(maxOrder-1)) == 0
(mo,~) = findmax(mos);
mo = convert(Int64,mo[1]);
ns = 0:mo;
# 'pochhammer' undefined in Julia
function poch(x,n) prod(x+(0:(n-1)) ); end
# binomial is only implemented for integers in version 0.3.1
function binom(x,n) prod(x:-1:(x-n+1) )/prod(1:n); end
# Will use gamma(1.0+n) instead of factorial(n) to avoid overflow
function brac(k) convert(FloatingPoint, (k == 0)) + convert(FloatingPoint, (k != 0))*prod(4*q^2 -(2*(1:k)-1).^2)/(2.0^(2*k)*(gamma(1.0+k) ) ); end

f = zeros(mo+1,1); # Coefficients in the expansion of \varphi(z)
for n = ns
    f[n+1] = poch(1/2,n)./(-r*2).^n./(1+2*n)./gamma(1.0+n);
end
# Or: f = poch(1/2,ns)./(-r*2).^ns./(1+2*ns)./gamma(1.0+ns);
g = zeros(max(mo*2+1,(maxOrder-1)),mo+1);
g[1,1] = 1; # Coefficients in the expansion of \varphi(z)^{-i}
for n = 1:mo
	g[1,n+1] = (-g[1,1:(n+1) ]*f[(n+1):-1:1])[1];
end

st = (1+1im)*zeros(max(mo*2+1,(maxOrder-1)),mo+1); # = \rho_{i,n,alpha+beta}
stp = (1+1im)*zeros(max(mo*2+1,(maxOrder-1)),mo+1); # = \rho_{i,n,alpha+beta+1}
stm = (1+1im)*zeros(max(mo*2+1,(maxOrder-1)),mo+1); # = \rho_{i,n,alpha+beta-1}
for n = ns
    tmp = 0;
    for j = 0:n
       tmp = tmp + binom(1/2,j)/(r*2)^j*cd[n-j+1];
    end
    st[1,n+1] = r*tmp + f[n+1]*(q+t);
    stp[1,n+1] = r*tmp + f[n+1]*(q+t+1);
    stm[1,n+1] = r*tmp + f[n+1]*(q+t-1);
end

for i = 2:max(mo*2+1,(maxOrder-1))
    for n = ns
        g[i,n+1] = sum(g[i-1,1:(n+1) ].*g[1,(n+1):-1:1] );
        st[i,n+1] = sum(st[i-1,1:(n+1) ].*st[1,(n+1):-1:1] );
        stp[i,n+1] = sum(stp[i-1,1:(n+1) ].*stp[1,(n+1):-1:1] );
        stm[i,n+1] = sum(stm[i-1,1:(n+1) ].*stm[1,(n+1):-1:1] );
    end
end

OmOdd = (1+1im)*zeros(mo+1,1); # = H_{n,alpha+beta}^{odd}
OmEven = (1+1im)*zeros(mo+1,1); # = H_{n,alpha+beta}^{even}
XiOdd = (1+1im)*zeros(mo+1,1); # = H_{n,alpha+beta+1}^{odd}
XiEven = (1+1im)*zeros(mo+1,1); # = H_{n,alpha+beta+1}^{even}
ThOdd = (1+1im)*zeros(mo+1,1); # = H_{n,alpha+beta-1}^{odd}
ThEven = (1+1im)*zeros(mo+1,1); # = H_{n,alpha+beta-1}^{even}
for n = ns
    if ( (mod(maxOrder,2)==0) || (n != mo) )
        OmEven[n+1] = OmEven[n+1] + st[1,n+1];
        XiEven[n+1] = XiEven[n+1] + stp[1,n+1];
        ThEven[n+1] = ThEven[n+1] + stm[1,n+1];
        for j = 1:n
            OmOdd[n+1] = OmOdd[n+1] + (r*2)^j*st[2*j,n-j+1]/gamma(1.0+2*j);
            XiOdd[n+1] = XiOdd[n+1] + (r*2)^j*stp[2*j,n-j+1]/gamma(1.0+2*j);
            ThOdd[n+1] = ThOdd[n+1] + (r*2)^j*stm[2*j,n-j+1]/gamma(1.0+2*j);
            OmEven[n+1] = OmEven[n+1] + (r*2)^j*st[2*j+1,n-j+1]/gamma(2.0+2*j);
            XiEven[n+1] = XiEven[n+1] + (r*2)^j*stp[2*j+1,n-j+1]/gamma(2.0+2*j);
            ThEven[n+1] = ThEven[n+1] + (r*2)^j*stm[2*j+1,n-j+1]/gamma(2.0+2*j);
        end
	else # This splitting is needed to avoid wrong OmEven[n+1] etc when n=mo and maxOrder is odd
        for j = 1:n
            OmOdd[n+1] = OmOdd[n+1] + (r*2)^j*st[2*j,n-j+1]/gamma(1.0+2*j);
            XiOdd[n+1] = XiOdd[n+1] + (r*2)^j*stp[2*j,n-j+1]/gamma(1.0+2*j);
            ThOdd[n+1] = ThOdd[n+1] + (r*2)^j*stm[2*j,n-j+1]/gamma(1.0+2*j);
        end
    end
end
OmOdd[1] = 1;
XiOdd[1] = 1; # Overwrite because else wrong
ThOdd[1] = 1;

OmO = (1+1im)*zeros(mo+1,1); # = X_{n,alpha+beta}^{odd}
OmE = (1+1im)*zeros(mo+1,1); # = X_{n,alpha+beta}^{even}
XiO = (1+1im)*zeros(mo+1,1); # = X_{n,alpha+beta+1}^{odd}
XiE = (1+1im)*zeros(mo+1,1); # = X_{n,alpha+beta+1}^{even}
ThO = (1+1im)*zeros(mo+1,1); # = X_{n,alpha+beta-1}^{odd}
ThE = (1+1im)*zeros(mo+1,1); # = X_{n,alpha+beta-1}^{even}
for n = ns
    for j = 0:n
        OmO[n+1] = OmO[n+1] + binom(-1/2,j)/(r*2)^j*OmOdd[n-j+1]/sqrt(r*2);
        XiO[n+1] = XiO[n+1] + binom(-1/2,j)/(r*2)^j*XiOdd[n-j+1]/sqrt(r*2);
        ThO[n+1] = ThO[n+1] + binom(-1/2,j)/(r*2)^j*ThOdd[n-j+1]/sqrt(r*2);
        if ( (mod(maxOrder,2)==0) || (n != mo) )
            OmE[n+1] = OmE[n+1] + binom(-1/2,j)/(r*2)^j*OmEven[n-j+1];
            XiE[n+1] = XiE[n+1] + binom(-1/2,j)/(r*2)^j*XiEven[n-j+1];
            ThE[n+1] = ThE[n+1] + binom(-1/2,j)/(r*2)^j*ThEven[n-j+1];
        end
    end
end

Ts = (1+1im)*zeros(2,2,mo+1);  # = T_{k,n}^{odd} or T_{k,n}^{even} depending on k, overwritten each new k
WVc = (1+1im)*zeros(2,2,maxOrder-1,mo+1);
for k = 1:(maxOrder-1)
    a = (q^2 + k/2-1/4)/k;
    b = -1im*r*(k-1/2);
    if (mosDefault)
        mo = ceil(Integer,  (maxOrder-1-k)/2)-1+ceil(Integer, k/2);
    else
        mo = mos[k];
    end
    Ts[:,:,:] = 0;
    if mod(k,2) == 1
# Actually, for m = (-ceil(Integer, k/2)):(ceil(Integer, (maxOrder-1-k)/2)-1) but it seems to be safer not to use negative indices so shift by ceil(Integer, k/2)+1 and easier to let n go to "mo" always iso ceil(Integer, (maxOrder-1-k)/2)+ceil(Integer, k/2)
        for n = 0:mo
	    # Or a*binom(-1/2,n)*(0im+r)^(n+1/2)/2^(n+1/2)*(2*n+1)/(2*n-1), or a*(-r*binom(-1/2,n)/(0im+r*2)^(n+1/2) -binom(-1/2,n-1)/(0im+r*2)^(n-1/2))
	    Ts[:,:,n+1] = [ (-a*binom(n-3/2,n)*(-r*2.0)^(-n)*(2*n+1.0)*sqrt(0im+r/2.0) +1im*b*OmO[n+1] ) 	Dinf^2*(1im*a*binom(-1/2,n)*(r*2.0)^(-n)/sqrt(0im+r*2) +r*b*XiO[n+1] )	;	(1im*a*binom(-1/2,n)*(r*2.0)^(-n)/sqrt(0im+r*2.0) +r*b*ThO[n+1] )/Dinf^2	 (a*binom(n-3/2,n)*(-r*2.0)^(-n)*(2*n+1.0)*sqrt(0im+r/2.0) -1im*b*OmO[n+1] )];
            WVc[:,:,k,n+1] = brac(k-1)/(r*2.0)^(3*k/2.0)*sum(repeat(reshape(g[k,1:(n+1) ],1,1,convert(Int64,n+1)),outer=[2,2,1]).*Ts[:,:,(n+1):-1:1],3);
        end
    else
        for n = 0:mo
	    Ts[:,:,n+1] = [(a*convert(FloatingPoint, (n == 0)) +1im*r*b*OmE[n+1])     Dinf^2*b*XiE[n+1] 	; 	b*ThE[n+1]/Dinf^2	(a*convert(FloatingPoint, (n == 0)) -1im*b*r*OmE[n+1] )];
            WVc[:,:,k,n+1] = brac(k-1)/(r*2.0)^(3*k/2.0)*sum(repeat(reshape(g[k,1:(n+1) ],1,1,convert(Int64,n+1)),outer=[2,2,1]).*Ts[:,:,(n+1):-1:1],3);
            if isW # The extra matrix of the theorem:
                WVc[:,:,k,n+1] = WVc[:,:,k,n+1] + brac(k-1)/(r*2.0)^(3*k/2.0)*(-(4*q^2+2*k-1)*g[k,n+1]/2/k*eye(2) );
            end
        end
    end
end
return WVc
end

# Get the U-matrices to construct the asymptotic expansion of R using the procedure with the convolutions as explained in the paper. Optionally, specify method or number of matrices to compute or get the Q-s.
# Input
#   alpha, beta  - Parts of the weight w(x) = (1-x).^alpha.*(1+x).^beta.*h(x)
#   Dinf         - Limit of the Szego function
#   c,d          - Coefficients needed for the U-matrices
#   maxOrder     - The maximum order of the error
#   [method       - 'UW' (default) to get the U-s through the W-s, 'UQW' to get the U- & Q-s through the W-s and 'UQV' to get the U- & Q-s through the V-s]
#   [ns          - If method contains 'Q', specify how many Q-s to compute so compute Q_{k,0:(ns(k)-1)}^{right/left}, also for W/V]
# Output
#   Uright       - Coefficient matrices of R_k(z) for (z-1)^(-m)
#   Uleft        - Coefficient matrices of R_k(z) for (z+1)^(-m)
#   [Qright      - Coefficient matrices of R_k^{right}(z) for (z-1)^n]
#   [Qleft       - Coefficient matrices of R_k^{left}(z) for (z+1)^n]
function UQ(alpha,beta,Dinf,c,d,maxOrder,method = "UW",ns= ceil(Integer,  (maxOrder-1-(1:(maxOrder-1)))/2)-1+ceil(Integer, (1:(maxOrder-1))/2) )

tmp = ceil(Integer,  (maxOrder-1-(1:(maxOrder-1)))/2)-1+ceil(Integer, (1:(maxOrder-1))/2)
existed = norm(ns[:]-tmp) != 0;
if ~existed
    if ismatch(r"Q",method)
	ns = convert(Array{Int64,1},maximum([ (maxOrder-2-(1:(maxOrder-1))+1+ceil(Integer, (1:(maxOrder-1))/2) )     zeros(1,maxOrder-1)[:]],2)[:] );
    else
        # Could use the following to specify the necessary (exact) number of W- or V-matrices to compute, but not really needed...
#         ns = ceil(Integer,  (maxOrder-1-(1:(maxOrder-1)))/2)-1+ceil(Integer, (1:(maxOrder-1))/2);
    end
end

Uright = (1+1im)*zeros(2,2,maxOrder-1,ceil(Integer,  (maxOrder-1)/2) );
Uleft = (1+1im)*zeros(2,2,maxOrder-1,ceil(Integer,  (maxOrder-1)/2));
# About half of these tensors will not be used
if ismatch(r"Q",method)
    tmp = ceil(Integer,  (maxOrder-1-(1:(maxOrder-1)))/2)-1+ceil(Integer, (1:(maxOrder-1))/2)
    if norm(ns[:]-tmp) != 0
	(mn,ignore) = findmax(ns);
	mn = convert(Int64,mn);
        Qright = (1+1im)*zeros(2,2,maxOrder-1,mn+1 );
        Qleft = (1+1im)*zeros(2,2,maxOrder-1,mn+1 );
    else
        Qright = (1+1im)*zeros(2,2,maxOrder-1,convert(Int64,ceil( (maxOrder-1)/2) ));
        Qleft = (1+1im)*zeros(2,2,maxOrder-1,convert(Int64,ceil( (maxOrder-1)/2) ));
    end
end

function poch(x,n) prod(x+(0:(n-1) ) ); end
# 'pochhammer' undefined in Julia

if ismatch(r"W",method)
    tmp = ceil(Integer,  (maxOrder-1-(1:(maxOrder-1)))/2)-1+ceil(Integer, (1:(maxOrder-1))/2)
    if norm(ns[:]-tmp) != 0
	Wr = WV(alpha,beta,Dinf,c,maxOrder,1,true,ns[:]+2+ceil(Integer,  (1:(maxOrder-1)) /2) );
	# Wr = W_{k,m}^{right}, k=1..maxOrder-1, m given by ns; Wl analogous.
	Wl = WV(beta,alpha,Dinf,d,maxOrder,(1im)^2,true,ns[:]+2+ceil(Integer,  (1:(maxOrder-1)) /2));
    else
	Wr = WV(alpha,beta,Dinf,c,maxOrder,1,true);
	Wl = WV(beta,alpha,Dinf,d,maxOrder,(1im)^2,true);
    end
    for k = 1:(maxOrder-1)
        # Actually, for m = (-ceil(Integer, k/2)):1, but better no negative indices so shift by ceil(Integer, k/2)+1
        for m = 1:ceil(Integer, k/2)
            Uright[:,:,k,m] = Wr[:,:,k,ceil(Integer, k/2)+1-m];
            Uleft[:,:,k,m] = Wl[:,:,k,ceil(Integer, k/2)+1-m];
            for j=1:(k-1)
                for l = max(m-ceil(Integer, j/2),1):ceil(Integer, (k-j)/2)
                    Uright[:,:,k,m] = Uright[:,:,k,m] + Uright[:,:,k-j,l]*Wr[:,:,j,ceil(Integer, j/2)+1+l-m];
                    Uleft[:,:,k,m] = Uleft[:,:,k,m] + Uleft[:,:,k-j,l]*Wl[:,:,j,ceil(Integer, j/2)+1+l-m];
                end
            end
            for j=1:(k-1)
                for n = 0:(ceil(Integer, j/2)-m)
                    for i = 1:ceil(Integer, (k-j)/2)
                        Uright[:,:,k,m] = Uright[:,:,k,m] + poch(1-i-n,n)/2.0^(i+n)/gamma(1.0+n)*Uleft[:,:,k-j,i]*Wr[:,:,j,ceil(Integer, j/2)+1-n-m];
                        Uleft[:,:,k,m] = Uleft[:,:,k,m] + poch(1-i-n,n)/(-2.0)^(i+n)/gamma(1.0+n)*Uright[:,:,k-j,i]*Wl[:,:,j,ceil(Integer, j/2)+1-n-m];
                    end
                end
            end
        end
    end
    if method == "UW"
	return (Uright,Uleft);
    end
elseif ismatch(r"V",method)
    tmp = ceil(Integer,  (maxOrder-1-(1:(maxOrder-1)))/2)-1+ceil(Integer, (1:(maxOrder-1))/2)
    if norm(ns[:]-tmp) != 0
		Vr = WV(alpha,beta,Dinf,c,maxOrder,1,false,ns[:]+2+ceil(Integer,  (1:(maxOrder-1)) /2) );
		Vl = WV(beta,alpha,Dinf,d,maxOrder,(1im)^2,false,ns[:]+2+ceil(Integer,  (1:(maxOrder-1)) /2) );
	else
		Vr = WV(alpha,beta,Dinf,c,maxOrder,1,false);
		Vl = WV(beta,alpha,Dinf,d,maxOrder,(1im)^2,false);
	end
else
    error(string("Wrong method format, got ", method));
end

if "UQW" == method
    # Already did the U's, calculate the Q-s from them:
    for k = 1:(maxOrder-1)
        for n = 0:ns[k]
            Qright[:,:,k,n+1] = -Wr[:,:,k,ceil(Integer, k/2)+1+n];
            Qleft[:,:,k,n+1] = -Wl[:,:,k,ceil(Integer, k/2)+1+n];
            for i = 1:ceil(Integer, k/2)
                Qright[:,:,k,n+1] = Qright[:,:,k,n+1] + poch(1-i-n,n)/2.0^(i+n)/gamma(1.0+convert(Int64,n) )*Uleft[:,:,k,i];
                Qleft[:,:,k,n+1] = Qleft[:,:,k,n+1] + poch(1-i-n,n)/(-2.0)^(i+n)/gamma(1.0+convert(Int64,n) )*Uright[:,:,k,i];
            end
            for j = 1:(k-1)
                for i = -ceil(Integer, j/2):n
                    for l = 1:ceil(Integer, (k-j)/2)
                        Qright[:,:,k,n+1] = Qright[:,:,k,n+1] - poch(1-l-n+i,n-i)/2.0^(n-i+l)/gamma(1.0+convert(Int64,n-i) )*Uleft[:,:,k-j,l]*Wr[:,:,j,i+1+ceil(Integer, j/2) ];
                        Qleft[:,:,k,n+1] = Qleft[:,:,k,n+1] - poch(1-l-n+i,n-i)/(-2.0)^(n-i+l)/gamma(1.0+convert(Int64,n-i) )*Uright[:,:,k-j,l]*Wl[:,:,j,i+1+ceil(Integer, j/2) ];
                    end
                end
                for i = (n+1):(n+ceil(Integer, (k-j)/2) )
                    Qright[:,:,k,n+1] = Qright[:,:,k,n+1] -Uright[:,:,k-j,i-n]*Wr[:,:,j,i+1+ceil(Integer, j/2) ];
                    Qleft[:,:,k,n+1] = Qleft[:,:,k,n+1] -Uleft[:,:,k-j,i-n]*Wl[:,:,j,i+1+ceil(Integer, j/2) ];
                end
            end
        end
    end
elseif "UQV"== method
    for k = 1:(maxOrder-1)
        for m = 1:ceil(Integer, k/2)
            Uright[:,:,k,m] = Vr[:,:,k,ceil(Integer, k/2)+1-m];
            Uleft[:,:,k,m] = Vl[:,:,k,ceil(Integer, k/2)+1-m];
            for j=1:(k-1)
                for l = 0:(ceil(Integer, j/2)-m)
                    Uright[:,:,k,m] = Uright[:,:,k,m] + Qright[:,:,k-j,l+1]*Vr[:,:,j,ceil(Integer, j/2)+1-l-m];
                    Uleft[:,:,k,m] = Uleft[:,:,k,m] + Qleft[:,:,k-j,l+1]*Vl[:,:,j,ceil(Integer, j/2)+1-l-m];
                end
            end
        end
        for n = 0:ns[k]
            Qright[:,:,k,n+1] = -Vr[:,:,k,ceil(Integer, k/2)+1+n];
            Qleft[:,:,k,n+1] = -Vl[:,:,k,ceil(Integer, k/2)+1+n];
            for i = 1:ceil(Integer, k/2)
                Qright[:,:,k,n+1] = Qright[:,:,k,n+1] + poch(1-i-n,n)/2.0^(i+n)/gamma(1.0+convert(Int64,n) )*Uleft[:,:,k,i];
                Qleft[:,:,k,n+1] = Qleft[:,:,k,n+1] + poch(1-i-n,n)/(-2.0)^(i+n)/gamma(1.0+convert(Int64,n) )*Uright[:,:,k,i];
            end
            for j = 1:(k-1)
                for l = 0:(ceil(Integer, j/2)+n)
                    Qright[:,:,k,n+1] = Qright[:,:,k,n+1] -Qright[:,:,k-j,l+1]*Vr[:,:,j,n-l+1+ceil(Integer, j/2) ];
                    Qleft[:,:,k,n+1] = Qleft[:,:,k,n+1] -Qleft[:,:,k-j,l+1]*Vl[:,:,j,n-l+1+ceil(Integer, j/2) ];
                end
            end
        end
    end
else
    error(string("Invalid method, got ", method));
end
return (Uright,Uleft,Qright,Qleft)
end


text:
println("testing periodic notion of distance")
X1 = [ JVecF(0.0,0.0,0.0), JVecF(0.5,0.5,0.5)]
at = Atoms(:H, X1)
set_pbc!(at, (true, false, false))
set_cell!(at, [1.0 0.2 0.3; 0.0 1.0 0.1; 0.0 0.0 1.0])
X2 = [ X1[1], X1[2] + JVecF(1.0, 0.2, 0.3)]
X3 = [ X1[1] - JVecF(0.1, 0.0, 0.0), X1[2]]
X4 = [ X1[1], X1[2] + JVecF(0.0, 1.0, 0.1)]

@testset "dist" begin
   @test JuLIP.dist(at, X1, X2) < 1e-14
   @test JuLIP.dist(at, X1, X3) ≈ 0.1
   @test JuLIP.dist(at, X1, X4) ≈ norm(JVecF(0.0, 1.0, 0.1))
end

ext: jl


text:
using Autotrack
using Test

@testset "Autotrack.jl" begin
    # Write your tests here.
end


text:
#--------------------------------------------------------------------------------------------------
using Test

#--------------------------------------------------------------------------------------------------

@testset "Hensel" begin
    @testset "Linear mapping tests" begin
        include("linear_mapping_tests.jl")
    end
    @testset "Functions on [a,b] tests" begin
        include("functions_on_ab_tests.jl")
    end
    @testset "Functions on [0,1] tests" begin
        include("functions_on_01_tests.jl")
    end
    @testset "LinMap tests" begin
        include("LinMap_tests.jl")
    end
    @testset "FunEnv tests" begin
        include("FunEnv_tests.jl")
    end
    @testset "Additional functions tests" begin
        include("additional_functions_tests.jl")
    end
end

#--------------------------------------------------------------------------------------------------

max_stars_repo_name: ayarodionov/Hensel

max_issues_repo_name: ayarodionov/Hensel

max_forks_repo_name: ayarodionov/Hensel


text:
using LinearAlgebra, ForwardDiff
using LinearOperators, Krylov

function newton_step(B,gx)
    pn, ks = Krylov.cg(B,-gx)
    if ks.solved == false
        @error ks
    end
    return pn
end

function dogleg(gx,B,Δ)
    pn = newton_step(B,gx)
    if norm(pn) <= Δ
        return pn
    end
    p = -(norm(gx)^2/(gx'*(B*gx)))*gx # Cauchy Step
    if norm(p) >= Δ
        return (p/norm(p))*Δ
    end
    proj = boundary_l2(p,pn-p,Δ)
    return proj
end

# Dogleg step calculation using l_\infty ball
function dogleg_box(x,gx,B,Δ)
    lb,ub = get_bounds(x,Δ)
    #@info lb,ub
    pn = newton_step(B,gx)
    if in_bounds(lb,Δ,pn)
        return pn
    end
    p = -(norm(gx)^2/(gx'*(B*gx)))*gx # Cauchy Step
    if in_bounds(lb,Δ,p) == false
        #@info "Scaled"
        t = step_to_bound(x,p/norm(p),lb,Δ)
        return (p/norm(p))*t
    end
    #@info "Dogleg"
    t = step_to_bound(x,pn-p,lb,Δ)
    return p + t * (pn-p)
end

# Distance to l_\infty bound
function step_to_bound(x,p,lb,ub)
    dist = max.(lb ./ p, ub ./ p)
    return minimum(dist)
end

# Test if vector is within the l_\infty ball
function in_bounds(lb,ub,x)
    return all(x .>= lb) && all(x .<= ub)
end

function get_bounds(x,Δ)
    lb = max.(-Δ,eps() .-x)
    ub = Δ * ones(size(x))
    return lb,ub
end

function boundary_l2(p,q,Δ)
    a = norm(q)^2
    b = 2*p'*q
    c = norm(p)^2 - Δ^2
    α = (-b + sqrt(b^2 - a*c))/(2*a)
    return p + α * q
end

function tr(f,∇f,x₀;maxiter=1000,tol=1e-6,Δ₀=1.0,freq=100)
    x = x₀
    fx = f(x₀)
    gx = ∇f(x₀)
    Δ = Δ₀
    #B = LinearOperators.LSR1Operator(length(x))
    B = LinearOperators.LBFGSOperator(length(x))
    iter = 1
    for i = 1:maxiter

        p = dogleg_box(x,gx,B,Δ) # Step calculation

        # Step update
        x̄ = x + p
		fx̄ = f(x̄)
		gx̄ = ∇f(x̄)

        # Quality evaluation
        pred = -gx'*p-0.5*p'*(B*p)
        ared = fx-fx̄
        ρ = ared/pred

        # Trust region modification
        if ρ > 0.75
            Δ = min(1e10,1.5*Δ)
        elseif ρ < 0.25
            Δ *= 0.25
        end

        # SR1 update
        y = gx̄-gx
        y_Bs = y - B*p
        if abs(p'*y_Bs) >= 1e-8 * norm(p) * norm(y_Bs)
            push!(B,p,gx̄-gx)
        end

        if ρ > 0.25
            x = x̄
            fx=fx̄
            gx = gx̄
        end

        if norm(gx) <= tol
            break
        end
        if iter%freq == 0
            @info "$i: \tx=$x,\tfx=$fx,\tgx=$(norm(gx)),\tΔ=$Δ"
        end
        iter += 1
    end
    return x,fx,norm(gx),iter
end

f(x) = (1-x[1])^2 + 100*(x[2]-x[1]^2)^2
∇f(x) = ForwardDiff.gradient(f,x)
x₀ = [10.1;10.2]
@time opt,fx,gx,iter = tr(f,∇f,x₀;maxiter=1000,freq=1)

max_stars_repo_name: dvillacis/BPLDenoising

max_issues_repo_name: dvillacis/BPLDenoising

max_forks_repo_name: dvillacis/BPLDenoising


text:
module TestMolecules

using Test
using LinearAlgebra
using Molecules
using StaticArrays: SVector
using Molecules: Bond, Carbon, Hydrogen, Lead, Molecule, Orbital,
    angularmomentum, chemicalpotential, countorbitals, hamiltonian, onsiteenergy, selfenergy

##################################################
# Setup
##################################################

atom1 = Hydrogen(.0, .0, .0, orbital"1s")
atom2 = Carbon(1.0, .0, .0, orbital"2s", orbital"2p_x", orbital"2p_y", orbital"2p_z")
atom3 = Hydrogen(1.0, .0, 1.0, orbital"1s")
atom4 = Carbon(1.0, .0, 1.0, orbital"2s", orbital"2p_x", orbital"2p_y", orbital"2p_z")

bond12 = Bond(atom1, atom2)
bond23 = Bond(atom2, atom3)
mol1 = Molecule([atom1, atom2, atom3], [bond12, bond23])

h1 = hamiltonian(atom1)
h2 = hamiltonian(atom2)
h3 = hamiltonian(atom3)
t12 = hamiltonian(atom1, atom2)
t23 = hamiltonian(atom2, atom3)
t13 = hamiltonian(atom1, atom3)

H = [h1 t12 0t13
     t12' h2 t23
     0t13' t23' h3]

Lx = [0.0  0.0  0.0  0.0
      0.0  0.0  0.0  0.0
      0.0  0.0  0.0 -im
      0.0  0.0  im  0.0]

Ly = [0.0  0.0  0.0  0.0
      0.0  0.0  0.0  im
      0.0  0.0  0.0  0.0
      0.0 -im  0.0  0.0]

Lz = [0.0  0.0  0.0  0.0
      0.0  0.0 -im  0.0
      0.0  im  0.0  0.0
      0.0  0.0  0.0  0.0]

vec = SVector(randn(3)...)

z12 = fill(0.0im, countorbitals(atom1), countorbitals(atom2))
z13 = fill(0.0im, countorbitals(atom1), countorbitals(atom3))
z23 = fill(0.0im, countorbitals(atom2), countorbitals(atom3))

atom6 = Hydrogen(-1.0, 0.0, 0.0, orbital"1s")
mol2 = Molecule([atom6], Bond[])
lead1 = Lead(mol2, 100.0)

μ1, Δ1 = chemicalpotential(mol1)
Σ1 = selfenergy(μ1 + Δ1, lead1, mol1, Set([Bond(atom1, atom6)]))
H1 = hamiltonian(mol1)

##################################################
# Tests
##################################################

@testset "Orbitals" begin
    @test orbital"1s" isa Orbital{1,:s,Nothing}
    @test orbital"1s" isa Orbital"*s_{*}"
    @test orbital"2p_x" ≡ orbital"2p_{x}"
    @test_throws LoadError eval(@macroexpand orbital"p_x")
    @test_throws LoadError eval(@macroexpand orbital"2p_x^2")
    @test Orbital"1s" ≡ Orbital{1,:s,Nothing}
    @test Orbital"2p_y" ≡ Orbital{2,:p,:y}
end

@testset "Atoms" begin
    @test countorbitals(atom1) == 1
    @test countorbitals(atom2) == 4
end

@testset "Molecules" begin
    @test countorbitals(mol1) == 6
end

# @testset "Energies" begin
#     # Onsite energies are negative (electrons are bound)

#     @test all(x -> x ≤ 0.0, hamiltonian(atom1))

#     t12 = hamiltonian(atom2, atom4)

#     # Check that energies match the expected bonding/anti--bonding
#     # behaviour

#     # s - *
#     @test t12[1,1] < 0.0
#     @test t12[1,2] ≈ 0.0
#     @test t12[1,3] ≈ 0.0
#     # s - p_z
#     @test t12[1,4] > 0.0
#     @test t12[4,1] < 0.0

#     # p_x - *
#     @test t12[2,1] ≈ 0.0
#     @test t12[2,2] < 0.0
#     @test t12[2,3] ≈ 0.0
#     @test t12[2,4] ≈ 0.0
# end

@testset "Angular momentum" begin
    @test angularmomentum(:z, orbital"1p_x", orbital"1p_y") == -im
    @test angularmomentum(:z, orbital"1p_y", orbital"1p_x") == im
    @test angularmomentum(:x, atom1) == fill(0.0, 1, 1)
    @test angularmomentum(:x, atom2) == Lx
    @test angularmomentum(:y, atom2) == Ly
    @test angularmomentum(:z, atom2) == Lz
    @test angularmomentum(vec, atom2) == vec[1]*Lx + vec[2]*Ly + vec[3]*Lz
    @test angularmomentum(vec, mol1) == [angularmomentum(vec, atom1) z12 z13
                                        z12' angularmomentum(vec, atom2) z23
                                        z13' z23' angularmomentum(vec, atom3)]
end

# @testset "Leads" begin
#     @test size(H1) == size(Σ1)
#     @test Σ1[1,1] ≠ 0.0eV
#     Σ1′ = copy(Σ1)
#     Σ1′[1, 1] = 0.0eV
#     @test all(x -> x == 0.0eV, Σ1′)
# end

end #module


text:
module HessianTest

import Calculus

using Base.Test
using ForwardDiff

include(joinpath(dirname(@__FILE__), "utils.jl"))

#############################
# rosenbrock hardcoded test #
#############################

x = [0.1, 0.2, 0.3]
v = rosenbrock(x)
g = [-9.4, 15.6, 52.0]
h = [-66.0  -40.0    0.0;
     -40.0  130.0  -80.0;
       0.0  -80.0  200.0]

for c in (Chunk{1}(), Chunk{2}(), Chunk{3}())

    # single-threaded #
    #-----------------#
    @test_approx_eq h ForwardDiff.hessian(rosenbrock, x, c)

    out = similar(x, 3, 3)
    ForwardDiff.hessian!(out, rosenbrock, x, c)
    @test_approx_eq out h

    out = HessianResult(x)
    ForwardDiff.hessian!(out, rosenbrock, x, c)
    @test_approx_eq ForwardDiff.value(out) v
    @test_approx_eq ForwardDiff.gradient(out) g
    @test_approx_eq ForwardDiff.hessian(out) h

    # multithreaded #
    #---------------#
    if ForwardDiff.IS_MULTITHREADED_JULIA
        @test_approx_eq h ForwardDiff.hessian(rosenbrock, x, c; multithread = true)

        out = similar(x, 3, 3)
        ForwardDiff.hessian!(out, rosenbrock, x, c; multithread = true)
        @test_approx_eq out h

        out = HessianResult(x)
        ForwardDiff.hessian!(out, rosenbrock, x, c; multithread = true)
        @test_approx_eq ForwardDiff.value(out) v
        @test_approx_eq ForwardDiff.gradient(out) g
        @test_approx_eq ForwardDiff.hessian(out) h
    end
end

########################
# test vs. Calculus.jl #
########################

for f in VECTOR_TO_NUMBER_FUNCS
    v = f(X)
    g = ForwardDiff.gradient(f, X)
    h = ForwardDiff.hessian(f, X)
    # finite difference approximation error is really bad for Hessians...
    @test_approx_eq_eps h Calculus.hessian(f, X) 0.01
    for c in CHUNK_SIZES
        for usecache in (true, false)
            println("  ...testing $f with (chunk size = $c) and (usecache = $usecache)")
            chunk = Chunk{c}()

            # single-threaded #
            #-----------------#
            out = ForwardDiff.hessian(f, X, chunk; usecache = usecache)
            @test_approx_eq out h

            out = similar(X, length(X), length(X))
            ForwardDiff.hessian!(out, f, X, chunk; usecache = usecache)
            @test_approx_eq out h

            out = HessianResult(X)
            ForwardDiff.hessian!(out, f, X, chunk; usecache = usecache)
            @test_approx_eq ForwardDiff.value(out) v
            @test_approx_eq ForwardDiff.gradient(out) g
            @test_approx_eq ForwardDiff.hessian(out) h

            # multithreaded #
            #---------------#
            if ForwardDiff.IS_MULTITHREADED_JULIA
                out = ForwardDiff.hessian(f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq out h

                out = similar(X, length(X), length(X))
                ForwardDiff.hessian!(out, f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq out h

                out = HessianResult(X)
                ForwardDiff.hessian!(out, f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq ForwardDiff.value(out) v
                @test_approx_eq ForwardDiff.gradient(out) g
                @test_approx_eq ForwardDiff.hessian(out) h
            end
        end
    end
end

end # module


text:
@testset "test Bool" begin
    t = TRUE()
    f = FALSE()
    @test repr(t) == "true"
    @test repr(f) == "false"
end


text:
# COMPAT

if VERSION < v"1.1"
    isnothing(o) = (o === nothing)
    export isnothing
    eachcol(A::AbstractArray) = (view(A, :, j) for j ∈ axes(A, 2))
end

if VERSION <= v"1.1"
    import Base:startswith
    Base.startswith(s::AbstractString, r::Regex) = (rr = Regex("^"*r.pattern); occursin(rr, s))
end

#######################################

abstract type Backend end

struct GLE <: Backend
    io::IOBuffer
end
GLE() = GLE(IOBuffer())

struct Gnuplot <: Backend
    io::IOBuffer
end
Gnuplot() = Gnuplot(IOBuffer())

# write to buffer (or buffer encapsulating object) with form (s |> b)
|>(s::String, b::Backend)   = write(b.io, s, " ")
|>(s::String, tio::Tuple)   = @. |>(s, tio)

# read buffer encapsulated by `b`
take!(b::Backend) = take!(b.io)

|>(bin::Backend, bout::Backend) = write(bout.io, take!(bin))

#######################################

isdef(el) = (el !== nothing)

# check if object `obj` has at least one field that is not "Nothing"
# this is useful when dealing with objects with lots of "Optional" fields
isanydef(obj) = any(isdef, (getfield(obj, f) for f ∈ fieldnames(typeof(obj))))

# see cla! (clear axes)
function reset!(obj::T; exclude=Vector{Symbol}(), inits...) where T
    # create a new object of the same type, assumes there is
    # a constructor that accepts empty input
    fresh = T(; inits...)
    for fn ∈ fieldnames(T)
        fn ∈ exclude && continue
        # set all fields to the field value given by the default
        # constructor (but keeps the address of the parent object)
        setfield!(obj, fn, deepcopy(getfield(fresh, fn)))
    end
    return obj
end

#######################################

rmin(x,::Nothing)  = x
rmin(::Nothing, y) = y
rmin(x,y) = min(x,y)

rmax(x,::Nothing)  = x
rmax(::Nothing, y) = y
rmax(x,y) = max(x,y)

#######################################

fl(::Missing) = missing
fl(::Nothing) = nothing
fl(x::Real)   = Float64(x)
fl(t::Tuple)  = fl.(t)
fl(v::AbstractVecOrMat) = fl.(v)

#######################################

# return a number with 3 digits accuracy, useful in col2str
round3d(x::Real) = round(x, digits=3)

# takes a colorant and transform it to a standard string rgba(...)
function col2str(col::Colorant; str=false)
    crgba = convert(RGBA, col)
    r, g, b, α = round3d.([crgba.r, crgba.g, crgba.b, crgba.alpha])
    s  = "rgba($r,$g,$b,$α)"
    # used by str(markerstyle), remove things that confuse gle
    sr = replace(s, r"[\(\),\.]"=>"_")
    return ifelse(str, sr, s)
end

# unroll a vector into a string with the elements separated by a space
vec2str(v::Vector{T}) where T<:Real = prod("$vi " for vi ∈ v)
vec2str(v::Vector{String}) = prod("\"$vi\" " for vi ∈ v)
# separated vec to str for things like d1,d2,d3 etc
function svec2str(v::Vector{String}, sep=",")
    length(v) > 1 || return v[1]
    return prod(vi*sep for vi∈v[1:end-1])*v[end]
end
svec2str(v::Base.Generator, sep=",") = svec2str(collect(v), sep)

#######################################

nvec(n::Int, T) = [T() for _ ∈ 1:n]

#######################################

function csv_writer(path::String, z, hasmissing::Bool)
    dirpath = splitdir(path)[1]
    isdir(dirpath) || mkpath(dirpath)
    if hasmissing
        tempio = IOBuffer()
        writedlm(tempio, z)
        # NOTE assumes it's fine to materialize the buffer with huge arrays
        # it's probably not ideal but in general should be fine, huge arrays are more
        # likely to happen with 3D objects (mesh) which are somewhat less likely to have missings
        temps = String(take!(tempio))
        write(path, replace(temps, r"missing|NaN|Inf"=>"?"))
    else
        writedlm(path, z)
    end
    return nothing
end

#######################################

"""
    tex_str(s)

Allows use of \$ and \\ while also allowing interpolation via `##name`. This is useful
for labels that contain tex-like expressions that may be generated in a sequence.

### Example

```julia
 i=5
t"\$x+##i\$" # amounts to "\$x+5\$"
```
"""
macro tex_str(s)
    m = match(r"##([_\p{L}](?:[\p{L}\d_]*))", s)
    m === nothing && return s
    v = Symbol(m.captures[1])
    esc(:(replace($s, r"##([_\p{L}](?:[\p{L}\d_]*))"=>string(eval($v)))))
end

@eval const $(Symbol("@t_str")) = $(Symbol("@tex_str"))

@eval const $(Symbol("@c_str")) = $(Symbol("@colorant_str"))

#######################################

struct NotImplementedError <: Exception
    msg::String
    NotImplementedError(s) = new("[$s] hasn't been implemented.")
end

struct UnknownOptionError <: Exception
    msg::String
    UnknownOptionError(s, o) = new("[$s] is not recognised as a valid option name for $(typeof(o)).")
end

struct OptionValueError <: Exception
    msg::String
    OptionValueError(s, v) = new("[$s] value given ($v) did not meet the expected format.")
end


text:
struct BarFramework
    vertices::Vector{Int64}
    edges::Vector{Tuple{Int64,Int64}}
    edgelengths::Dict{Tuple{Int64,Int64}, Number}
    dimension::Int64

    # function BarFramework(vertices, edges, edgelengths, dimension)
    #     new(vertices, edges, edgelengths, dimension)
    # end
end

struct RealizationsResult
    barframework::BarFramework
    HCSystem::HomotopyContinuation.System
    HCResult::HomotopyContinuation.Result
    realsolutions::Vector{Array{Float64,1}}
end


text:
# utility functions for the CUSPARSE wrappers

"""
convert `SparseChar` {`N`,`T`,`C`} to `cusparseOperation_t`
"""
function cusparseop(trans::SparseChar)
    if trans == 'N'
        return CUSPARSE_OPERATION_NON_TRANSPOSE
    end
    if trans == 'T'
        return CUSPARSE_OPERATION_TRANSPOSE
    end
    if trans == 'C'
        return CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE
    end
    throw(ArgumentError("unknown cusparse operation $trans"))
end

"""
convert `SparseChar` {`G`,`S`,`H`,`T`} to `cusparseMatrixType_t`
"""
function cusparsetype(mattype::SparseChar)
    if mattype == 'G'
        return CUSPARSE_MATRIX_TYPE_GENERAL
    end
    if mattype == 'T'
        return CUSPARSE_MATRIX_TYPE_TRIANGULAR
    end
    if mattype == 'S'
        return CUSPARSE_MATRIX_TYPE_SYMMETRIC
    end
    if mattype == 'H'
        return CUSPARSE_MATRIX_TYPE_HERMITIAN
    end
    throw(ArgumentError("unknown cusparse matrix type $mattype"))
end

"""
convert `SparseChar` {`U`,`L`} to `cusparseFillMode_t`
"""
function cusparsefill(uplo::SparseChar)
    if uplo == 'U'
        return CUSPARSE_FILL_MODE_UPPER
    end
    if uplo == 'L'
        return CUSPARSE_FILL_MODE_LOWER
    end
    throw(ArgumentError("unknown cusparse fill mode $uplo"))
end

"""
convert `SparseChar` {`U`,`N`} to `cusparseDiagType_t`
"""
function cusparsediag(diag::SparseChar)
    if diag == 'U'
        return CUSPARSE_DIAG_TYPE_UNIT
    end
    if diag == 'N'
        return CUSPARSE_DIAG_TYPE_NON_UNIT
    end
    throw(ArgumentError("unknown cusparse diag mode $diag"))
end

"""
convert `SparseChar` {`Z`,`O`} to `cusparseIndexBase_t`
"""
function cusparseindex(index::SparseChar)
    if index == 'Z'
        return CUSPARSE_INDEX_BASE_ZERO
    end
    if index == 'O'
        return CUSPARSE_INDEX_BASE_ONE
    end
    throw(ArgumentError("unknown cusparse index base"))
end

"""
convert `SparseChar` {`R`,`C`} to `cusparseDirection_t`
"""
function cusparsedir(dir::SparseChar)
    if dir == 'R'
        return CUSPARSE_DIRECTION_ROW
    end
    if dir == 'C'
        return CUSPARSE_DIRECTION_COLUMN
    end
    throw(ArgumentError("unknown cusparse direction $dir"))
end

"""
check that the dimensions of matrix `X` and vector `Y` make sense for a multiplication
"""
function chkmvdims(X, n, Y, m)
    if length(X) != n
        throw(DimensionMismatch("X must have length $n, but has length $(length(X))"))
    elseif length(Y) != m
        throw(DimensionMismatch("Y must have length $m, but has length $(length(Y))"))
    end
end

"""
check that the dimensions of matrices `X` and `Y` make sense for a multiplication
"""
function chkmmdims( B, C, k, l, m, n )
    if size(B) != (k,l)
        throw(DimensionMismatch("B has dimensions $(size(B)) but needs ($k,$l)"))
    elseif size(C) != (m,n)
        throw(DimensionMismatch("C has dimensions $(size(C)) but needs ($m,$n)"))
    end
end

"""
form a `cusparseMatDescr` from a `CuSparseMatrix`
"""
function getDescr( A::CuSparseMatrix, index::SparseChar )
    cuind = cusparseindex(index)
    typ   = CUSPARSE_MATRIX_TYPE_GENERAL
    fill  = CUSPARSE_FILL_MODE_LOWER
    if ishermitian(A)
        typ  = CUSPARSE_MATRIX_TYPE_HERMITIAN
        fill = cusparsefill(A.uplo)
    elseif issymmetric(A)
        typ  = CUSPARSE_MATRIX_TYPE_SYMMETRIC
        fill = cusparsefill(A.uplo)
    end
    cudesc = cusparseMatDescr(typ, fill,CUSPARSE_DIAG_TYPE_NON_UNIT, cuind)
end

function getDescr( A::Symmetric, index::SparseChar )
    cuind = cusparseindex(index)
    typ  = CUSPARSE_MATRIX_TYPE_SYMMETRIC
    fill = cusparsefill(A.uplo)
    cudesc = cusparseMatDescr(typ, fill,CUSPARSE_DIAG_TYPE_NON_UNIT, cuind)
end

function getDescr( A::Hermitian, index::SparseChar )
    cuind = cusparseindex(index)
    typ  = CUSPARSE_MATRIX_TYPE_HERMITIAN
    fill = cusparsefill(A.uplo)
    cudesc = cusparseMatDescr(typ, fill,CUSPARSE_DIAG_TYPE_NON_UNIT, cuind)
end


text:
dir_path = @__DIR__
push!(LOAD_PATH, dir_path * "/sim") # simulation
push!(LOAD_PATH, dir_path * "/vis") # visualization
push!(LOAD_PATH, dir_path * "/adv") # adversary policies
push!(LOAD_PATH, dir_path * "/dis") # discretization
push!(LOAD_PATH, dir_path * "/pomdp")
using sim
using vis
using adv
using dis
using pomdp

using Serialization
using Printf

include("update_info.jl")

function discretize_adv(world::World)
  pts_per_acc = 3
  pts_per_ste = 3
  min_acc = -15
  max_acc = 5
  min_ste = -1e-1
  max_ste = 1e-1
  ctrl_d = Discretization([pts_per_acc, pts_per_ste],
                          [min_acc, min_ste],
                          [max_acc, max_ste])
  return ctrl_d
end

# Scenario: Racer scenario
function main()
  # load the graphical context (OpenGL handle, the graphical window, etc.)
  context = vis.setup()
  vis_scaling = 0.01

  # make the road
  path = make_figure8_path()
  road_width = 20.0
  road = Road(path, road_width,
              vis.make_road(context, path.X, path.Y, road_width))
  road.road.T = vis.scale_mat(vis_scaling, vis_scaling)

  # make the world
  world = World(road, [], vis_scaling)

	len_rd = path.S[end]

	v_track = len_rd / 15
	p_track = 0.0
	ctrl_d = discretize_adv(world)

  # make the agents
  x01 = [len_rd / 2 - 100.0; 25.0; 0]
  agent1 = Agent(1, copy(x01), vis.make_car(context, [0.0, 1.0, 0.0]))
	agent1.controller! = pomdp.adv_controller!
	agent1.custom = [v_track, p_track, pomdp.NOTHING, ctrl_d, pomdp.WEAK]

  x02 = [len_rd - 100.0; 25.0; 0.0]
  agent2 = Agent(2, copy(x02), vis.make_car(context, [1.0, 0.0, 0.0]))
	agent2.controller! = pomdp.adv_controller!
	agent2.custom = [v_track, p_track, pomdp.NOTHING, ctrl_d, pomdp.STRONG]

  x03 = [len_rd - 150.0; 25.0; 0.0]
  agent3 = Agent(3, copy(x03), vis.make_car(context, [1.0, 0.0, 0.0]))
	agent3.controller! = pomdp.adv_controller!
	agent3.custom = [v_track, p_track, pomdp.NOTHING, ctrl_d, pomdp.MEDIUM]

  push!(world.agents, agent1)
  push!(world.agents, agent2)
	push!(world.agents, agent3)
  # ------------------------------------------------------------------------- #

  # make diagnostics render objects
  info1 = vis.InfoBox(context, 0.75, 0.75, vis_scaling)
  info2 = vis.InfoBox(context, -0.75, -0.25, vis_scaling)

	v1 = vis.make_vector_xy(context, 0.0, 0.0, 0.0, 0.0)
	v2 = vis.make_vector_xy(context, 0.0, 0.0, 0.0, 0.0)

  # main loop for rendering and simulation
  window = true
  h = 1e-2
  t0 = time_ns()
  oldt = (time_ns() - t0) / 1e9

	t_prev_replan = 0.0
	plan_ex_time = 0.25

  while window
    t = (time_ns() - t0) / 1e9

    to_visualize = []
    if world.road.road != nothing
      push!(to_visualize, world.road.road)
    end

    if mod(agent2.x[1], road.path.S[end]) < 150.0 &&
      mod(agent2.x[1], path.S[end]) > 50.0
      agent1.x = copy(x01)
      agent2.x = copy(x02)

      println("Resetting")
    end

    for agent in world.agents

			c_v = [NaN, NaN]
			if t > t_prev_replan + plan_ex_time
				(agent.custom[3], c_v) = adv.replan_adv(agent, world)

				if agent.id == 1
					(o1, _) = predict_collision(agent1.x, agent2.x, world)
					(o2, _) = predict_collision(agent2.x, agent1.x, world)
					print("[AGENT 1]")
					println(o1.ctype)
					print("[AGENT 2]")
					println(o2.ctype)

					println("---------------------------")
				end

			end

      ## advance one frame in time
      advance!(agent.dynamics!, agent.x, Pair(agent, world), oldt, t, h)
			if agent.id == 1
      	update_info(info1, agent, world, t)
      else
      	update_info(info2, agent, world, t)
      end

      ## visualize
      (x, y, sx, sy, dx, u) = diagnostic(agent, world, t)
      agent.is_braking = dx[1] * u[1] < 0
			if agent.id == 1
					vis.update_vector_xy!(v1, vis_scaling * x, vis_scaling * y,
													vis_scaling * c_v[1], vis_scaling * c_v[2])
			elseif agent.id == 2
					vis.update_vector_xy!(v2, vis_scaling * x, vis_scaling * y,
													vis_scaling * c_v[1], vis_scaling * c_v[2])
			end

      ## get additional information
      dp = dx[3]
      th = atan(sy, sx)
      th_car = atan(dp, agent.x[2]) + atan(sy, sx)

      if agent.car != nothing
        ## update the car
        vis.car_lights!(agent.car, agent.is_braking)
        agent.car.T = (vis.scale_mat(world.vis_scaling, world.vis_scaling) *
                       vis.translate_mat(x, y) *
                       vis.rotate_mat(th_car - pi / 2))
        push!(to_visualize, agent.car)
      end

    end

		if t > t_prev_replan + plan_ex_time
			t_prev_replan = t
		end

    push!(to_visualize, info1)
    push!(to_visualize, info2)
		push!(to_visualize, v1)
		push!(to_visualize, v2)

    window = vis.visualize(context, to_visualize)

    oldt = t
  end
end


text:
using Plots

backtrg = [[0.0; 0.0],[10.0; 0.0],[5.0; 8.66],[0.0; 0.0]]

lasttringles = [[[0.0; 0.0],[10.0; 0.0],[5.0; 8.66],[0.0; 0.0]]]
for i in 1:8
    appendTringles = []
    for points in lasttringles
        delta = points[1]
        Trgpoints = (points - fill(delta,4)) / 2
        movecor1 = Trgpoints[2]
        movecor2 = Trgpoints[3]
        Trgpoints += fill(delta,4)
        append!(appendTringles, [Trgpoints])
        append!(appendTringles, [Trgpoints + fill(movecor1, 4)])
        append!(appendTringles, [Trgpoints + fill(movecor2, 4)])
    end
    lasttringles = appendTringles
end

plot(hcat(backtrg...)[1,:], hcat(backtrg...)[2,:], legend = false, border=:none, dpi = 500)

for points in lasttringles
    plot!(hcat(points...)[1,:], hcat(points...)[2,:], legend = false, fill = (0, 0.5, :green), dpi = 500)
end
plot!(dpi = 500)
savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\A-few-fractals-in-Julia\\Fractals-Fig\\ST.png")

max_stars_repo_name: shahmari/A-few-fractals-in-Julia

max_issues_repo_name: shahmari/A-few-fractals-in-Julia

max_forks_repo_name: shahmari/A-few-fractals-in-Julia


text:
using Test, Sockets, LanguageServer, CSTParser, SymbolServer, SymbolServer.Pkg, StaticLint, JSON
using LanguageServer: Document, get_text, get_offset, get_line_offsets, get_position_at, get_open_in_editor, set_open_in_editor, is_workspace_file, applytextdocumentchanges
import JSONRPC
const LS = LanguageServer
const Range = LanguageServer.Range

# TODO Replace this with a proper mock endpoint
JSONRPC.send(::Nothing, ::Any, ::Any) = nothing

@testset "LanguageServer" begin

    @testset "document" begin
        include("test_document.jl")
    end
    @testset "communication" begin
        include("test_communication.jl")
    end
    @testset "intellisense" begin
        include("test_intellisense.jl")
    end
    @testset "edit" begin
        include("test_edit.jl")
    end
    @testset "actions" begin
        include("test_actions.jl")
    end
    @testset "paths" begin
        include("test_paths.jl")
    end

end


text:
## -*-Julia-*-
## Test suite for Julia's PulseAudio module
using PulseAudio
using Base.Test

@testset "basic sink construction" begin
    sink = PulseAudioSink()
    # test defaults
    @test nchannels(sink) == 2
    @test samplerate(sink) == 48000
    @test eltype(sink) == Float32
    close(sink)
end


text:
module EnsembleKalmanInversions

using Distributions
using OffsetArrays
using ProgressBars
using Random
using Printf
using LinearAlgebra
using Suppressor: @suppress
using Statistics
using EnsembleKalmanProcesses.EnsembleKalmanProcessModule
using EnsembleKalmanProcesses.ParameterDistributionStorage

using EnsembleKalmanProcesses.EnsembleKalmanProcessModule: sample_distribution

using ..InverseProblems: n_ensemble, observation_map, forward_map, tupify_parameters

#####
##### Priors
#####

function lognormal_with_mean_std(mean, std)
    k = std^2 / mean^2 + 1
    μ = log(mean / sqrt(k))
    σ = sqrt(log(k))
    return LogNormal(μ, σ)
end

struct ConstrainedNormal{FT}
    # θ is the original constrained paramter, θ̃ is the unconstrained parameter ~ N(μ, σ)
    # θ = lower_bound + (upper_bound - lower_bound)/（1 + exp(θ̃)）
    μ :: FT
    σ :: FT
    lower_bound :: FT
    upper_bound :: FT
end

# Scaling factor to give the parameter a magnitude of one
sf(prior) = 1 / abs(prior.μ)

# Model priors are sometimes constrained; EKI deals with unconstrained, Normal priors.
convert_prior(prior::LogNormal) = Normal(sf(prior) * prior.μ, sf(prior) * prior.σ)
convert_prior(prior::Normal) = Normal(sf(prior) * prior.μ, sf(prior) * prior.σ)
convert_prior(prior::ConstrainedNormal) = Normal(prior.μ, prior.σ)

# Convert parameters to unconstrained for EKI
forward_parameter_transform(prior::LogNormal, parameter) = log(parameter^sf(prior))
forward_parameter_transform(prior::Normal, parameter) = parameter * sf(prior)
forward_parameter_transform(cn::ConstrainedNormal, parameter) =
    log((cn.upper_bound - parameter) / (cn.upper_bound - cn.lower_bound))

# Convert parameters from unconstrained (EKI) to constrained
inverse_parameter_transform(prior::LogNormal, parameter) = exp(parameter / sf(prior))
inverse_parameter_transform(prior::Normal, parameter) = parameter / sf(prior)
inverse_parameter_transform(cn::ConstrainedNormal, parameter) =
    cn.lower_bound + (cn.upper_bound - cn.lower_bound) / (1 + exp(parameter))

# Convenience vectorized version
inverse_parameter_transform(priors::NamedTuple, parameters::Vector) =
    NamedTuple(name => inverse_parameter_transform(priors[name], parameters[i])
               for (i, name) in enumerate(keys(priors)))

# Convert covariance from unconstrained (EKI) to constrained
inverse_covariance_transform(::Tuple{Vararg{LogNormal}}, parameters, covariance) =
    Diagonal(exp.(parameters)) * covariance * Diagonal(exp.(parameters))

inverse_covariance_transform(::Tuple{Vararg{Normal}}, parameters, covariance) = covariance

function inverse_covariance_transform(cn::Tuple{Vararg{ConstrainedNormal}}, parameters, covariance)
    upper_bound = [cn[i].upper_bound for i = 1:length(cn)]
    lower_bound = [cn[i].lower_bound for i = 1:length(cn)]
    dT = Diagonal(@. -(upper_bound - lower_bound) * exp(parameters) / (1 + exp(parameters))^2)
    return dT * covariance * dT'
end

mutable struct EnsembleKalmanInversion{I, P, E, M, O, F, S, R}
    inverse_problem :: I
    parameter_distribution :: P
    ensemble_kalman_process :: E
    mapped_observations :: M
    noise_covariance :: O
    inverting_forward_map :: F
    iteration :: Int
    iteration_summaries :: S
    resampler :: R
end

"""
    parameter_ensemble(eki::EnsembleKalmanInversion)

Return a `Vector` of parameter sets (in physical / constrained space) for each ensemble member.
"""
function parameter_ensemble(eki::EnsembleKalmanInversion)
    priors = eki.inverse_problem.free_parameters.priors
    return parameter_ensemble(eki.ensemble_kalman_process, priors)
end

function parameter_ensemble(ensemble_kalman_process, priors)
    unconstrained_parameters = get_u_final(ensemble_kalman_process) # (N_params, ensemble_size) array
    ensemble_size = size(unconstrained_parameters, 2)
    return [inverse_parameter_transform(priors, unconstrained_parameters[:, n]) for n = 1:ensemble_size]
end

Base.show(io::IO, eki::EnsembleKalmanInversion) =
    print(io, "EnsembleKalmanInversion", '\n',
              "├── inverse_problem: ", typeof(eki.inverse_problem).name.wrapper, '\n',
              "├── parameter_distribution: ", typeof(eki.parameter_distribution).name.wrapper, '\n',
              "├── ensemble_kalman_process: ", typeof(eki.ensemble_kalman_process), '\n',
              "├── mapped_observations: ", summary(eki.mapped_observations), '\n',
              "├── noise_covariance: ", summary(eki.noise_covariance), '\n',
              "├── inverting_forward_map: ", typeof(eki.inverting_forward_map).name.wrapper, '\n',
              "├── iteration: $(eki.iteration)", '\n',
              "└── resampler: $(typeof(eki.resampler))")

construct_noise_covariance(noise_covariance::AbstractMatrix, y) = noise_covariance

function construct_noise_covariance(noise_covariance::Number, y)
    # Independent noise for synthetic observations
    n_obs = length(y)
    return noise_covariance * Matrix(I, n_obs, n_obs)
end

"""
    EnsembleKalmanInversion(inverse_problem; noise_covariance=1e-2)

 Return an object that interfaces with [EnsembleKalmanProcesses.jl](https://github.com/CliMA/EnsembleKalmanProcesses.jl)
and uses Ensemble Kalman Inversion to iteratively "solve" the inverse problem:

```math
y = G(θ) + η,
```

for the parameters ``θ``, where ``y`` is a "normalized" vector of observations,
``G(θ)`` is a forward map that predicts the observations, and ``η ∼ N(0, Γ_y)`` is zero-mean
random noise with covariance matrix ``Γ_y`` representing uncertainty in the observations.

By "solve", we mean that the iteration finds the parameter values ``θ`` that minimizes the
distance between ``y`` and ``G(θ)``.

The "forward map output" `G` can have many interpretations. The specific statistics that `G` computes
have to be selected for each use case to provide a concise summary of the complex model solution that
contains the values that we would most like to match to the corresponding truth values `y`. For example,
in the context of an ocean-surface boundary layer parametrization, this summary could be a vector of
concatenated `u`, `v`, `b`, `e` profiles at all or some time steps of the CATKE solution.

(For more details on the Ensemble Kalman Inversion algorithm refer to the
[EnsembleKalmanProcesses.jl Documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/ensemble_kalman_inversion/).)

Arguments
=========

- `inverse_problem :: InverseProblem`: Represents an inverse problem representing the comparison between
                                       synthetic observations generated by
                                       [Oceananigans.jl](https://clima.github.io/OceananigansDocumentation/stable/)
                                       and model predictions, also generated by Oceananigans.jl.

- `noise_covariance` (`AbstractMatrix` or `Number`): normalized covariance representing observational
                                                     uncertainty. If `noise_covariance isa Number` then
                                                     it's converted to an identity matrix scaled by
                                                     `noise_covariance`.

- `resampler`: controls resampling procedure when the forward map contains NaNs. See `NaNResampler`.
"""
function EnsembleKalmanInversion(inverse_problem; noise_covariance=1e-2, resampler=NaNResampler())

    free_parameters = inverse_problem.free_parameters
    original_priors = free_parameters.priors

    transformed_priors = [Parameterized(convert_prior(prior)) for prior in original_priors]
    no_constraints = [[no_constraint()] for _ in transformed_priors]

    parameter_distribution = ParameterDistribution(transformed_priors,
                                                   no_constraints,
                                                   collect(string.(free_parameters.names)))

    ek_process = Inversion()
    initial_ensemble = sample_distribution(parameter_distribution, n_ensemble(inverse_problem))

    # Build EKP-friendly observations "y" and the covariance matrix of observational uncertainty "Γy"
    y = dropdims(observation_map(inverse_problem), dims=2) # length(forward_map_output) column vector
    Γy = construct_noise_covariance(noise_covariance, y)

    # The closure G(θ) maps (N_params, ensemble_size) array to (length(forward_map_output), ensemble_size)
    function inverting_forward_map(θ)
        θ = parameter_ensemble(ensemble_kalman_process, original_priors)

        return forward_map(inverse_problem, θ)
    end

    ensemble_kalman_process = EnsembleKalmanProcess(initial_ensemble, y, Γy, ek_process)

    eki = EnsembleKalmanInversion(inverse_problem,
                                  parameter_distribution,
                                  ensemble_kalman_process,
                                  y,
                                  Γy,
                                  inverting_forward_map,
                                  0,
                                  OffsetArray([], -1),
                                  resampler)

    return eki
end

"""
    UnscentedKalmanInversion(inverse_problem, prior_mean, prior_cov;
                             noise_covariance = 1e-2, α_reg = 1, update_freq = 0)

Return an object that interfaces with [EnsembleKalmanProcesses.jl](https://github.com/CliMA/EnsembleKalmanProcesses.jl)
and uses Unscented Kalman Inversion to iteratively "solve" the inverse problem:

```math
y = G(θ) + η,
```

for the parameters ``θ``, where ``y`` is a "normalized" vector of observations,
``G(θ)`` is a forward map that predicts the observations, and ``η ∼ N(0, Γ_y)`` is zero-mean
random noise with covariance matrix ``Γ_y`` representing uncertainty in the observations.

By "solve", we mean that the iteration finds the parameter values ``θ`` that minimizes the
distance between ``y`` and ``G(θ)``.

(For more details on the Unscented Kalman Inversion algorithm refer to the
[EnsembleKalmanProcesses.jl Documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/unscented_kalman_inversion/).)

Arguments
=========

- `inverse_problem :: InverseProblem`: an inverse problem representing the comparison between
                                       synthetic observations generated by
                                       [Oceananigans.jl](https://clima.github.io/OceananigansDocumentation/stable/)
                                       and model predictions also generated by Oceananigans.jl.

- `prior_mean :: Vector{Float64}`: prior mean

- `prior_cov :: Matrix{Float64}`: prior covariance

- `noise_covariance :: Float64`: observation error covariance

- `α_reg :: Float64`: regularization parameter toward the prior mean (0 < `α_reg` ≤ 1);
                      default `α_reg=1` implies no regularization

- `update_freq :: IT`: set to 0 when the inverse problem is not identifiable (default), namely the
                       inverse problem has multiple solutions, the covariance matrix will represent
                       only the sensitivity of the parameters, instead of posterior covariance information;
                       set to 1 (or anything > 0) when the inverse problem is identifiable, and
                       the covariance matrix will converge to a good approximation of the
                       posterior covariance with an uninformative prior
"""
function UnscentedKalmanInversion(inverse_problem, prior_mean, prior_cov;
                                  noise_covariance = 1e-2, α_reg = 1, update_freq = 0, resampler = nothing)

    free_parameters = inverse_problem.free_parameters
    original_priors = free_parameters.priors

    transformed_priors = [Parameterized(convert_prior(prior)) for prior in original_priors]
    no_constraints = [[no_constraint()] for _ in transformed_priors]
    parameter_distribution = ParameterDistribution(transformed_priors, no_constraints, collect(string.(free_parameters.names)))

    # Build EKP-friendly observations "y" and the covariance matrix of observational uncertainty "Γy"
    y = dropdims(observation_map(inverse_problem), dims=2) # length(forward_map_output) column vector
    Γy = construct_noise_covariance(noise_covariance, y)

    # The closure G(θ) maps (N_params, ensemble_size) array to (length(forward_map_output), ensemble_size)
    function G(θ)
        batch_size = size(θ, 2)
        inverted_parameters = [inverse_parameter_transform.(values(original_priors), θ[:, i]) for i = 1:batch_size]
        return forward_map(inverse_problem, inverted_parameters)
    end

    ensemble_kalman_process = EnsembleKalmanProcess(y, Γy, Unscented(prior_mean, prior_cov, α_reg, update_freq))

    eki = EnsembleKalmanInversion(inverse_problem,
                                  parameter_distribution,
                                  ensemble_kalman_process,
                                  y,
                                  Γy,
                                  G,
                                  0,
                                  OffsetArray([], -1),
                                  resampler)

    return eki
end

"""
    UnscentedKalmanInversionPostprocess(eki)

Returns
=======

- `mean :: Matrix{Float64}`: `N_iterations` × `N_parameters` mean matrix
- `cov :: Vector{Matrix{Float64}}`: `N_iterations` vector of `N_parameters` × `N_parameters` covariance matrix
- `std :: Matrix{Float64}`: `N_iterations` × `N_parameters` standard deviation matrix
- `err :: Vector{Float64}`: `N_iterations` error array
"""
function UnscentedKalmanInversionPostprocess(eki)
    original_priors = eki.inverse_problem.free_parameters.priors
    θ_mean_raw = hcat(eki.ensemble_kalman_process.process.u_mean...)
    θθ_cov_raw = eki.ensemble_kalman_process.process.uu_cov

    θ_mean = similar(θ_mean_raw)
    θθ_cov = similar(θθ_cov_raw)
    θθ_std_arr = similar(θ_mean_raw)

    for i = 1:size(θ_mean, 2) # number of iterations
        θ_mean[:, i] = inverse_parameter_transform.(values(original_priors), θ_mean_raw[:, i])
        θθ_cov[i] = inverse_covariance_transform(values(original_priors), θ_mean_raw[:, i], θθ_cov_raw[i])

        for j = 1:size(θ_mean, 1) # number of parameters
            θθ_std_arr[j, i] = sqrt(θθ_cov[i][j, j])
        end
    end

    return θ_mean, θθ_cov, θθ_std_arr, eki.ensemble_kalman_process.err
end

"""
    struct IterationSummary{P, M, C, V, E}

Container with information about each iteration of the Ensemble Kalman Process.
"""
struct IterationSummary{P, M, C, V, E}
    parameters :: P # constrained
    ensemble_mean :: M # constrained
    ensemble_cov :: C # constrained
    ensemble_var :: V
    mean_square_errors :: E
end

"""
    IterationSummary(eki, parameters, forward_map)

Return the summary for Ensemble Kalman Process `eki` with free `parameters` and `forward_map`.
"""
function IterationSummary(eki, parameters, forward_map)
    N_observations, N_ensemble = size(forward_map)
    original_priors = eki.inverse_problem.free_parameters.priors

    ensemble_mean = mean(parameters, dims=2)
    constrained_ensemble_mean = inverse_parameter_transform.(values(original_priors), ensemble_mean)
    constrained_ensemble_mean = tupify_parameters(eki.inverse_problem, constrained_ensemble_mean)

    ensemble_covariance = cov(parameters, dims=2)
    constrained_ensemble_covariance = inverse_covariance_transform(values(original_priors), parameters, ensemble_covariance)
    constrained_ensemble_variance = tupify_parameters(eki.inverse_problem, diag(constrained_ensemble_covariance))

    constrained_parameters = inverse_parameter_transform.(values(original_priors), parameters)

    constrained_parameters = [tupify_parameters(eki.inverse_problem, constrained_parameters[:, i])
                              for i = 1:size(constrained_parameters, 2)]

    mean_square_errors = [
        mapreduce((x, y) -> (x - y)^2, +, eki.mapped_observations, view(forward_map, :, m)) / N_observations
        for m = 1:N_ensemble
    ]

    return IterationSummary(constrained_parameters,
                            constrained_ensemble_mean,
                            constrained_ensemble_covariance,
                            constrained_ensemble_variance,
                            mean_square_errors)
end

function IterationSummary(eki, parameters)
    original_priors = eki.inverse_problem.free_parameters.priors

    ensemble_mean = mean(parameters, dims=2)
    constrained_ensemble_mean = inverse_parameter_transform.(values(original_priors), ensemble_mean)
    constrained_ensemble_mean = tupify_parameters(eki.inverse_problem, constrained_ensemble_mean)

    ensemble_covariance = cov(parameters, dims=2)
    constrained_ensemble_covariance = inverse_covariance_transform(values(original_priors), parameters, ensemble_covariance)
    constrained_ensemble_variance = tupify_parameters(eki.inverse_problem, diag(constrained_ensemble_covariance))

    constrained_parameters = inverse_parameter_transform.(values(original_priors), parameters)

    constrained_parameters = [tupify_parameters(eki.inverse_problem, constrained_parameters[:, i])
                              for i = 1:size(constrained_parameters, 2)]

    return IterationSummary(constrained_parameters,
                            constrained_ensemble_mean,
                            constrained_ensemble_covariance,
                            constrained_ensemble_variance,
                            fill(NaN, size(constrained_parameters, 2)))
end

"""
    IterationSummary(eki)

Return the summary for Ensemble Kalman Process `eki` before any iteration.
"""
function IterationSummary(eki)
    parameters = get_u_final(eki.ensemble_kalman_process) # (N_params, ensemble_size) array
    return IterationSummary(eki, parameters)
end

function Base.show(io::IO, is::IterationSummary)
    print(io, "IterationSummary(ensemble = ", length(is.mean_square_errors), ")", '\n',
              "                      ", param_str.(keys(is.ensemble_mean))..., '\n',
              "       ensemble_mean: ", param_str.(values(is.ensemble_mean))..., '\n',
              "   ensemble_variance: ", param_str.(values(is.ensemble_var))..., '\n',
              particle_str.(1:length(is.parameters), is.mean_square_errors, is.parameters)...)
    return nothing
end

quick_summary(iter, is) = println("Iter $iter ", is.ensemble_mean)

function param_str(p::Symbol)
    p_str = string(p)
    length(p_str) > 9 && (p_str = p_str[1:9])
    return @sprintf("% 9s | ", p_str)
end

param_str(p::Number) = @sprintf("%1.3e | ", p)

particle_str(particle, error, parameters) =
    @sprintf("% 7s particle % 3d: ", " ", particle) *
    string(param_str.(values(parameters))...) *
    @sprintf("error = %.3e", error) * "\n"

"""
    sample(eki, θ, G, n)

Generate `n` new particles sampled from a multivariate Normal distribution parameterized
by the ensemble mean and covariance computed based on the `N_θ` × `N_ensemble` ensemble
array `θ`, under the condition that all `n` particles lead to forward map outputs that
are "stable" (don't include `NaNs`). `G` is the inverting forward map computed on
ensemble `θ`.

Return an `N_θ` × `n` array of new particles, along with the inverting forward
map output corresponding to the new particles.
"""
function sample(eki, θ, G, n)

    n_params, ens_size = size(θ)
    G_length = size(G, 1)

    ens_dist = eki.resampler.distribution(θ, G)

    found_θ = zeros((n_params, 0))
    found_G = zeros((G_length, 0))

    while size(found_θ, 2) < n
        θ_sample = rand(ens_dist, ens_size)
        G_sample = eki.inverting_forward_map(θ_sample)

        nan_values = nan_cols(G_sample)
        success_columns = findall(Bool.(1 .- nan_values))

        found_θ = hcat(found_θ, θ_sample[:, success_columns])
        found_G = hcat(found_G, G_sample[:, success_columns])
    end

    return found_θ[:, 1:n], found_G[:, 1:n]
end

nan_cols(G) = vec(mapslices(any, isnan.(G); dims=1))

"""
    iterate!(eki::EnsembleKalmanInversion; iterations=1)

Iterate the ensemble Kalman inversion problem `eki` forward by `iterations`.
"""
function iterate!(eki::EnsembleKalmanInversion; iterations = 1, show_progress = true)

    eki.iteration == 0 && (iterations += 1; eki.iteration -= 1)

    iterator = show_progress ? ProgressBar(1:iterations) : 1:iterations

    for _ in iterator

        eki.iteration += 1

        θ = get_u_final(eki.ensemble_kalman_process) # (N_params, ensemble_size) array
        G = eki.inverting_forward_map(θ) # (len(G), ensemble_size)

        eki.iteration != 0 && update_ensemble!(eki.ensemble_kalman_process, G)

        resample!(eki.resampler, G, θ, eki)

        # Save the parameter values and mean square error between forward map
        # and observations at the current iteration
        summary = IterationSummary(eki, θ, G)
        push!(eki.iteration_summaries, summary)
    end

    # Return ensemble mean (best guess for optimal parameters)
    best_parameters = eki.iteration_summaries[end].ensemble_mean

    return tupify_parameters(eki.inverse_problem, best_parameters)
end

#####
##### Resampling
#####

abstract type EnsembleDistribution end

function ensemble_dist(θ)
    μ = [mean(θ, dims=2)...]
    Σ = cov(θ, dims=2)
    return MvNormal(μ, Σ)
end

struct FullEnsembleDistribution <: EnsembleDistribution end

(::FullEnsembleDistribution)(θ, G) = ensemble_dist(θ)

struct SuccessfulEnsembleDistribution <: EnsembleDistribution end

(::SuccessfulEnsembleDistribution)(θ, G) = ensemble_dist(θ[:, findall(Bool.(1 .- nan_cols(G)))])

resample!(::Nothing, args...) = nothing

struct NaNResampler{D}
    abort_fraction :: Float64
    distribution :: D
end

NaNResampler(; abort_fraction=0.0, distribution=FullEnsembleDistribution()) = NaNResampler(abort_fraction, distribution)

function resample!(resampler::NaNResampler, G, θ, eki)

    # `ensemble_size` vector of bits indicating, for each ensemble member, whether the forward map contained `NaN`s
    nan_values = nan_cols(G)
    nan_columns = findall(Bool.(nan_values)) # indices of columns (particles) with `NaN`s
    nan_count = length(nan_columns)
    nan_fraction = nan_count / size(θ, 2)

    if nan_fraction > resampler.abort_fraction
        error("The forward map for $(100nan_fraction)% of particles included NaNs. Consider \n" *
              "    1. Increasing the `abort_fraction` for `resampler = NaNResampler(abort_fraction=0.0)` \n" *
              "        when constructing `EnsembleKalmanInversion`; \n" *
              "    2. Reducing the time-step for `InverseProblem.simulation`, \n" *
              "    3. Evolving `InverseProblem.simulation` for less time \n" *
              "    4. Narrowing `FreeParameters` priors.")
    end

    if nan_fraction > 0
        @warn "The forward map for $nan_count particles ($(100nan_fraction)%) included NaNs. Resampling
                $nan_count particles from a multivariate Normal distribution parameterized by the
                ensemble mean and covariance."

        found_θ, found_G = sample(eki, θ, G, nan_count)
        view(θ, :, nan_columns) .= found_θ
        view(G, :, nan_columns) .= found_G

        new_process = EnsembleKalmanProcess(θ,
                                            eki.mapped_observations,
                                            eki.noise_covariance,
                                            eki.ensemble_kalman_process.process)

        eki.ensemble_kalman_process = new_process
    end

    return G
end

end # module


text:
using Test
using Random
using DataValues
# import DataValues: unsafe_getindex_notnull, unsafe_getvalue_notnull

@testset "DataValueArray: Indexing" begin

    x = DataValueArray{Int}((5, 2))

    for i in eachindex(x)
        x[i] = i
    end

    for i in eachindex(x)
        y = x[i]
        @test isa(y, DataValue{Int})
        @test get(y) === i
    end

    for i in eachindex(x)
        x[i] = DataValue{Int}()
    end

    for i in eachindex(x)
        y = x[i]
        @test isa(y, DataValue{Int})
        @test isna(y)
    end

    _values = rand(10, 10)
    _isna = rand(Bool, 10, 10)
    X = DataValueArray(_values, _isna)

# Base.getindex{T, N}(X::DataValueArray{T, N}, I::DataValue{Int}...)
    @test_throws DataValueException getindex(X, DataValue{Int}(), DataValue{Int}())
    if _isna[1]
        @test isna(getindex(X, DataValue(1)))
    else
        @test isequal(getindex(X, DataValue(1)), DataValue(_values[1]))
    end

# Scalar getindex
    for i = 1:100
        if _isna[i]
            @test isna(X[i])
        else
            @test isequal(X[i], DataValue(_values[i]))
        end
    end
    for i in 1:10, j in 1:10
        if _isna[i, j]
            @test isna(X[i, j])
        else
            @test isequal(X[i, j], DataValue(_values[i, j]))
        end
    end

# getindex with AbstractVectors
    rg = 2:9
    v = X[rg]
    for i = 1:length(rg)
        if _isna[rg[i]]
            @test isna(v[i])
        else
            @test isequal(v[i], DataValue(_values[rg[i]]))
        end
    end

    v = X[rg, 9]
    for i = 1:length(rg)
        if _isna[rg[i], 9]
            @test isna(v[i])
        else
            @test isequal(v[i], DataValue(_values[rg[i], 9]))
        end
    end

    rg2 = 5:7
    v = X[rg, rg2]
    for j in 1:length(rg2), i in 1:length(rg)
        if _isna[rg[i], rg2[j]]
            @test isna(v[i, j])
        else
            @test isequal(v[i, j], DataValue(_values[rg[i], rg2[j]]))
        end
    end

# range indexing
    Z_values = reshape(collect(1:125), (5, 5, 5))
    Z = DataValueArray(Z_values)

    @test isequal(Z[1, 1:4, 1], DataValueArray([1, 6, 11, 16]))

# getindex with AbstractVector{Bool}
    b = bitrand(10, 10)
    rg = (LinearIndices(b))[findall(b)]
    v = X[b]
    for i = 1:length(rg)
        if _isna[rg[i]]
            @test isna(v[i])
        else
            @test isequal(v[i], DataValue(_values[rg[i]]))
        end
    end

# getindex with DataValueVector with null entries throws error
# DA TODO Decide whether I want to keep this
# @test_throws DataValueException X[DataValueArray([1, 2, 3, NA])]

# getindex with DataValueVector and non-null entries
# DA TODO Decide whether I want this or not
# @test isequal(X[DataValueArray([1, 2, 3])], X[[1, 2, 3]])

# indexing with DataValues

    n = rand(1:5)
    siz = [ rand(2:5) for i in n ]
    A = rand(siz...)
    M = rand(Bool, siz...)
    Z = DataValueArray(A, M)
    i = rand(1:length(Z))
    @test isequal(Z[DataValue(i)], Z[i])
    I = [ rand(1:size(Z, i)) for i in 1:n ]
    NI = [ DataValue(i) for i in I ]
    @test isequal(Z[NI...], Z[I...])

# ----- test setindex! -----#

# setindex! with scalar indices
    _values = rand(10, 10)
    for i = 1:100
        X[i] = _values[i]
    end
    @test isequal(X, DataValueArray(_values))

    _values = rand(10, 10)
    for i in 1:10, j in 1:10
        X[i, j] = _values[i, j]
    end
    @test isequal(X, DataValueArray(_values))

    _values = rand(10, 10)
    for i in 1:10, j in 1:10
        X[i, j] = DataValue(_values[i, j])
    end
    @test isequal(X, DataValueArray(_values))

# ----- test nullify! -----#
    _isna = bitrand(10, 10)
    for i = 1:100
        if _isna[i]
            X[i] = NA
        end
    end

# setindex! with scalar and vector indices
    rg = 2:9
    _values[rg] .= 1.0
    X[rg] .= 1.0
    for i = 1:length(rg)
        @test isequal(X[rg[i]], DataValue(1.0))
    end

# setindex! with NA and vector indices
    rg = 5:13
    _isna[rg] .= true
# TODO This should be changed to .=, but currently crashes on 0.7
    X[rg] .= Ref(NA)
    for i = 1:length(rg)
        @test isna(X[rg[i]])
    end

# setindex! with vector and vector indices
    rg = 12:67
    _values[rg] .= rand(length(rg))
    X[rg] .= _values[rg]
    for i = 1:length(rg)
        @test isequal(X[rg[i]], DataValue(_values[rg[i]]))
    end

# ----- test UNSAFE INDEXING -----#

# DA TODO Disabled for now
# X = DataValueArray([1, 2, 3, 4, 5], [true, false, false, false, false])

# @test isequal(unsafe_getindex_notnull(X, 1), DataValue(1))
# @test isequal(unsafe_getindex_notnull(X, 2), DataValue(2))
# @test isequal(unsafe_getvalue_notnull(X, 1), 1)
# @test isequal(unsafe_getvalue_notnull(X, 2), 2)

# ----- test Base.checkbounds -----#

    X = DataValueArray([1:10...])
    b = vcat(false, fill(true, 9))

# Base.checkindex(::Type{Bool}, inds::UnitRange, i::DataValue)
# DA TODO Decide whether I want these
# @test_throws DataValueException checkindex(Bool, 1:1, DataValue{Int}())
# @test checkindex(Bool, 1:10, DataValue(1)) == true
# @test isequal(X[DataValue(1)], DataValue(1))

# Base.checkindex{N}(::Type{Bool}, inds::UnitRange, I::DataValueArray{Bool, N})
# DA TODO Decide whether I want these
# @test checkindex(Bool, 1:5, DataValueArray([true, false, true, false, true]))
    @test isequal(X[b], DataValueArray([2:10...]))

# Base.checkindex{T<:Real}(::Type{Bool}, inds::UnitRange, I::DataValueArray{T})
# DA TODO Decide whether I want these
# @test checkindex(Bool, 1:10, DataValueArray([1:10...]))
# @test checkindex(Bool, 1:10, DataValueArray([10, 11])) == false
# @test_throws BoundsError checkbounds(X, DataValueArray([10, 11]))

# ---- test Base.to_index -----#

# Base.to_index(X::DataValueArray)
# DA TODO Decide whether I want these
# @test Base.to_index(X) == [1:10...]
# push!(X, DataValue{Int}())
# @test_throws DataValueException Base.to_index(X)

end


text:
### Views

"""
    unblock(block_sizes, inds, I)

Returns the indices associated with a block as a `BlockSlice`.
"""
function unblock(A, inds, I)
    B = first(I)
    BlockSlice(B,inds[1][B])
end
# Allow `ones(2)[Block(1)[1:1], Block(1)[1:1]]` which is
# similar to `ones(2)[1:1, 1:1]`.
unblock(A, ::Tuple{}, I) = BlockSlice(first(I),Base.OneTo(1))

to_index(::Block) = throw(ArgumentError("Block must be converted by to_indices(...)"))
to_index(::BlockIndex) = throw(ArgumentError("BlockIndex must be converted by to_indices(...)"))
to_index(::BlockIndexRange) = throw(ArgumentError("BlockIndexRange must be converted by to_indices(...)"))
to_index(::BlockRange) = throw(ArgumentError("BlockRange must be converted by to_indices(...)"))

@inline to_indices(A, inds, I::Tuple{Block{1}, Vararg{Any}}) =
    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
@inline to_indices(A, inds, I::Tuple{BlockRange{1,R}, Vararg{Any}}) where R =
    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
@inline to_indices(A, inds, I::Tuple{BlockIndex{1}, Vararg{Any}}) where R =
    (inds[1][I[1]], to_indices(A, _maybetail(inds), tail(I))...)
@inline to_indices(A, inds, I::Tuple{BlockIndexRange{1,R}, Vararg{Any}}) where R =
    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)

# splat out higher dimensional blocks
# this mimics view of a CartesianIndex
@inline to_indices(A, inds, I::Tuple{Block, Vararg{Any}}) =
    to_indices(A, inds, (Block.(I[1].n)..., tail(I)...))
@inline to_indices(A, inds, I::Tuple{BlockIndex, Vararg{Any}}) =
    to_indices(A, inds, (BlockIndex.(I[1].I, I[1].α)..., tail(I)...))
@inline to_indices(A, inds, I::Tuple{BlockIndexRange, Vararg{Any}}) =
    to_indices(A, inds, (BlockIndexRange.(Block.(I[1].block.n), tuple.(I[1].indices))..., tail(I)...))
@inline to_indices(A, inds, I::Tuple{BlockRange, Vararg{Any}}) =
    to_indices(A, inds, (BlockRange.(tuple.(I[1].indices))..., tail(I)...))

# In 0.7, we need to override to_indices to avoid calling linearindices
@inline to_indices(A, I::Tuple{BlockIndexRange, Vararg{Any}}) = to_indices(A, axes(A), I)
@inline to_indices(A, I::Tuple{Block, Vararg{Any}}) = to_indices(A, axes(A), I)
@inline to_indices(A, I::Tuple{BlockRange, Vararg{Any}}) = to_indices(A, axes(A), I)

reindex(idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
        subidxs::Tuple{BlockSlice{<:BlockIndexRange}, Vararg{Any}}) =
    (@_propagate_inbounds_meta; (BlockSlice(BlockIndexRange(Block(idxs[1].block.indices[1][Int(subidxs[1].block.block)]),
                                                            subidxs[1].block.indices),
                                            idxs[1].indices[subidxs[1].indices]),
                                reindex(tail(idxs), tail(subidxs))...))
reindex(idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
        subidxs::Tuple{BlockSlice{<:Block}, Vararg{Any}}) =
    (@_propagate_inbounds_meta; (BlockSlice(idxs[1].block[Int(subidxs[1].block)],
                                            idxs[1].indices[subidxs[1].indices]),
                                reindex(tail(idxs), tail(subidxs))...))

# _splatmap taken from Base:
_splatmap(f, ::Tuple{}) = ()
_splatmap(f, t::Tuple) = (f(t[1])..., _splatmap(f, tail(t))...)

# De-reference blocks before creating a view to avoid taking `global2blockindex`
# path in `AbstractBlockStyle` broadcasting.
@inline function Base.unsafe_view(
        A::BlockArray{<:Any, N},
        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
    @_propagate_inbounds_meta
    B = view(A, map(block, I)...)
    return view(B, _splatmap(x -> x.block.indices, I)...)
end

@inline function Base.unsafe_view(
        A::PseudoBlockArray{<:Any, N},
        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
    @_propagate_inbounds_meta
    return view(A.blocks, map(x -> x.indices, I)...)
end

@inline function Base.unsafe_view(
        A::ReshapedArray{<:Any, N, <:AbstractBlockArray{<:Any, M}},
        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N, M}
    @_propagate_inbounds_meta
    # Note: assuming that I[M+1:end] are verified to be singletons
    return reshape(view(A.parent, I[1:M]...), Val(N))
end

@inline function Base.unsafe_view(
        A::Array{<:Any, N},
        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
    @_propagate_inbounds_meta
    return view(A, map(x -> x.indices, I)...)
end

# make sure we reindex correctrly
function Base._maybe_reindex(V, I::Tuple{BlockSlice{<:BlockIndexRange{1}}, Vararg{Any}}, ::Tuple{})
    @_inline_meta
    @inbounds idxs = to_indices(V.parent, reindex(V.indices, I))
    view(V.parent, idxs...)
end

# BlockSlices map the blocks and the indices
# this is loosely based on Slice reindex in subarray.jl
reindex(idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
        subidxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}}) =
    (@_propagate_inbounds_meta; (BlockSlice(BlockRange(idxs[1].block.indices[1][Int.(subidxs[1].block)]),
                                            idxs[1].indices[subidxs[1].block]),
                                reindex(tail(idxs), tail(subidxs))...))

function reindex(idxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}},
        subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}})
        (idxs[1], reindex(tail(idxs), tail(subidxs))...)
end

#################
# support for pointers
#################

const BlockOrRangeIndex = Union{RangeIndex, BlockSlice}

## BlockSlice1 is a convenience for views
const BlockSlice1 = BlockSlice{Block{1,Int}}

block(A::BlockSlice) = block(A.block)
block(A::Block) = A

# unwind BLockSlice1 for AbstractBlockArray
@inline Base.view(block_arr::AbstractBlockArray{<:Any,N}, blocks::Vararg{BlockSlice1, N}) where N =
    view(block_arr, map(block,blocks)...)

@inline Base.view(V::SubArray{<:Any,N,<:AbstractBlockArray,<:NTuple{N,BlockSlice{<:BlockRange{1}}}}, block::Block{N}) where N =
    view(parent(V), map((b, i) -> b.block[i], parentindices(V), block.n)...)

function Base.view(A::Adjoint{<:Any,<:BlockArray}, b::Block{2})
    k, j = b.n
    view(parent(A), Block(j), Block(k))'
end
function Base.view(A::Transpose{<:Any,<:BlockArray}, b::Block{2})
    k, j = b.n
    transpose(view(parent(A), Block(j), Block(k)))
end

Base.view(A::AdjOrTrans{<:Any,<:BlockArray}, K::Block{1}, J::Block{1}) = view(A, Block(Int(K), Int(J)))

unsafe_convert(::Type{Ptr{T}}, V::SubArray{T,N,PseudoBlockArray{T,N,AT},<:Tuple{Vararg{BlockOrRangeIndex}}}) where {T,N,AT} =
    unsafe_convert(Ptr{T}, V.parent) + (Base.first_index(V)-1)*sizeof(T)

# support for strided array interface for subblocks. Typically
# these aren't needed as `view(A, Block(K))` will return the block itself,
# not a `SubArray`, but currently `view(view(A, Block.(1:3)), Block(2))` does
# not.

strides(A::SubArray{<:Any,N,<:BlockArray,<:NTuple{N,BlockSlice1}}) where N =
    strides(view(parent(A), block.(parentindices(A))...))

for Adj in (:Transpose, :Adjoint)
    @eval strides(A::SubArray{<:Any,N,<:$Adj{<:Any,<:BlockArray},<:NTuple{N,BlockSlice1}}) where N =
        strides(view(parent(A), block.(parentindices(A))...))
end

unsafe_convert(::Type{Ptr{T}}, V::SubArray{T, N, <:BlockArray, <:NTuple{N, BlockSlice1}}) where {T,N} =
    unsafe_convert(Ptr{T}, view(parent(V), block.(parentindices(V))...))

# The default blocksize(V) is slow for views as it calls axes(V), which
# allocates. Here we work around this.

_sub_blocksize() = ()
_sub_blocksize(ind::BlockSlice{<:BlockRange{1}}, inds...) = tuple(length(ind.block),_sub_blocksize(inds...)...)
_sub_blocksize(ind::BlockSlice{<:Block{1}}, inds...) = tuple(1,_sub_blocksize(inds...)...)
_sub_blocksize(ind::AbstractVector, inds...) = tuple(blocksize(ind,1),_sub_blocksize(inds...)...)
_sub_blocksize(ind::Integer, inds...) = _sub_blocksize(inds...)
blocksize(V::SubArray) = _sub_blocksize(parentindices(V)...)
blocksize(V::SubArray, i::Int) = _sub_blocksize(parentindices(V)[i])[1]

function hasmatchingblocks(V::SubArray{<:Any,2,<:Any,<:NTuple{2,BlockSlice{<:BlockRange{1}}}})
    a,b = axes(parent(V))
    kr,jr = parentindices(V)
    KR,JR = (kr.block),(jr.block)
    length(KR) == length(JR) || return false
    for (K,J) in zip(KR,JR)
        length(a[K]) == length(b[J]) || return false
    end
    true
end


text:
# this is a little Monte Carlo exercise that illustrates that
# the OLS estimator is unbiased when we have strong exogeneity
using Plots
reps = 1000 # number of Monte Carlo reps.
n = 20 # sample size
sig = 3.0  # st. dev. of errors

x = [ones(n,1) randn(n,1)]  # x is fixed over repeated samples
beta = [1.0, 2.0] # true beta
PopRegLine = x*beta

e = sig*randn(n,reps) # reps will be in columns
y = PopRegLine .+ e

betas = inv(x'x)*x'y

betas[2,:] .-=  2.0
histogram(betas[2,:], label = "")
gui()
#savefig("Unbiased.svg")

max_stars_repo_name: nilshg/Econometrics

max_issues_repo_name: nilshg/Econometrics

max_forks_repo_name: nilshg/Econometrics


text:
export l2

doc"""
    l2(x::Var, lambda::Float64)

L2 regularization.

```math
y = \frac{\lambda}{2}\left\Vert \mathbf{x} \right\Vert ^{2}
```

```julia
x = Var(rand(Float32,10,5))
y = l2(x, 0.01)
```
"""
function l2(x::Var, lambda::Float64)
    isnothing(x.data) && return Var(nothing,l2,(x,lambda))
    T = eltype(x)
    y = mapreduce(x -> x*x, +, x.data) * T(lambda) / 2
    Var([y], (l2,x,lambda))
end

function addgrad!(y::Var, ::typeof(l2), x::Var, lambda::Float64)
    T = eltype(y)
    isvoid(y.grad) || ∇l2!(y.grad, x.data, x.grad, T(lambda))
end

function ∇l2!(gy::Vector{T}, x::Array{T}, gx::Array{T}, lambda::T) where T
    @inbounds for i = 1:length(x)
        gx[i] += gy[1] * lambda * x[i]
    end
end


text:
# Generative model using the change of variables formula
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using LinearAlgebra, InvertibleNetworks, PyPlot, Flux, Random
import Flux.Optimise.update!

# Random seed
Random.seed!(11)

####################################################################################################

# Define network
nx = 1
ny = 1
n_in = 2
n_hidden = 64
batchsize = 64
depth = 8
AN = Array{ActNorm}(undef, depth)
L = Array{CouplingLayerHINT}(undef, depth)
Params = Array{Parameter}(undef, 0)

# Create layers
for j=1:depth
    AN[j] = ActNorm(n_in; logdet=true)
    L[j] = CouplingLayerHINT(n_in, n_hidden; logdet=true, permute="lower", k1=1, k2=1, p1=0, p2=0)

    # Collect parameters
    global Params = cat(Params, get_params(AN[j]); dims=1)
    global Params = cat(Params, get_params(L[j]); dims=1)
end

# Forward pass
function forward(X)
    logdet = 0f0
    for j=1:depth
        X_, logdet1 = AN[j].forward(X)
        X, logdet2 = L[j].forward(X_)
        logdet += (logdet1 + logdet2)
    end
    return X, logdet
end

# Backward pass
function backward(ΔX, X)
    for j=depth:-1:1
        ΔX_, X_ = L[j].backward(ΔX, X)
        ΔX, X = AN[j].backward(ΔX_, X_)
    end
    return ΔX, X
end

####################################################################################################

# Negative log-prior for banana distribution
c = [1f0, 4f0]
function neglogprior(X::Array{Float32, 4})
    z1 = c[1]*X[:, :, 1:1, :]
    z2 = X[:, :, 2:2, :]/c[1]-c[2]*c[1]^2*(X[:, :, 1:1, :].^2 .+1)
    return sum(0.5f0*(z1.^2+z2.^2)), cat(c[1]*z1-2*c[2]*c[1]^2*z2.*X[:, :, 1:1, :], z2/c[1]; dims = 3)
end

# Negative log-likelihood (Gaussian)
σ2 = (0.2f0)^2
Ydat = sample_banana(1)+sqrt(σ2)*randn(Float32, 1, 1, 2, 1)
function negloglike(X::Array{Float32, 4}; Y::Array{Float32, 4} = Ydat, σ2::Float32 = σ2)
    return -log_likelihood(Y.-X)/σ2, ∇log_likelihood(Y.-X)/σ2
end

# Negative log-posterior
function neglogpost(X::Array{Float32, 4})
    flike,  glike  = negloglike(X)
    fprior, gprior = neglogprior(X)
    return flike+fprior, glike+gprior
end

# Loss
function loss(Z::Array{Float32, 4})
    X, logdet = forward(Z)
    fpost, gpost = neglogpost(X)
    f = fpost-logdet
    ΔX = gpost
    backward(ΔX, X)
    return f
end

# Training
maxiter = 2000
opt = Flux.ADAM(1f-3)
lr_step = 100
lr_decay_fn = Flux.ExpDecay(1f-3, .9, lr_step, 0.)
fval = zeros(Float32, maxiter)

for j=1:maxiter

    # Evaluate objective and gradients
    Z = randn(Float32, 1, 1, 2, batchsize)
    fval[j] = loss(Z)
    mod(j, 10) == 0 && (print("Iteration: ", j, "; f = ", fval[j], "\n"))

    # Update params
    for p in Params
        update!(opt, p.data, p.grad)
        update!(lr_decay_fn, p.data, p.grad)
    end
    clear_grad!(Params)

end

####################################################################################################

# Testing (posterior sampling)
test_size = 500
Z = randn(Float32, 1, 1, 2, test_size)
Xpost = forward(Z)[1]

# Prior sampling
Xprior = sample_banana(test_size)

# Plot
fig = figure(figsize = [8, 8])
title("Prior and posterior samples")
plot(Xprior[1, 1, 1, :], Xprior[1, 1, 2, :], "y.", label = "Prior")
plot(Xpost[1, 1, 1, :], Xpost[1, 1, 2, :], "b.", label = "Posterior")
plot(Ydat[1], Ydat[2], "r*", label = "Data")
legend()
ax = gca()
ax.set_xlim([-3.5, 3.5]); ax.set_ylim([0, 20])


text:
__precompile__()

module FEDVR

function __init__()
    @warn "The FEDVR.jl package has been deprecated in favour of JuliaApproximation/CompactBases.jl"
    nothing
end

using SparseArrays
using RecipesBase

include("grid.jl")
include("lagrange.jl")
include("basis.jl")
include("derivatives.jl")
include("projections.jl")

end # module


text:
@doc raw"""
    RandomSubset(m::Int)

## Arguments
- `m::Int`: Number of inducing points

Uniform sampling of a subset of `m` points of the data.
"""
struct RandomSubset <: OffIPSA
    m::Int
    function RandomSubset(m::Int)
        m > 0 || throw(ArgumentError("The number of inducing points m should be positive"))
        return new(m)
    end
end

"""
     inducingpoints([rng::AbstractRNG], alg::RandomSubset, X::AbstractVector; [weights::Vector=nothing])
     inducingpoints([rng::AbstractRNG], alg::RandomSubset, X::AbstractMatrix; obsdim=1, [weights::Vector=nothing])

Select inducing points by taking a Random Subset. Optionally accepts a weight vector for the inputs `X`.
"""
function inducingpoints(
    rng::AbstractRNG, alg::RandomSubset, X::AbstractVector; weights=nothing, kwargs...
)
    nSamples = length(X)
    alg.m > nSamples && error(
        "Number of inducing points larger than the input collection size ($(alg.m) > $(length(X))",
    )
    alg.m == nSamples && return X

    samp = if isnothing(weights)
        sample(rng, 1:nSamples, alg.m; replace=false)
    else
        wsample(rng, 1:nSamples, weights, alg.m; replace=false)
    end
    return X[samp]
end


text:
function lind_jump(lind::LindbladLiouvillian, u, p, s::Real)
    l = length(lind)
    prob = Float64[]
    ops = Vector{Matrix{ComplexF64}}()
    for (γfun, Lfun) in zip(lind.γ, lind.L)
        L = Lfun(s)
        γ = γfun(s)
        push!(prob, γ * norm(L * u))
        push!(ops, L)
    end
    sample(ops, Weights(prob))
end

function ame_jump(D::DaviesGenerator, u, ω_ba, v, s)
    lvl = size(ω_ba, 1)
    sys_dim = size(v, 1)
    num_noise = length(D.coupling)
    prob_dim = (lvl * (lvl - 1) + 1) * num_noise
    γm = D.γ.(ω_ba)
    prob = Array{Float64,1}(undef, prob_dim)
    tag = Array{Tuple{Int,Int,Int},1}(undef, prob_dim)
    idx = 1
    ϕb = abs2.(v' * u)
    σab = [v' * op * v for op in D.coupling(s)]
    for i = 1:num_noise
        γA = γm .* abs2.(σab[i])
        for b = 1:lvl
            for a = 1:b - 1
                prob[idx] = γA[a, b] * ϕb[b]
                tag[idx] = (i, a, b)
                idx += 1
                prob[idx] = γA[b, a] * ϕb[a]
                tag[idx] = (i, b, a)
                idx += 1
            end
        end
        prob[idx] = transpose(diag(γA)) * ϕb
        tag[idx] = (i, 0, 0)
        idx += 1
    end
    choice = sample(tag, Weights(prob))
    if choice[2] == 0
        res = zeros(ComplexF64, sys_dim, sys_dim)
        for i in range(1, stop=sys_dim)
            res += sqrt(γm[1, 1]) * σab[choice[1]][i, i] * v[:, i] * v[:, i]'
        end
    else
        res =
            sqrt(γm[choice[2], choice[3]]) *
            σab[choice[1]][choice[2], choice[3]] *
            v[:, choice[2]] *
            v[:, choice[3]]'
    end
    res
end

# TODO: Better implemention of ame_jump function
"""
$(SIGNATURES)

Calculate the jump operator for the `DiffEqLiouvillian` at time `t`.
"""
function lindblad_jump(Op::DiffEqLiouvillian{true,false}, u, p, t::Real)
    s = p(t)
    w, v = Op.H.EIGS(Op.H, s, Op.lvl)
    ω_ba = transpose(w) .- w
    resample([ame_jump(x, u, ω_ba, v, s) for x in Op.opensys_eig], u)
end

function lindblad_jump(Op::DiffEqLiouvillian{false,false}, u, p, t::Real)
    s = p(t)
    jump_ops = resample([lind_jump(x, u, p, s) for x in Op.opensys], u)
end

function resample(Ls, u)
    if length(Ls) == 1
        Ls[1]
    else
        prob = [norm(L * u) for L in Ls]
        sample(Ls, Weights(prob))
    end
end


text:
# if `b` is a builder, then `b(model, rng, shape...)` is called to make a
# new chain, where `shape` is the return value of this method:
function MLJFlux.shape(model::NeuralNetworkClassifier, X, y)
    levels = MLJModelInterface.classes(y[1])
    n_output = length(levels)
    n_input = Tables.schema(X).names |> length
    return (n_input, n_output)
end

# builds the end-to-end Flux chain needed, given the `model` and `shape`:
MLJFlux.build(model::NeuralNetworkClassifier, rng, shape) =
    Flux.Chain(build(model.builder, rng, shape...),
               model.finaliser)

# returns the model `fitresult` (see "Adding Models for General Use"
# section of the MLJ manual) which must always have the form `(chain,
# metadata)`, where `metadata` is anything extra neede by `predict` may
# require:
MLJFlux.fitresult(model::NeuralNetworkClassifier, chain, y) =
    (chain, MLJModelInterface.classes(y[1]))

function MLJModelInterface.predict(model::NeuralNetworkClassifier,
                                   fitresult,
                                   Xnew)
    chain, levels = fitresult
    X = reformat(Xnew)
    probs = vcat([chain(tomat(X[:,i]))' for i in 1:size(X, 2)]...)
    return MLJModelInterface.UnivariateFinite(levels, probs)
end

MLJModelInterface.metadata_model(NeuralNetworkClassifier,
                                 input=Table(Continuous),
                                 target=AbstractVector{<:Finite},
                                 path="MLJFlux.NeuralNetworkClassifier",
                                 descr="A neural network model for making "*
                                 "probabilistic predictions of a "*
                                 "`Multiclass` or `OrderedFactor` target, "*
                                 "given a table of `Continuous` features. ")


text:
import jInv.ForwardShare.getSensMatVec

function getSensMatVec(v::Vector,m::Vector,pFor::FWIparam)

    # extract pointers
    Mesh  = pFor.Mesh
    omega = pFor.omega
    gamma = pFor.gamma
    Q     = pFor.Sources
    P     = pFor.Receivers
    U     = pFor.Fields
    LU    = pFor.Ainv

    nsrc = size(Q,2)
    nfreq = length(omega)

    # allocate space for matvec product
    Jv   = zeros(size(P,2),nsrc,nfreq)

    # derivative of mass matrix
    An2cc = getNodalAverageMatrix(Mesh)
    dM   = repeat(An2cc'*((1 .- 1im*vec(gamma)).*v),1,nsrc)
    for i=1:nfreq
        R   = U[:,:,i].*dM
        Lam = LU[i]\R
        Jv[:,:,i] = real(omega[i]^2*P'*Lam)
    end

    return vec(Jv)
end