text:
Vtr = (ℂ^3,
        (ℂ^4)',
        ℂ^5,
        ℂ^6,
        (ℂ^7)')
Vℤ₂ = (ℂ[ℤ₂](0=>1, 1=>1),
        ℂ[ℤ₂](0=>1, 1=>2)',
        ℂ[ℤ₂](0=>3, 1=>2)',
        ℂ[ℤ₂](0=>2, 1=>3),
        ℂ[ℤ₂](0=>2, 1=>5))
Vℤ₃ = (ℂ[ℤ₃](0=>1, 1=>2, 2=>2),
        ℂ[ℤ₃](0=>3, 1=>1, 2=>1),
        ℂ[ℤ₃](0=>2, 1=>2, 2=>1)',
        ℂ[ℤ₃](0=>1, 1=>2, 2=>3),
        ℂ[ℤ₃](0=>1, 1=>3, 2=>3)')
VU₁ = (ℂ[U₁](0=>1, 1=>2, -1=>2),
        ℂ[U₁](0=>3, 1=>1, -1=>1),
        ℂ[U₁](0=>2, 1=>2, -1=>1)',
        ℂ[U₁](0=>1, 1=>2, -1=>3),
        ℂ[U₁](0=>1, 1=>3, -1=>3)')
VCU₁ = (ℂ[CU₁]((0,0)=>1, (0,1)=>2, 1=>1),
        ℂ[CU₁]((0,0)=>3, (0,1)=>0, 1=>1),
        ℂ[CU₁]((0,0)=>1, (0,1)=>0, 1=>2)',
        ℂ[CU₁]((0,0)=>2, (0,1)=>2, 1=>1),
        ℂ[CU₁]((0,0)=>2, (0,1)=>1, 1=>2)')
VSU₂ = (ℂ[SU₂](0=>3, 1//2=>1),
        ℂ[SU₂](0=>2, 1=>1),
        ℂ[SU₂](1//2=>1, 1=>1)',
        ℂ[SU₂](0=>2, 1//2=>2),
        ℂ[SU₂](0=>1, 1//2=>1, 3//2=>1)')

for (G,V) in ((Trivial, Vtr), (ℤ₂, Vℤ₂), (ℤ₃, Vℤ₃), (U₁, VU₁), (CU₁, VCU₁), (SU₂, VSU₂))
    println("------------------------------------")
    println("Tensors with symmetry: $G")
    println("------------------------------------")
    ti = time()
    V1, V2, V3, V4, V5 = V
    @testset TimedTestSet "Basic tensor properties" begin
        W = V1 ⊗ V2 ⊗ V3 ⊗ V4 ⊗ V5
        for T in (Int, Float32, Float64, ComplexF32, ComplexF64, BigFloat)
            t = Tensor(zeros, T, W)
            @test @inferred(hash(t)) == hash(deepcopy(t))
            @test eltype(t) == T
            @test norm(t) == 0
            @test codomain(t) == W
            @test space(t) == (W ← one(W))
            @test domain(t) == one(W)
        end
    end
    @testset TimedTestSet "Tensor Dict conversion" begin
    W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
        for T in (Int, Float32, ComplexF64)
            t = TensorMap(rand, T, W)
            d = convert(Dict, t)
            @test t == convert(TensorMap, d)
        end
    end
    @testset TimedTestSet "Basic linear algebra" begin
        W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
        for T in (Float32, ComplexF64)
            t = TensorMap(rand, T, W)
            @test eltype(t) == T
            @test space(t) == W
            @test space(t') == W'
            @test dim(t) == dim(space(t))
            @test codomain(t) == codomain(W)
            @test domain(t) == domain(W)
            @test isa(@inferred(norm(t)), real(T))
            @test norm(t)^2 ≈ dot(t,t)
            α = rand(T)
            @test norm(α*t) ≈ abs(α)*norm(t)
            @test norm(t+t, 2) ≈ 2*norm(t, 2)
            @test norm(t+t, 1) ≈ 2*norm(t, 1)
            @test norm(t+t, Inf) ≈ 2*norm(t, Inf)
            p = 3*rand(Float64)
            @test norm(t+t, p) ≈ 2*norm(t, p)
            @test norm(t) ≈ norm(t')

            t2 = TensorMap(rand, T, W)
            β = rand(T)
            @test dot(β*t2,α*t) ≈ conj(β)*α*conj(dot(t,t2))
            @test dot(t2,t) ≈ conj(dot(t, t2))
            @test dot(t2,t) ≈ conj(dot(t2', t'))
            @test dot(t2,t) ≈ dot(t', t2')

            i1 = @inferred(isomorphism(Matrix{T}, V1 ⊗ V2, V2 ⊗ V1))
            i2 = @inferred(isomorphism(Matrix{T}, V2 ⊗ V1, V1 ⊗ V2))
            @test i1 * i2 == @inferred(id(Matrix{T}, V1 ⊗ V2))
            @test i2 * i1 == @inferred(id(Matrix{T}, V2 ⊗ V1))

            w = @inferred(isometry(Matrix{T}, V1 ⊗ (oneunit(V1) ⊕ oneunit(V1)), V1))
            @test dim(w) == 2*dim(V1←V1)
            @test w'*w == id(Matrix{T}, V1)
            @test w*w' == (w*w')^2
        end
    end
    @testset TimedTestSet "Basic linear algebra: test via conversion" begin
        W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5
        for T in (Float32, ComplexF64)
            t = TensorMap(rand, T, W)
            t2 = TensorMap(rand, T, W)
            @test norm(t, 2) ≈ norm(convert(Array,t), 2)
            @test dot(t2,t) ≈ dot(convert(Array,t2), convert(Array, t))
            α = rand(T)
            @test convert(Array, α*t) ≈ α*convert(Array,t)
            @test convert(Array, t+t) ≈ 2*convert(Array,t)
        end
    end
    @testset TimedTestSet "Permutations: test via inner product invariance" begin
        W = V1 ⊗ V2 ⊗ V3 ⊗ V4 ⊗ V5
        t = Tensor(rand, ComplexF64, W);
        t′ = Tensor(rand, ComplexF64, W);
        for k = 0:5
            for p in permutations(1:5)
                p1 = ntuple(n->p[n], StaticLength(k))
                p2 = ntuple(n->p[k+n], StaticLength(5-k))
                t2 = @inferred permute(t, p1, p2)
                @test norm(t2) ≈ norm(t)
                t2′= permute(t′, p1, p2)
                @test dot(t2′,t2) ≈ dot(t′,t) ≈ dot(transpose(t2′), transpose(t2))
            end
        end
    end
    @testset TimedTestSet "Permutations: test via conversion" begin
        W = V1 ⊗ V2 ⊗ V3 ⊗ V4 ⊗ V5
        t = Tensor(rand, ComplexF64, W);
        for k = 0:5
            for p in permutations(1:5)
                p1 = ntuple(n->p[n], StaticLength(k))
                p2 = ntuple(n->p[k+n], StaticLength(5-k))
                t2 = permute(t, p1, p2)
                a2 = convert(Array, t2)
                @test a2 ≈ permutedims(convert(Array, t), (p1...,p2...))
                @test convert(Array, transpose(t2)) ≈ permutedims(a2, (5,4,3,2,1))
            end
        end
    end
    @testset TimedTestSet "Full trace: test self-consistency" begin
        t = Tensor(rand, ComplexF64, V1 ⊗ V2' ⊗ V2 ⊗ V1');
        t2 = permute(t, (1,2), (4,3))
        s = @inferred tr(t2)
        @test conj(s) ≈ tr(t2')
        @tensor s2 = t[a,b,b,a]
        @tensor t3[a,b] := t[a,c,c,b]
        @tensor s3 = t3[a,a]
        @test s ≈ s2
        @test s ≈ s3
    end
    @testset TimedTestSet "Partial trace: test self-consistency" begin
        t = Tensor(rand, ComplexF64, V1 ⊗ V2' ⊗ V3 ⊗ V2 ⊗ V1' ⊗ V1);
        @tensor t2[a,b] := t[c,d,b,d,c,a]
        @tensor t4[a,b,c,d] := t[d,e,b,e,c,a]
        @tensor t5[a,b] := t4[a,b,c,c]
        @test t2 ≈ t5
    end
    @testset TimedTestSet "Trace: test via conversion" begin
        t = Tensor(rand, ComplexF64, V1 ⊗ V2' ⊗ V3 ⊗ V2 ⊗ V1' ⊗ V1);
        @tensor t2[a,b] := t[c,d,b,d,c,a]
        @tensor t3[a,b] := convert(Array, t)[c,d,b,d,c,a]
        @test t3 ≈ convert(Array, t2)
    end
    @testset TimedTestSet "Trace and contraction" begin
        t1 = Tensor(rand, ComplexF64, V1 ⊗ V2 ⊗ V3);
        t2 = Tensor(rand, ComplexF64, V2' ⊗ V4 ⊗ V1');
        t3 = t1 ⊗ t2
        @tensor ta[a,b] := t1[x,y,a]*t2[y,b,x]
        @tensor tb[a,b] := t3[x,y,a,y,b,x]
        @test ta ≈ tb
    end
    @testset TimedTestSet "Tensor contraction: test via conversion" begin
        A1 = TensorMap(randn, ComplexF64, V1'*V2', V3')
        A2 = TensorMap(randn, ComplexF64, V3*V4, V5)
        rhoL = TensorMap(randn, ComplexF64, V1, V1)
        rhoR = TensorMap(randn, ComplexF64, V5, V5)' # test adjoint tensor
        H = TensorMap(randn, ComplexF64, V2*V4, V2*V4)
        @tensor HrA12[a, s1, s2, c] := rhoL[a, a'] * conj(A1[a', t1, b]) *
            A2[b, t2, c'] * rhoR[c', c] * H[s1, s2, t1, t2]

        @tensor HrA12array[a, s1, s2, c] := convert(Array, rhoL)[a, a'] *
            conj(convert(Array, A1)[a', t1, b]) *
            convert(Array, A2)[b, t2, c'] *
            convert(Array, rhoR)[c', c] *
            convert(Array, H)[s1, s2, t1, t2]

        @test HrA12array ≈ convert(Array, HrA12)
    end
    @testset TimedTestSet "Multiplication and inverse: test compatibility" begin
        W1 = V1 ⊗ V2 ⊗ V3
        W2 = V4 ⊗ V5
        for T in (Float64, ComplexF64)
            t1 = TensorMap(rand, T, W1, W1)
            t2 = TensorMap(rand, T, W2, W2)
            t = TensorMap(rand, T, W1, W2)
            @test t1*(t1\t) ≈ t
            @test (t/t2)*t2 ≈ t
            @test t1\one(t1) ≈ inv(t1)
            @test one(t1)/t1 ≈ pinv(t1)
            @test_throws SpaceMismatch inv(t)
            @test_throws SpaceMismatch t2\t
            @test_throws SpaceMismatch t/t1
            tp = pinv(t)*t
            @test tp ≈ tp*tp
        end
    end
    @testset TimedTestSet "Multiplication and inverse: test via conversion" begin
        W1 = V1 ⊗ V2 ⊗ V3
        W2 = V4 ⊗ V5
        for T in (Float32, Float64, ComplexF32, ComplexF64)
            t1 = TensorMap(rand, T, W1, W1)
            t2 = TensorMap(rand, T, W2, W2)
            t = TensorMap(rand, T, W1, W2)
            d1 = dim(W1)
            d2 = dim(W2)
            At1 = reshape(convert(Array, t1), d1, d1)
            At2 = reshape(convert(Array, t2), d2, d2)
            At = reshape(convert(Array, t), d1, d2)
            @test reshape(convert(Array, t1*t), d1, d2) ≈ At1*At
            @test reshape(convert(Array, t1'*t), d1, d2) ≈ At1'*At
            @test reshape(convert(Array, t2*t'), d2, d1) ≈ At2*At'
            @test reshape(convert(Array, t2'*t'), d2, d1) ≈ At2'*At'

            @test reshape(convert(Array, inv(t1)), d1, d1) ≈ inv(At1)
            @test reshape(convert(Array, pinv(t)), d2, d1) ≈ pinv(At)

            if T == Float32 || T == ComplexF32
                continue
            end

            @test reshape(convert(Array, t1\t), d1, d2) ≈ At1\At
            @test reshape(convert(Array, t1'\t), d1, d2) ≈ At1'\At
            @test reshape(convert(Array, t2\t'), d2, d1) ≈ At2\At'
            @test reshape(convert(Array, t2'\t'), d2, d1) ≈ At2'\At'

            @test reshape(convert(Array, t2/t), d2, d1) ≈ At2/At
            @test reshape(convert(Array, t2'/t), d2, d1) ≈ At2'/At
            @test reshape(convert(Array, t1/t'), d1, d2) ≈ At1/At'
            @test reshape(convert(Array, t1'/t'), d1, d2) ≈ At1'/At'
        end
    end
    @testset TimedTestSet "Factorization" begin
        W = V1 ⊗ V2 ⊗ V3 ⊗ V4 ⊗ V5
        for T in (Float32, ComplexF64)
            # Test both a normal tensor and an adjoint one.
            ts = (Tensor(rand, T, W), Tensor(rand, T, W)')
            for t in ts
                @testset "leftorth with $alg" for alg in (TensorKit.QR(), TensorKit.QRpos(), TensorKit.QL(), TensorKit.QLpos(), TensorKit.Polar(), TensorKit.SVD(), TensorKit.SDD())
                    Q, R = @inferred leftorth(t, (3,4,2),(1,5); alg = alg)
                    QdQ = Q'*Q
                    @test QdQ ≈ one(QdQ)
                    @test Q*R ≈ permute(t, (3,4,2),(1,5))
                end
                @testset "leftnull with $alg" for alg in (TensorKit.QR(), TensorKit.SVD(), TensorKit.SDD())
                    N = @inferred leftnull(t, (3,4,2),(1,5); alg = alg)
                    NdN = N'*N
                    @test NdN ≈ one(NdN)
                    @test norm(N'*permute(t, (3,4,2),(1,5))) < 100*eps(norm(t))
                end
                @testset "rightorth with $alg" for alg in (TensorKit.RQ(), TensorKit.RQpos(), TensorKit.LQ(), TensorKit.LQpos(), TensorKit.Polar(), TensorKit.SVD(), TensorKit.SDD())
                    L, Q = @inferred rightorth(t, (3,4),(2,1,5); alg = alg)
                    QQd = Q*Q'
                    @test QQd ≈ one(QQd)
                    @test L*Q ≈ permute(t, (3,4),(2,1,5))
                end
                @testset "rightnull with $alg" for alg in (TensorKit.LQ(), TensorKit.SVD(), TensorKit.SDD())
                    M = @inferred rightnull(t, (3,4),(2,1,5); alg = alg)
                    MMd = M*M'
                    @test MMd ≈ one(MMd)
                    @test norm(permute(t, (3,4),(2,1,5))*M') < 100*eps(norm(t))
                end
                @testset "tsvd with $alg" for alg in (TensorKit.SVD(), TensorKit.SDD())
                    U, S, V = @inferred tsvd(t, (3,4,2),(1,5); alg = alg)
                    UdU = U'*U
                    @test UdU ≈ one(UdU)
                    VVd = V*V'
                    @test VVd ≈ one(VVd)
                    @test U*S*V ≈ permute(t, (3,4,2),(1,5))
                end
            end
            @testset "empty tensor" begin
                t = TensorMap(randn, T, V1 ⊗ V2, typeof(V1)())
                @testset "leftorth with $alg" for alg in (TensorKit.QR(), TensorKit.QRpos(), TensorKit.QL(), TensorKit.QLpos(), TensorKit.Polar(), TensorKit.SVD(), TensorKit.SDD())
                    Q, R = @inferred leftorth(t; alg = alg)
                    @test Q == t
                    @test dim(Q) == dim(R) == 0
                end
                @testset "leftnull with $alg" for alg in (TensorKit.QR(), TensorKit.SVD(), TensorKit.SDD())
                    N = @inferred leftnull(t; alg = alg)
                    @test N'*N ≈ id(domain(N))
                    @test N*N' ≈ id(codomain(N))
                end
                @testset "rightorth with $alg" for alg in (TensorKit.RQ(), TensorKit.RQpos(), TensorKit.LQ(), TensorKit.LQpos(), TensorKit.Polar(), TensorKit.SVD(), TensorKit.SDD())
                    L, Q = @inferred rightorth(copy(t'); alg = alg)
                    @test Q == t'
                    @test dim(Q) == dim(L) == 0
                end
                @testset "rightnull with $alg" for alg in (TensorKit.LQ(), TensorKit.SVD(), TensorKit.SDD())
                    M = @inferred rightnull(copy(t'); alg = alg)
                    @test M*M' ≈ id(codomain(M))
                    @test M'*M ≈ id(domain(M))
                end
                @testset "tsvd with $alg" for alg in (TensorKit.SVD(), TensorKit.SDD())
                    U, S, V = @inferred tsvd(t; alg = alg)
                    @test U == t
                    @test dim(U) == dim(S) == dim(V)
                end
            end

            t = Tensor(rand, T, V1 ⊗ V1' ⊗ V2 ⊗ V2')
            @testset "eig and isposdef" begin
                D, V = eigen(t, (1,3), (2,4))
                VdV = V'*V
                VdV = (VdV + VdV')/2
                @test isposdef(VdV)
                t2 = permute(t, (1,3), (2,4))
                @test t2*V ≈ V*D
                @test !isposdef(t2) # unlikely for non-hermitian map
                t2 = (t2 + t2');
                D, V = eigen(t2)
                VdV = V'*V
                @test VdV ≈ one(VdV)
                λ = minimum(minimum(real(diag(b))) for (c,b) in blocks(D))
                @test isposdef(t2) == isposdef(λ)
                @test isposdef(t2 - λ*one(t2) + 0.1*one(t2))
                @test !isposdef(t2 - λ*one(t2) - 0.1*one(t2))
            end
        end
    end
    @testset TimedTestSet "Tensor truncation" begin
        for T in (Float32, ComplexF64)
            for p in (1, 2, 3, Inf)
                # Test both a normal tensor and an adjoint one.
                ts = (TensorMap(randn, T, V1 ⊗ V2 ⊗ V3, V4 ⊗ V5),
                      TensorMap(randn, T, V4 ⊗ V5, V1 ⊗ V2 ⊗ V3)')
                for t in ts
                    U₀, S₀, V₀, = tsvd(t)
                    t = rmul!(t, 1/norm(S₀, p))
                    U, S, V, ϵ = @inferred tsvd(t; trunc = truncerr(5e-1), p = p)
                    # @show p, ϵ
                    # @show domain(S)
                    # @test min(space(S,1), space(S₀,1)) != space(S₀,1)
                    U′, S′, V′, ϵ′ = tsvd(t; trunc = truncerr(nextfloat(ϵ)), p = p)
                    @test (U, S, V, ϵ) == (U′, S′, V′, ϵ′)
                    U′, S′, V′, ϵ′ = tsvd(t; trunc = truncdim(dim(domain(S))), p = p)
                    @test (U, S, V, ϵ) == (U′, S′, V′, ϵ′)
                    U′, S′, V′, ϵ′ = tsvd(t; trunc = truncspace(space(S,1)), p = p)
                    @test (U, S, V, ϵ) == (U′, S′, V′, ϵ′)
                    # results with truncationcutoff cannot be compared because they don't take degeneracy into account, and thus truncate differently
                    U, S, V, ϵ = tsvd(t; trunc = truncbelow(1/dim(domain(S₀))), p = p)
                    # @show p, ϵ
                    # @show domain(S)
                    # @test min(space(S,1), space(S₀,1)) != space(S₀,1)
                    U′, S′, V′, ϵ′ = tsvd(t; trunc = truncspace(space(S,1)), p = p)
                    @test (U, S, V, ϵ) == (U′, S′, V′, ϵ′)
                end
            end
        end
    end
    @testset TimedTestSet "Tensor functions" begin
        W = V1 ⊗ V2
        for T in (Float64, ComplexF64)
            t = TensorMap(randn, T, W, W)
            s = dim(W)
            expt = @inferred exp(t)
            @test reshape(convert(Array, expt), (s,s)) ≈
                    exp(reshape(convert(Array, t), (s,s)))

            @test (@inferred sqrt(t))^2 ≈ t
            @test reshape(convert(Array, sqrt(t^2)), (s,s)) ≈
                    sqrt(reshape(convert(Array, t^2), (s,s)))

            @test exp(@inferred log(expt)) ≈ expt
            @test reshape(convert(Array, log(expt)), (s,s)) ≈
                    log(reshape(convert(Array, expt), (s,s)))

            @test (@inferred cos(t))^2 + (@inferred sin(t))^2 ≈ id(W)
            @test (@inferred tan(t)) ≈ sin(t)/cos(t)
            @test (@inferred cot(t)) ≈ cos(t)/sin(t)
            @test (@inferred cosh(t))^2 - (@inferred sinh(t))^2 ≈ id(W)
            @test (@inferred tanh(t)) ≈ sinh(t)/cosh(t)
            @test (@inferred coth(t)) ≈ cosh(t)/sinh(t)

            t1 = sin(t)
            @test sin(@inferred asin(t1)) ≈ t1
            t2 = cos(t)
            @test cos(@inferred acos(t2)) ≈ t2
            t3 = sinh(t)
            @test sinh(@inferred asinh(t3)) ≈ t3
            t4 = cosh(t)
            @test cosh(@inferred acosh(t4)) ≈ t4
            t5 = tan(t)
            @test tan(@inferred atan(t5)) ≈ t5
            t6 = cot(t)
            @test cot(@inferred acot(t6)) ≈ t6
            t7 = tanh(t)
            @test tanh(@inferred atanh(t7)) ≈ t7
            t8 = coth(t)
            @test coth(@inferred acoth(t8)) ≈ t8
        end
    end
    @testset TimedTestSet "Tensor product: test via norm preservation" begin
        for T in (Float32, ComplexF64)
            t1 = TensorMap(rand, T, V2 ⊗ V3 ⊗ V1, V1 ⊗ V2)
            t2 = TensorMap(rand, T, V2 ⊗ V1 ⊗ V3, V1 ⊗ V1)
            t = @inferred (t1 ⊗ t2)
            @test norm(t) ≈ norm(t1) * norm(t2)
        end
    end
    @testset TimedTestSet "Tensor product: test via conversion" begin
        for T in (Float32, ComplexF64)
            t1 = TensorMap(rand, T, V2 ⊗ V3 ⊗ V1, V1)
            t2 = TensorMap(rand, T, V2 ⊗ V1 ⊗ V3, V2)
            t = @inferred (t1 ⊗ t2)
            d1 = dim(codomain(t1))
            d2 = dim(codomain(t2))
            d3 = dim(domain(t1))
            d4 = dim(domain(t2))
            At = convert(Array, t)
            @test reshape(At, (d1, d2, d3, d4)) ≈
                    reshape(convert(Array, t1), (d1, 1, d3, 1)) .*
                    reshape(convert(Array, t2), (1, d2, 1, d4))
        end
    end
    @testset TimedTestSet "Tensor product: test via tensor contraction" begin
        for T in (Float32, ComplexF64)
            t1 = Tensor(rand, T, V2 ⊗ V3 ⊗ V1)
            t2 = Tensor(rand, T, V2 ⊗ V1 ⊗ V3)
            t = @inferred (t1 ⊗ t2)
            @tensor t′[1, 2, 3, 4, 5, 6] := t1[1,2,3]*t2[4,5,6]
            @test t ≈ t′
        end
    end
    tf = time()
    printstyled("Finished tensor tests with symmetry $G in ",
                string(round(tf-ti; sigdigits=3)),
                " seconds."; bold = true, color = Base.info_color())
    println()
end


text:
using StatsBase
using Base.Test

n = 5000

# 1D integer counts

x = rand(1:5, n)
w = weights(rand(n))

c = counts(x, 5)
@test size(c) == (5,)
c0 = Int[countnz(x .== i) for i in 1 : 5]
@test c == c0
@test counts(x .+ 1, 2:6) == c0
@test_approx_eq proportions(x, 1:5) (c0 ./ n)

c = counts(x)
@test size(c) == (5,)
c0 = Int[countnz(x .== i) for i in 1 : 5]
@test c == c0
@test counts(x .+ 1, 2:6) == c0
@test_approx_eq proportions(x) (c0 ./ n)

c = counts(x, 5, w)
@test size(c) == (5,)
c0 = Float64[sum(w.values[x .== i]) for i in 1 : 5]
@test_approx_eq c c0
@test_approx_eq counts(x .+ 1, 2:6, w) c0
@test_approx_eq proportions(x, 1:5, w) (c0 ./ sum(w))

c = counts(x, w)
@test size(c) == (5,)
c0 = Float64[sum(w.values[x .== i]) for i in 1 : 5]
@test_approx_eq c c0
@test_approx_eq counts(x .+ 1, 2:6, w) c0
@test_approx_eq proportions(x, w) (c0 ./ sum(w))

# 2D integer counts

x = rand(1:4, n)
y = rand(1:5, n)
w = weights(rand(n))

c = counts(x, y, (4, 5))
@test size(c) == (4, 5)
c0 = Int[countnz((x .== i) & (y .== j)) for i in 1 : 4, j in 1 : 5]
@test c == c0
@test counts(x .+ 2, y .+ 3, (3:6, 4:8)) == c0
@test_approx_eq proportions(x, y, (1:4, 1:5)) (c0 ./ n)

c = counts(x, y)
@test size(c) == (4, 5)
c0 = Int[countnz((x .== i) & (y .== j)) for i in 1 : 4, j in 1 : 5]
@test c == c0
@test counts(x .+ 2, y .+ 3, (3:6, 4:8)) == c0
@test_approx_eq proportions(x, y,) (c0 ./ n)

c = counts(x, y, (4, 5), w)
@test size(c) == (4, 5)
c0 = Float64[sum(w.values[(x .== i) & (y .== j)]) for i in 1 : 4, j in 1 : 5]
@test_approx_eq c c0
@test_approx_eq counts(x .+ 2, y .+ 3, (3:6, 4:8), w) c0
@test_approx_eq proportions(x, y, (1:4, 1:5), w) (c0 ./ sum(w))

c = counts(x, y, w)
@test size(c) == (4, 5)
c0 = Float64[sum(w.values[(x .== i) & (y .== j)]) for i in 1 : 4, j in 1 : 5]
@test_approx_eq c c0
@test_approx_eq counts(x .+ 2, y .+ 3, (3:6, 4:8), w) c0
@test_approx_eq proportions(x, y, w) (c0 ./ sum(w))

# count map

x = ["a", "b", "a", "a", "b", "c"]
w = [1.0, 1.5, 2.0, 2.5, 3.0, 3.5]

cm = countmap(x)
@test cm["a"] == 3
@test cm["b"] == 2
@test cm["c"] == 1
pm = proportionmap(x)
@test_approx_eq pm["a"] (1/2)
@test_approx_eq pm["b"] (1/3)
@test_approx_eq pm["c"] (1/6)

cm = countmap(x, weights(w))
@test cm["a"] == 5.5
@test cm["b"] == 4.5
@test cm["c"] == 3.5
pm = proportionmap(x, weights(w))
@test_approx_eq pm["a"] (5.5 / 13.5)
@test_approx_eq pm["b"] (4.5 / 13.5)
@test_approx_eq pm["c"] (3.5 / 13.5)

ext: jl


text:
module HELICS

__precompile__(true)

import Libdl

# Load in `deps.jl`, complaining if it does not exist
const depsjl_path = joinpath(@__DIR__, "..", "deps", "deps.jl")

if !isfile(depsjl_path)
    error("libhelicsSharedLib cannot be loaded. Please run Pkg.build(\"HELICS\").")
end

include(depsjl_path)

# Module initialization function
function __init__()
    check_deps()
    atexit(helicsCloseLibrary)
end

module Lib

    using CEnum

    import ..HELICS
    const libhelicsSharedLib = HELICS.libhelicsSharedLib

    const HELICS_EXPORT = nothing
    const HELICS_NO_EXPORT = nothing

    include("ctypes.jl")

    include("common.jl")

    include("manual.jl")

    include("lib.jl")

end

const HELICS_LIBRARY_VERSION = VersionNumber(split(Lib.HELICS_VERSION_STRING)[1])
const HELICS_PACKAGE_VERSION = VersionNumber(match(r"version\s*=\s*\"(.+)\"", read(joinpath(@__DIR__, "..", "Project.toml"), String)).captures[1])

include("wrapper.jl")

include("utils.jl")

include("api.jl")

end # module


text:
abstract type AbstractAggregateOperators end

struct MyExceptionTree <: Exception
    msg::String
end

"""
getAggHamSystemSmall(agg)

Get Hamiltonian of the system, ``H_S`` only for ``\\mathcal{H}_S``.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function getAggHamSystemSmall(
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited,
    groundEnergy::Bool = true,
)
    Ham_sys = zeros(Float64, (aggCore.molCount + 1, aggCore.molCount + 1))

    agg_shifts = getShifts(aggCore)
    agg_frequencies = getFrequencies(aggCore)
    reorganisation_energies = []
    for mol_i = 1:aggCore.molCount
        mol = aggCore.molecules[mol_i]
        reorganisation_energy = 0.0
        for mode_i = 1:length(mol.modes)
            reorganisation_energy +=
                agg_frequencies[mol_i][mode_i] * agg_shifts[mol_i][mode_i]^2 / 2.0
        end
        push!(reorganisation_energies, reorganisation_energy)
    end

    E_agg = zeros(Float64, (2, aggCore.molCount))
    E_agg[1, :] = map(mol -> mol.E[1], aggCore.molecules)
    E_agg[2, :] = map(mol -> mol.E[2], aggCore.molecules)
    for elInd in aggTools.elIndices
        ind = OpenQuantumSystems.elIndOrder(elInd)
        E_state = 0
        for mol_i = 1:aggCore.molCount
            E_state += E_agg[elInd[mol_i], mol_i]
        end
        if vib_basis == :ground_ground
            # add reorganisation energy
            if ind > 1
                E_state += reorganisation_energies[ind-1]
            end
        end
        Ham_sys[ind, ind] = E_state
    end

    Ham_sys[:, :] += aggCore.coupling[:, :]
    if !groundEnergy
        E0 = Ham_sys[1, 1]
        for i = 1:aggTools.bSystemSize
            Ham_sys[i, i] -= E0
        end
    end

    return DenseOperator(aggTools.basisSystem, aggTools.basisSystem, Ham_sys)
end

"""
getAggHamSystemBig(agg)

Get Hamiltonian of the system, ``H_S`` for ``\\mathcal{H}_S \\otimes \\mathcal{H}_B``.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function getAggHamSystemBig(
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited,
    groundEnergy::Bool = true,
)
    Ham_sys = getAggHamSystemSmall(aggCore, aggTools; vib_basis = vib_basis, groundEnergy = groundEnergy)
    Ham = zeros(Float64, (aggTools.bSize, aggTools.bSize))
    for I = 1:aggTools.bSize
        elind1, vibind1 = aggTools.indices[I]
        elOrder1 = elIndOrder(elind1)
        for J = 1:aggTools.bSize
            elind2, vibind2 = aggTools.indices[J]
            elOrder2 = elIndOrder(elind2)

            # Ham[I, J] = Ham_sys.data[elOrder1, elOrder2] * aggTools.FCfactors[I, J]
            if vibind1 == vibind2
                Ham[I, J] = Ham_sys.data[elOrder1, elOrder2]
            end
        end
    end
    if !groundEnergy
        E0 = Ham[1, 1]
        for I = 1:aggTools.bSize
            Ham[I, I] -= E0
        end
    end

    return DenseOperator(aggTools.basis, aggTools.basis, Ham)
end

"""
getAggHamBathSmall(agg)

Get Hamiltonian of the bath, ``H_B`` only for ``\\mathcal{H}_B``.

"""
function getAggHamBathSmall(
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited,
    groundEnergy::Bool = true,
)
    Ham_bath = zeros(Float64, (aggTools.bBathSize, aggTools.bBathSize))

    elind = fill(1, (aggCore.molCount + 1))
    Ham_sys = getAggHamSystemSmall(aggCore, aggTools; vib_basis = vib_basis, groundEnergy = true)
    for I = 1:aggTools.bBathSize
        vibind = aggTools.vibIndices[I]
        Ham_bath[I, I] = getAggStateEnergy(aggCore, elind, vibind) - Ham_sys.data[1, 1]
    end

    if !groundEnergy
        E0 = Ham_bath[1, 1]
        for i = 1:aggTools.bBathSize
            Ham_bath[i, i] -= E0
        end
    end
    return DenseOperator(aggTools.basisBath, aggTools.basisBath, Ham_bath)
end

"""
getAggHamBathBig(agg)

Get Hamiltonian of the bath, ``H_B`` only for ``\\mathcal{H}_S \\otimes \\mathcal{H}_B``.

"""
function getAggHamBathBig(
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    groundEnergy::Bool = true,
)

    Ham_bath = getAggHamBathSmall(aggCore, aggTools)
    one_sys = OneDenseOperator(aggTools.basisSystem)
    # TODO: optimise types in future
    one_sys = DenseOperator(aggTools.basisSystem, aggTools.basisSystem, real(one_sys.data))
    Ham_bath = tensor(Ham_bath, one_sys)
    if !groundEnergy
        E0 = Ham_bath.data[1, 1]
        for I = 1:aggTools.bSize
            Ham_bath.data[I, I] -= E0
        end
    end
    return DenseOperator(aggTools.basis, aggTools.basis, Ham_bath.data)
end

"""
getAggHamSystemBath(agg, aggTools.indices,
\tfranckCondonFactors; groundEnergy = false)
getAggHamSystemBath(agg, aggTools.indices; groundEnergy = false)
getAggHamSystemBath(agg; groundEnergy = false)

Get system and bath Hamiltonian of the [`Aggregate`](@ref) in a form of DenseOperator.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggTools.indices`: Aggregate indices generated by [`getIndices`](@ref).
* `franckCondonFactors`: Franck-Condon factors generated by [`getFranckCondonFactors`](@ref).
"""
function getAggHamSystemBath(
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited,
    groundEnergy::Bool = true,
)
    Ham_B = getAggHamBathBig(aggCore, aggTools; groundEnergy = true)
    Ham_S = getAggHamSystemBig(aggCore, aggTools; vib_basis = vib_basis, groundEnergy = true)
    Ham_0 = Ham_B + Ham_S
    if !groundEnergy
        E0 = Ham_0.data[1, 1]
        for I = 1:aggTools.bSize
            Ham_0.data[I, I] -= E0
        end
    end
    return Ham_0
end

"""
getAggHamInteraction(agg, aggTools.indices, franckCondonFactors)
getAggHamInteraction(agg, aggTools.indices)
getAggHamInteraction(agg)

Get interation Hamiltonian of the [`Aggregate`](@ref).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggTools.indices`: Aggregate indices generated by [`getIndices`](@ref).
* `franckCondonFactors`: Franck-Condon factors generated by [`getFranckCondonFactors`](@ref).
"""
function getAggHamInteraction(aggCore::AggregateCore, aggTools::AggregateTools; vib_basis::Symbol=:ground_excited)
    if vib_basis ∉ (:ground_ground, :ground_excited)
        throw(ArgumentError("Optional argument vib_basis has to be selected from (:ground_ground, :ground_excited)"))
    end
    Ham_sys = getAggHamSystemSmall(aggCore, aggTools; vib_basis=vib_basis, groundEnergy = true)

    Ham_I = zeros(Float64, (aggTools.bSize, aggTools.bSize))
    agg_shifts = getShifts(aggCore)
    agg_frequencies = getFrequencies(aggCore)

    for I = 1:aggTools.bSize
        elind1, vibind1 = aggTools.indices[I]
        elOrder1 = elIndOrder(elind1)
        if elOrder1 == 1
            continue
        end
        for J = 1:aggTools.bSize
            elind2, vibind2 = aggTools.indices[J]
            elOrder2 = elIndOrder(elind2)
            if elOrder2 == 1
                continue
            end
            if elOrder1 != elOrder2
                if vib_basis == :ground_excited
                    Ham_I[I, J] = Ham_sys.data[elOrder1, elOrder2] * aggTools.FCfactors[I, J]
                    if vibind1 == vibind2
                        Ham_I[I, J] -= Ham_sys.data[elOrder1, elOrder2]
                    end
                end
            else
                # permits only one vib change
                # check all moleciles and all modes, if they differ in max one vib state
                # calculate the values of \Delta V
                #=
                diff_num = 0 # in molecules
                mol_j = 0
                mode_j = 0
                for mol_i = 1:aggCore.molCount
                    mol = aggCore.molecules[mol_i]
                    for mode_i = 1:length(mol.modes)
                        diff_vib = abs(vibind1[mol_i][mode_i] - vibind2[mol_i][mode_i])
                        diff_num += diff_vib
                        if diff_vib == 1 && diff_num == 1
                            mol_j = mol_i
                            mode_j = mode_i
                        end
                    end
                end

                if diff_num != 1 || mol_j != elOrder1 - 1
                    continue
                end
                vib_n = vibind1[mol_j][mode_j]
                vib_m = vibind2[mol_j][mode_j]
                coeff = agg_coeffs[mol_j][mode_j]
                Ham_I[I, J] = -coeff * sqrt(min(vib_n, vib_m))
                =#
                # permits any molecule change
                #=
                coeff = 1.
                mol_i = elOrder1 - 1
                mol = aggCore.molecules[mol_i]
                for mode_i = 1:length(mol.modes)
                    diff_vib = abs(vibind1[mol_i][mode_i] - vibind2[mol_i][mode_i])
                    if diff_vib == 1
                        vib_n = vibind1[mol_i][mode_i]
                        vib_m = vibind2[mol_i][mode_i]
                        coeff *= agg_coeffs[mol_i][mode_i] * sqrt(min(vib_n, vib_m))
                    end
                end
                if coeff != 1.
                    Ham_I[I, J] = - coeff
                end
                =#
                # somehing smells here, eigenvals of Ham are not the same in both vib bath bases
                # total hours spent: 11
                # What are the correct elements of ΔV?
                # Battle not with monsters, lest ye become a monster, and if you gaze into the abyss, the abyss gazes also into you.
                #      --Nietzsche
                if vib_basis == :ground_ground
                    coeff = 0.
                    diff_vib_sum = 0
                    for mol_i = 1:aggCore.molCount
                        mol = aggCore.molecules[mol_i]
                        for mode_i = 1:length(mol.modes)
                            diff_vib = abs(vibind1[mol_i][mode_i] - vibind2[mol_i][mode_i])
                            diff_vib_sum += diff_vib
                            if diff_vib == 1 && mol_i == elOrder1 - 1
                                vib_n = vibind1[mol_i][mode_i]
                                vib_m = vibind2[mol_i][mode_i]
                                coeff += agg_frequencies[mol_i][mode_i] * agg_shifts[mol_i][mode_i] / Float64(min(vib_n, vib_m))^0.5 / 2.0
                                # println(coeff, " ", vibind1, " ", vibind2, " ", mol_i, " ", mode_i, " ", min(vib_n, vib_m))
                            end # agg_frequencies[mol_i][mode_i] * agg_shifts[mol_i][mode_i]
                        end
                    end
                    if diff_vib_sum == 1
                        Ham_I[I, J] = coeff  # -
                    end
                end
            end
        end
    end

    return DenseOperator(aggTools.basis, aggTools.basis, Ham_I)
end

"""
getAggHamiltonian(agg, aggTools.indices, franckCondonFactors;
\t, groundEnergy = true)
getAggHamiltonian(agg, aggTools.indices; groundEnergy = true)
getAggHamiltonian(agg; groundEnergy = true)

Get Hamiltonian of the [`Aggregate`](@ref).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggTools.indices`: Aggregate indices generated by [`getIndices`](@ref).
* `franckCondonFactors`: Franck-Condon factors generated by [`getFranckCondonFactors`](@ref).
"""
function getAggHamiltonian(
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    groundEnergy::Bool = true,
    vib_basis::Symbol=:ground_excited
)
    Ham_I = getAggHamInteraction(aggCore, aggTools; vib_basis = vib_basis)
    Ham_0 = getAggHamSystemBath(aggCore, aggTools; groundEnergy = groundEnergy)
    return Ham_0 + Ham_I
end

### Sparse versions
#=
"""
getAggHamiltonianSparse(agg, aggTools.indices)
getAggHamiltonianSparse(agg)

Sparse version of [`getAggHamiltonian`](@ref).

"""
function getAggHamiltonianSparse(
agg::AggregateCore{T,C1,C2},
aggTools.indices::Any,
franckCondonFactors::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
if aggTools.indices === nothing
    aggTools.indices = getIndices(agg)
end
aggTools.bSize = length(aggTools.indices)
aggCore.molCount = length(agg.molecules)
if franckCondonFactors === nothing
    franckCondonFactors = getFranckCondonFactors(agg, aggTools.indices)
end
Ham = spzeros(C2, aggTools.bSize, aggTools.bSize)
for I = 1:aggTools.bSize
    elind1, vibind1 = aggTools.indices[I]
    elOrder1 = elIndOrder(elind1)
    for J = 1:aggTools.bSize
        elind2, vibind2 = aggTools.indices[J]
        elOrder2 = elIndOrder(elind2)
        if I == J
            val = getAggStateEnergy(agg, elind1, vibind1)
            if val != 0.0
                Ham[I, J] = val
            end
        else
            if elind1 != elind2
                val = agg.coupling[elOrder1, elOrder2] * franckCondonFactors[I, J]
                if val != 0.0
                    Ham[I, J] = val
                end
            end
        end
    end
end
E0 = Ham[1, 1]
for I = 1:aggTools.bSize
    Ham[I, I] -= E0
end
b = GenericBasis([aggTools.bSize])
return SparseOperator(b, b, Ham)
end

getAggHamiltonianSparse(
agg::AggregateCore{T,C1,C2}
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamiltonianSparse(
agg::AggregateCore{T,C1,C2},
nothing,
nothing
)

getAggHamiltonianSparse(
agg::AggregateCore{T,C1,C2},
aggTools.indices::Any
) where {T<:Integer,C1<:ComputableType,C2<:ComputableType} = getAggHamiltonianSparse(
agg::AggregateCore{T,C1,C2},
aggTools.indices,
nothing
)
=#

# TODO: add docs and proper types
struct AggregateOperators{O_sys,O_bath,O} <: AbstractAggregateOperators
    Ham_sys::O_sys
    Ham_bath::O_bath
    Ham_S::O
    Ham_B::O
    Ham_0::O
    Ham_I::O
    Ham::O
    groundEnergy::Bool
    vib_basis::Symbol
    function AggregateOperators{O_sys,O_bath,O}(
        Ham_sys::O_sys,
        Ham_bath::O_bath,
        Ham_S::O,
        Ham_B::O,
        Ham_0::O,
        Ham_I::O,
        Ham::O,
        groundEnergy::Bool,
        vib_basis::Symbol,
    ) where {O_sys<:Operator,O_bath<:Operator,O<:Operator}
        new(Ham_sys, Ham_bath, Ham_S, Ham_B, Ham_0, Ham_I, Ham, groundEnergy, vib_basis)
    end
end

function AggregateOperators(
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    groundEnergy::Bool = true,
    vib_basis::Symbol=:ground_excited
)

    Ham_sys = getAggHamSystemSmall(aggCore, aggTools; vib_basis = vib_basis, groundEnergy = true)
    Ham_bath = getAggHamBathSmall(aggCore, aggTools; groundEnergy = true)
    Ham_S = getAggHamSystemBig(aggCore, aggTools; vib_basis = vib_basis, groundEnergy = groundEnergy)
    Ham_B = getAggHamBathBig(aggCore, aggTools; groundEnergy = groundEnergy)
    Ham_0 = getAggHamSystemBath(aggCore, aggTools; groundEnergy = groundEnergy)
    Ham_I = getAggHamInteraction(aggCore, aggTools; vib_basis = vib_basis)
    Ham = getAggHamiltonian(aggCore, aggTools; groundEnergy = groundEnergy, vib_basis = vib_basis)

    O_sys = typeof(Ham_sys)
    O_bath = typeof(Ham_bath)
    O = typeof(Ham_S)

    return AggregateOperators{O_sys,O_bath,O}(
        Ham_sys,
        Ham_bath,
        Ham_S,
        Ham_B,
        Ham_0,
        Ham_I,
        Ham,
        groundEnergy,
        vib_basis
    )
end

Base.:(==)(x::AggregateOperators, y::AggregateOperators) =
    x.Ham_sys == y.Ham_sys &&
    x.Ham_bath == y.Ham_bath &&
    x.Ham_S == y.Ham_S &&
    x.Ham_B == y.Ham_B &&
    x.Ham_0 == y.Ham_0 &&
    x.Ham_I == y.Ham_I &&
    x.Ham == y.Ham &&
    x.groundEnergy == y.groundEnergy &&
    x.vib_basis == y.vib_basis

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems


text:
using WeavePynb
using Mustache

mmd(fname) = mmd_to_html(fname, BRAND_HREF="../toc.html", BRAND_NAME="Calculus with Julia")
## uncomment to generate just .md files
#mmd(fname) = mmd_to_md(fname, BRAND_HREF="../toc.html", BRAND_NAME="Calculus with Julia")

fnames = [
          "double_triple_integrals",
          "line_integrals",
          "div_grad_curl",
          "stokes_theorem",
          "review"
]

function process_file(nm, twice=false)
    include("$nm.jl")
    mmd_to_md("$nm.mmd")
    markdownToHTML("$nm.md")
    twice && markdownToHTML("$nm.md")
end

process_files(twice=false) = [process_file(nm, twice) for nm in fnames]

"""
## TODO integral_vector_calculus
* line integrals needs an image (lasso?)
* more questions
"""


text:
module Visualize_SolidFileMesh

using Modia3D
using Modia3D.StaticArrays

vmat1 = Modia3D.Material(color="LightBlue" , transparency=0.5)    # material of SolidFileMesh
vmat2 = deepcopy(vmat1)                                           # material of convex decomposition of SolidFileMesh
vmat2.transparency = 0.7

font  = Modia3D.Font(fontFamily="Arial", bold=true, charSize=0.2)
rtext = [0.0, 1.0, 0.0]
filenamePascal = joinpath(Modia3D.path, "objects", "pascal", "pascal.obj")

@assembly FileMeshes begin
  world = Modia3D.Object3D(visualizeFrame=false)
  solidFileMesh = Modia3D.SolidFileMesh(filenamePascal,0.2)
  fileMesh1 = Modia3D.Object3D(world, Modia3D.Solid(                       solidFileMesh, nothing, vmat1       ); r=[0.0,0.0,0.0])
  fileMesh2 = Modia3D.Object3D(world, Modia3D.SolidWithConvexDecomposition(solidFileMesh, nothing, vmat1, vmat2); r=[3.0,0.0,0.0])

end

Modia3D.visualizeAssembly!( FileMeshes(sceneOptions=Modia3D.SceneOptions(visualizeFrames=true,defaultFrameLength=0.7,enableContactDetection=false) ) )

println("... success of Visualize_SolidFileMesh.jl!")

end


text:
function equilibrate!(M, aff, opt)
    max_iters=opt.equilibration_iters
    lb = opt.equilibration_lb
    ub = opt.equilibration_ub

    @timeit "eq init" begin
        α = (aff.n / (aff.m + aff.p)) ^ .25
        β = ((aff.m + aff.p) / aff.n ) ^ .25
        α2, β2 = α^2, β^2
        γ = .1

        u, v           = zeros(aff.m + aff.p), zeros(aff.n)
        u_, v_         = zeros(aff.m + aff.p), zeros(aff.n)
        u_grad, v_grad = zeros(aff.m + aff.p), zeros(aff.n)
        row_norms, col_norms = zeros(aff.m + aff.p), zeros(aff.n)
        E = LinearAlgebra.Diagonal(u)
        D = LinearAlgebra.Diagonal(v)
        M_ = copy(M)

        rows_M_ = SparseArrays.rowvals(M_)
    end

    for iter in 1:max_iters
        @timeit "update diags" begin
            E.diag .= exp.(u)
            D.diag .= exp.(v)
        end
        @timeit "M_" begin
            LinearAlgebra.mul!(M_, M, D)
            LinearAlgebra.mul!(M_, E, M_)
        end

        step_size = 2. / (γ * (iter + 1.))

        # u gradient step
        @timeit "norms" begin
            fill!(row_norms, 0.0)
            fill!(col_norms, 0.0)
            for col in 1:aff.n
                for line in SparseArrays.nzrange(M_, col)
                    row_norms[rows_M_[line]] += abs2(M_[rows_M_[line], col])
                    col_norms[col] += abs2(M_[rows_M_[line], col])
                end
            end
        end
        @timeit "grad" begin
            @. u_grad = row_norms - α2 + γ * u
            @. v_grad = col_norms - β2 + γ * v
        end
        @timeit "proj " begin
            # u
            @. u -= step_size * u_grad
            box_project!(u, lb, ub)
            # v
            @. v -= step_size * v_grad
            sum_v = sum(v)
            v .= sum_v / aff.n
            box_project!(v, 0., ub)
        end
        # Update averages.
        @timeit "update" begin
            @. u_ = 2 * u / (iter + 2) + iter * u_ / (iter + 2)
            @. v_ = 2 * v / (iter + 2) + iter * v_ / (iter + 2)
        end
    end

    @timeit "update diags" begin
        E.diag .= exp.(u_)
        D.diag .= exp.(v_)
    end

    return E, D
end

function box_project!(y, lb, ub)
    y .= min.(ub, max.(y, lb))
    nothing
end


text:
# -*- encoding: utf-8 -*-

export NobsMatrixElement, Observation, DestripingData, CleanedTOD
export update_nobs!, update_nobs_matrix!, observation, compute_nobs_matrix!
export update_binned_map!, reset_maps!, compute_z_and_subgroup!, compute_z_and_group!
export compute_residuals!, array_dot, calc_stopping_factor, apply_offset_to_baselines!
export calculate_cleaned_map!, destripe!

import LinearAlgebra: I, Symmetric, cond, diagm, dot
import RunLengthArrays: RunLengthArray, runs, values
import Healpix
import Logging: @debug, @info, @warn
import Logging: ConsoleLogger, global_logger, Info, Debug
import Statistics: mean, std
import Printf: @sprintf

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

# This implementation of the destriping algorithm is based on the
# paper «Destriping CMB temperature and polarization maps» by
# Kurki-Suonio et al. 2009, A&A 506, 1511–1539 (2009),
# https://dx.doi.org/10.1051/0004-6361/200912361
#
# It is important to have that paper at hand while reading this code,
# as many functions and variable defined here use the same letters and
# symbols of that paper. We refer to it in code comments and
# docstrings as "KurkiSuonio2009".

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

include("mapmaking_nobsmatrixelement.jl")
include("mapmaking_observations.jl")
include("mapmaking_destripingdata.jl")
include("mapmaking_cleanedtod.jl")
include("mapmaking_preconditioners.jl")
include("mapmaking_covmatrix.jl")

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

@doc raw"""
    compute_nobs_matrix!(nobs_matrix::Vector{NobsMatrixElement}, observations::Vector{Observation{T}}; comm = nothing)

Initialize all the elements in `nobs_matrix` so that they are the
matrices M_p in Eq. (10) of KurkiSuonio2009. The TOD is taken from the
list of observations in the parameter `observations`.

"""
function compute_nobs_matrix!(
    nobs_matrix::Vector{NobsMatrixElement{T}},
    observations::Vector{Observation{T}};
    comm = nothing,
) where {T <: Real}

    for submatrix in nobs_matrix
        submatrix.m .= 0.0
    end

    for idx in 1:length(observations)
        update_nobs_matrix!(
            nobs_matrix,
            observations[idx].psi_angle,
            observations[idx].sigma_squared,
            observations[idx].pixidx,
            observations[idx].flagged,
            comm = comm,
        )
    end

    for idx in 1:length(nobs_matrix)
        update_nobs!(nobs_matrix[idx])
    end
end

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

function update_binned_map!(
    vec,
    skymap::Healpix.PolarizedMap{T, O},
    hitmap::Healpix.Map{T, O},
    obs::Observation{T};
    comm = nothing,
    unseen = missing,
) where {T <: Real, O <: Healpix.Order}

    # We use "zip" to iterate over those arrays that might not be
    # plain Julia Arrays
    for (idx, vecsamp, sigmasamp, mapidx) in zip(1:length(vec),
                                                 vec,
                                                 obs.sigma_squared,
                                                 obs.pixidx)

        sin_term, cos_term = sincos(2 * obs.psi_angle[idx])

        skymap.i[mapidx] += vecsamp / sigmasamp
        skymap.q[mapidx] += vecsamp * cos_term / sigmasamp
        skymap.u[mapidx] += vecsamp * sin_term / sigmasamp
        hitmap[mapidx] += 1 / sigmasamp
    end

    if use_mpi() && comm != nothing
        MPI.Allreduce!(MPI.IN_PLACE, skymap.i.pixels, MPI.SUM, comm)
        MPI.Allreduce!(MPI.IN_PLACE, skymap.q.pixels, MPI.SUM, comm)
        MPI.Allreduce!(MPI.IN_PLACE, skymap.u.pixels, MPI.SUM, comm)
        MPI.Allreduce!(MPI.IN_PLACE, hitmap.pixels, MPI.SUM, comm)
    end
end

function update_binned_map!(
    skymap::Healpix.PolarizedMap{T, O},
    hitmap::Healpix.Map{T, O},
    obs::Observation{T};
    comm = nothing,
    unseen = missing,
) where {T <: Real, O <: Healpix.Order}

    update_binned_map!(
        obs.tod,
        skymap,
        hitmap,
        obs,
        comm = comm,
        unseen = unseen,
    )
end

function update_binned_map!(
    vec_list,
    skymap::Healpix.PolarizedMap{T, O},
    hitmap::Healpix.Map{T, O},
    obs_list::Vector{Observation{T}};
    comm = nothing,
    unseen = missing,
) where {T <: Real, O <: Healpix.Order}

    @assert length(vec_list) == length(obs_list)

    for (cur_vec, cur_obs) in zip(vec_list, obs_list)
        update_binned_map!(
            cur_vec,
            skymap,
            hitmap,
            cur_obs,
            comm = comm,
            unseen = unseen,
        )
    end
end

function update_binned_map!(
    skymap::Healpix.PolarizedMap{T, O},
    hitmap::Healpix.Map{T, O},
    obs_list::Vector{Observation{T}};
    comm = nothing,
    unseen = missing,
) where {T <: Real, O <: Healpix.Order}

    for cur_obs in obs_list
        update_binned_map!(
            skymap,
            hitmap,
            cur_obs,
            comm = comm,
            unseen = unseen,
        )
    end
end

@doc raw"""
    update_binned_map!(vec, skymap::PolarizedMap{T, O}, hitmap::Healpix.Map{T, O}, obs::Observation{T}; comm=nothing, unseen=NaN) where {T <: Real, O <: Healpix.Order}
    update_binned_map!(skymap::PolarizedMap{T, O}, hitmap::Healpix.Map{T, O}, obs::Observation{T}; comm=nothing, unseen=NaN) where {T <: Real, O <: Healpix.Order}
    update_binned_map!(skymap::PolarizedMap{T, O}, hitmap::Healpix.Map{T, O}, obs_list::Vector{Observation{T}}; comm=nothing, unseen=NaN) where {T <: Real, O <: Healpix.Order}

Apply Eqs. (18), (19), and (20) to compute the values of I, Q, and U
from a TOD, and save the result in `skymap` and `hitmap`.

The three versions differ for the source of the TOD calibrated data:

- if the `vec` parameter is present (first form), it will be used
  instead of `obs.tod` (second form)
- instead of `obs` (a single [`Observation`](@ref) object), you can
  pass an array `obs_list` (third form). All the elements in this
  array will be used to update `skymap` and `hitmap`.

"""
update_binned_map!

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

function compute_z_and_subgroup!(
    dest_baselines,
    vec,
    baseline_lengths,
    destriping_data::DestripingData{T, O},
    obs::Observation{T};
    comm = nothing,
    unseen = missing,
) where {T, O <: Healpix.Order}

    # WARNING: unlike compute_z_and_group!, this does *not* update
    # `destriping_data`!

    length(baseline_lengths) > 0 || return
    @assert sum(baseline_lengths) == length(obs.tod)
    @assert length(baseline_lengths) == length(dest_baselines)

    iqu = Array{Float64}(undef, 3)

    @inbounds @simd for idx in eachindex(dest_baselines)
        dest_baselines[idx] = 0.0
    end

    baseline_idx = 1
    samples_in_baseline = 0
    for (idx, sigmasamp, mapsamp) in zip(eachindex(obs.pixidx),
                                         obs.sigma_squared,
                                         obs.pixidx)
        iqu[1] = destriping_data.skymap.i[mapsamp]
        iqu[2] = destriping_data.skymap.q[mapsamp]
        iqu[3] = destriping_data.skymap.u[mapsamp]

        iqu .= destriping_data.nobs_matrix[mapsamp].invm * iqu
        sin_term, cos_term = sincos(2 * obs.psi_angle[idx])
        signal = (iqu[1] + iqu[2] * cos_term + iqu[3] * sin_term)
        dest_baselines[baseline_idx] += (vec[idx] - signal) / sigmasamp

        samples_in_baseline += 1
        if samples_in_baseline == baseline_lengths[baseline_idx]
            baseline_idx += 1
            samples_in_baseline = 0
        end
    end
end

function compute_z_and_subgroup!(
    dest_baselines,
    baseline_lengths,
    destriping_data::DestripingData{T, O},
    obs::Observation{T};
    comm = nothing,
    unseen = missing,
) where {T, O <: Healpix.Order}

    # WARNING: unlike compute_z_and_group!, this does *not* update
    # `destriping_data`!

    compute_z_and_group!(
        dest_baselines,
        obs.tod,
        baseline_lengths,
        destriping_data,
        obs,
        comm = comm,
        unseen = unseen,
    )
end

@doc raw"""
    compute_z_and_subgroup!(dest_baselines, vec, baseline_lengths, destriping_data::DestripingData{T, O}, obs::Observation{T}; comm=nothing, unseen=NaN) where {T, O <: Healpix.Order}
    compute_z_and_subgroup!(dest_baselines, baseline_lengths, destriping_data::DestripingData{T, O}, obs::Observation{T}; comm=nothing, unseen=NaN) where {T, O <: Healpix.Order}

Compute the application of matrix ``A`` in Eq. (25) in
KurkiSuonio2009. The result is a set of baselines, and it is saved in
`dest_baselines` (an array of elements whose type is `T`). The two
forms differ only in the presence of the `vec` parameter: if it is not
present, the value of `obs.tod` will be used.

This function only operates on single observations. If you want to
apply it to an array of observations, use
[`compute_z_and_group!`](@ref).

"""
compute_z_and_subgroup!

function compute_z_and_group!(
    dest_baselines,
    vectors,
    baseline_lengths,
    destriping_data::DestripingData{T, O},
    obs_list::Vector{Observation{T}};
    comm = nothing,
    unseen = missing,
) where {T, O <: Healpix.Order}

    @assert length(obs_list) == length(vectors)
    @assert length(obs_list) == length(baseline_lengths)
    @assert length(obs_list) == length(dest_baselines)

    for (idx, cur_obs) in enumerate(obs_list)
        update_binned_map!(
            vectors[idx],
            destriping_data.skymap,
            destriping_data.hitmap,
            cur_obs,
            comm = comm,
            unseen = unseen,
        )
    end

    for (idx, cur_obs, cur_baseline_length) in zip(1:length(obs_list), obs_list, baseline_lengths)
        @assert length(cur_baseline_length) == length(dest_baselines[idx])

        compute_z_and_subgroup!(
            dest_baselines[idx],
            vectors[idx],
            cur_baseline_length,
            destriping_data,
            cur_obs,
            comm = comm,
            unseen = unseen,
        )
    end
end

function compute_z_and_group!(
    dest_baselines,
    baseline_lengths,
    destriping_data::DestripingData{T, O},
    obs_list::Vector{Observation{T}};
    comm = nothing,
    unseen = missing,
) where {T, O <: Healpix.Order}

    @assert length(obs_list) == length(baseline_lengths)
    @assert length(obs_list) == length(dest_baselines)

    for cur_obs in obs_list
        update_binned_map!(
            cur_obs.tod,
            destriping_data.skymap,
            destriping_data.hitmap,
            cur_obs,
            comm = comm,
            unseen = unseen,
        )
    end

    for (idx, cur_obs, cur_baseline_length) in zip(1:length(obs_list), obs_list, baseline_lengths)
        @assert length(cur_baseline_length) == length(dest_baselines[idx])

        compute_z_and_subgroup!(
            dest_baselines[idx],
            cur_obs.tod,
            cur_baseline_length,
            destriping_data,
            cur_obs,
            comm = comm,
            unseen = unseen,
        )
    end
end

@doc raw"""
    compute_z_and_group!(dest_baselines, vectors, baseline_lengths, destriping_data::DestripingData{T, O}, obs_list::Vector{Observation{T}}; comm=nothing, unseen=NaN) where {T, O <: Healpix.Order}
    compute_z_and_group!(dest_baselines, baseline_lengths, destriping_data::DestripingData{T, O}, obs_list::Vector{Observation{T}}; comm=nothing, unseen=NaN) where {T, O <: Healpix.Order}

Compute the application of matrix ``A`` in Eq. (25) in
KurkiSuonio2009. The result is a set of baselines, and it is saved in
`dest_baselines` (an array of elements whose type is `T`). The two
forms differ only in the presence of the `vec` parameter: if it is not
present, the value of `obs.tod` will be used.

"""
compute_z_and_group!

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

@doc raw"""
    compute_residuals!(residuals, baseline_lengths, destriping_data, obs_list; comm=nothing, unseen=missing)

This function is used to compute the value of ``A x - b`` at the
beginning of the Conjugated-Gradient algorithm implemented in
[`destripe!`](@ref).

"""
function compute_residuals!(
    residuals,
    baseline_lengths,
    destriping_data::DestripingData{T, O},
    obs_list::Vector{Observation{T}};
    comm = nothing,
    unseen = missing,
) where {T, O <: Healpix.Order}

    @assert length(residuals) == length(obs_list)
    @assert length(residuals) == length(baseline_lengths)

    reset_maps!(destriping_data)

    compute_z_and_group!(
        residuals,
        baseline_lengths,
        destriping_data,
        obs_list,
        comm = comm,
        unseen = unseen,
    )

    left_side = [Array{T}(undef, length(residuals[idx])) for idx in 1:length(residuals)]

    reset_maps!(destriping_data)

    compute_z_and_group!(
        left_side,
        destriping_data.baselines,
        baseline_lengths,
        destriping_data,
        obs_list,
        comm = comm,
        unseen = unseen,
    )
end

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

function array_dot(
    x::Array{RunLengthArray{N, T}},
    y;
    comm = nothing,
) where {N <: Integer, T <: Real}

    @assert length(x) == length(y)

    result = 0.0
    for (curx, cury) in zip(x, y)
        result += dot(values(curx), cury)
    end

    if use_mpi() && comm != nothing
        result = MPI.Allreduce(result, +, comm)
    end

    result
end

function array_dot(x, y; comm = nothing)
    @assert length(x) == length(y)

    result = 0.0
    for (curx, cury) in zip(x, y)
        result += dot(curx, cury)
    end

    if use_mpi() && comm != nothing
        result = MPI.Allreduce(result, +, comm)
    end

    result
end

@doc raw"""
    array_dot(x, y; comm = nothing)

Compute the dot product between `x` and `y`, assuming that both are
arrays of arrays.

"""
array_dot

@doc raw"""
    calc_stopping_factor(r; comm=nothing)

Calculate the stopping factor required by the Conjugated Gradient
method to determine if the iteration must be stopped or not. The
parameter `r` is the vector of the residuals.

"""
calc_stopping_factor(r; comm = nothing) = maximum(abs.(Iterators.flatten(r)))

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

@doc raw"""
    apply_offset_to_baselines!(baselines)

Add a constant offset to all the baselines such that their sum is zero.

"""
function apply_offset_to_baselines!(
    baselines::Array{RunLengthArray{Int, T}, 1};
    comm = nothing,
) where {T <: Real}

    baseline_sum = 0.0
    baseline_num = 0
    for cur_baseline in baselines
        baseline_sum += sum(values(cur_baseline))
        baseline_num += length(values(cur_baseline))
    end

    if use_mpi() && comm != nothing
        baseline_sum = MPI.Allreduce(baseline_sum, +, comm)
        baseline_num = MPI.Allreduce(baseline_num, +, comm)
    end

    baseline_sum /= baseline_num

    for cur_baseline in baselines
        cur_baseline.values .-= baseline_sum
    end
end

@doc raw"""
    calculate_cleaned_map!(obs_list, destriping_data; comm=nothing, unseen=missing)

Provided a set of baselines in `baselines`, this function cleans the
TOD in `obs_list` and computes the I/Q/U maps, which are saved in
`destriping_data`.

"""
function calculate_cleaned_map!(
    obs_list::Vector{Observation{T}},
    destriping_data::DestripingData{T, O};
    comm = nothing,
    unseen = missing,
) where {T <: Real, O <: Healpix.Order}

    cleaned_tod = CleanedTOD{Float64}([], [])
    reset_maps!(destriping_data)

    for (cur_obs, cur_baseline) in zip(obs_list, destriping_data.baselines)
        cleaned_tod.tod = cur_obs.tod
        cleaned_tod.baselines = cur_baseline

        update_binned_map!(
            cleaned_tod,
            destriping_data.skymap,
            destriping_data.hitmap,
            cur_obs,
            comm = comm,
            unseen = unseen,
        )
    end

    # Now solve for the I, Q, U parameters
    iqu = Array{Float64}(undef, 3)
    for curpix in eachindex(destriping_data.nobs_matrix)
        iqu[1] = destriping_data.skymap.i[curpix]
        iqu[2] = destriping_data.skymap.q[curpix]
        iqu[3] = destriping_data.skymap.u[curpix]

        iqu .= destriping_data.nobs_matrix[curpix].invm * iqu

        destriping_data.skymap.i[curpix] = iqu[1]
        destriping_data.skymap.q[curpix] = iqu[2]
        destriping_data.skymap.u[curpix] = iqu[3]
    end
end

@doc raw"""
    destripe!(obs_list, destriping_data::DestripingData{T,O}; comm, unseen, callback) where {T <: Real, O <: Healpix.Order}

Apply the destriping algorithm to the TOD in `obs_list`. The result is
returned in `destriping_data`, which is of type
[`DestripingData`](@ref) and contains the baselines, I/Q/U maps,
weight map, etc.

The parameters must satisfy the following constraints:

- `obs_list` must be an array of `N` [`Observation`](@ref) objects;
- `destriping_data` must be a [`DestripingData`](@ref) object. It does
  not need to have been initialized using [`reset_maps!`](@ref).

The optional keyword have the following meaning and default value:

- `comm` is a MPI communicator object, or `nothing` is MPI is not
  used.
- `unseen` is the value to be used for pixels in the map that have not
  been observed. The default is `missing`; other sensible values are
  `nothing` and `NaN`.
- `callback` is either `nothing` (the default) or a function that is
  called before the CG algorithm starts, and then once every
  iteration. It can be used to monitor the progress of the destriper,
  or to save intermediate results in a database or a file. The
  function receives `destriping_data` as its parameter

The function uses Julia's `Logging` module. Therefore, you can enable
the output of diagnostic messages as usual; consider the following
example:

```julia
global_logger(ConsoleLogger(stderr, Debug))

# This command is going to produce a lot of output…
destripe!(...)
```

"""
function destripe!(
    obs_list::Vector{Observation{T}},
    destriping_data::DestripingData{T, O};
    comm = nothing,
    unseen = missing,
    callback = nothing,
) where {T <: Real, O <: Healpix.Order}

    @assert length(obs_list) == length(destriping_data.baselines)

    # This variable is not going to change during the execution of the function
    baseline_lengths = [runs(destriping_data.baselines[idx])
                        for idx in 1:length(obs_list)]

    r = [Array{T}(undef, length(values(destriping_data.baselines[idx])))
         for idx in 1:length(obs_list)]
    rnext = deepcopy(r)

    compute_nobs_matrix!(destriping_data.nobs_matrix, obs_list, comm = comm)

    compute_residuals!(
        r,
        baseline_lengths,
        destriping_data,
        obs_list,
        comm = comm,
        unseen = unseen,
    )
    k = 0

    destriping_data.stopping_factors = T[]

    stopping_factor = calc_stopping_factor(r, comm = comm)
    push!(destriping_data.stopping_factors, stopping_factor)
    stopping_factor < destriping_data.threshold && return

    precond = if destriping_data.use_preconditioner
        jacobi_preconditioner(
            [runs(b) for b in destriping_data.baselines],
            T
        )
    else
        IdentityPreconditioner()
    end

    z = deepcopy(r)
    apply!(precond, z)

    old_z_dot_r = array_dot(z, r, comm = comm)
    p = [RunLengthArray{Int, T}(runs(destriping_data.baselines[idx]), z[idx])
         for idx in 1:length(destriping_data.baselines)]
    best_stopping_factor = stopping_factor
    best_a = deepcopy(destriping_data.baselines)

    ap = [Array{T}(undef, length(r[idx])) for idx in 1:length(r)]
    while true
        k += 1
        @debug "CG iteration $k/$(destriping_data.max_iterations), stopping factor = $stopping_factor"

        if k >= destriping_data.max_iterations
            @warn "The destriper did not converge after $k iterations"
            destriping_data.baselines = best_a
            break
        end

        reset_maps!(destriping_data)
        compute_z_and_group!(
            ap,
            p,
            baseline_lengths,
            destriping_data,
            obs_list,
            comm = comm,
            unseen = unseen,
        )
        γ = old_z_dot_r / array_dot(p, ap, comm = comm)
        for obsidx in 1:length(obs_list)
            destriping_data.baselines[obsidx].values .+= γ * p[obsidx].values
            rnext[obsidx] .= r[obsidx] - γ * ap[obsidx]
        end
        # Remove the mean value from the baselines
        apply_offset_to_baselines!(destriping_data.baselines, comm = comm)

        stopping_factor = calc_stopping_factor(r, comm = comm)
        push!(destriping_data.stopping_factors, stopping_factor)
        if stopping_factor < best_stopping_factor
            best_a = deepcopy(destriping_data.baselines)
            best_stopping_factor = stopping_factor
        end
        if stopping_factor < destriping_data.threshold
            @info "Destriper converged after $k iterations, stopping factor ($(stopping_factor)) < threshold ($(destriping_data.threshold))"
            break
        end

        z = deepcopy(rnext)
        apply!(precond, z)
        new_z_dot_r = array_dot(z, rnext, comm = comm)

        for obsidx in 1:length(obs_list)
            r[obsidx] .= rnext[obsidx]
            p[obsidx].values .= z[obsidx] + (new_z_dot_r / old_z_dot_r) * p[obsidx].values
        end

        (callback === nothing) || callback(destriping_data)
        old_z_dot_r = new_z_dot_r
    end

    @debug "Producing the map"
    calculate_cleaned_map!(obs_list, destriping_data, comm = comm, unseen = unseen)
end


text:
module PropagatorUnitTests

using Test
using UnitfulAstrodynamics
using UnitfulAstrodynamics.TwoBody.Systems: Earth

@testset "Propagators" begin

    # Twobody Orbit
    r̅ = [0.0, 11681, 0.0] * u"km"
    v̅ = [5.134, 4.226, 2.787] * u"km/s"
    orbit = CartesianOrbit(r̅, v̅, Earth, 0u"s")

    # Propagate twobody
    sols = propagate(orbit, 5u"s"; save_everystep=false)
    @test isapprox(sols[end], kepler(orbit, 5u"s"; tol=1e-14), atol=1e-6)

end

end


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

using Markdown
using InteractiveUtils

# ╔═╡ c788534a-ee16-4972-91ce-b0b3ed2799fa
md"""
# Quick Start Guide

This guide is for people who:

- doesn't know much Julia but knows Python (or other programming languages)
- knows quantum computing basics
- want to try out programming quantum circuits quickly with Julia

the contents are organized as following

1. Julia Basics
2. Basic Concepts of Yao
3. How to write a Quantum Circuit Born Machine using Yao
"""

# ╔═╡ 01ebf798-c068-4115-9cb9-381f5d8e8b35
md"""
## Julia Basics

You can jump to next section if you already know how to use Julia. We will only be introduce some basic concepts of Julia and differences comparing to Python in this section, if you wish to learn this language more seriously please refer to the learning materials on the official website: [julialang.org/learning/](https://julialang.org/learning/). This short Julia tutorial is based on [MIT computational thinking course](https://computationalthinking.mit.edu/).

If you haven't installed Julia, please refer to the front page of Julia: [julialang.org](https://julialang.org/)

If you wish to run this notebook locally, you will need to install [Pluto](https://github.com/fonsp/Pluto.jl):

1. open your [Julia interactive session (known as REPL)](https://docs.julialang.org/en/v1/manual/getting-started/)
2. press `]` key in the REPL to use the package mode, then type the following command

```julia
pkg> add Pluto
```

you will now see Julia's package manager start downloading this package. After the download is finished, press `backspace` and run the following command.

```julia
julia> import Pluto

julia> Pluto.run()
```

now it should open a web page for you in the browser, choose the downloaded notebook file `quick-start.jl`, you are good to go!
"""

# ╔═╡ Cell order:
# ╟─c788534a-ee16-4972-91ce-b0b3ed2799fa
# ╟─01ebf798-c068-4115-9cb9-381f5d8e8b35


text:
function implicitUseLockModel(lockType,conf,data)

	Val = ""
	if lockType == "waitTime"
		Val = useLockModel([0.125 0.0001 conf], data, "TPCC").TimeSpentWaiting
	elseif lockType == "numberOfLocks"
		Val = useLockModel([0.125 0.0001 conf], data, "TPCC").LocksBeingHeld
	elseif lockType == "numberOfConflicts"
		Val = useLockModel([0.125 0.0001 conf], data, "TPCC").totalWaits
	else
		error(string("Invalid lockType:", lockType))
	end
	return Val
end

max_stars_repo_name: barzan/dbseer

max_issues_repo_name: barzan/dbseer

max_forks_repo_name: barzan/dbseer


text:
module TestSubDataFrame

using Test, DataFrames

@testset "copy - SubDataFrame" begin
    df = DataFrame(x = 1:10, y = 1.0:10.0)
    sdf = view(df, 1:2, 1:1)
    @test sdf isa SubDataFrame
    @test copy(sdf) isa DataFrame
    @test sdf == copy(sdf)
    @test view(sdf, :, :) === sdf
    @test view(sdf, :, r"") == sdf
end

@testset "view -- DataFrame" begin
    df = DataFrame(x = 1:10, y = 1.0:10.0)
    @test view(df, 1, :) == DataFrameRow(df, 1, :)
    @test view(df, UInt(1), :) == DataFrameRow(df, 1, :)
    @test view(df, BigInt(1), :) == DataFrameRow(df, 1, :)
    @test view(df, UInt(1):UInt(1), :) == SubDataFrame(df, 1:1, :)
    @test view(df, BigInt(1):BigInt(1), :) == SubDataFrame(df, 1:1, :)
    @test view(df, 1:2, :) == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), :) == first(df, 2)
    @test view(df, [1, 2], :) == first(df, 2)

    @test view(df, 1, r"") == DataFrameRow(df, 1, :)
    @test view(df, UInt(1), r"") == DataFrameRow(df, 1, :)
    @test view(df, BigInt(1), r"") == DataFrameRow(df, 1, :)
    @test view(df, UInt(1):UInt(1), r"") == SubDataFrame(df, 1:1, :)
    @test view(df, BigInt(1):BigInt(1), r"") == SubDataFrame(df, 1:1, :)
    @test view(df, 1:2, r"") == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), r"") == first(df, 2)
    @test view(df, [1, 2], r"") == first(df, 2)

    @test view(df, 1, r"") == DataFrameRow(df, 1, r"")
    @test view(df, UInt(1), r"") == DataFrameRow(df, 1, r"")
    @test view(df, BigInt(1), r"") == DataFrameRow(df, 1, r"")
    @test view(df, UInt(1):UInt(1), r"") == SubDataFrame(df, 1:1, r"")
    @test view(df, BigInt(1):BigInt(1), r"") == SubDataFrame(df, 1:1, r"")

    @test view(df, 1, :x) == view(df[!, :x], 1)
    @test view(df, 1, :x) isa SubArray
    @test size(view(df, 1, :x)) == ()
    @test view(df, 1:2, :x) == df[!, :x][1:2]
    @test view(df, 1:2, :x) isa SubArray
    @test view(df, vcat(trues(2), falses(8)), :x) == view(df[!, :x], vcat(trues(2), falses(8)))
    @test view(df, [1, 2], :x) == view(df[!, :x], [1, 2])
    @test view(df, 1, 1) == view(df[!, 1], 1)
    @test view(df, 1, 1) isa SubArray
    @test size(view(df, 1, 1)) == ()
    @test view(df, 1:2, 1) == df[!, 1][1:2]
    @test view(df, 1:2, 1) isa SubArray
    @test view(df, vcat(trues(2), falses(8)), 1) == view(df[!, 1], vcat(trues(2), falses(8)))
    @test view(df, [1, 2], 1) == view(df[!, 1], [1,2])
    @test view(df, 1:2, 1) == df[!, 1][1:2]
    @test view(df, 1:2, 1) isa SubArray

    @test view(df, 1, [:x, :y]) == DataFrameRow(df[:, [:x, :y]], 1, :)
    @test view(df, 1, [:x, :y]) == DataFrameRow(df, 1, [:x, :y])
    @test view(df, 1:2, [:x, :y]) == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), [:x, :y]) == first(df, 2)
    @test view(df, [1, 2], [:x, :y]) == first(df, 2)

    @test view(df, 1, r"[xy]") == DataFrameRow(df[:, [:x, :y]], 1, :)
    @test view(df, 1, r"[xy]") == DataFrameRow(df, 1, [:x, :y])
    @test view(df, 1:2, r"[xy]") == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), r"[xy]") == first(df, 2)
    @test view(df, [1, 2], r"[xy]") == first(df, 2)

    @test view(df, 1, [1, 2]) == DataFrameRow(df[:, 1:2], 1, :)
    @test view(df, 1, [1, 2]) == DataFrameRow(df, 1, 1:2)
    @test view(df, 1:2, [1, 2]) == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), [1, 2]) == first(df, 2)
    @test view(df, [1, 2], [1, 2]) == first(df, 2)

    @test view(df, 1, trues(2)) == DataFrameRow(df[:, trues(2)], 1, :)
    @test view(df, 1, trues(2)) == DataFrameRow(df, 1, trues(2))
    @test view(df, 1:2, trues(2)) == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), trues(2)) == first(df, 2)
    @test view(df, [1, 2], trues(2)) == first(df, 2)

    @test view(df, Integer[1, 2], :) == first(df, 2)
    @test view(df, UInt[1, 2], :) == first(df, 2)
    @test view(df, BigInt[1, 2], :) == first(df, 2)
    @test view(df, Union{Int, Missing}[1, 2], :) == first(df, 2)
    @test view(df, Union{Integer, Missing}[1, 2], :) == first(df, 2)
    @test view(df, Union{UInt, Missing}[1, 2], :) == first(df, 2)
    @test view(df, Union{BigInt, Missing}[1, 2], :) == first(df, 2)

    @test view(df, :, :) == df
    @test view(df, 1, :) == DataFrameRow(df, 1, :)
    @test view(df, :, 1) == df[:, 1]
    @test view(df, :, 1) isa SubArray

    @test view(df, 1, r"") == DataFrameRow(df, 1, :)
    @test view(df, :, 1) == df[:, 1]
    @test view(df, :, 1) isa SubArray

    @test size(view(df, :, r"a")) == (0, 0)
    @test size(view(df, 1:2, r"a")) == (0, 0)
    @test size(view(df, 2, r"a")) == (0,)

    @test_throws ArgumentError view(df, :, [missing, 1])
    @test_throws ArgumentError view(df, [missing, 1], :)
    @test_throws ArgumentError view(df, [missing, 1], r"")
end

@testset "view -- SubDataFrame" begin
    df = view(DataFrame(x = 1:10, y = 1.0:10.0), 1:10, :)

    @test view(df, 1, :) == DataFrameRow(df, 1, :)
    @test view(df, UInt(1), :) == DataFrameRow(df, 1, :)
    @test view(df, BigInt(1), :) == DataFrameRow(df, 1, :)
    @test view(df, 1:2, :) == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), :) == first(df, 2)
    @test view(df, [1, 2], :) == first(df, 2)

    @test view(df, 1, r"") == DataFrameRow(df, 1, :)
    @test view(df, UInt(1), r"") == DataFrameRow(df, 1, :)
    @test view(df, BigInt(1), r"") == DataFrameRow(df, 1, :)
    @test view(df, 1:2, r"") == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), r"") == first(df, 2)
    @test view(df, [1, 2], r"") == first(df, 2)

    @test view(df, 1, r"") == DataFrameRow(df, 1, r"")
    @test view(df, UInt(1), r"") == DataFrameRow(df, 1, r"")
    @test view(df, BigInt(1), r"") == DataFrameRow(df, 1, r"")

    @test view(df, 1, :x) == view(df[!, :x], 1)
    @test view(df, 1, 1) isa SubArray
    @test size(view(df, 1, 1)) == ()
    @test view(df, 1:2, :x) == view(df[!, :x], 1:2)
    @test view(df, vcat(trues(2), falses(8)), :x) == view(df[!, :x], vcat(trues(2), falses(8)))
    @test view(df, [1, 2], :x) == view(df[!, :x], [1, 2])

    @test view(df, 1, 1) == view(df[!, 1], 1)
    @test view(df, 1, 1) isa SubArray
    @test size(view(df, 1, 1)) == ()
    @test view(df, 1:2, 1) == view(df[!, :x], 1:2)
    @test view(df, vcat(trues(2), falses(8)), 1) == view(df[!, :x], vcat(trues(2), falses(8)))
    @test view(df, [1, 2], 1) == view(df[!, :x], [1, 2])

    @test view(df, 1, [:x, :y]) == DataFrameRow(df[:, [:x,:y]], 1, :)
    @test view(df, 1, [:x, :y]) == DataFrameRow(df, 1, [:x,:y])
    @test view(df, 1:2, [:x, :y]) == first(df, 2)

    @test view(df, 1, r"[xy]") == DataFrameRow(df[:, [:x,:y]], 1, :)
    @test view(df, 1, r"[xy]") == DataFrameRow(df, 1, [:x,:y])
    @test view(df, 1:2, r"[xy]") == first(df, 2)

    @test view(df, vcat(trues(2), falses(8)), [:x, :y]) == first(df, 2)
    @test view(df, [1, 2], [:x, :y]) == first(df, 2)
    @test view(df, 1, [1, 2]) == DataFrameRow(df[:, 1:2], 1, :)
    @test view(df, 1, [1, 2]) == DataFrameRow(df, 1, 1:2)
    @test view(df, 1:2, [1, 2]) == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), [1, 2]) == first(df, 2)
    @test view(df, [1, 2], [1, 2]) == first(df, 2)
    @test view(df, 1, trues(2)) == DataFrameRow(df[:, trues(2)], 1, :)
    @test view(df, 1, trues(2)) == DataFrameRow(df, 1, trues(2))
    @test view(df, 1:2, trues(2)) == first(df, 2)
    @test view(df, vcat(trues(2), falses(8)), trues(2)) == first(df, 2)
    @test view(df, [1, 2], trues(2)) == first(df, 2)
    @test view(df, Integer[1, 2], :) == first(df, 2)
    @test view(df, UInt[1, 2], :) == first(df, 2)
    @test view(df, BigInt[1, 2], :) == first(df, 2)
    @test view(df, Union{Int, Missing}[1, 2], :) == first(df, 2)
    @test view(df, Union{Integer, Missing}[1, 2], :) == first(df, 2)
    @test view(df, Union{UInt, Missing}[1, 2], :) == first(df, 2)
    @test view(df, Union{BigInt, Missing}[1, 2], :) == first(df, 2)

    @test view(df, :, :) == df
    @test view(df, 1, :) == DataFrameRow(df, 1, :)
    @test view(df, :, 1) == df[:, 1]
    @test view(df, :, 1) isa SubArray

    @test view(df, :, r"") == df
    @test view(df, 1, r"") == DataFrameRow(df, 1, :)

    @test_throws ArgumentError view(df, :, [missing, 1])
    @test_throws ArgumentError view(df, [missing, 1], :)
    @test_throws ArgumentError view(df, [missing, 1], r"")
    @test_throws ArgumentError view(df, :, true)
end

@testset "getproperty, setproperty! and propertynames" begin
    x = collect(1:10)
    y = collect(1.0:10.0)
    df = view(DataFrame(:x=>x, :y=>y, copycols=false), 2:6, :)

    @test Base.propertynames(df) == names(df)

    @test df.x == 2:6
    @test df.y == 2:6
    @test_throws ArgumentError df.z

    df[:, :x] = 1:5
    @test df.x == 1:5
    @test x == [1; 1:5; 7:10]
    df[:, :y] .= 1
    @test df.y == [1, 1, 1, 1, 1]
    @test y == [1; 1; 1; 1; 1; 1; 7:10]
end

@testset "index" begin
    y = 1.0:10.0
    df = view(DataFrame(y=y), 2:6, :)
    df2 = view(DataFrame(x=y, y=y), 2:6, 2:2)
    @test DataFrames.index(df) == DataFrames.index(df2)
    @test haskey(DataFrames.index(df2), :y)
    @test !haskey(DataFrames.index(df2), :x)
    @test haskey(DataFrames.index(df2), 1)
    @test !haskey(DataFrames.index(df2), 2)
    @test !haskey(DataFrames.index(df2), 0)
    @test_throws ArgumentError haskey(DataFrames.index(df2), true)
    @test keys(DataFrames.index(df2)) == [:y]

    x = DataFrame(ones(5,4))
    df = view(x, 2:3, 2:3)
    @test names(df) == names(x)[2:3]
    df = view(x, 2:3, [4,2])
    @test names(df) == names(x)[[4,2]]
end

@testset "deleterows!" begin
    y = 1.0:10.0
    df = view(DataFrame(y=y), 2:6, :)
    @test_throws ArgumentError deleterows!(df, 1)
end

@testset "parent" begin
    df = DataFrame(a=Union{Int, Missing}[1, 2, 3, 1, 2, 2],
                   b=[2.0, missing, 1.2, 2.0, missing, missing],
                   c=["A", "B", "C", "A", "B", missing])
    @test parent(view(df, [4, 2], :)) === df
    @test parent(view(df, [4, 2], r"")) === df
    @test parentindices(view(df, [4, 2], :)) == ([4,2], Base.OneTo(3))
    @test parentindices(view(df, [4, 2], r"")) == ([4,2], [1,2,3])
    @test parent(view(df, [4, 2], 1:3)) === df
    @test parentindices(view(df, [4, 2], 1:3)) == ([4, 2], Base.OneTo(3))
end

@testset "duplicate column" begin
    df = DataFrame([11:16 21:26 31:36 41:46])
    sdf = view(df, [3,1,4], [3,3,3])
    @test names(sdf) == fill(:x3, 3)
    @test sdf[!, 1] == [33, 31, 34]
    @test sdf[!, 1] === sdf[!, 2] === sdf[!, 3]
    @test sdf.x3[1] == 33
    sdf.x3[1] = 333
    @test df.x3[3] == 333
    @test_throws ArgumentError sdf.x1
end

@testset "conversion to DataFrame" begin
    df = DataFrame([11:16 21:26 31:36 41:46])
    sdf = view(df, [3,1,4], [3,2,1])
    df2 = DataFrame(sdf)
    @test df2 isa DataFrame
    @test df2 == df[[3,1,4], [3,2,1]]
    @test all(x -> x isa Vector{Int}, eachcol(df2))
    df2 = convert(DataFrame, sdf)
    @test df2 isa DataFrame
    @test df2 == df[[3,1,4], [3,2,1]]
    @test all(x -> x isa Vector{Int}, eachcol(df2))

    df = DataFrame(x=1:4, y=11:14, z=21:24)
    sdf = @view df[2:3, [2]]
    df2 = DataFrame(sdf)
    @test size(df2) == (2, 1)
    @test df2.y isa Vector{Int}
    @test df2.y == [12, 13]
    df2 = DataFrame(sdf, copycols=false)
    @test size(df2) == (2, 1)
    @test df2.y isa SubArray{Int,1,Vector{Int},Tuple{UnitRange{Int}},true}
    @test df2.y == [12, 13]
end

end # module


text:
############################################################
## joAbstractLinearOperator - outer constructors
############################################################

############################################################
## joMatrix

"""
joMatrix outer constructor

    joMatrix(array::AbstractMatrix;
             DDT::DataType=eltype(array),
             RDT::DataType=promote_type(eltype(array),DDT),
             name::String="joMatrix")

Look up argument names in help to joMatrix type.

# Example
- joMatrix(rand(4,3)) # implicit domain and range
- joMatrix(rand(4,3);DDT=Float32) # implicit range
- joMatrix(rand(4,3);DDT=Float32,RDT=Float64)
- joMatrix(rand(4,3);name="my matrix") # adding name

# Notes
- if DDT:<Real for complex matrix then imaginary part will be neglected for transpose/adjoint operator
- if RDT:<Real for complex matrix then imaginary part will be neglected for forward/conjugate operator

"""
function joMatrix(array::AbstractMatrix{EDT};
    DDT::DataType=EDT,RDT::DataType=promote_type(EDT,DDT),name::String="joMatrix") where {EDT}

        (typeof(array)<:DArray || typeof(array)<:SharedArray) && @warn "Creating joMatrix from non-local array like $(typeof(array)) is likely going to have adverse impact on JOLI's health. Please, avoid it."

        return joMatrix{DDT,RDT}(name,size(array,1),size(array,2),
            v1->jo_convert(RDT,array*v1,false),
            v2->jo_convert(DDT,transpose(array)*v2,false),
            v3->jo_convert(DDT,adjoint(array)*v3,false),
            v4->jo_convert(RDT,conj(array)*v4,false),
            v5->jo_convert(DDT,array\v5,false),
            v6->jo_convert(RDT,transpose(array)\v6,false),
            v7->jo_convert(RDT,adjoint(array)\v7,false),
            v8->jo_convert(DDT,conj(array)\v8,false)
            )
end

############################################################
## joLinearFunction

export joLinearFunctionAll, joLinearFunction_T, joLinearFunction_A,
       joLinearFunctionFwd, joLinearFunctionFwd_T, joLinearFunctionFwd_A

# outer constructors

"""
joLinearFunction outer constructor

    joLinearFunctionAll(m::Integer,n::Integer,
        fop::Function,fop_T::Function,fop_A::Function,fop_C::Function,
        iop::Function,iop_T::Function,iop_A::Function,iop_C::Function,
        DDT::DataType,RDT::DataType=DDT;
        fMVok::Bool=false,iMVok::Bool=false,
        name::String="joLinearFunctionAll")

Look up argument names in help to joLinearFunction type.

# Notes
- the developer is responsible for ensuring that used functions take/return correct DDT/RDT

"""
joLinearFunctionAll(m::Integer,n::Integer,
    fop::Function,fop_T::Function,fop_A::Function,fop_C::Function,
    iop::Function,iop_T::Function,iop_A::Function,iop_C::Function,
    DDT::DataType,RDT::DataType=DDT;
    fMVok::Bool=false,iMVok::Bool=false,
    name::String="joLinearFunctionAll") =
        joLinearFunction{DDT,RDT}(name,m,n,
            fop,fop_T,fop_A,fop_C,fMVok,
            iop,iop_T,iop_A,iop_C,iMVok
            )
"""
joLinearFunction outer constructor

    joLinearFunction_T(m::Integer,n::Integer,
        fop::Function,fop_T::Function, iop::Function,iop_T::Function,
        DDT::DataType,RDT::DataType=DDT;
        fMVok::Bool=false,iMVok::Bool=false,
        name::String="joLinearFunction_T")

Look up argument names in help to joLinearFunction type.

# Notes
- the developer is responsible for ensuring that used functions take/return correct DDT/RDT

"""
joLinearFunction_T(m::Integer,n::Integer,
    fop::Function,fop_T::Function, iop::Function,iop_T::Function,
    DDT::DataType,RDT::DataType=DDT;
    fMVok::Bool=false,iMVok::Bool=false,
    name::String="joLinearFunction_T") =
        joLinearFunction{DDT,RDT}(name,m,n,
            fop,
            fop_T,
            v3->conj(fop_T(conj(v3))),
            v4->conj(fop(conj(v4))),
            fMVok,
            iop,
            iop_T,
            v7->conj(iop_T(conj(v7))),
            v8->conj(iop(conj(v8))),
            iMVok
            )
"""
joLinearFunction outer constructor

    joLinearFunction_A(m::Integer,n::Integer,
        fop::Function,fop_A::Function, iop::Function,iop_A::Function,
        DDT::DataType,RDT::DataType=DDT;
        fMVok::Bool=false,iMVok::Bool=false,
        name::String="joLinearFunction_A")

Look up argument names in help to joLinearFunction type.

# Notes
- the developer is responsible for ensuring that used functions take/return correct DDT/RDT

"""
joLinearFunction_A(m::Integer,n::Integer,
    fop::Function,fop_A::Function, iop::Function,iop_A::Function,
    DDT::DataType,RDT::DataType=DDT;
    fMVok::Bool=false,iMVok::Bool=false,
    name::String="joLinearFunction_A") =
        joLinearFunction{DDT,RDT}(name,m,n,
            fop,
            v2->conj(fop_A(conj(v2))),
            fop_A,
            v4->conj(fop(conj(v4))),
            fMVok,
            iop,
            v6->conj(iop_A(conj(v6))),
            iop_A,
            v8->conj(iop(conj(v8))),
            iMVok
            )
"""
joLinearFunction outer constructor

    joLinearFunctionFwd(m::Integer,n::Integer,
        fop::Function,fop_T::Function,fop_A::Function,fop_C::Function,
        DDT::DataType,RDT::DataType=DDT;
        fMVok::Bool=false,
        name::String="joLinearFunctionFwd")

Look up argument names in help to joLinearFunction type.

# Notes
- the developer is responsible for ensuring that used functions take/return correct DDT/RDT

"""
joLinearFunctionFwd(m::Integer,n::Integer,
    fop::Function,fop_T::Function,fop_A::Function,fop_C::Function,
    DDT::DataType,RDT::DataType=DDT;
    fMVok::Bool=false,
    name::String="joLinearFunctionFwd") =
        joLinearFunction{DDT,RDT}(name,m,n,
            fop,fop_T,fop_A,fop_C,fMVok,
            @joNF, @joNF, @joNF, @joNF, false
            )
"""
joLinearFunction outer constructor

    joLinearFunctionFwd_T(m::Integer,n::Integer,
        fop::Function,fop_T::Function,
        DDT::DataType,RDT::DataType=DDT;
        fMVok::Bool=false,
        name::String="joLinearFunctionFwd_T")

Look up argument names in help to joLinearFunction type.

# Notes
- the developer is responsible for ensuring that used functions take/return correct DDT/RDT

"""
joLinearFunctionFwd_T(m::Integer,n::Integer,
    fop::Function,fop_T::Function,
    DDT::DataType,RDT::DataType=DDT;
    fMVok::Bool=false,
    name::String="joLinearFunctionFwd_T") =
        joLinearFunction{DDT,RDT}(name,m,n,
            fop,
            fop_T,
            v3->conj(fop_T(conj(v3))),
            v4->conj(fop(conj(v4))),
            fMVok,
            @joNF, @joNF, @joNF, @joNF, false
            )
"""
joLinearFunction outer constructor

    joLinearFunctionFwd_A(m::Integer,n::Integer,
        fop::Function,fop_A::Function,
        DDT::DataType,RDT::DataType=DDT;
        fMVok::Bool=false,
        name::String="joLinearFunctionFwd_A")

Look up argument names in help to joLinearFunction type.

# Notes
- the developer is responsible for ensuring that used functions take/return correct DDT/RDT

"""
joLinearFunctionFwd_A(m::Integer,n::Integer,
    fop::Function,fop_A::Function,
    DDT::DataType,RDT::DataType=DDT;
    fMVok::Bool=false,
    name::String="joLinearFunctionFwd_A") =
        joLinearFunction{DDT,RDT}(name,m,n,
            fop,
            v2->conj(fop_A(conj(v2))),
            fop_A,
            v4->conj(fop(conj(v4))),
            fMVok,
            @joNF, @joNF, @joNF, @joNF, false
            )

############################################################
## joAddSolver{Any,All}

# exported in joAbstractOperator/constructors.jl

# joAddSolver - outer constructor for adding solver to operator
export joAddSolverAny
"""
joAddSolver outer constructor

    joAddSolverAny(A::joAbstractLinearOperator{DDT,RDT},solver::Function)

Create joLinearOperator with added solver for \\(jo,[m]vec),
same for each form of the operator

# Example (for all forms of O)

    O=joAddSolverAny(O,(s,x)->my_solver(s,x))

"""
joAddSolverAny(A::joAbstractLinearOperator{DDT,RDT},slvr::Function) where {DDT,RDT} =
        joLinearOperator{DDT,RDT}("("*A.name*"+solver)",A.m,A.n,
            v1->A*v1,
            v2->transpose(A)*v2,
            v3->adjoint(A)*v3,
            v4->conj(A)*v4,
            v5->slvr(A,v5),
            v6->slvr(transpose(A),v6),
            v7->slvr(adjoint(A),v7),
            v8->slvr(conj(A),v8)
            )
export joAddSolverAll
"""
joAddSolver outer constructor

    joAddSolverAll(A::joAbstractLinearOperator{DDT,RDT},
        solver::Function,solver_T::Function,solver_A::Function,solver_C::Function)

Create joLinearOperator with added specific solver(s) for \\(jo,[m]vec),
distinct for each form of the operator.

# Examples

    O=joAddSolverAll(O,
        (s,x)->my_solver(s,x),
        (s,x)->my_solver_T(s,x),
        (s,x)->my_solver_A(s,x),
        (s,x)->my_solver_C(s,x))

    O=joAddSolverAll(O,
        (s,x)->my_solver(s,x),
        @joNF,
        (s,x)->my_solver_A(s,x),
        @joNF)

    O=joAddSolverAll(O,
        (s,x)->my_solver(s,x),
        @joNF,
        @joNF,
        @joNF)

"""
joAddSolverAll(A::joAbstractLinearOperator{DDT,RDT},
    slvr::Function,slvr_T::Function,slvr_A::Function,slvr_C::Function) where {DDT,RDT} =
        joLinearOperator{DDT,RDT}("("*A.name*"+solver)",A.m,A.n,
            v1->A*v1,
            v2->transpose(A)*v2,
            v3->adjoint(A)*v3,
            v4->conj(A)*v4,
            v5->slvr(A,v5),
            v6->slvr_T(transpose(A),v6),
            v7->slvr_A(adjoint(A),v7),
            v8->slvr_C(conj(A),v8)
            )


text:
export prox_l2s!

function prox_l2s!(x::Vector{TF}, rho::TF, m::Vector{TF}) where {TF<:Real}
  x .= (x .* rho .+ m) ./ (rho .+ 1.0);
  return x
end

max_stars_repo_name: slimgroup/SetIntersectionProjection

max_issues_repo_name: slimgroup/SetIntersectionProjection

max_forks_repo_name: slimgroup/SetIntersectionProjection


text:
"""
    LiftingLineSectionControl <: AbstractModel

Lifting line section control input model with state variable ``\\delta``, zero
inputs, and zero parameters.  Two-dimensional control surface models may be
extended to three dimensional models by coupling with this model.  Note that this
model has no rate equations of its own since its state variables are defined as
functions of the 3D system's control inputs.
"""
struct LiftingLineSectionControl <: AbstractModel end

"""
    LiftingLineSectionControl()

Initialize an object of type [`LiftingLineSectionControl`](@ref)
"""

# --- Traits --- #

number_of_states(::Type{<:LiftingLineSectionControl}) = 1
number_of_inputs(::Type{<:LiftingLineSectionControl}) = 0
number_of_parameters(::Type{<:LiftingLineSectionControl}) = 0

inplaceness(::Type{<:LiftingLineSectionControl}) = OutOfPlace()

# --- Internal Methods for Couplings --- #


text:
using Dates

length(ARGS) != 1 && error("Need to provide the kernel number as an argument.")

num = lpad(ARGS[1], 4, "0")
file = "naif$num.tls"
url = "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/lsk/" * file
download(url, file)

const MJD = 2400000.5

t = Vector{Int}()
leapseconds = Vector{Float64}()
re = r"(?<dat>[0-9]{2}),\s+@(?<date>[0-9]{4}-[A-Z]{3}-[0-9])"
lines = open(readlines, file)
for line in lines
    s = string(line)
    if occursin(re, s)
        m = match(re, s)
        push!(leapseconds, parse(Float64, m["dat"]))
        push!(t, datetime2julian(DateTime(m["date"], "y-u-d")) - MJD)
    end
end

open("leap_seconds.jl", "w") do f
    write(f, """
    # Automatically generated by $(basename(@__FILE__)), do not edit!

    const LS_EPOCHS = [
        $(join(t, ",\n    ")),
    ]

    const LEAP_SECONDS = [
        $(join(leapseconds, ",\n    ")),
    ]
    """)
end


text:
# This is an example which uses much of the functionality provided by
# Interplanetary.jl. It is based on a homework assignment, but that is
# no guarentee that it's correct.
#
# The objective is to explore optimal trajectories from Earth to
# asteroid 2015PDC - a fictional asteroid created for a hypothetical
# asteroid impact scenario held at the 2015 Planetary Defense
# Conference (see http://neo.jpl.nasa.gov/pdc15/). This is
# accomplished via a Pork Chop Plot - a contour plot showing the value
# of some objective function (mass at arrival, delta V, etc.) over a
# grid of Earth departure dates and time of flight values.
#
# This particular example searches for the minimum-delta V solution. I
# also included some of my experiments with using Julia's built-in
# parallel computing capabilities, because it seemed like a good place
# for parallelization.

using DelimitedFiles, Dates, LinearAlgebra

function meshgrid(x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat}
    return (repeat(x, 1, length(y)),
             repeat(y',length(x),1))
end

"""
This function creates a Pork Chop Plot given a range of departure
dates, time of flight for a given grid point (trajectory), and the
figure of merit value for a given trajectory.
"""
function createPCC( departure_date_range::Vector{T},
                   TOF_range::Vector{T},
                   FOM_array::Matrix{T}; levels = [],
                   FOM_label::String = "") where {T<:AbstractFloat}
    (depart_grid, TOF_grid) = meshgrid(departure_date_range, TOF_range)

    if isempty(levels)
        cplot = contourf(depart_grid, TOF_grid, FOM_array)
    else
        cplot = contourf(depart_grid, TOF_grid, FOM_array, levels)
    end

    # formatting
    ylocs,ylabels = yticks()
    xlocs,xlabels = xticks()
    cb = colorbar()
    cb.set_label(FOM_label)

    xticks(xlocs,Date.(Dates.julian2datetime.(mjd2jd.(xlocs))),rotation=30)
    #yticks(ylocs,Date(julian2datetime(mjd2jd(ylocs))))
    xlabel("Departure date")
    ylabel("TOF [days]")
    # need to set axis limits
end

# read ephemeris files for Earth and Asteroid 2015PDC
Earth_ephem = readdlm("data/Earth_2015-04-01_to_2022-09-03_step_1d.orb")
ast_2015_PDC_ephem = readdlm("data/ast_2015_PDC_2015-04-01_to_2022-09-03_step_1d.orb")

#@everywhere begin # need to add cores before running this: addprocs(n)
using Interplanetary, PyPlot
departure_lb_mjd = jd2mjd( Dates.datetime2julian( DateTime("2015-10-01") ) )
departure_ub_mjd = jd2mjd( Dates.datetime2julian( DateTime("2022-02-26") ) )

# set up ranges for grid search
departure_step = 5. #days
TOF_step = 5. #days

departure_date_range = collect(departure_lb_mjd:departure_step:departure_ub_mjd)
TOF_range = collect(90.:TOF_step:1500.)

# now loop over all combinations

dV_total = Array{Float64}(undef,length(departure_date_range),length(TOF_range)) # for non-parallel
#dV_total = SharedArray{Float64}(length(departure_date_range),length(TOF_range))

rPark = RE + 185. #parking orbit radius
#end

#@time @sync @parallel for i = 1:length(departure_date_range)
for i = 1:length(departure_date_range)
    for j = 1:length(TOF_range)
        dep_date = departure_date_range[i]
        arr_date = dep_date + TOF_range[j]
        if arr_date <= ast_2015_PDC_ephem[end,1]
            (dv1,dv2,c3,rla,dla) = lambert_transfer(muEarth,
                                        TOF_range[j]*86400.,
                                        mjd2jd(departure_date_range[i]),
                                        rPark,
                                        vec(Earth_ephem[Earth_ephem[:,1].==dep_date,2:7]),
                                        vec(ast_2015_PDC_ephem[ast_2015_PDC_ephem[:,1].==arr_date,2:7]))
            dV_total[i,j] = dv1+dv2
        end#if
    end#for
end#for

# create the contour plot
# sdata() makes a SharedArray into a regular array
createPCC(departure_date_range,
          TOF_range,sdata(dV_total),
          levels = collect(5.:2:30),
          FOM_label = "Delta V [km/s]")

# find the minimum dV solution
# linear index of minimum, excluding zeros:
ind = findall(dV_total .== minimum(dV_total[dV_total.>0]))[1]

subs = CartesianIndices(dV_total)[ind] # convert to array subscripts
opt_arrival_date = departure_date_range[subs[1][1]] + TOF_range[subs[2][1]]

# plot the optimal solution on the pork chop plot
plot(departure_date_range[subs[1][1]],TOF_range[subs[2][1]],"r*",markersize=10)

# rerun the Lambert targeter for the optimal solution
(dv1,dv2,c3,rla,dla) = lambert_transfer(muEarth,
                            TOF_range[subs[2][1]]*86400.,
                            mjd2jd(departure_date_range[subs[1][1]]),
                            rPark,
                            vec(Earth_ephem[Earth_ephem[:,1].==departure_date_range[subs[1][1]],2:7]),
                            vec(ast_2015_PDC_ephem[ast_2015_PDC_ephem[:,1].==opt_arrival_date,2:7]))

# print results to screen
println("Optimal solution:")
println("Departure date: ", Date.(Dates.julian2datetime.(mjd2jd(departure_date_range[subs[1][1]]))))
println("TOF (days): ",TOF_range[subs[2][1]])
println("Earth departure ΔV [km/s]: ", dv1)
println("Earth departure C3 [km^2/s^2]: ", c3)
println("DLA (deg): ", rad2deg(dla))
println("Arrival ΔV [km/s]: ", dv2)
println("Total solution ΔV [km/s]: ", dV_total[subs][1])


text:
function test_zygote(f, xs...; kws...)
    config = ZygoteRuleConfig()
    test_rrule(config, f, xs...; kws..., rrule_f = rrule_via_ad)
end


text:
module RulesDotProductIn1Test

using Test
using ReactiveMP
using Random

import ReactiveMP: @test_rules
import LinearAlgebra: dot

@testset "rules:typeof(dot):in1" begin

     @testset "Belief Propagation: (m_out::UnivariateNormalDistributionsFamily, m_in2::PointMass)" begin

        @test_rules [ with_float_conversions = true ] typeof(dot)(:in1, Marginalisation) [
            (input = (m_out = NormalMeanVariance(2.0, 2.0), m_in2 = PointMass(-1.0), meta = NoCorrection()), output = NormalWeightedMeanPrecision(-1.0, 0.5)),
            (input = (m_out = NormalMeanPrecision(1.0, 1.0), m_in2 = PointMass(-2.0), meta = NoCorrection()), output = NormalWeightedMeanPrecision(-2.0, 4.0)),
            (input = (m_out = NormalWeightedMeanPrecision(2.0, 1.0), m_in2 = PointMass(-1.0), meta = NoCorrection()), output = NormalWeightedMeanPrecision(-2.0, 1.0)),
        ]

        @test_rules [ with_float_conversions = true ] typeof(dot)(:in1, Marginalisation) [
            (input = (m_out = NormalMeanVariance(2.0, 2.0), m_in2 = PointMass(-1.0), meta = TinyCorrection()), output = NormalWeightedMeanPrecision(-1.0, 0.5)),
            (input = (m_out = NormalMeanPrecision(1.0, 1.0), m_in2 = PointMass(-2.0), meta = TinyCorrection()), output = NormalWeightedMeanPrecision(-2.0, 4.0)),
            (input = (m_out = NormalWeightedMeanPrecision(2.0, 1.0), m_in2 = PointMass(-1.0), meta = TinyCorrection()), output = NormalWeightedMeanPrecision(-2.0, 1.0)),
        ]

        @test_rules [ with_float_conversions = true ] typeof(dot)(:in1, Marginalisation) [
            (input = (m_out = NormalMeanVariance(2.0, 2.0), m_in2 = PointMass(0.0), meta = NoCorrection()), output = NormalWeightedMeanPrecision(0.0, 0.0)),
            (input = (m_out = NormalMeanPrecision(1.0, 1.0), m_in2 = PointMass(0.0), meta = NoCorrection()), output = NormalWeightedMeanPrecision(0.0, 0.0)),
            (input = (m_out = NormalWeightedMeanPrecision(2.0, 1.0), m_in2 = PointMass(0.0), meta = NoCorrection()), output = NormalWeightedMeanPrecision(0.0, 0.0)),
        ]

        @test_rules [ with_float_conversions = true ] typeof(dot)(:in1, Marginalisation) [
            (input = (m_out = NormalMeanVariance(2.0, 2.0), m_in2 = PointMass(0.0), meta = TinyCorrection()), output = NormalWeightedMeanPrecision(0.0, tiny)),
            (input = (m_out = NormalMeanPrecision(1.0, 1.0), m_in2 = PointMass(0.0), meta = TinyCorrection()), output = NormalWeightedMeanPrecision(0.0, tiny)),
            (input = (m_out = NormalWeightedMeanPrecision(2.0, 1.0), m_in2 = PointMass(0.0), meta = TinyCorrection()), output = NormalWeightedMeanPrecision(0.0, tiny)),
        ]

    end

    @testset "Belief Propagation: (m_out::UnivariateNormalDistributionsFamily, m_in2::PointMass)" begin

        @test_rules [ with_float_conversions = true ] typeof(dot)(:in1, Marginalisation) [
            (input = (m_out = NormalMeanVariance(2.0, 1.0), m_in2 = PointMass([-1.0, 2.0]), meta = NoCorrection()), output = MvNormalWeightedMeanPrecision([-2.0, 4.0], [1.0 -2.0; -2.0 4.0])),
            (input = (m_out = NormalMeanPrecision(1.0, inv(2.0)), m_in2 = PointMass([1.0, 1.0]), meta = NoCorrection()), output = MvNormalWeightedMeanPrecision([0.5, 0.5], [0.5 0.5; 0.5 0.5])),
            (input = (m_out = NormalWeightedMeanPrecision(2.0, 1.0), m_in2 = PointMass([-2.0, 3.0]), meta = NoCorrection()), output = MvNormalWeightedMeanPrecision([-4.0, 6.0], [4.0 -6.0; -6.0 9.0])),
        ]

        @test_rules [ with_float_conversions = false ] typeof(dot)(:in1, Marginalisation) [
            (input = (m_out = NormalMeanVariance(2.0, 1.0), m_in2 = PointMass([-1.0, 2.0]), meta = TinyCorrection()), output = MvNormalWeightedMeanPrecision([-2.0, 4.0], [1.0 + tiny -2.0; -2.0 4.0 + tiny])),
            (input = (m_out = NormalMeanPrecision(1.0, inv(2.0)), m_in2 = PointMass([1.0, 1.0]), meta = TinyCorrection()), output = MvNormalWeightedMeanPrecision([0.5, 0.5], [0.5 + tiny 0.5; 0.5 0.5 + tiny])),
            (input = (m_out = NormalWeightedMeanPrecision(2.0, 1.0), m_in2 = PointMass([-2.0, 3.0]), meta = TinyCorrection()), output = MvNormalWeightedMeanPrecision([-4.0, 6.0], [4.0 + tiny -6.0; -6.0 9.0 + tiny])),
        ]
    end

end

end