text:
## gradient methods

"""
    gradient(inFLD::MeshArray,Γ::NamedTuple)

Compute spatial derivatives. Other methods:
```
gradient(inFLD::MeshArray,Γ::NamedTuple,doDIV::Bool)
gradient(inFLD::MeshArray,iDXC::MeshArray,iDYC::MeshArray)
```
"""
function gradient(inFLD::MeshArray,Γ::NamedTuple)
(dFLDdx, dFLDdy)=gradient(inFLD,Γ,true)
return dFLDdx, dFLDdy
end

function gradient(inFLD::MeshArray,Γ::NamedTuple,doDIV::Bool)

exFLD=exchange(inFLD,1)
dFLDdx=similar(inFLD)
dFLDdy=similar(inFLD)

for a=1:inFLD.grid.nFaces
  (s1,s2)=size(exFLD.f[a])
  tmpA=view(exFLD.f[a],2:s1-1,2:s2-1)
  tmpB=tmpA-view(exFLD.f[a],1:s1-2,2:s2-1)
  tmpC=tmpA-view(exFLD.f[a],2:s1-1,1:s2-2)
  if doDIV
    dFLDdx.f[a]=tmpB./Γ.DXC.f[a]
    dFLDdy.f[a]=tmpC./Γ.DYC.f[a]
  else
    dFLDdx.f[a]=tmpB
    dFLDdy.f[a]=tmpC
  end
end

return dFLDdx, dFLDdy
end

function gradient(inFLD::MeshArray,iDXC::MeshArray,iDYC::MeshArray)

exFLD=exchange(inFLD,1)
dFLDdx=similar(inFLD)
dFLDdy=similar(inFLD)

for a=1:inFLD.grid.nFaces
  (s1,s2)=size(exFLD.f[a])
  tmpA=view(exFLD.f[a],2:s1-1,2:s2-1)
  tmpB=tmpA-view(exFLD.f[a],1:s1-2,2:s2-1)
  tmpC=tmpA-view(exFLD.f[a],2:s1-1,1:s2-2)
  dFLDdx.f[a]=tmpB.*iDXC.f[a]
  dFLDdy.f[a]=tmpC.*iDYC.f[a]
end

return dFLDdx, dFLDdy
end

##

"""
    curl(u::MeshArray,v::MeshArray,Γ::NamedTuple)

Compute curl of a velocity field.
"""
function curl(u::MeshArray,v::MeshArray,Γ::NamedTuple)

	uvcurl=similar(Γ.XC)
	fac=exchange(1.0 ./Γ.RAZ,1)
	(U,V)=exchange(u,v,1)
    (DXC,DYC)=exchange(Γ.DXC,Γ.DYC,1)
	[DXC[i].=abs.(DXC[i]) for i in eachindex(U)]
	[DYC[i].=abs.(DYC[i]) for i in eachindex(V)]

	for i in eachindex(U)
    ucur=U[i][2:end,:]
    vcur=V[i][:,2:end]
    tmpcurl=ucur[:,1:end-1]-ucur[:,2:end]
    tmpcurl=tmpcurl-(vcur[1:end-1,:]-vcur[2:end,:])
    tmpcurl=tmpcurl.*fac[i][1:end-1,1:end-1]

		##still needed:
		##- deal with corners
		##- if putCurlOnTpoints

		tmpcurl=1/4*(tmpcurl[1:end-1,2:end]+tmpcurl[1:end-1,1:end-1]+
					tmpcurl[2:end,2:end]+tmpcurl[2:end,1:end-1])

		uvcurl[i]=tmpcurl
	end

	return uvcurl
end

## mask methods

function mask(fld::MeshArray)
fldmsk=mask(fld,NaN)
return fldmsk
end

"""
    mask(fld::MeshArray, val::Number)

Replace non finite values with val. Other methods:
```
mask(fld::MeshArray)
mask(fld::MeshArray, val::Number, noval::Number)
```
"""
function mask(fld::MeshArray, val::Number)
  fldmsk=similar(fld)
  for a=1:fld.grid.nFaces
    tmp1=copy(fld.f[a])
    replace!(x -> !isfinite(x) ? val : x, tmp1 )
    fldmsk.f[a]=tmp1
  end
  return fldmsk
end

function mask(fld::MeshArray, val::Number, noval::Number)
  fldmsk=similar(fld)
  for a=1:fld.grid.nFaces
    tmp1=copy(fld.f[a])
    replace!(x -> x==noval ? val : x, tmp1  )
    fldmsk.f[a]=tmp1
  end
  return fldmsk
end

## convergence methods

"""
    convergence(uFLD::MeshArray,vFLD::MeshArray)

Compute convergence of a vector field
"""
function convergence(uFLD::MeshArray,vFLD::MeshArray)

#important note:
#  Normally uFLD, vFLD should not contain any NaN;
#  if otherwise then something this may be needed:
#  uFLD=mask(uFLD,0.0); vFLD=mask(vFLD,0.0);

CONV=similar(uFLD)

(tmpU,tmpV)=exch_UV(uFLD,vFLD)
for a=1:tmpU.grid.nFaces
  (s1,s2)=size(uFLD.f[a])
  tmpU1=view(tmpU.f[a],1:s1,1:s2)
  tmpU2=view(tmpU.f[a],2:s1+1,1:s2)
  tmpV1=view(tmpV.f[a],1:s1,1:s2)
  tmpV2=view(tmpV.f[a],1:s1,2:s2+1)
  CONV.f[a]=tmpU1-tmpU2+tmpV1-tmpV2
end

return CONV
end

## smooth function

"""
    smooth(FLD::MeshArray,DXCsm::MeshArray,DYCsm::MeshArray,Γ::NamedTuple)

Smooth out scales below DXCsm / DYCsm via diffusion
"""
function smooth(FLD::MeshArray,DXCsm::MeshArray,DYCsm::MeshArray,Γ::NamedTuple)

#important note:
#input FLD should be land masked (NaN/1) by caller if needed

#get land masks (NaN/1):
mskC=fill(1.0,FLD) + 0.0 * mask(FLD)
(mskW,mskS)=gradient(FLD,Γ,false)
mskW=fill(1.0,FLD) + 0.0 * mask(mskW)
mskS=fill(1.0,FLD) + 0.0 * mask(mskS)

#replace NaN with 0. in FLD and land masks:
FLD=mask(FLD,0.0)
mskC=mask(mskC,0.0)
mskW=mask(mskW,0.0)
mskS=mask(mskS,0.0)

#get inverse grid spacing:
iDXC=similar(FLD)
iDYC=similar(FLD)
for a=1:FLD.grid.nFaces
  iDXC.f[a]=1.0./Γ.DXC.f[a]
  iDYC.f[a]=1.0./Γ.DYC.f[a]
end

#Before scaling the diffusive operator ...
tmp0=DXCsm*iDXC*mskW;
tmp00=maximum(tmp0);
tmp0=DYCsm*iDYC*mskS;
tmp00=max(tmp00,maximum(tmp0));

#... determine a suitable time period:
nbt=ceil(1.1*2*tmp00^2);
dt=1.;
T=nbt*dt;
#println("nbt="*"$nbt")

#diffusion operator times DYG / DXG
KuxFac=mskW*DXCsm*DXCsm/T/2.0*Γ.DYG
KvyFac=mskS*DYCsm*DYCsm/T/2.0*Γ.DXG

#time steping factor:
dtFac=dt*mskC/Γ.RAC

#loop:
for it=1:nbt
  (dTdxAtU,dTdyAtV)=gradient(FLD,iDXC,iDYC);
  tmpU=similar(FLD)
  tmpV=similar(FLD)
  for a=1:FLD.grid.nFaces
      tmpU.f[a]=dTdxAtU.f[a].*KuxFac.f[a];
      tmpV.f[a]=dTdyAtV.f[a].*KvyFac.f[a];
  end
  tmpC=convergence(tmpU,tmpV);
  for a=1:FLD.grid.nFaces
      FLD.f[a]=FLD.f[a]-dtFac.f[a].*tmpC.f[a];
  end
end

#Apply land mask (NaN/1) to end result:
mskC=mask(mskC,NaN,0.0)
FLD=mskC*FLD

return FLD

end

## ThroughFlow function

"""
    ThroughFlow(VectorField,IntegralPath,Γ::NamedTuple)

Compute transport through an integration path
"""
function ThroughFlow(VectorField,IntegralPath,Γ::NamedTuple)

    #Note: vertical intergration is not always wanted; left for user to do outside

    U=VectorField["U"]
    V=VectorField["V"]

    nd=ndims(U)
    #println("nd=$nd and d=$d")

    n=fill(1,4)
    tmp=size(U)
    n[1:nd].=tmp[1:nd]

    haskey(VectorField,"factors") ? f=VectorField["factors"] : f=Array{String,1}(undef,0)
    haskey(VectorField,"dimensions") ? d=VectorField["dimensions"] : d=Array{String,1}(undef,nd)

    #a bit of a hack to distinguish gcmfaces v gcmarray indexing:
    isdefined(U,:fIndex) ? ndoffset=1 : ndoffset=0
    length(d)!=nd+ndoffset ? error("inconsistent specification of dims") : nothing

    trsp=Array{Float64}(undef,1,n[3],n[4])
    do_dz=sum(f.=="dz")
    do_dxory=sum(f.=="dxory")

    for i3=1:n[3]
        #method 1: quite slow
        #mskW=IntegralPath.mskW
        #do_dxory==1 ? mskW=mskW*Γ.DYG : nothing
        #do_dz==1 ? mskW=Γ.DRF[i3]*mskW : nothing
        #mskS=IntegralPath.mskS
        #do_dxory==1 ? mskS=mskS*Γ.DXG : nothing
        #do_dz==1 ? mskS=Γ.DRF[i3]*mskS : nothing
        #
        #method 2: less slow
        tabW=IntegralPath.tabW
        tabS=IntegralPath.tabS
        for i4=1:n[4]
            #method 1: quite slow
            #trsp[1,i3,i4]=sum(mskW*U[:,:,i3,i4])+sum(mskS*V[:,:,i3,i4])
            #
            #method 2: less slow
            trsp[1,i3,i4]=0.0
            for k=1:size(tabW,1)
                (a,i1,i2,w)=tabW[k,:]
                do_dxory==1 ? w=w*Γ.DYG.f[a][i1,i2] : nothing
                do_dz==1 ? w=w*Γ.DRF[i3] : nothing
                isdefined(U,:fIndex) ? u=U.f[a,i3,i4][i1,i2] : u=U.f[a][i1,i2,i3,i4]
                trsp[1,i3,i4]=trsp[1,i3,i4]+w*u
            end
            for k=1:size(tabS,1)
                (a,i1,i2,w)=tabS[k,:]
                do_dxory==1 ? w=w*Γ.DXG.f[a][i1,i2] : nothing
                do_dz==1 ? w=w*Γ.DRF[i3] : nothing
                isdefined(V,:fIndex) ? v=V.f[a,i3,i4][i1,i2] : v=V.f[a][i1,i2,i3,i4]
                trsp[1,i3,i4]=trsp[1,i3,i4]+w*v
            end
        end
    end

    nd+ndoffset<4 ? trsp=dropdims(trsp,dims=3) : nothing
    nd+ndoffset<3 ? trsp=dropdims(trsp,dims=2) : nothing
    nd+ndoffset==2 ? trsp=trsp[1] : nothing

    return trsp
end

## LatitudeCircles function

"""
    LatitudeCircles(LatValues,Γ::NamedTuple)

Compute integration paths that follow latitude circles
"""
function LatitudeCircles(LatValues,Γ::NamedTuple)
    LatitudeCircles=Array{NamedTuple}(undef,length(LatValues))
    for j=1:length(LatValues)
        mskCint=1*(Γ.YC .>= LatValues[j])
        mskC,mskW,mskS=edge_mask(mskCint)
        LatitudeCircles[j]=(lat=LatValues[j],
        tabC=MskToTab(mskC),tabW=MskToTab(mskW),tabS=MskToTab(mskS))
    end
    return LatitudeCircles
end

function MskToTab(msk::MeshArray)
  n=Int(sum(msk .!= 0)); k=0
  tab=Array{Int,2}(undef,n,4)
  for i in eachindex(msk)
    a=msk[i]
    b=findall( a .!= 0)
    for ii in eachindex(b)
      k += 1
      tab[k,:]=[i,b[ii][1],b[ii][2],a[b[ii]]]
    end
  end
  return tab
end

"""
    edge_mask(mskCint::MeshArray)

Compute edge mask (mskC,mskW,mskS) from domain interior mask (mskCint).
This is used in `LatitudeCircles` and `Transect`.
"""
function edge_mask(mskCint::MeshArray)
  mskCint=1.0*mskCint

  #treat the case of blank tiles:
  #mskCint[findall(RAC.==0)].=NaN

  mskC=similar(mskCint)
  mskW=similar(mskCint)
  mskS=similar(mskCint)

  mskCint=exchange(mskCint,1)

  for i in eachindex(mskCint)
      tmp1=mskCint[i]
      # tracer mask:
      tmp2=tmp1[2:end-1,1:end-2]+tmp1[2:end-1,3:end]+
      tmp1[1:end-2,2:end-1]+tmp1[3:end,2:end-1]
      mskC[i]=1((tmp2.>0).&(tmp1[2:end-1,2:end-1].==0))
      # velocity masks:
      mskW[i]=tmp1[2:end-1,2:end-1] - tmp1[1:end-2,2:end-1]
      mskS[i]=tmp1[2:end-1,2:end-1] - tmp1[2:end-1,1:end-2]
  end

  #treat the case of blank tiles:
  #mskC[findall(isnan.(mskC))].=0.0
  #mskW[findall(isnan.(mskW))].=0.0
  #mskS[findall(isnan.(mskS))].=0.0

  return mskC,mskW,mskS
end

##

"""
    Transect(name,lons,lats,Γ)

Compute integration paths that follow a great circle between two geolocations give by `lons`, `lats`.
"""
function Transect(name,lons,lats,Γ)
  x0,y0,z0,R=rotate_points(lons,lats)
  x,y,z=rotate_XCYC(Γ,R)
  mskCint=1.0*(z.>0)
  mskCedge,mskWedge,mskSedge=edge_mask(mskCint)
  mskCedge,mskWedge,mskSedge=shorter_paths!((x,y,z),(x0,y0,z0),(mskCedge,mskWedge,mskSedge))
  tabC=MskToTab(mskCedge)
  tabW=MskToTab(mskWedge)
  tabS=MskToTab(mskSedge)
  return (name=name,tabC=tabC,tabW=tabW,tabS=tabS)
end

##

nanmean(x) = mean(filter(!isnan,x))
nanmean(x,y) = mapslices(nanmean,x,dims=y)

"""
    UVtoUEVN(u,v,G::NamedTuple)

1. Interpolate to grid cell centers (uC,vC)
2. Convert to `Eastward/Northward` components (uE,vN)
"""
function UVtoUEVN(u::MeshArray,v::MeshArray,G::NamedTuple)
    u[findall(G.hFacW[:,1].==0)].=NaN
    v[findall(G.hFacS[:,1].==0)].=NaN

    (u,v)=exch_UV(u,v); uC=similar(u); vC=similar(v)
    for iF=1:u.grid.nFaces
        tmp1=u[iF][1:end-1,:]; tmp2=u[iF][2:end,:]
        uC[iF]=reshape(nanmean([tmp1[:] tmp2[:]],2),size(tmp1))
        tmp1=v[iF][:,1:end-1]; tmp2=v[iF][:,2:end]
        vC[iF]=reshape(nanmean([tmp1[:] tmp2[:]],2),size(tmp1))
    end

    return uC.*G.AngleCS-vC.*G.AngleSN, uC.*G.AngleSN+vC.*G.AngleCS
end

"""
    UVtoTransport(U,V,G::NamedTuple)

Convert e.g. velocity (m/s) to transport (m^3/s) by multiplying by `DRF` and `DXG`,`DYG`.
"""
function UVtoTransport(U::MeshArray,V::MeshArray,G::NamedTuple)
  uTr=deepcopy(U)
  vTr=deepcopy(V)
  for i in eachindex(U)
      tmp1=U[i]; tmp1[(!isfinite).(tmp1)] .= 0.0
      tmp1=V[i]; tmp1[(!isfinite).(tmp1)] .= 0.0
      uTr[i]=G.DRF[i[2]]*U[i].*G.DYG[i[1]]
      vTr[i]=G.DRF[i[2]]*V[i].*G.DXG[i[1]]
  end
  return uTr,vTr
end


text:
##########################################################
# Rafał Nowak
#
# Instytut Informatyki
# Uniwersytet Wrocławski
#
#
# Pomocniczne materiały do wykładu z analizy numerycznej
#
##########################################################

wartoscDokladna_1e5 = 1.6449240668982262698057485033127;
wartoscDokladna_1e6 = 1.6449330668487264363057484999794;
wartoscDokladna_1e7 = 1.6449339668482314364722484999793;
wartoscDokladna_1e8 = 1.6449340568482264864724149999793;
wartoscDokladna_1e9 = 1.6449340658482264369724151664793;

real = Float32

function algorytm1(n)
  suma = zero(real);
  for k = 1:1:n
    suma = suma + one(real)/convert(real, k)/convert(real,k);
  end
  return suma
end

function algorytm2(n)
  suma = zero(real);
  for k = n:-1:1
    suma = suma + one(real)/convert(real, k)/convert(real,k);
  end
  return suma
end

function doTest(x, algorytm, n)
  println(algorytm, " n=", n)
  @printf("Wynik dokładny  = %.16f\n", x);
  y = algorytm(n);
  @printf("Wynik obliczony = %.16f   błąd względny = %.3e\n", y, abs(x-y)/x);
end

doTest(wartoscDokladna_1e5, algorytm1, 1e5);
doTest(wartoscDokladna_1e5, algorytm2, 1e5);

doTest(wartoscDokladna_1e6, algorytm1, 1e6);
doTest(wartoscDokladna_1e6, algorytm2, 1e6);

@printf("Precyzja arytmetyki = %.3e\n", 0.5*eps(real));

ext: jl

max_stars_repo_name: Magikis/Uniwersity

max_issues_repo_name: Magikis/Uniwersity

max_forks_repo_name: Magikis/Uniwersity


text:
# Identification of EMIC and mirror modes in two frequency bands from single satellite data.
#
# EMIC: δB⟂ vs. δv_⟂
# Mirror: δB∥ vs. δn
#
# The moving box window size is decided empirically.
#
# Usage:
#   julia -t 4 demo_wave_satellite_mt.jl
# or
#   JULIA_NUM_THREADS=4 julia demo_wave_satellite_mt.jl
#
# Hongyang Zhou, hyzhou@umich.edu

using Statistics: mean
using LinearAlgebra: ⋅, normalize!, norm
using DSP, PyPlot, Glob
using Vlasiator
using Vlasiator: μ₀, mᵢ, RE
using Polyester

"Extract variables at `loc` from VLSV `files`."
function extract_vars(files, loc)
   nfiles = length(files)

   t = zeros(Float32, nfiles)
   n = zeros(Float32, nfiles)
   v = zeros(Float32, 3, nfiles)
   b = zeros(Float32, 3, nfiles)
   e = zeros(Float32, 3, nfiles)

   id = load(files[1]) do meta
      getcell(meta, loc)
   end

   # Extract data from each frame
   @batch for i = eachindex(files)
      meta = load(files[i])
      t[i] = meta.time
      n[i] = readvariable(meta, "proton/vg_rho", id)[1]
      v[:,i] = readvariable(meta, "proton/vg_v", id)
      b[:,i] = readvariable(meta, "vg_b_vol", id)
      e[:,i] = readvariable(meta, "vg_e_vol", id)
   end

   t, n, v, e, b
end

function moving_average(g::AbstractVector{<:AbstractFloat}, n::Int)
   nbackward = div(n,2)
   nforward = isodd(n) ? nbackward : nbackward - 1
   len = length(g)
   g_avg = similar(g)
   @inbounds @batch for i = 1:len
      lo = max(1, i - nbackward)
      hi = min(len, i + nforward)
      g_avg[i] = mean(@view g[lo:hi])
   end
   g_avg
end

function moving_average(g::AbstractMatrix{<:AbstractFloat}, n::Int)
   nbackward = div(n,2)
   nforward = isodd(n) ? nbackward : nbackward - 1
   len = size(g,2)
   g_avg = similar(g)
   for icomp = axes(g,1)
      @inbounds @batch for i = 1:len
         lo = max(1, i - nbackward)
         hi = min(len, i + nforward)
         g_avg[icomp,i] = mean(@view g[icomp, lo:hi])
      end
   end
   g_avg
end

function detrend(v::AbstractVector{<:AbstractFloat}; nbox=length(v)÷12)
   v̄ = moving_average(v, nbox)
   dv = v .- v̄
end

function detrend(v::AbstractMatrix{<:AbstractFloat}; nbox=size(v,2)÷12)
   v̄ = similar(v)
   for i in axes(v,1)
      v̄[i,:] = @views moving_average(v[i,:], nbox)
   end
   dv = v .- v̄
end

function align_yaxis(ax1, ax2)
   y_lims = [ax.get_ylim() for ax in [ax1, ax2]]

   # normalize both axes
   y_mags = ntuple(i -> y_lims[i][2] - y_lims[i][1], Val(2))
   y_lims_normalized = ntuple(i -> y_lims[i] ./ y_mags[i], Val(2))

   # find combined range
   y_new_lims_normalized = (min(y_lims_normalized[1][1],y_lims_normalized[2][1]),
      max(y_lims_normalized[1][2], y_lims_normalized[2][2]))

   # denormalize combined range to get new axes
   new_lim1 = y_new_lims_normalized .* y_mags[1]
   new_lim2 = y_new_lims_normalized .* y_mags[2]
   ax1.set_ylim(new_lim1)
   ax2.set_ylim(new_lim2)
   return
end

################# Main

files = glob("bulk*.vlsv", ".")

# virtual satellite location
loc = [12Vlasiator.RE, 0, 0]

nbox = 40 # moving box size, [# of indices]
fs = 2.0 # sampling rate, [Hz]

println("Running with $(Threads.nthreads()) threads...")
println("Number of files: $(length(files))")
println("Extracting location: $loc [m]")
println("Sampling rate: $fs [Hz]")
println("Moving box window size: $(nbox / fs) [s]")

t, n, v, e, b = extract_vars(files, loc)

bmag = [sqrt(b[1,i]^2 + b[2,i]^2 + b[3,i]^2) for i in eachindex(n)]

# Detrend before filtering to remove the lowest frequency changes
dn = detrend(n; nbox)
dbmag = detrend(bmag; nbox)

n̄ = moving_average(n, nbox)
b̄mag = moving_average(bmag, nbox)

v̄a = @. b̄mag / sqrt(μ₀ * mᵢ * n̄)

dv = detrend(v; nbox) # [m/s]
db = detrend(b; nbox) # [T]

# Decompose into parallel and perpendicular vector components
b̂₀ = moving_average(b, nbox)
for i in axes(b̂₀, 2)
   normalize!(@view b̂₀[:,i])
end

# δB∥
db_par = [db[:,i] ⋅ b̂₀[:,i] for i in axes(db, 2)]
# δB⟂
db_perp = [norm(db[:,i] .- db_par[i] .* b̂₀[:,i]) for i in axes(db, 2)]
# δv∥
dv_par = [dv[:,i] ⋅ b̂₀[:,i] for i in axes(dv, 2)]
# δv⟂
dv_perp = [norm(dv[:,i] .- dv_par[i] .* b̂₀[:,i]) for i in axes(dv, 2)]

designmethod = Butterworth(5)

responsetype = Highpass(0.1; fs)
dn_high = filtfilt(digitalfilter(responsetype, designmethod), dn)
dbpar_high = filtfilt(digitalfilter(responsetype, designmethod), db_par)
dbperp_high = filtfilt(digitalfilter(responsetype, designmethod), db_perp)
dvperp_high = filtfilt(digitalfilter(responsetype, designmethod), dv_perp)

responsetype = Bandpass(0.02, 0.067; fs)
dn_low = filtfilt(digitalfilter(responsetype, designmethod), dn)
dbpar_low = filtfilt(digitalfilter(responsetype, designmethod), db_par)

color1 = "tab:blue"
color2 = "tab:red"

fig, axs = plt.subplots(3, 1; figsize=(13, 9), sharex=true, constrained_layout=true)

fig.suptitle("Virtual satellite at $(string(round.(loc./RE, digits=2))) "*L"R_E";
   fontsize="x-large")

axs[1].plot(t, dn_low ./ n̄, color1, label=L"\delta n, 0.02-0.067\, Hz")
ax12 = axs[1].twinx()
ax12.plot(t, dbpar_low ./ b̄mag, color=color2, label=L"\delta B, 0.02-0.067\, Hz")

axs[2].plot(t, (dn_high ./ n̄), color1, label=L"\delta n, 0.1-1.0\, Hz")
ax22 = axs[2].twinx()
ax22.plot(t, (dbpar_high ./ b̄mag), color=color2, label=L"\delta B, 0.1-1.0\, Hz")

axs[3].plot(t, dvperp_high ./ v̄a, color1, label=L"\delta v_\perp, 0.1-1.0\, Hz")
ax32 = axs[3].twinx()
ax32.plot(t, dbperp_high ./ b̄mag, color=color2, label=L"\delta B, 0.1-1.0\, Hz")

axs[3].set_xlabel("time [s]"; fontsize="large")

n_str = (L"\delta n /n", L"\delta n / n", L"\delta v_\perp / V_A")

for (i, a) in enumerate(axs)
   a.hlines(0.0, t[1], t[end]; colors="k", linestyle="dashed", alpha=0.3)
   a.tick_params(axis="y", labelcolor=color1)
   a.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
   a.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
   a.grid(true)
   a.set_ylabel(n_str[i], color=color1, fontsize=14)
end

axtwin = (ax12, ax22, ax32)
b_str = ( L"\delta B_\parallel / B_0", L"\delta B_\parallel / B_0", L"\delta B_\perp / B_0")

for (i, a) in enumerate(axtwin)
   a.tick_params(axis="y", labelcolor=color2)
   a.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
   a.set_ylabel(b_str[i], color=color2, fontsize="large")
   align_yaxis(axs[i], a)
end

AnchoredText = matplotlib.offsetbox.AnchoredText

at = AnchoredText(
   "[0.02, 0.067] Hz band", prop=Dict("size"=>"medium"), frameon=true, loc="lower left")
at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
axs[1].add_artist(at)

at = AnchoredText(
   "[0.1, 1.0] Hz band", prop=Dict("size"=>"medium"), frameon=true, loc="lower left")
at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
axs[2].add_artist(at)

at = AnchoredText(
   "[0.1, 1.0] Hz band", prop=Dict("size"=>"medium"), frameon=true, loc="lower left")
at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
axs[3].add_artist(at)

savefig("virtual_satellite_wave.png"; dpi=300)


text:
const CALLBACK_DEBUG_SIGNAL = Ref(false)

function debug_callback(args...)
    CALLBACK_DEBUG_SIGNAL[] = true
    default_debug_callback(args...)
end

const debug_callback_c = @cfunction(debug_callback, UInt32, (DebugUtilsMessageSeverityFlagEXT, DebugUtilsMessageTypeFlagEXT, Ptr{vk.VkDebugUtilsMessengerCallbackDataEXT}, Ptr{Cvoid}))
const API_VERSION = v"1.2"
const VALIDATION_LAYER = "VK_LAYER_KHRONOS_validation"

const INSTANCE_LAYERS = [
]
const INSTANCE_EXTENSIONS = [
]
const DEVICE_EXTENSIONS = [
]
const ENABLED_FEATURES = PhysicalDeviceFeatures(
)

let available_layers = unwrap(enumerate_instance_layer_properties())
    if VALIDATION_LAYER ∈ getproperty.(available_layers, :layer_name)
        push!(INSTANCE_LAYERS, VALIDATION_LAYER)
        true
    else
        @warn "Validation layer not found."
        false
    end
end

const WITH_DEBUG = let available_extensions = unwrap(enumerate_instance_extension_properties())
    if "VK_EXT_debug_utils" ∈ getproperty.(available_extensions, :extension_name)
        push!(INSTANCE_EXTENSIONS, "VK_EXT_debug_utils")
        true
    else
        @warn "VK_EXT_debug_utils not supported"
        false
    end
end

@testset "Vulkan tests" begin
    include("init.jl")

    @testset "Utilities" begin
        @testset "Function pointers" begin
            @test function_pointer(instance, "vkEnumeratePhysicalDevices") == get_instance_proc_addr("vkEnumeratePhysicalDevices"; instance)
            @test function_pointer(device, "vkQueueWaitIdle") == get_device_proc_addr(device, "vkQueueWaitIdle")
        end
    end

    @testset "Debugging" begin
        @testset "Validation" begin
            if WITH_DEBUG
                @test CALLBACK_DEBUG_SIGNAL[]
            end
        end
    end

    @testset "Bitmask flags" begin
        buffer = Buffer(device, 24, BUFFER_USAGE_INDIRECT_BUFFER_BIT, SHARING_MODE_EXCLUSIVE, [0]; flags=BUFFER_CREATE_PROTECTED_BIT & BUFFER_CREATE_SPARSE_ALIASED_BIT)
        buffer_2 = Buffer(device, 24, BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, SHARING_MODE_EXCLUSIVE, [0]; flags=0)
        @test buffer isa Buffer
        @test buffer_2 isa Buffer
    end

    @testset "Buffers" begin
        include("buffers.jl")
    end
end

GC.gc()


text:
"""
```julia
@enum(BoundConstraint_Enum,
      ACTIVE_LB = -1,
      INACTIVE_BC = 0,
      ACTIVE_UB = +1)
```

An enum to store bound constraint state.
"""
@enum(BoundConstraint_Enum,
      ACTIVE_LB = -1,
      INACTIVE_BC = 0,
      ACTIVE_UB = +1)


text:
#Based on http://www.aei.mpg.de/~gielen/diffgeo.pdf

using Equations

Ea=Tensor(:E,0,:a,2) 			#1.1
u=Tensor(:u,:a,0,1)*Ea			#1.2

eq1_3=@equ BraKet(ω,α*v+β*w)=α*BraKet(ω,v)+β*BraKet(ω,w)	#1.3
ω=Tensor(:ω,0,:a,1)*Tensor(:E,:a,0,2)				#1.4

eq1_5=@equ BraKet(Tensor(E,a,0,2),Tensor(E,0,b,2))=Delta(a,b)	#1.5
eq1_6=@equ BraKet(αη+βλ,u)=α*BraKet(η,u)+β*BraKet(λ,u)		#1.6
eq1_7=@equ BraKet(ω,u)=Tensor(ω,0,a,1)*Tensor(u,a,0,1)		#1.7

eq1_8=@equ BraKet(D(f),X)=X*f					#1.8
xj=Tensor(:x,:j,0,1)
eq1_9=@equ BraKet(DerOp(Tensor(x,i,0,1)),D(xj))=Delta(i,j)	#1.9
eq1_10=@equ BraKet(DerOp(Tensor(x,i,0,1)),D(xj))=DerOp(Tensor(x,i,0,1))*xj
Xj=Tensor(:X,:j,0,1)
X=Xj*DerOp(xj)
eq1_11=@equ BraKet(D(f),X)=X*f

#If
eq1_12=@equ BraKet(D(f),X)=0
#Then f is a constant in the direction of the vector X.

eq1_13=@equ D(f)=Der(f,xi)*D(xi)
eq1_13&(@equ f=Der(f,xj)*D(xj))

eq1_14=@equ T=Tensor(T,[a1,a2,a3],[b1,b2])*Tensor(E,0,a1)*Tensor(E,0,a2)*Tensor(E,0,a3)*Tensor(E,b1,0)*Tensor(E,b2,0)
eq1_15=@equ Ea´=Tensor(χ,a,a´)*Tensor(E,0,a)
eq1_16=@equ Tensor(E,a´,0)=Tensor(Φ,a´,a)*Tensor(E,0,a)
eq1_17=@equ Delta(a´,b´)=Tensor(Φ,a´,a)*Tensor(χ,a,b´) #Tensor(χ,b,b´)*Delta(a,b)=Tensor(χ,a,b´)

eq1_19=@equ Tensor(T,a´,b´)*Tensor(χ,a,a´)*Tensor(Φ,b´,b)=Tensor(T,a,b)

eq1_22=@equ Tensor(E,a,0)∧Tensor(E,b,0)=Tensor(E,a,0)⊗Tensor(E,b,0)-Tensor(E,b,0)⊗Tensor(E,a,0)

eq1_24=@equ P=Form(1/p!,Tensor(P,0,[a1,:...,ap],p),Tensor(E,a1,0)∧:...∧Tensor(E,ap,0),p)
eq1_25=@equ Q=Form(1/q!,Tensor(Q,0,[b1,:...,bq],q),Tensor(E,b1,0)∧:...∧Tensor(E,bq,0),q)

eq1_28=@equ P∧Q=-1^(p*q)*Q∧P

eq1_32=@equ P=Form(1/p!,Tensor(P,0,[a1,:...,ap],p),Tensor(dx,a1,0)∧:...∧Tensor(dx,ap,0),p)
eq1_33=@equ dP=Form(1/p!,Der(Tensor(P,0,[a1,:...,ap],p),Tensor(x,b,0)),Tensor(dx,b,0)∧Tensor(dx,a1,0)∧:...∧Tensor(dx,ap,0),p)

eq1_36=@equ Tensor(A,a,a´)=Der(Tensor(x,a,0),Tensor(x´,a´,0))
eq1_37=@equ Tensor(P,0,[a´1,:...,a´p])=Tensor(A,a1,a´1)*:...*Tensor(A,ap,a´p)*Tensor(P,0,[a1,:...,ap])

eq1_45=@equ D(P∧Q)=D(P)∧Q+(-1)^p*P∧D(Q)

eq1_68=@equ D(F)=0
eq1_69=@equ F=Tensor(F,0,[a,b])*D(Tensor(x,a,0))∧D(Tensor(x,b,0))/2

eq1_77=@equ ds^2=-dt^2+dρ^2+ρ^2*dθ^2+dz^2

eq1_80=@equ F=-Bθ*dρ∧dz
eq1_81=@equ D(F)=0
eq1_82=@equ Tensor("*F",0,[a,b])=Tensor(ε,[c,d],[a,b])*Tensor(F,0,[c,d])/2

eq1_87=@equ A=Tensor(A,0,a)*D(Tensor(x,a,0))
eq1_88=@equ Tensor(η,0,[α,β])=-2Trace(Tensor(T,0,α)*Tensor(T,0,β))


text:
using KNITRO
using Test

const KN_VERBOSE = false

@testset "Test C API" begin
    include("knitroapi.jl")
end

@testset "Test examples" begin
    for file in filter(f -> endswith(f, ".jl"), readdir(joinpath(dirname(@__FILE__), "..", "examples")))
        if occursin("mps_reader", file)
            continue
        end
        @testset "Test example $file" begin
            include(joinpath(dirname(@__FILE__), "..", "examples", file))
        end
    end
end

@testset "Test MathOptInterface" begin
    include("MOI_wrapper.jl")
end

try
    @testset "Test C API License" begin
        include("knitroapi_licman.jl")
    end
catch e
    @warn("License tests failed, but this might be due to License Manager" *
    " not being supported by your license.")
    println("The error catched was:\n")
    println("$e\n")
    println("See table above for more details.")
end


text:
function test_primal_problems()
    n  = rand(50:100)
    m  = rand(50:100)
    X1 = rand(Normal(-1.5,1), n, 2)
    X2 = rand(Normal(1.5,1), m, 2)
    X  = vcat(X1, X2)
    y  = vcat(ones(n), zeros(m)) .== 1

    τ    = 0.2*rand() + 0.3
    K    = rand(1:10)
    C    = 0.5*rand() + 1
    ϑ1   = 0.4*rand() + 0.5
    ϑ2   = 0.5*rand() + 1
    data = Primal(X, y);

    @testset "PatMat with $surrogate loss" for surrogate in [Hinge, Quadratic]
        l1    = surrogate(ϑ1);
        l2    = surrogate(ϑ2);
        model = PatMat(τ, C, l1, l2);

        test_primal(model, data, 100000, ADAM(0.0001))
    end

    @testset "PatMatNP with $surrogate loss" for surrogate in [Hinge, Quadratic]
        l1    = surrogate(ϑ1);
        l2    = surrogate(ϑ2);
        model = PatMatNP(τ, C, l1, l2);

        test_primal(model, data, 100000, ADAM(0.0001))
    end

    @testset "TopPushK with $surrogate loss" for surrogate in [Hinge, Quadratic]
        l     = surrogate(ϑ1);
        model = TopPushK(K, C, l);

        test_primal(model, data, 100000, ADAM())
    end

    @testset "TopPush with $surrogate loss" for surrogate in [Hinge, Quadratic]
        l     = surrogate(ϑ1);
        model = TopPush(C, l);

        test_primal(model, data, 100000, ADAM(0.001))
    end
end

function test_primal(model::AbstractModel, data::Primal, maxiter::Integer, optimizer::Any; atol::Real = 1e-4)

    sol1 = solve(Gradient(maxiter = maxiter, optimizer = optimizer, verbose = false), model, data)
    L1   = AccuracyAtTopKernels.objective(model, data, sol1)

    sol2 = solve(General(), model, data)
    L2   = AccuracyAtTopKernels.objective(model, data, sol2)

    @test L1 ≈ L2 atol = atol
end


text:
###########################################
# Concrete implementations of AbstractWord:
#   Word and SubWord
#

"""
    Word{T} <: AbstractWord{T}
Word as written in an alphabet storing only pointers to letters of an Alphabet.

Note that the negative values in `genptrs` field represent the inverse of letter.
If type is not specified in the constructor it will default to `Int16`.
"""
struct Word{T} <: AbstractWord{T}
    ptrs::Vector{T}

    function Word{T}(v::AbstractVector{<:Integer}, check=true) where {T}
        check && @assert all(x -> x > 0, v) "All entries of a Word must be positive integers"
        return new{T}(v)
    end
end

# setting the default type to UInt16
Word(x::AbstractVector{<:Integer}) = Word{UInt16}(x)
Word(w::AbstractWord{T}) where T = Word{T}(w, false)

struct SubWord{T, V<:SubArray{T,1}} <: AbstractWord{T}
    ptrs::V
end

Base.view(w::Union{Word, SubWord}, u::UnitRange{Int}) =
    KnuthBendix.SubWord(view(w.ptrs, u))

# AbstractWord Interface:

Base.@propagate_inbounds function Base.getindex(w::Union{Word, SubWord}, n::Integer)
    @boundscheck checkbounds(w, n)
    return @inbounds w.ptrs[n]
end

Base.@propagate_inbounds function Base.setindex!(w::Union{Word, SubWord}, v::Integer, n::Integer)
    @boundscheck checkbounds(w, n)
    @assert v > 0 "All entries of a Word must be positive integers"
    return @inbounds w.ptrs[n] = v
end
Base.length(w::Union{Word, SubWord}) = length(w.ptrs)

Base.push!(w::Word, n::Integer) = (@assert n > 0; push!(w.ptrs, n); w)
Base.pushfirst!(w::Word, n::Integer) = (@assert n > 0; pushfirst!(w.ptrs, n); w)
Base.pop!(w::Word) = (pop!(w.ptrs))
Base.popfirst!(w::Word) = (popfirst!(w.ptrs))

Base.append!(w::Word, v::AbstractVector{<:Integer}) = (append!(w.ptrs, v); w)
Base.prepend!(w::Word, v::AbstractVector{<:Integer}) = (prepend!(w.ptrs, v); w)
Base.append!(w::Word, v::Union{Word, SubWord}) = append!(w, v.ptrs)
Base.prepend!(w::Word, v::Union{Word, SubWord}) = prepend!(w, v.ptrs)

Base.resize!(w::Word, n::Integer) = (resize!(w.ptrs, n); w)

Base.:*(w::Union{Word{S}, SubWord{S}}, v::Union{Word{T}, SubWord{T}}) where {S,T} =
    (TT = promote_type(S, T); Word{TT}(TT[w.ptrs; v.ptrs], false))

Base.similar(w::Union{Word, SubWord}, ::Type{S}) where {S} = Word{S}(Vector{S}(undef, length(w)), false)


text:
using ICapeThermoUniversalConstants,ICapeThermoCompounds,ECapeExceptions,PhysicalPropertyCalculator
import CapeOpen: MaterialObject, PropertyPackage
import ICapeThermoCompounds.TempPropData

#ICapeThermoUniversalConstants
universalconstantlist=getuniversalconstantlist()
println("universalconstantlist=",universalconstantlist)
@test_throws ECapeInvalidArgument (getuniversalconstant("avogadroConstan"))
@test getuniversalconstant("avogadroConstant") == 6.0221419947e23

#ICapeThermoCompounds
constproplist=getconstproplist()
@test length(constproplist) == 17
println("constproplist=",constproplist)
tdependentproplist=gettdependentproplist()
@test length(tdependentproplist) == 11
println("tdependentproplist=",tdependentproplist)
pdependentproplist=getpdependentproplist()
@test length(pdependentproplist) == 1
@test getnumcompounds() == 345
compIds,formulae,names,boilTemps,molwts,casnos=getcompoundlist();
comps=UInt16[findfirst(names,name) for name in ["Air","Water","Nitrogen"]];
propvals=Vector{Union{ASCIIString,Float64}}();
@test_throws ECapeThrmPropertyNotAvailable getcompoundconstant!(["casRegistryNumber","chemicalFormula","criticalPressure","liquidDensityAt25C"],comps,propvals)
allconstants=Vector{Union{ASCIIString,Float64}}()
@test_throws ECapeThrmPropertyNotAvailable getcompoundconstant!(constproplist,compIds,allconstants)
@test length(allconstants) == length(constproplist)*length(compIds)
propvals=Vector{Float64}()
getpdependentproperty!(["boilingPointTemperature"],80000.,comps,propvals)
propvals=Vector{Float64}()
@test_throws ECapeThrmPropertyNotAvailable gettdependentproperty!(tdependentproplist,300.,compIds,propvals)


text:
using Test
using Suppressor
using LinearAlgebra
using RecipesBase
using Triangulate # load before StartUpDG b/c of @require
using StartUpDG

include("triangulate_tests.jl")
include("geometric_term_tests.jl")
include("reference_elem_tests.jl")
include("sbp_tests.jl")
include("MeshData_tests.jl")
include("boundary_util_tests.jl")
include("misc_tests.jl")


text:
using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators, DomainSets, Distributions

@parameters t x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dt = Differential(t)
t_min= 0.
t_max = 2.0
x_min = 0.
x_max = 2.
y_min = 0.
y_max = 2.

nx_emis = 5
ny_emis = 5
nt_emis = 5
dx_emis = (x_max - x_min)/nx_emis
dy_emis = (y_max - y_min)/ny_emis
#dt_emis = (t_max - t_min)/nt_emis
#add t variable later

#Created random emission, later to be changed by real emission data.
emis = rand(LogNormal(1,3), nx_emis, ny_emis)

#Locate/ align emis into coordinate we're using
function emissions(x,y)
    i = Int(ceil((x - x_min)/dx_emis))
    j = Int(ceil((y - y_min)/dy_emis))
    println(x, " ",y," ", i," ", j)
    return emis[i,j]
end

@register emissions(x, y)

# 3D PDE (diffusion eq + emissions)
eq  = Dt(u(t,x,y)) ~ Dxx(u(t,x,y)) + Dyy(u(t,x,y)) + emissions(x, y)

# Initial and boundary conditionn
bcs = #[u(t_min,x,y) ~ analytic_sol_func(t_min,x,y),
      [u(t_min,x,y) ~ 0,
       u(t,x_min,y) ~ 0,
       u(t,x_max,y) ~ 0,
       u(t,x,y_min) ~ 0,
       u(t,x,y_max) ~ 0]

# Space and time domains
domains = [t ∈ IntervalDomain(t_min,t_max),
           x ∈ IntervalDomain(x_min,x_max),
           y ∈ IntervalDomain(y_min,y_max)]
pdesys = PDESystem([eq],bcs,domains,[t,x,y],[u(t,x,y)])

# Method of lines discretization
dx = 0.1; dy = 0.1
discretization = MOLFiniteDifference([x=>dx,y=>dy],t)
prob = ModelingToolkit.discretize(pdesys,discretization)
sol = solve(prob,saveat = 0.1, Tsit5())

# Plotting
using Plots
xs,ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
#= for i ∈ length(sol.t)
    u_sol = reshape(sol.u[i], length(xs)-2,length(ys)-2)
    p = plot(xs[2:length(xs)-1], ys[2:length(ys)-1], u_sol, linetype=:contourf,title = "solution")

end
plottogether = plot(p) =#

#Ploting 1 ~ 4 index of the solution
u_sol = reshape(sol.u[1], length(xs)-2,length(ys)-2)
p1 = plot(xs[2:length(xs)-1], ys[2:length(ys)-1], u_sol, linetype=:contourf,title = "solution")
u_sol = reshape(sol.u[2], length(xs)-2,length(ys)-2)
p2 = plot(xs[2:length(xs)-1], ys[2:length(ys)-1], u_sol, linetype=:contourf,title = "solution")
u_sol = reshape(sol.u[3], length(xs)-2,length(ys)-2)
p3 = plot(xs[2:length(xs)-1], ys[2:length(ys)-1], u_sol, linetype=:contourf,title = "solution")
u_sol = reshape(sol.u[4], length(xs)-2,length(ys)-2)
p4 = plot(xs[2:length(xs)-1], ys[2:length(ys)-1], u_sol, linetype=:contourf,title = "solution")
plottogether = plot(p1,p2,p3,p4)

#Animation (didn't work)
anim = @animate for i ∈ 1:length(sol.t)
    u_sol = reshape(sol.u[i], length(xs)-2,length(ys)-2)
    plot(xs[2:length(xs)-1], ys[2:length(ys)-1], u_sol, linetype=:contourf,title = "solution")

end
gif(anim, "anim_fps15.gif", fps = 15)


text:
function get_name(obj::TCBScale)
    return jcall(obj, "getName", JString, ())
end

function offset_from_tai(obj::TCBScale, arg0::AbsoluteDate)
    return jcall(obj, "offsetFromTAI", jdouble, (AbsoluteDate,), arg0)
end

function offset_from_tai(obj::TCBScale, arg0::FieldAbsoluteDate)
    return jcall(obj, "offsetFromTAI", RealFieldElement, (FieldAbsoluteDate,), arg0)
end

function to_string(obj::TCBScale)
    return jcall(obj, "toString", JString, ())
end


text:
"""
    Represents a tangent vector, with momentum `p=2π/N*k`, living in the tangent space of a puMPState `state`.
"""
mutable struct puMPSTvec{T}
    state::puMPState{T}
    B::MPSTensor{T}
    k::Float64 #p=2π/N*k
end

Base.copy(Tvec::puMPSTvec) = puMPSTvec(Tvec.state, copy(Tvec.B), Tvec.k, Tvec.N) #does not copy state
Base.eltype(Tvec::puMPSTvec{T}) where {T} = T

MPS.bond_dim(Tvec::puMPSTvec) = bond_dim(Tvec.state)
MPS.phys_dim(Tvec::puMPSTvec) = phys_dim(Tvec.state)
state(Tvec::puMPSTvec) = Tvec.state
tvec_tensor(Tvec::puMPSTvec) = Tvec.B
mps_tensors(Tvec::puMPSTvec) = (mps_tensor(Tvec.state), tvec_tensor(Tvec))
num_sites(Tvec::puMPSTvec) = num_sites(Tvec.state)
spin(Tvec::puMPSTvec) = Tvec.k
momentum(Tvec::puMPSTvec) = 2π/num_sites(Tvec) * spin(Tvec)

set_tvec_tensor!(Tvec::puMPSTvec{T}, B::MPSTensor{T}) where {T} = Tvec.B = B

"""
    Base.Vector(M::puMPState{T}) where T

Computes the full (dense) state vector from the (compressed) puMPS form.
Since puMPS typically represent state vectors that are far too big to fit
in memory, this should only be used for very small systems!
"""
function Base.Vector(Tvec::puMPSTvec{T}) where T
    N = num_sites(Tvec)
    d = phys_dim(Tvec)
    D = bond_dim(Tvec)
    A, B = mps_tensors(Tvec)
    p = momentum(Tvec)

    N == 0 && return Vector{T}()

    onlyAs = A
    res = B * cis(p)
    for j in 2:N
        @tensor res[l,p1,p2,r] := res[l,p1,ri] * A[ri,p2,r]
        Bp = B * cis(p*j)
        @tensor res[l,p1,p2,r] += onlyAs[l,p1,ri] * Bp[ri,p2,r]
        res = reshape(res, (D, d^j, D))
        if j < N
            @tensor onlyAs[l,p1,p2,r] := onlyAs[l,p1,ri] * A[ri,p2,r]
            onlyAs = reshape(onlyAs, (D, d^j, D))
        end
    end
    # res now has D^2 d^N elements!
    # TODO: Break the system into two parts so that the largest
    # object has max(D^2 * d^(N/2), d^N).
    @tensor res[p] := res[v,p,v]

    res
end

"""
    LinearAlgebra.norm(Tvec::puMPSTvec{T}) where {T}

Computes the norm of a puMPState tangent vector `Tvec`.
"""
function LinearAlgebra.norm(Tvec::puMPSTvec{T}) where {T}
    N = num_sites(Tvec)
    A, B = mps_tensors(Tvec)
    p = momentum(Tvec)

    #conj(B) at fixed location
    TBs = TM_dense(B, B)
    TA = TM_dense(A, B)

    TMres = similar(TA)
    work = workvec_applyTM_l(A, A, TA)
    for n in 2:N
        TBs = applyTM_l!(TBs, A, A, TBs, work)

        TAB = applyTM_l!(TMres, B, A, TA, work)
        axpy!(cis(p*n), TAB, TBs) #Add to complete terms (containing both B's)

        TA = applyTM_l!(TA, A, A, TA, work)
    end

    sqrt(N*tr(TBs))
end

"""
    LinearAlgebra.normalize!(Tvec::puMPSTvec{T}) where {T} = rmul!(Tvec.B, 1.0/norm(Tvec))

Normalizes a puMPState tangent vector `Tvec` in place.
"""
LinearAlgebra.normalize!(Tvec::puMPSTvec{T}) where {T} = rmul!(Tvec.B, 1.0/norm(Tvec))

"""
    LinearAlgebra.dot(Tvec::puMPSTvec, psi::AbstractVector)

Computes the inner product of a puMPSTvec with a state vector (dense representation).
"""
function LinearAlgebra.dot(Tvec::puMPSTvec, psi::AbstractVector)
    N = num_sites(Tvec)
    D = bond_dim(Tvec)
    d = phys_dim(Tvec)
    length(psi) == d^N || DimensionMismatch("")

    A, B = mps_tensors(Tvec)
    p = momentum(Tvec)

    #Project psi into momentum sector of Tvec
    psi_tmp = Vector{eltype(Tvec)}(undef, length(psi))
    Ppsi = similar(psi_tmp)
    copyto!(Ppsi, psi)
    rmul!(Ppsi, cis(-p))
    for n=1:N-1
        psi_r = reshape(psi, (d^(n), d^(N-n)))
        copyto!(psi_tmp, transpose(psi_r))
        axpy!(cis(-(n+1)*p), psi_tmp, Ppsi)
    end

    Ppsi_r = reshape(Ppsi, (d, d^(N-1)))
    @tensor con[r,phys,l] := conj(B[l,p1,r]) * Ppsi_r[p1,phys]
    for n=2:N
        con = reshape(con, (D, d, d^(N-n), D)) #r,p1,phys,l
        @tensor con[r,phys,l] := conj(A[li,p1,r]) * con[li,p1,phys,l]
    end
    con = reshape(con, (D,D))

    tr(con)
end

LinearAlgebra.dot(psi::AbstractVector, Tvec::puMPSTvec) = conj(dot(Tvec, psi))

"""
    expect(Tvec::puMPSTvec{T}, O::MPO_PBC_uniform{T}) where {T}

The expectation value of a translation-invariant operator specified as an MPO
with periodic boundary conditions, with respect to the puMPState tangent vector `Tvec`.
"""
function expect(Tvec::puMPSTvec{T}, O::MPO_PBC_uniform{T}) where {T}
    A, B = mps_tensors(Tvec)
    OB, OM = O

    N = num_sites(Tvec)
    p = momentum(Tvec)

    #conj(B) at fixed location...
    TBs = TM_dense_MPO(B, B, OB)
    TA = TM_dense_MPO(A, B, OB)

    TMres = similar(TA)
    work = workvec_applyTM_MPO_l(A, A, OM, TA)
    for n in 2:N
        TBs = applyTM_MPO_l!(TBs, A, A, OM, TBs, work)

        TAB = applyTM_MPO_l!(TMres, B, A, OM, TA, work)
        axpy!(cis(p*(n-1)), TAB, TBs) #Add to complete terms (containing both B's)

        n < N && (TA = applyTM_MPO_l!(TA, A, A, OM, TA, work))
    end

    N*tr(TBs)
end

#Slower version of the above
function expect_using_overlap(Tvec::puMPSTvec{T}, O::MPO_PBC_uniform{T}) where {T}
    OB, OM = O
    expect(Tvec, MPOTensor{T}[OB, (OM for j in 1:num_sites(Tvec)-1)...])
end

"""
    expect_global_product(Tvec::puMPSTvec{T}, O::Matrix{TO}) where {T,TO}

The expectation value of a global product of on-site operators `O` with
respect to the puMPState tangent vector `Tvec`.
"""
function expect_global_product(Tvec::puMPSTvec{T}, O::Matrix{TO}) where {T,TO}
    Or = copy(reshape(transpose(O), (1,size(O,2),1,size(O,1))))
    O_MPO = (Or,Or)
    expect(Tvec, O_MPO)
end

"""
    expect(Tvec::puMPSTvec{T}, O::MPO_open{T}) where {T}

The expectation value of an arbitrary MPO with respect to the puMPState tangent
vector `Tvec`.

Note: This is not the most efficient implementation, since it uses `overlap()` and
      does not take advantage of the two tangent vectors being the same.
"""
function expect(Tvec::puMPSTvec{T}, O::MPO_open{T}) where {T}
    overlap(Tvec, O, Tvec)
end

"""
    overlap(Tvec2::puMPSTvec{T}, O::MPO_open{T}, Tvec1::puMPSTvec{T};
        TAA_all::Vector{Array}=Array[], TBAs_all::Vector{Array}=Array[]) where {T}

Computes the physical-space inner product `<Tvec2 | O | Tvec1>`,
where `Tvec1` and `Tvec2` are puMPState tangent vectors and `O` is an operator in MPO form.
The tangent vectors need not live in the same puMPState tangent space, since the inner product
is taken in the physical space. They must have the same number of physical sites.

The operator MPO may have between 1 and `N = num_sites()` tensors.
In case it has `N` tensors, it may have a bond around the circle, between sites `N` and 1.
"""
function overlap(Tvec2::puMPSTvec{T}, O::MPO_open{T}, Tvec1::puMPSTvec{T};
    TAA_all::Vector{Array}=Array[], TBAs_all::Vector{Array}=Array[]) where {T}

    A1, B1 = mps_tensors(Tvec1)
    A2, B2 = mps_tensors(Tvec2)

    N = num_sites(Tvec1)
    @assert num_sites(Tvec2) == N

    p1 = momentum(Tvec1)
    p2 = momentum(Tvec2)

    #Fix site numbering with O at sites 1..length(O)

    TAA = length(TAA_all) > 0 ? TAA_all[1] : (length(O) > 0 ? TM_dense_MPO(A1, A2, O[1]) : TM_dense(A1, A2))

    TBBs = length(O) > 0 ? TM_dense_MPO(B1, B2, O[1]) : TM_dense(B1, B2)
    rmul!(TBBs, cis(p1-p2))

    TABs = length(O) > 0 ? TM_dense_MPO(A1, B2, O[1]) : TM_dense(A1, B2)
    rmul!(TABs, cis(-p2))

    if length(TBAs_all) > 0
        TBAs = TBAs_all[1]
    else
        TBAs = length(O) > 0 ? TM_dense_MPO(B1, A2, O[1]) : TM_dense(B1, A2)
        rmul!(TBAs, cis(p1))
    end

    work = Vector{T}()
    for n in 2:length(O)
        work = workvec_applyTM_MPO_l!(work, A1, A2, O[n], TAA)
        TMres = res_applyTM_MPO_l(A1, A2, O[n], TAA)

        TBBs = applyTM_MPO_l!(similar(TMres), A1, A2, O[n], TBBs, work)

        axpy!(cis(n*(p1-p2)), applyTM_MPO_l!(TMres, B1, B2, O[n], TAA, work), TBBs)
        axpy!(cis(n*(-p2)), applyTM_MPO_l!(TMres, A1, B2, O[n], TBAs, work), TBBs)
        axpy!(cis(n*(p1)), applyTM_MPO_l!(TMres, B1, A2, O[n], TABs, work), TBBs)

        if n < N
            TABs = applyTM_MPO_l!(similar(TMres), A1, A2, O[n], TABs, work)
            axpy!(cis(n*(-p2)), applyTM_MPO_l!(TMres, A1, B2, O[n], TAA, work), TABs)

            if length(TBAs_all) > 0
                TBAs = TBAs_all[n]
            else
                TBAs = applyTM_MPO_l!(similar(TMres), A1, A2, O[n], TBAs, work)
                axpy!(cis(n*(p1)), applyTM_MPO_l!(TMres, B1, A2, O[n], TAA, work), TBAs)
            end

            TAA = length(TAA_all) > 0 ? TAA_all[n] : applyTM_MPO_l!(similar(TMres), A1, A2, O[n], TAA, work)
        end
    end

    if length(O) < N
        TMres = similar(TAA)
        for n in max(2, length(O)+1):N
            work = workvec_applyTM_l!(work, A1, A2, TAA)

            TBBs = applyTM_l!(TBBs, A1, A2, TBBs, work)

            axpy!(cis(n*(p1-p2)), applyTM_l!(TMres, B1, B2, TAA, work), TBBs)
            axpy!(cis(n*(-p2)), applyTM_l!(TMres, A1, B2, TBAs, work), TBBs)
            axpy!(cis(n*(p1)), applyTM_l!(TMres, B1, A2, TABs, work), TBBs)

            if n < N
                TABs = applyTM_l!(TABs, A1, A2, TABs, work)
                axpy!(cis(n*(-p2)), applyTM_l!(TMres, A1, B2, TAA, work), TABs)

                if length(TBAs_all) > 0
                    TBAs = TBAs_all[n]
                else
                    TBAs = applyTM_l!(TBAs, A1, A2, TBAs, work)
                    axpy!(cis(n*(p1)), applyTM_l!(TMres, B1, A2, TAA, work), TBAs)
                end

                TAA = length(TAA_all) > 0 ? TAA_all[n] : applyTM_l!(TAA, A1, A2, TAA, work)
            end
        end
    end

    tr(TBBs)
end

function fidelity(Tvec2::puMPSTvec{T}, Tvec1::puMPSTvec{T}) where {T}
    N = num_sites(Tvec1)
    @assert num_sites(Tvec2) == N
    overlap(Tvec2, MPOTensor{T}[], Tvec1)
end

"""
    overlap_precomp_samestate(O::MPO_open{T}, Tvec1::puMPSTvec{T}) where {T}

Precompute some data for `overlap()`, useful for cases where many overlaps are computed
with tangent vectors living in the same tangent space.
"""
function overlap_precomp_samestate(O::MPO_open{T}, Tvec1::puMPSTvec{T}) where {T}
    A1, B1 = mps_tensors(Tvec1)
    p1 = momentum(Tvec1)
    N = num_sites(Tvec1)

    TAA_all = Array[]
    TBAs_all = Array[]

    TAA = TM_dense_MPO(A1, A1, O[1])
    push!(TAA_all, TAA)
    TBAs = TM_dense_MPO(B1, A1, O[1])
    rmul!(TBAs, cis(p1))
    push!(TBAs_all, TBAs)

    work = Vector{T}()
    for n in 2:min(length(O), N-1)
        work = workvec_applyTM_MPO_l!(work, A1, A1, O[n], TAA)
        TMres = res_applyTM_MPO_l(A1, A1, O[n], TAA)

        TBAs = applyTM_MPO_l!(similar(TMres), A1, A1, O[n], TBAs, work)
        axpy!(cis(n*(p1)), applyTM_MPO_l!(TMres, B1, A1, O[n], TAA, work), TBAs)
        push!(TBAs_all, TBAs)

        TAA = applyTM_MPO_l!(similar(TMres), A1, A1, O[n], TAA, work)
        push!(TAA_all, TAA)
    end

    if length(O) < N
        TMres = similar(TAA)
        for n in length(O)+1:N-1
            work = workvec_applyTM_l!(work, A1, A1, TAA)

            TBAs = applyTM_l!(TBAs, A1, A1, TBAs, work)
            axpy!(cis(n*(p1)), applyTM_l!(TMres, B1, A1, TAA, work), TBAs)
            push!(TBAs_all, TBAs)

            TAA = applyTM_l!(TAA, A1, A1, TAA, work)
            push!(TAA_all, TAA)
        end
    end

    TAA_all, TBAs_all
end

"""
    op_in_basis(Tvec_basis::Vector{puMPSTvec{T}}, O::MPO_open{T}) where {T}

Compute the matrix elements of the operator `O` in the basis of tangent vectors `Tvec_basis`.

CAUTION: Possibly broken?
"""
function op_in_basis(Tvec_basis::Vector{puMPSTvec{T}}, O::MPO_open{T}) where {T}
    Nvec = length(Tvec_basis)

    M = state(Tvec_basis[1])
    @assert all(state(tv) === M for tv in Tvec_basis)

    O_in_basis = zeros(T, (Nvec, Nvec))
    for j in 1:Nvec
        #Increases memory usage by approx. 2*N*D^4
        TAA_all, TBAs_all = overlap_precomp_samestate(O, Tvec_basis[j])
        for k in 1:Nvec
            O_in_basis[k,j] = overlap(Tvec_basis[k], O, Tvec_basis[j]; TAA_all=TAA_all, TBAs_all=TBAs_all)
        end
    end
    O_in_basis
end

"""
    tangent_space_metric(M::puMPState{T}, ks::Vector{<:Real}, blkTMs::Vector{MPS_MPO_TM{T}}, L::Integer=1) where {T}

Computes the physical metric induced on the tangent space of the puMPState `M` for the tangent-space
momentum sectors specified in `ks`. The momenta are `ps = 2π/N .* ks` where `N = num_sites(M)`.
"""
function tangent_space_metric(M::puMPState{T}, ks::Vector{<:Real}, blkTMs::Vector{MPS_MPO_TM{T}}, L::Integer=1) where {T}
    #Time cost O(N*d^2*D^6).. could save a factor of dD by implementing the sparse mat-vec operation, but then we'd add a factor Nitr
    #space cost O(N)
    A = mps_tensor(M)
    N = num_sites(M)
    d = phys_dim(M)
    D = bond_dim(M)
    ps = 2π/N .* ks

    Gs = Array{T,6}[zeros(T, (D,d,D, D,d,D)) for j in 1:length(ks)]

    Gpart = zeros(T, (D,d,D, D,d,D))

    #Add I on the physical index. This is the same-site term.
    let E = Matrix{T}(I,d^L,d^L), blk = TM_convert(blkTMs[N-L])
        GC.enable(false)
        @tensor Gpart[V1b,Pb,V2b, V1t,Pt,V2t] = E[Pt,Pb] * blk[V2t,V2b, V1t,V1b]
        GC.enable(true)
        GC.gc(false)
        for j in 1:length(ks)
            axpy!(N, Gpart, Gs[j])
        end
    end

    #This deals with the terms where the gaps overlap or are neighbours.
    Ablk = A
    @show L
    for j in 1:L
        E = Matrix{T}(I,d^(L-j),d^(L-j))
        blk = TM_convert(blkTMs[N-L-j])

        GC.enable(false)
        @tensor Gpart_pre[V1b,Pb,V2b, V1t,Pt,V2t] := Ablk[vt,Pb,V1t] * conj(Ablk[V2b,Pt,vb]) * blk[V2t,vb, vt,V1b]
        GC.enable(true)
        GC.gc(false)

        Gpart_r = reshape(Gpart, (D,d^j,d^(L-j),D, D,d^(L-j),d^j,D))
        GC.enable(false)
        @tensor Gpart_r[V1b,Pb1,Pb2,V2b, V1t,Pt1,Pt2,V2t] = Gpart_pre[V1b,Pb1,V2b, V1t,Pt2,V2t] * E[Pb2,Pt1]
        GC.enable(true)
        GC.gc(false)

        for k in 1:length(ps)
            axpy!(N*cis(ps[k]*L), Gpart, Gs[k])
        end

        GC.enable(false)
        @tensor Gpart_pre[V1b,Pb,V2b, V1t,Pt,V2t] = Ablk[V2t,Pb,vt] * blk[vt,V2b,V1t,vb] * conj(Ablk[vb,Pt,V1b])
        GC.enable(true)
        GC.gc(false)

        Gpart_r = reshape(Gpart, (D,d^(L-j),d^j,D, D,d^j,d^(L-j),D))
        GC.enable(false)
        @tensor Gpart_r[V1b,Pb2,Pb1,V2b, V1t,Pt2,Pt1,V2t] = Gpart_pre[V1b,Pb1,V2b, V1t,Pt2,V2t] * E[Pb2,Pt1]
        GC.enable(true)
        GC.gc(false)

        for k in 1:length(ps)
            axpy!(N*cis(ps[k]*(N-L)), Gpart, Gs[k])
        end

        if j < L
            @tensor Ablk[V1, P1,P2, V2] := Ablk[V1, P1, v] * A[v, P2, V2]
            Ablk = reshape(Ablk, (D, d^j, D))
        end
    end

    #Now we handle terms where the gaps are separated: We need two block transfer matrices.
    left_T = zeros(T, (D, D,D, D,d))
    right_T = zeros(T, (D,d,D, D,D))
    for i in 2L:N-2L
        blk = TM_convert(blkTMs[i-1])
        GC.enable(false)
        @tensor left_T[V1b, V2t,V2b, V1t,Pt] = Ablk[V1t,Pt,vt] * blk[vt,V1b, V2t,V2b]
        GC.enable(true)
        GC.gc(false)

        blk = TM_convert(blkTMs[N-i-1])
        GC.enable(false)
        @tensor right_T[V1t,Pb,V1b, V2t,V2b] = conj(Ablk[V1b,Pb,vb]) * blk[V1t,vb, V2t,V2b]
        GC.enable(true)
        GC.gc(false)

        GC.enable(false)
        @tensor Gpart[V1b,Pb,V2b, V1t,Pt,V2t] = left_T[V2b, V1t,vb, vt,Pb] * right_T[V2t,Pt,vb, vt,V1b] #complete loop, cost O(d^2 * D^6)
        GC.enable(true)
        GC.gc(false)
        for j in 1:length(ps)
            axpy!(N*cis(ps[j]*i), Gpart, Gs[j])
        end
    end

    Gs
end

"""
    tangent_space_MPO!(op_effs::Vector{Array{T,6}}, M::puMPState{T}, op::MPO_open{T}, ks::Vector{<:Real}) where {T}

Compute the effective operator on the tangent space parameters, in the momentum sectors specified in `ks`, of the physical operator
`op`, given as an MPO. This version adds the result to existing tangent-space operators `op_effs`.
It is assumed that the operator is translation invariant (even if the MPO representation is not), or is part
of a sum of operators that are translation invariant.
This assumption is used as follows: We fix the first tensor of the MPO to always be located on the same site as the
`B` tensor of one of the tangent vectors. We then sum over all locations of the `B` tensor of the other tangent vector.
This removes N-1 redundant calculations.
"""
function tangent_space_MPO!(op_effs::Vector{Array{T,6}}, M::puMPState{T}, op::MPO_open{T}, ks::Vector{<:Real}) where {T}
    A = mps_tensor(M)
    N = num_sites(M)
    ps = 2π/N .* ks

    #First construct block transfer matrices. A conjugate gap at site 1, where the Hamiltonian begins.
    #This is essentially an MPO TM, but with a physical index on the left instead of an MPO index.
    op_1 = op[1]
    op_n = op[2]
    GC.enable(false)
    @tensor blk_l[V1t,P,M1,V1b, V2t,M2,V2b] := ((((A[V1t,p1,vt] * op_1[M1,p1,m,P])
                                             * A[vt,p2t,V2t]) * op_n[m,p2t,M2,p2b]) * conj(A[V1b,p2b,V2b]))
    GC.enable(true)
    GC.gc(false)

    #We also need block TM's for the latter part of the Hamiltonian. We will add the gap on the left later.
    #We precompute these right blocks, since we will use the intermediates
    blks_r = MPS_MPO_TM{T}[TM_dense_MPO(A, A, op[N])] #site N
    work = Vector{T}()
    for n in N-1:-1:2
        work = workvec_applyTM_MPO_r!(work, A, A, op[n], blks_r[end])
        res = res_applyTM_MPO_r(A, A, op[n], blks_r[end])
        push!(blks_r, applyTM_MPO_r!(res, A, A, op[n], blks_r[end], work))
    end

    #Same-site term
    blk_r = pop!(blks_r) #block for N-1 sites
    op_n = op[1]
    GC.enable(false)
    @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] := op_n[m2,Pt,m1,Pb] * blk_r[V2t,m1,V2b, V1t,m2,V1b]
    GC.enable(true)
    GC.gc(false)
    for j in 1:length(ks)
        axpy!(N, part, op_effs[j])
    end

    #Nearest-neighbour term with conjugate gap on the left, gap on the right. Conjugate gap is site 1.
    blk_r = pop!(blks_r) #block for N-2 sites
    op_1 = op[1]; op_n = op[2]
    GC.enable(false)
    @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = (op_1[m3,pt,m1,Pb] * (A[vt,pt,V1t] *
                                                (op_n[m1,Pt,m2,pb] * (conj(A[V2b,pb,vb]) * blk_r[V2t,m2,vb, vt,m3,V1b]))))
    GC.enable(true)
    GC.gc(false)
    for j in 1:length(ks)
        axpy!(N*cis(ps[j]), part, op_effs[j])
    end

    #Terms separated by one or more sites
    for n in 3:N-1 #if conj(gap) is at site 1, gap is at site n
        blk_r = pop!(blks_r) #N-n sites
        op_n = op[n]
        GC.enable(false)
        @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = blk_l[vt,Pb,m3,V2b, V1t,m1,vb1] *
                                                  (op_n[m1,Pt,m2,pb] * (conj(A[vb1,pb,vb2]) * blk_r[V2t,m2,vb2, vt,m3,V1b]))
        GC.enable(true)
        GC.gc(false)
        for j in 1:length(ks)
            axpy!(N*cis(ps[j]*(n-1)), part, op_effs[j])
        end

        #Reshape blk_l into an MPO transfer matrix
        blsz = size(blk_l)
        blk_l_r = reshape(blk_l, (blsz[1],blsz[2]*blsz[3],blsz[4:end]...))

        #If the dimension of o[n] is not the same everywhere, we need to allocate a new result array here.
        #Otherwise we can modify blk_l_r in place.
        res = size(op[n],3) == blsz[6] ? blk_l_r : res_applyTM_MPO_l(A, A, op[n], blk_l_r)

        work = workvec_applyTM_MPO_l!(work, A, A, op[n], blk_l_r)
        applyTM_MPO_l!(res, A, A, op[n], blk_l_r, work) #extend blk_l by multiplying from the right by an MPO TM
        blk_l = reshape(res, (blsz[1],blsz[2],blsz[3],blsz[4], blsz[5],size(op[n],3),blsz[7]))
    end

    #Nearest-neighbour term with gap on the left, conjugate gap on the right. Conjugate gap is site 1.

    op_N = op[N]
    GC.enable(false)
    @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = (blk_l[V2t,Pb,m2,V2b, V1t,m1,vb] * conj(A[vb,p,V1b])) * op_N[m1,Pt,m2,p]
    GC.enable(true)
    GC.gc(false)
    for j in 1:length(ks)
        axpy!(N*cis(ps[j]*(N-1)), part, op_effs[j])
    end

    op_effs
end

#This will return the transfer matrix (TM) from site nL to site nR, given a boundary Hamiltonian centred at the
#boundary between sites N and 1.
function blockTM_MPO_boundary(M::puMPState{T}, op_b::MPO_open{T}, blkTMs::Vector{MPS_MPO_TM{T}}, nL::Integer, nR::Integer) where {T}
    A = mps_tensor(M)
    N = num_sites(M)
    Nop = length(op_b)

    @assert N > Nop
    @assert nL <= nR

    #Split op_b into left (starting at site 1) and right (ending at site N) parts
    op_b_L = op_b[Nop÷2+1:end]
    op_b_R = op_b[1:Nop÷2]

    #Illustration for four-site boundary Hamiltonian:
    #
    # <-1-2-3-4-.-.-.-.-.-.-.-.-.-.-.-.-N->    state ket
    #   | | | |                     | | |
    # <-o-o | |        ...          | o-o->    operator
    #   | | | |                     | | |
    # <-1-2-3-4-.-.-.-.-.-.-.-.-.-.-.-.-N->    state bra
    #
    # Just the operator:
    #   | |                           | |
    # <-o-o                           o-o->
    #   | |                           | |
    # op_b_L                          op_b_R

    NopL = length(op_b_L)
    NopR = length(op_b_R)

    nL_in_op = nL - (N - NopR)
    nR_in_op = nR - (N - NopR)

    Nmid = min(nR, N-NopR) - max(nL, NopL + 1) + 1 #length of middle TM (no h_b terms)

    if Nmid > 0
        blk = blkTMs[Nmid]

        #These will add MPO terms if needed:
        #Note: The index ranges may be empty, in which case applyTM_MPO_x() returns blk unmodified.
        blk = applyTM_MPO_r(M, op_b_L[nL:end], blk)
        blk = applyTM_MPO_l(M, op_b_R[1:nR_in_op], blk)
    elseif nR <= length(op_b_L)
        blk = TM_dense_MPO(A, A, op_b_L[nR])
        blk = applyTM_MPO_r(M, op_b_L[nL:nR-1], blk)
    else
        blk = TM_dense_MPO(A, A, op_b_R[nL_in_op])
        blk = applyTM_MPO_l(M, op_b_R[nL_in_op+1:nR_in_op], blk)
    end

    blk
end

#This will return the boundary MPO tensor for site n. Where the MPO does not act on site n,
#return a zero-length result.
function op_term_MPO_boundary(M::puMPState{T}, op_b::MPO_open{T}, n::Integer) where {T}
    A = mps_tensor(M)
    N = num_sites(M)
    Nop = length(op_b)

    @assert N > Nop

    op_b_L = op_b[Nop÷2+1:end]
    op_b_R = op_b[1:Nop÷2]

    NopL = length(op_b_L)
    NopR = length(op_b_R)

    if n <= length(op_b_L)
        return op_b_L[n]
    elseif n >= N-NopR+1
        return op_b_R[n-(N-NopR)]
    else
        return MPOTensor{T}(undef, (0,0,0,0)) #length zero MPO tensor
    end
end

"""
    tangent_space_MPO_boundary!(op_effs::Vector{Array{T,6}}, M::puMPState{T}, op_b::MPO_open{T},
        ks::Vector{<:Real}, blkTMs::Vector{MPS_MPO_TM{T}}) where {T}

Compute the boundary part of the effective MPO in the momentum sectors specified in `ks`.
See `tangent_space_metric_and_MPO()`.

This is very similar to the OBC MPO case, except that we now have a
short MPO `op_b` centred on the boundary between sites N and 1.
"""
function tangent_space_MPO_boundary!(op_effs::Vector{Array{T,6}}, M::puMPState{T}, op_b::MPO_open{T},
                                        ks::Vector{<:Real}, blkTMs::Vector{MPS_MPO_TM{T}}) where {T}
    A = mps_tensor(M)
    N = num_sites(M)
    ps = 2π/N .* ks

    Nh = length(op_b)
    @assert iseven(Nh)

    #Check this is really OBC
    @assert size(op_b[1], 1) == 1
    @assert size(op_b[end], 3) == 1

    #On-site term
    blk = blockTM_MPO_boundary(M, op_b, blkTMs, 2, N)
    op_b_1 = op_term_MPO_boundary(M, op_b, 1)
    GC.enable(false)
    @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] := blk[V2t,m2,V2b, V1t,m1,V1b] * op_b_1[m1,Pt,m2,Pb]
    GC.enable(true)
    GC.gc(false)
    for j in 1:length(ks)
        axpy!(N, part, op_effs[j])
    end

    #NN-term 1, with conjugate gap (site 1) then gap
    blk = blockTM_MPO_boundary(M, op_b, blkTMs, 3, N)
    op_b_n = op_term_MPO_boundary(M, op_b, 2)
    GC.enable(false)
    if length(op_b_n) == 0
        @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = op_b_1[m2,pt,m1,Pb] * (A[vt,pt,V1t] * (conj(A[V2b,Pt,vb])
                * blk[V2t,m1,vb, vt,m2,V1b]))
    else
        @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = op_b_1[m3,pt,m1,Pb] *
        (A[vt,pt,V1t] * (op_b_n[m1,Pt,m2,pb] * (conj(A[V2b,pb,vb]) * blk[V2t,m2,vb, vt,m3,V1b])))
    end
    GC.enable(true)
    GC.gc(false)
    for j in 1:length(ks)
        axpy!(N*cis(ps[j]), part, op_effs[j])
    end

    #gap at site n
    for n in 3:N-1
        blkL = blockTM_MPO_boundary(M, op_b, blkTMs, 2, n-1)
        GC.enable(false)
        @tensor blkL_cgap[V1t,M1,P,V1b, V2t,M2,V2b] := op_b_1[M1,p,m,P] * (A[V1t,p,vt] * blkL[vt,m,V1b, V2t,M2,V2b])
        GC.enable(true)
        GC.gc(false)

        blkR = blockTM_MPO_boundary(M, op_b, blkTMs, n+1, N)
        op_b_n = op_term_MPO_boundary(M, op_b, n)
        GC.enable(false)
        if length(op_b_n) == 0
            @tensor blkR_gap[V1t,P,M1,V1b, V2t,M2,V2b] := conj(A[V1b,P,vb]) * blkR[V1t,M1,vb, V2t,M2,V2b]
        else
            #Note: This is never called for a nearest-neighbour boundary Hamiltonian
            @tensor blkR_gap[V1t,P,M1,V1b, V2t,M2,V2b] := op_b_n[M1,P,m,p] * (conj(A[V1b,p,vb]) * blkR[V1t,m,vb, V2t,M2,V2b])
        end
        GC.enable(true)
        GC.gc(false)

        GC.enable(false)
        @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = blkL_cgap[vt,m2,Pb,V2b, V1t,m1,vb] * blkR_gap[V2t,Pt,m1,vb, vt,m2,V1b]
        GC.enable(true)
        GC.gc(false)
        for j in 1:length(ks)
            axpy!(N*cis(ps[j]*(n-1)), part, op_effs[j])
        end
    end

    #NN-term 2, with gap (site N) then conjugate gap (site 1)
    blk = blockTM_MPO_boundary(M, op_b, blkTMs, 2, N-1)
    op_b_n = op_term_MPO_boundary(M, op_b, N)
    GC.enable(false)
    if length(op_b_n) == 0
        @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = ((op_b_1[m2,pt,m1,Pb] * (A[V2t,pt,vt] * blk[vt,m1,V2b, V1t,m2,vb])) *
            conj(A[vb,Pt,V1b]))
    else
        @tensor part[V1b,Pb,V2b, V1t,Pt,V2t] = ((op_b_1[m3,pt,m1,Pb] * (A[V2t,pt,vt] * blk[vt,m1,V2b, V1t,m2,vb])) *
            conj(A[vb,pb,V1b])) * op_b_n[m2,Pt,m3,pb]
    end
    GC.enable(true)
    GC.gc(false)
    for j in 1:length(ks)
        axpy!(N*cis(ps[j]*(N-1)), part, op_effs[j])
    end

    op_effs
end

function tangent_space_MPO!(op_effs::Vector{Array{T,6}}, M::puMPState{T}, op::MPO_PBC_uniform{T}, ks::Vector{<:Real}) where {T}
    N = num_sites(M)

    op_boundary, op_bulk = op
    op_full = MPOTensor{T}[op_boundary, (op_bulk for j in 1:N-1)...]

    tangent_space_MPO!(op_effs, M, op_full, ks)
end

function tangent_space_MPO(M::puMPState{T}, op::MPO_PBC_uniform{T}, ks::Vector{<:Real}) where {T}
    D = bond_dim(M)
    d = phys_dim(M)
    N = num_sites(M)

    op_effs = Array{T,6}[zeros(T, (D,d,D, D,d,D)) for j in 1:length(ks)]
    tangent_space_MPO!(op_effs, M, op, ks)
end

function tangent_space_MPO!(op_effs::Vector{Array{T,6}}, M::puMPState{T}, op::MPO_open_uniform{T}, ks::Vector{<:Real}) where {T}
    N = num_sites(M)
    opL, opM, opR = op

    op_full = MPOTensor{T}[opL, (opM for j in 1:N-2)..., opR]

    tangent_space_MPO!(op_effs, M, op_full, ks)
end

"""
    tangent_space_MPO(M::puMPState{T}, op::Union{MPO_PBC_uniform_split{T}, MPO_PBC_split{T}}, ks::Vector{<:Real}) where {T}

Given an MPO representation of an operator `op`, split into a large OBC MPO and a boundary
MPO, this computes its representation on the tangent space of the puMPState `M` for the tangent
vector momenta (of the ket tangent vector) specified in `ks`.
The sum of the OBC and boundary parts of the operator must be translation invariant (up to a phase), which we exploit
to reduce the cost of the computation.
"""
function tangent_space_MPO(M::puMPState{T}, op::Union{MPO_PBC_uniform_split{T}, MPO_PBC_split{T}}, ks::Vector{<:Real}) where {T}
    A = mps_tensor(M)
    D = bond_dim(M)
    d = phys_dim(M)
    N = num_sites(M)

    op_bulk, op_boundary = op

    op_effs = Array{T,6}[zeros(T, (D,d,D, D,d,D)) for j in 1:length(ks)]

    @time tangent_space_MPO_boundary!(op_effs, M, op_boundary, ks, blockTMs(M))
    blkTMs = nothing #free up this space
    @time tangent_space_MPO!(op_effs, M, op_bulk, ks)

    op_effs
end

function tangent_space_MPO(M::puMPState{T}, op::MPO_open{T}, ks::Vector{<:Real}) where {T}
    A = mps_tensor(M)
    D = bond_dim(M)
    d = phys_dim(M)
    N = num_sites(M)

    op_effs = Array{T,6}[zeros(T, (D,d,D, D,d,D)) for j in 1:length(ks)]
    @time tangent_space_MPO!(op_effs, M, op, ks)

    op_effs
end

function tangent_space_MPO(M::puMPState{T}, op::MPO_open_uniform{T}, ks::Vector{<:Real}) where {T}
    N = num_sites(M)

    op_full = MPOTensor{T}[op[1], (op[2] for j in 1:N-2)..., op[3]]

    tangent_space_MPO(M, op_full, ks)
end

"""
    tangent_space_metric_and_MPO(M::puMPState{T}, op::MPO_PBC_uniform{T}, ks::Vector{<:Real}, lambda_i::Matrix{T}) where {T}

Computes the tangent-space metric and effective operator given a physical-space operator as a periodic, uniform MPO,
for the momentum sectors specified in `ks`. The momenta are `ps = 2π/N .* ks` where `N = num_sites(M)`.
"""
function tangent_space_metric_and_MPO(M::puMPState{T}, op::MPO_PBC_uniform{T}, ks::Vector{<:Real}, lambda_i::Matrix{T}) where {T}
    @time Gs = tangent_space_metric(M, ks, blockTMs(M))
    @time op_effs = tangent_space_MPO(M, op, ks)

    Gs, op_effs
end

"""
    tangent_space_metric_and_MPO(M::puMPState{T}, op::Union{MPO_PBC_uniform_split{T}, MPO_PBC_split{T}},
        ks::Vector{<:Real}, lambda_i::Matrix{T}) where {T}

Computes the tangent-space metric and effective operator given a physical-space operator as a combination
of an open MPO, representing for example the Hamiltonian with open boundary conditions (OBC), and a boundary MPO.
The sum of the OBC and boundary parts of the operator must be translation invariant (up to a phase), which we exploit
to reduce the cost of the computation.
Metrics and effective Hamiltonians are computed for the momentum sectors specified in `ks`.
The momenta are `ps = 2π/N .* ks` where `N = num_sites(M)`.
"""
function tangent_space_metric_and_MPO(M::puMPState{T}, op::Union{MPO_PBC_uniform_split{T}, MPO_PBC_split{T}},
                                                      ks::Vector{<:Real}, lambda_i::Matrix{T}) where {T}
    A = mps_tensor(M)
    D = bond_dim(M)
    d = phys_dim(M)
    N = num_sites(M)

    bTMs = blockTMs(M) #requires N * D^4 bytes
    @time Gs = tangent_space_metric(M, ks, bTMs)

    op_bulk, op_boundary = op

    op_effs = Array{T,6}[zeros(T, (D,d,D, D,d,D)) for j in 1:length(ks)]
    @time tangent_space_MPO_boundary!(op_effs, M, op_boundary, ks, bTMs)
    blkTMs = nothing #free up this space
    @time tangent_space_MPO!(op_effs, M, op_bulk, ks)

    Gs, op_effs
end

function tspace_ops_to_center_gauge!(ops::Vector{Array{T,6}}, lambda_i::Matrix{T}) where {T}
    for op in ops
        GC.enable(false)
        @tensor op[V1b,Pb,V2b, V1t,Pt,V2t] = lambda_i[vb,V2b] * (op[V1b,Pb,vb, V1t,Pt,vt] * lambda_i[vt,V2t])
        GC.enable(true)
        GC.gc(false)
    end
    ops
end

"""
    excitations!(M::puMPState{T}, H::Union{MPO_PBC_uniform{T}, MPO_PBC_uniform_split{T}}, ks::AbstractVector{<:Real}, num_states::AbstractVector{<:Integer}; pinv_tol::Real=1e-10) where {T}

Computes eigenstates of the effective Hamiltonian obtained by projecting `H` onto the tangent space of the puMPState `M`.
This is done in the momentum sectors specified in `ks`, where each entry of `k = ks[j]` specified a momentum `k*2pi/N`,
where `N` is the number of sites.

The number of eigenstates to be computed for each momentum sector is specified in `num_states`.

The function returns a list of energies, a list of momenta (entries of `ks`), and a list of normalized tangent vectors.
"""
function excitations!(M::puMPState{T}, H::Union{MPO_PBC_uniform{T}, MPO_PBC_uniform_split{T}}, ks::AbstractVector{<:Real}, num_states::AbstractVector{<:Integer}; pinv_tol::Real=1e-10) where {T}
    M, lambda, lambda_i = canonicalize_left!(M)
    lambda_i = Matrix(lambda_i)

    @time Gs, Heffs = tangent_space_metric_and_MPO(M, H, ks, lambda_i)
    tspace_ops_to_center_gauge!(Gs, lambda_i)
    tspace_ops_to_center_gauge!(Heffs, lambda_i)

    excitations(M, Gs, Heffs, lambda_i, ks, num_states, pinv_tol=pinv_tol)
end

function excitations(M::puMPState{T}, Gs::Vector{Array{T,6}}, Heffs::Vector{Array{T,6}}, lambda_i::Matrix{T},
                        ks::Vector{Tk}, num_states::Vector{<:Integer}; pinv_tol::Real=1e-12) where {T,Tk<:Real}
    D = bond_dim(M)
    d = phys_dim(M)
    A = mps_tensor(M)
    Bshp = mps_tensor_shape(d, D)
    par_dim = prod(Bshp)

    ens = Vector{T}[]
    exs = Vector{puMPSTvec{T}}[]
    ks_rep = Vector{Tk}[]
    for j in 1:length(ks)
        println("k=$(ks[j])")
        G = reshape(Gs[j], (par_dim, par_dim))
        Heff = reshape(Heffs[j], (par_dim, par_dim))

        #Force Hermiticity
        G = (G + G') / 2
        Heff = (Heff + Heff') / 2

        # A dense pinv seems to avoid problems with spurious near-zero eigenvalues
        @time GiH = pinv(G, pinv_tol) * Heff

        @time ev, eV, info = eigsolve(GiH, vec(A), num_states[j], :LM)
        info.converged < num_states[j] && @warn "$(num_states[j] - info.converged) eigenvectors did not converge after $(info.numiter) iterations."

        exs_k = puMPSTvec{T}[]
        @time for v in eV
            Bnrm = sqrt(dot(v, G * v)) #We can use G to compute the norm
            rmul!(v, 1.0/Bnrm)
            Bc_mat = reshape(v, (D*d, D))
            Bl = Bc_mat * lambda_i #We assume G and Heff were in the center gauge.
            Bl = reshape(Bl, Bshp)
            push!(exs_k, puMPSTvec{T}(M, Bl, ks[j]))
        end

        push!(ens, ev)
        push!(exs, exs_k)
        push!(ks_rep, repeat(ks[j:j], inner=length(ev)))
    end

    vcat(ens...), vcat(ks_rep...), vcat(exs...)
end

"""
    Hn_in_basis(M::puMPState{T}, Hn_split::Tuple{Integer, MPO_PBC_split{T}}, Tvec_basis::Vector{puMPSTvec{T}}, ks::Vector{<:Real}) where {T}

Given an MPO representation of a Hamiltonian Fourier mode, split into a large OBC MPO and a boundary
MPO, `Hn_split`, computes its matrix elements in the basis of puMPState tangent vectors `Tvec_basis`
which are assumed to live in the tangent space of the puMPState `M`.
"""
function Hn_in_basis(M::puMPState{T}, Hn_split::Tuple{Integer, MPO_PBC_split{T}}, Tvec_basis::Vector{puMPSTvec{T}}, ks::Vector{<:Real}) where {T}
    #Hn is a Fourier mode of the Hamiltonian density with "momentum" n * 2pi/N
    #Each entry in ks, times 2pi/N is the momentum of one of the excitations (the ket).
    #Due to the way tangent_space_MPO() works, the momentum of the other excitation is automatically (ks[j] + n) * 2pi/N.
    ks = 1.0 * ks

    n = Hn_split[1]
    Hn_effs = tangent_space_MPO(M, Hn_split[2], ks)
    Ntvs = length(Tvec_basis)
    Tvspins = map(spin, Tvec_basis)

    Hn_in_basis = zeros(T, (Ntvs,Ntvs))
    for j in 1:Ntvs
        ind = findfirst(k->k==Tvspins[j], ks)
        if ind !== nothing
            Bj = tvec_tensor(Tvec_basis[j])
            Hn_eff = reshape(Hn_effs[ind], (length(Bj), length(Bj)))
            for k in 1:Ntvs
                if Tvspins[k] == Tvspins[j] + n
                    Bk = tvec_tensor(Tvec_basis[k])
                    Hn_in_basis[k,j] = dot(vec(Bk), Hn_eff * vec(Bj))
                end
            end
        end
    end

    Hn_in_basis
end

max_stars_repo_name: FuTen/puMPS

max_issues_repo_name: FuTen/puMPS

max_forks_repo_name: FuTen/puMPS


text:
module PortfolioOpt

using Distributions
using JuMP
using LinearAlgebra
using LinearAlgebra: dot
import Reexport

include("formulations.jl")
include("mean_variance.jl")
include("mean_variance_robust.jl")
include("mean_variance_dro.jl")
include("sample_based.jl")
include("sample_based_robust.jl")
include("sample_based_stochastic.jl")
include("simple_rules.jl")
include("utils.jl")
include("forecasts.jl")
include("test_utils/testutils.jl")

export AbstractPortfolioFormulation,
    AbstractMeanVariance,
    AbstractSampleBased,
    base_model,
    compute_solution,
    MeanVariance,
    mixed_signals_predict_return,
    portfolio_return!,
    portfolio_variance!,
    po_max_conditional_expectation_limit_predicted_return,
    po_max_predicted_return_limit_conditional_expectation,
    po_max_predicted_return_limit_return,
    po_max_return_limit_variance,
    po_max_utility_return,
    po_min_variance_limit_return,
    predicted_portfolio_return!,
    predicted_portfolio_variance!,
    readjust_volumes!,
    RobustBenTal,
    RobustBertsimas,
    RobustDelague,
    RobustBetina,
    SampleBased,
    # end-to-end
    mean_variance_noRf_analytical,
    max_sharpe,
    equal_weights

    Reexport.@reexport using JuMP

end


text:
function μ(E0::Real, K1::Real, K2::Real, cϕ::Real, sϕ::Real, cθ::Real, sθ::Real)
  n̂i = n̂(cϕ, sϕ, cθ, sθ);
  return ((K1 - K2) * E0 * cθ * n̂i) + (K2 * [0.0; 0.0; E0]);
end

max_stars_repo_name: grasingerm/polymer-stats

max_issues_repo_name: grasingerm/polymer-stats

max_forks_repo_name: grasingerm/polymer-stats


text:
function enclose(obj::Encloser, arg0::Iterable)
    return jcall(obj, "enclose", EnclosingBall, (Iterable,), arg0)
end


text:
import Base: hcat, vcat

module Locate
    import Arrow
    using Arrow: ALIGNMENT, DefaultOffset

    abstract type Abstract end

    function checkalignment(loc::Integer)
        if (loc-1) % ALIGNMENT ≠ 0
            throw(ArgumentError("Improperly aligned sub-buffer locator value: $loc"))
        end
    end

    containertype(::Type{J}, v) where {J} = Arrow.Primitive
    containertype(::Type{Union{J,Missing}}, v) where {J} = Arrow.NullablePrimitive
    containertype(::Type{String}, v) = Arrow.List
    containertype(::Type{Union{String,Missing}}, v) = Arrow.NullableList
    containertype(::Type{Bool}, v) = Arrow.BitPrimitive
    containertype(::Type{Union{Bool,Missing}}, v) = Arrow.NullableBitPrimitive

    length(v) = Base.length(v)

    function values end
    function valueslength end
    function bitmask end
    function offsets end

    struct Values{J} <: Abstract
        loc::Int64
        len::Int64
        Values{J}(loc::Integer, len::Integer) where {J} = (checkalignment(loc); new{J}(loc, len))
    end
    Values{J}(x) where {J} = Values{J}(values(x), valueslength(x))
    Values(::Type{String}, x) = Values{UInt8}(values(x), valueslength(x))
    Values(::Type{Union{String,Missing}}, x) = Values{UInt8}(values(x), valueslength(x))
    struct Bitmask <: Abstract
        loc::Int64
        Bitmask(loc::Integer) = (checkalignment(loc); new(loc))
    end
    Bitmask(x) = Bitmask(bitmask(x))
    struct Offsets{J<:Integer} <: Abstract
        loc::Int64
        Offsets{J}(loc::Integer) where {J} = (checkalignment(loc); new{J}(loc))
    end
    Offsets{J}(x) where {J} = Offsets{J}(offsets(x))
    Offsets(x) = Offsets{DefaultOffset}(offsets(x))
end

# type constructor methods
function locate(data::Vector{UInt8}, vals::Locate.Values{J}, len::Integer) where J
    Primitive{J}(data, vals.loc, len)
end
function locate(data::Vector{UInt8}, bmask::Locate.Bitmask, vals::Locate.Values{J}, len::Integer) where J
    NullablePrimitive{J}(data, bmask.loc, vals.loc, len)
end
function locate(data::Vector{UInt8}, ::Type{J}, offs::Locate.Offsets{K}, vals::Locate.Values{C},
                len::Integer) where {J,K<:Integer,C}
    List{J,K}(data, offs.loc, vals.loc, len, C, vals.len)
end
function locate(data::Vector{UInt8}, ::Type{Union{J,Missing}}, bmask::Locate.Bitmask,
                offs::Locate.Offsets{K}, vals::Locate.Values{C}, len::Integer) where {J,K<:Integer,C}
    NullableList{J,K}(data, bmask.loc, offs.loc, vals.loc, len, C, vals.len)
end

# method called by users
function locate(data::Vector{UInt8}, ::Type{J}, vec::T) where {J,T}
    locate(Locate.containertype(J, vec), data, J, vec)
end

# these call the locate methods above
function locate(::Type{Primitive}, data::Vector{UInt8}, ::Type{J}, vec::T) where {J,T}
    locate(data, Locate.Values{J}(vec), Locate.length(vec))
end
function locate(::Type{NullablePrimitive}, data::Vector{UInt8}, ::Type{Union{J,Missing}}, vec::T
               ) where {J,T}
    locate(data, Locate.Bitmask(vec), Locate.Values{J}(vec), Locate.length(vec))
end
function locate(::Type{List}, data::Vector{UInt8}, ::Type{J}, vec::T) where {J,T}
    locate(data, J, Locate.Offsets(vec), Locate.Values(J, vec), Locate.length(vec))
end
function locate(::Type{NullableList}, data::Vector{UInt8}, ::Type{J}, vec::T) where {J,T}
    locate(data, J, Locate.Bitmask(vec), Locate.Offsets(vec), Locate.Values(J, vec), Locate.length(vec))
end
# these are a little inflexible and need more specific declarations
function locate(::Type{BitPrimitive}, data::Vector{UInt8}, ::Type{Bool}, vec::T) where {T}
    vals = Locate.Values{Bool}(vec)
    BitPrimitive(data, vals.loc, Locate.length(vec))
end
function locate(::Type{NullableBitPrimitive}, data::Vector{UInt8}, ::Type{Union{Bool,Missing}},
                vec::T) where {T}
    bmask = Locate.Bitmask(vec)
    vals = Locate.Values{Bool}(vec)
    NullableBitPrimitive(data, bmask.loc, vals.loc, Locate.length(vec))
end


text:
using Documenter, Pkg, Vulkan

Pkg.develop(path="generator")

using VulkanGen

makedocs(;
    modules=[Vulkan, VulkanGen],
    format=Documenter.HTML(prettyurls = true),
    pages=[
        "Home" => "index.md",
        "Introduction" => "intro.md",
        "Design" => "design.md",
        "Features" => "features.md",
        "API" => "api.md",
        "Utility" => "utility.md",
        "Troubleshooting" => "troubleshooting.md",
        "Developer documentation" => "dev.md",
    ],
    repo="https://github.com/JuliaGPU/Vulkan.jl/blob/{commit}{path}#L{line}",
    sitename="Vulkan.jl",
    authors="serenity4 <cedric.bel@hotmail.fr>",
    strict=true,
    doctest=false,
    checkdocs=:exports,
    linkcheck=:true,
)

deploydocs(
    repo = "github.com/JuliaGPU/Vulkan.jl.git",
)


text:
"""

    Sgd(;lr=0.001)
    update!(w,g,p::Sgd)
    update!(w,g;lr=0.001)

Container for parameters of the Stochastic gradient descent (SGD)
optimization algorithm used by [`update!`](@ref).

SGD is an optimization technique to minimize an objective function by
updating its weights in the opposite direction of their gradient. The
learning rate (lr) determines the size of the step.  SGD updates the
weights with the following formula:

    w = w - lr * g

where `w` is a weight array, `g` is the gradient of the loss function
w.r.t `w` and `lr` is the learning rate.

SGD is used by default if no algorithm is specified in the two
argument version of `update!`[@ref].

"""
type Sgd
    lr::AbstractFloat
end

const SGDLR=0.001

Sgd(;lr=SGDLR)=Sgd(lr)

"""

    Momentum(;lr=0.001, gamma=0.9)
    update(w,g,p::Momentum)

Container for parameters of the Momentum optimization algorithm used
by [`update!`](@ref).

The Momentum method tries to accelerate SGD by adding a velocity term
to the update.  This also decreases the oscillation between successive
steps. It updates the weights with the following formulas:

    velocity = gamma * velocity + lr * g
    w = w - velocity

where `w` is a weight array, `g` is the gradient of the objective
function w.r.t `w`, `lr` is the learning rate, `gamma` is the momentum
parameter, `velocity` is an array with the same size and type of `w`
and holds the accelerated gradients.

Reference: [Qian,
N. (1999)](http://doi.org/10.1016/S0893-6080(98)00116-6). On the
momentum term in gradient descent learning algorithms.  Neural
Networks : The Official Journal of the International Neural Network
Society, 12(1), 145–151.

"""
type Momentum
    lr::AbstractFloat
    gamma::AbstractFloat
    velocity
end

Momentum(; lr=0.001, gamma=0.9)=Momentum(lr, gamma, nothing)

"""

    Adagrad(;lr=0.1, eps=1e-6)
    update(w,g,p::Adagrad)

Container for parameters of the Adagrad optimization algorithm used by
[`update!`](@ref).

Adagrad is one of the methods that adapts the learning rate to each of
the weights.  It stores the sum of the squares of the gradients to
scale the learning rate.  The learning rate is adapted for each weight
by the value of current gradient divided by the accumulated
gradients. Hence, the learning rate is greater for the parameters
where the accumulated gradients are small and the learning rate is
small if the accumulated gradients are large. It updates the weights
with the following formulas:

    G = G + g .^ 2
    w = w - g .* lr ./ sqrt(G + eps)

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `G` is an array with the same
size and type of `w` and holds the sum of the squares of the
gradients. `eps` is a small constant to prevent a zero value in the
denominator.

Reference: [Duchi, J., Hazan, E., & Singer,
Y. (2011)](http://jmlr.org/papers/v12/duchi11a.html). Adaptive
Subgradient Methods for Online Learning and Stochastic Optimization.
Journal of Machine Learning Research, 12, 2121–2159.

"""
type Adagrad
    lr::AbstractFloat
    eps::AbstractFloat
    G
end

Adagrad(; lr=0.1, eps=1e-6)=Adagrad(lr, eps, nothing)

"""

    Adadelta(;lr=0.01, rho=0.9, eps=1e-6)
    update(w,g,p::Adadelta)

Container for parameters of the Adadelta optimization algorithm used by
[`update!`](@ref).

Adadelta is an extension of Adagrad that tries to prevent the decrease
of the learning rates to zero as training progresses. It scales the
learning rate based on the accumulated gradients like Adagrad and
holds the acceleration term like Momentum. It updates the weights with
the following formulas:

    G = (1-rho) * g .^ 2 + rho * G
    update = g .* sqrt(delta + eps) ./ sqrt(G + eps)
    w = w - lr * update
    delta = rho * delta + (1-rho) * update .^ 2

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `G` is an array with the same
size and type of `w` and holds the sum of the squares of the
gradients. `eps` is a small constant to prevent a zero value in the
denominator.  `rho` is the momentum parameter and `delta` is an array
with the same size and type of `w` and holds the sum of the squared
updates.

Reference: [Zeiler,
M. D. (2012)](http://arxiv.org/abs/1212.5701). ADADELTA: An Adaptive
Learning Rate Method.

"""
type Adadelta
	lr::AbstractFloat
	rho::AbstractFloat
	eps::AbstractFloat
	G
	delta
end

Adadelta(; lr=0.01, rho=0.9, eps=1e-6)=Adadelta(lr, rho, eps, nothing, nothing)

"""

    Rmsprop(;lr=0.001, rho=0.9, eps=1e-6)
    update(w,g,p::Rmsprop)

Container for parameters of the Rmsprop optimization algorithm used by
[`update!`](@ref).

Rmsprop scales the learning rates by dividing the root mean squared of
the gradients. It updates the weights with the following formula:

    G = (1-rho) * g .^ 2 + rho * G
    w = w - lr * g ./ sqrt(G + eps)

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `G` is an array with the same
size and type of `w` and holds the sum of the squares of the
gradients. `eps` is a small constant to prevent a zero value in the
denominator.  `rho` is the momentum parameter and `delta` is an array
with the same size and type of `w` and holds the sum of the squared
updates.

Reference: [Tijmen Tieleman and Geoffrey Hinton
(2012)](https://dirtysalt.github.io/images/nn-class-lec6.pdf). "Lecture
6.5-rmsprop: Divide the gradient by a running average of its recent
magnitude."  COURSERA: Neural Networks for Machine Learning 4.2.

"""
type Rmsprop
	lr::AbstractFloat
	rho::AbstractFloat
	eps::AbstractFloat
	G
end

Rmsprop(; lr=0.001, rho=0.9, eps=1e-6)=Rmsprop(lr, rho, eps, nothing)

"""

    Adam(;lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8)
    update(w,g,p::Adam)

Container for parameters of the Adam optimization algorithm used by
[`update!`](@ref).

Adam is one of the methods that compute the adaptive learning rate. It
stores accumulated gradients (first moment) and the sum of the squared
of gradients (second).  It scales the first and second moment as a
function of time. Here is the update formulas:

    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * g .* g
    mhat = m ./ (1 - beta1 ^ t)
    vhat = v ./ (1 - beta2 ^ t)
    w = w - (lr / (sqrt(vhat) + eps)) * mhat

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `m` is an array with the same
size and type of `w` and holds the accumulated gradients. `v` is an
array with the same size and type of `w` and holds the sum of the
squares of the gradients. `eps` is a small constant to prevent a zero
denominator. `beta1` and `beta2` are the parameters to calculate bias
corrected first and second moments. `t` is the update count.

Reference: [Kingma, D. P., & Ba,
J. L. (2015)](https://arxiv.org/abs/1412.6980). Adam: a Method for
Stochastic Optimization. International Conference on Learning
Representations, 1–13.

"""
type Adam
    lr::AbstractFloat
    beta1::AbstractFloat
    beta2::AbstractFloat
    eps::AbstractFloat
    t::Int
    fstm
    scndm
end

Adam(; lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8)=Adam(lr, beta1, beta2, eps, 0, nothing, nothing)

"""

    update!(weights, gradients, params)
    update!(weights, gradients; lr=0.001)

Update the `weights` using their `gradients` and the optimization
algorithm parameters specified by `params`.  The 2-arg version
defaults to the [`Sgd`](@ref) algorithm with learning rate `lr`.  The
`weights` and possibly `params` are modified in-place.

`weights` can be an individual numeric array or a collection of arrays
represented by an iterator or dictionary.  In the individual case,
`gradients` should be a similar numeric array of `size(weights)` and
`params` should be a single object.  In the collection case, each
individual weight array should have a corresponding params object.
This way different weight arrays can have their own optimization
state, different learning rates, or even different optimization
algorithms running in parallel.  In the iterator case, `gradients` and
`params` should be iterators of the same length as `weights` with
corresponding elements.  In the dictionary case, `gradients` and
`params` should be dictionaries with the same keys as `weights`.

Individual optimization parameters can be one of the following types:
* [`Sgd`](@ref)`(;lr=0.001)`
* [`Momentum`](@ref)`(;lr=0.001, gamma=0.9)`
* [`Rmsprop`](@ref)`(;lr=0.001, rho=0.9, eps=1e-6)`
* [`Adagrad`](@ref)`(;lr=0.1, eps=1e-6)`
* [`Adadelta`](@ref)`(;lr=0.01, rho=0.9, eps=1e-6)`
* [`Adam`](@ref)`(;lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8)`

# Example:

    w = rand(d)                 # an individual weight array
    g = lossgradient(w)         # gradient g has the same shape as w
    update!(w, g)               # update w in-place with Sgd()
    update!(w, g; lr=0.1)       # update w in-place with Sgd(lr=0.1)
    update!(w, g, Sgd(lr=0.1))  # update w in-place with Sgd(lr=0.1)

    w = (rand(d1), rand(d2))    # a tuple of weight arrays
    g = lossgradient2(w)        # g will also be a tuple
    p = (Adam(), Sgd())         # p has params for each w[i]
    update!(w, g, p)            # update each w[i] in-place with g[i],p[i]

    w = Any[rand(d1), rand(d2)] # any iterator can be used
    g = lossgradient3(w)        # g will be similar to w
    p = Any[Adam(), Sgd()]      # p should be an iterator of same length
    update!(w, g, p)            # update each w[i] in-place with g[i],p[i]

    w = Dict(:a => rand(d1), :b => rand(d2)) # dictionaries can be used
    g = lossgradient4(w)
    p = Dict(:a => Adam(), :b => Sgd())
    update!(w, g, p)

"""
function update!{T<:AbstractFloat}(w::KorA{T}, g::KorA{T}, p::Sgd)
    axpy!(-p.lr, g, w)
end

function update!{T<:AbstractFloat}(w::KorA{T}, g::KorA{T}, p::Momentum)
    if p.velocity===nothing; p.velocity=zeros(w); end
    scale!(p.gamma, p.velocity)
    axpy!(p.lr, g, p.velocity)
    axpy!(-1, p.velocity, w)
end

function update!{T<:AbstractFloat}(w::KorA{T}, g::KorA{T}, p::Adam)
    if p.fstm===nothing; p.fstm=zeros(w); p.scndm=zeros(w); end
    p.t += 1
    scale!(p.beta1, p.fstm)
    axpy!(1-p.beta1, g, p.fstm)
    scale!(p.beta2, p.scndm)
    axpy!(1-p.beta2, g .* g, p.scndm)
    fstm_corrected = p.fstm / (1 - p.beta1 ^ p.t)
    scndm_corrected = p.scndm / (1 - p.beta2 ^ p.t)
    axpy!(-p.lr, (fstm_corrected ./ (sqrt(scndm_corrected) + p.eps)), w)
end

function update!{T<:AbstractFloat}(w::KorA{T}, g::KorA{T}, p::Adagrad)
    if p.G===nothing; p.G=zeros(w); end
    axpy!(1, g .* g, p.G)
    axpy!(-p.lr, g ./ sqrt(p.G + p.eps), w)
end

function update!{T<:AbstractFloat}(w::KorA{T}, g::KorA{T}, p::Adadelta)
    if p.G===nothing; p.G=zeros(w); p.delta=zeros(w); end
    scale!(p.rho, p.G)
    axpy!(1-p.rho, g .* g, p.G)
    dw = g .* sqrt(p.delta + p.eps) ./ sqrt(p.G + p.eps)
    scale!(p.rho, p.delta)
    axpy!(1-p.rho, dw .* dw , p.delta)
    axpy!(-p.lr, dw, w)
end

function update!{T<:AbstractFloat}(w::KorA{T}, g::KorA{T}, p::Rmsprop)
    if p.G===nothing; p.G=zeros(w); end
    scale!(p.rho, p.G)
    axpy!(1-p.rho, g .* g, p.G)
    axpy!(-p.lr, g ./ sqrt(p.G + p.eps), w)
end

# This takes care of arrays, tuples, iterators in general.
function update!(w,g,p)
    if !(length(w)==length(g)==length(p))
        error("weight, gradient, and optimization parameters not the same length.")
    end
    if isbits(eltype(w))
        error("Bad args: $((typeof(w),typeof(g),typeof(p)))")
    end
    for (wi,gi,pi) in zip(w,g,p)
        update!(wi,gi,pi)
    end
end

# We still need an extra method for Dict.
function update!(w::Associative,g::Associative,p::Associative)
    for k in keys(w)
        update!(w[k],g[k],p[k])
    end
end

# Two arg version for the simple default Sgd update.
function update!(w,g;lr=SGDLR)
    if !(length(w)==length(g))
        error("weight, gradient not the same length.")
    end
    sgd = Sgd(lr)
    for (wi,gi) in zip(w,g)
        update!(wi,gi,sgd)
    end
end

function update!(w::Associative,g::Associative;lr=SGDLR)
    sgd = Sgd(lr)
    for k in keys(w)
        update!(w[k],g[k],sgd)
    end
end

# To distinuish the two-arg update! for the numeric weight arrays:
update!{T<:AbstractFloat}(w::KorA{T},g::KorA{T};lr=SGDLR)=update!(w,g,Sgd(lr))