several fixes

docs/src/advanced/meanfield.md

@@ -41,7 +41,7 @@ julia> h(SA[])
 julia> g = greenfunction(h, GS.Spectrum());
 
 julia> ρ0 = densitymatrix(g[])(0.5, 0) ## density matrix at chemical potential `µ=0.5` and temperature `kBT = 0`  on all sites
-4×4 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+4×4 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
  0.138197+0.0im  0.223607+0.0im  0.223607+0.0im  0.138197+0.0im
  0.223607+0.0im  0.361803+0.0im  0.361803+0.0im  0.223607+0.0im
  0.223607+0.0im  0.361803+0.0im  0.361803+0.0im  0.223607+0.0im

docs/src/tutorial/greenfunctions.md

@@ -213,15 +213,15 @@ julia> g(0.2)
 GreenSolution{Float64,2,0}: Green function at arbitrary positions, but at a fixed energy
 
 julia> g(0.2)[1, 3]
-5×5 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+5×5 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
  -2.56906+0.000123273im  -4.28767+0.00020578im   -4.88512+0.000234514im  -4.28534+0.00020578im    -2.5664+0.000123273im
  -4.28767+0.00020578im   -7.15613+0.00034351im   -8.15346+0.000391475im  -7.15257+0.00034351im    -4.2836+0.000205781im
  -4.88512+0.000234514im  -8.15346+0.000391475im  -9.29002+0.000446138im  -8.14982+0.000391476im  -4.88095+0.000234514im
  -4.28534+0.00020578im   -7.15257+0.00034351im   -8.14982+0.000391476im  -7.14974+0.000343511im  -4.28211+0.000205781im
   -2.5664+0.000123273im   -4.2836+0.000205781im  -4.88095+0.000234514im  -4.28211+0.000205781im  -2.56469+0.000123273im
 
 julia> g(0.2)[siteselector(region = RP.circle(1, (0.5, 0))), 3]
-2×5 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+2×5 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
  0.0749214+3.15744e-8im   0.124325+5.27948e-8im   0.141366+6.01987e-8im   0.124325+5.27948e-8im  0.0749214+3.15744e-8im
  -0.374862+2.15287e-5im  -0.625946+3.5938e-5im   -0.712983+4.09561e-5im  -0.624747+3.59379e-5im   -0.37348+2.15285e-5im
 ```

docs/src/tutorial/observables.md

@@ -42,7 +42,7 @@ DensityMatrix: density matrix on specified sites using solver of type DensityMat
 
 julia> @time ρ(4)
   6.150548 seconds (57.84 k allocations: 5.670 GiB, 1.12% gc time)
-5×5 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+5×5 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
           0.5+0.0im          -7.34893e-10-3.94035e-15im  0.204478+1.9366e-14im   -7.34889e-10-1.44892e-15im  -5.70089e-10+5.48867e-15im
  -7.34893e-10+3.94035e-15im           0.5+0.0im          0.200693-2.6646e-14im   -5.70089e-10-1.95251e-15im  -7.34891e-10-2.13804e-15im
      0.204478-1.9366e-14im       0.200693+2.6646e-14im        0.5+0.0im              0.200693+3.55692e-14im      0.204779-4.27255e-14im
@@ -62,7 +62,7 @@ DensityMatrix: density matrix on specified sites with solver of type DensityMatr
 
 julia> @time ρ(4)
   0.001659 seconds (9 allocations: 430.906 KiB)
-5×5 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+5×5 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
           0.5+0.0im  -2.21437e-15+0.0im  0.204478+0.0im   2.67668e-15+0.0im   3.49438e-16+0.0im
  -2.21437e-15+0.0im           0.5+0.0im  0.200693+0.0im  -1.40057e-15+0.0im  -2.92995e-15+0.0im
      0.204478+0.0im      0.200693+0.0im       0.5+0.0im      0.200693+0.0im      0.204779+0.0im
@@ -81,7 +81,7 @@ DensityMatrix: density matrix on specified sites with solver of type DensityMatr
 
 julia> @time ρ(4)
   0.006580 seconds (3 allocations: 1.156 KiB)
-5×5 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+5×5 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
          0.5+0.0im  2.15097e-17+0.0im   0.20456+0.0im  2.15097e-17+0.0im   3.9251e-17+0.0im
  2.15097e-17+0.0im          0.5+0.0im  0.200631+0.0im  1.05873e-16+0.0im  1.70531e-18+0.0im
      0.20456+0.0im     0.200631+0.0im       0.5+0.0im     0.200631+0.0im      0.20482+0.0im

src/Quantica.jl

@@ -31,7 +31,7 @@ export sublat, bravais_matrix, lattice, sites, supercell, hamiltonian,
        greenfunction, selfenergy, attach,
        plotlattice, plotlattice!, plotbands, plotbands!, qplot, qplot!, qplotdefaults,
        conductance, josephson, ldos, current, transmission, densitymatrix,
-       OrbitalSliceArray, OrbitalSliceVector, OrbitalSliceMatrix, orbaxes,
+       OrbitalSliceArray, OrbitalSliceVector, OrbitalSliceMatrix, orbaxes, siteindexdict,
        serializer, serialize, serialize!, deserialize!, deserialize,
        meanfield
 

src/docstrings.jl

@@ -716,7 +716,7 @@ julia> h[sites(1), sites(2)]
  1.0+0.0im  0.0+0.0im
 
 julia> ph = h |> @hopping!((t; p = 3) -> p*t); ph[region = RP.square(1)]
-4×4 OrbitalSliceMatrix{SparseArrays.SparseMatrixCSC{ComplexF64, Int64}}:
+4×4 OrbitalSliceMatrix{ComplexF64,SparseArrays.SparseMatrixCSC{ComplexF64, Int64}}:
  0.0+0.0im  0.0+0.0im  0.0+0.0im  3.0+0.0im
  0.0+0.0im  0.0+0.0im  3.0+0.0im  0.0+0.0im
  0.0+0.0im  3.0+0.0im  0.0+0.0im  0.0+0.0im
@@ -1677,7 +1677,7 @@ julia> gω[ss] == gs(0.1; t = 2)
 true
 
 julia> summary(gω[ss])
-"14×14 OrbitalSliceMatrix{Matrix{ComplexF64}}"
+"14×14 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}"
 ```
 
 # See also
@@ -1738,14 +1738,14 @@ Equivalent to `diagonal(siteselector(; sites...); kernel)`
 julia> g = HP.graphene(orbitals = 2) |> attach(nothing, cells = (0,0)) |> greenfunction();
 
 julia> g(1)[diagonal(:)]                            # g(ω = 1) diagonal on all contact orbitals
-4×4 OrbitalSliceMatrix{LinearAlgebra.Diagonal{ComplexF64, Vector{ComplexF64}}}:
+4×4 OrbitalSliceMatrix{ComplexF64,LinearAlgebra.Diagonal{ComplexF64, Vector{ComplexF64}}}:
  -0.10919-0.0839858im       0.0+0.0im             0.0+0.0im             0.0+0.0im
       0.0+0.0im        -0.10919-0.0839858im       0.0+0.0im             0.0+0.0im
       0.0+0.0im             0.0+0.0im        -0.10919-0.0839858im       0.0+0.0im
       0.0+0.0im             0.0+0.0im             0.0+0.0im        -0.10919-0.0839858im
 
 julia> g(1)[diagonal(:, kernel = SA[1 0; 0 -1])]    # σz spin density of the above
-2×2 OrbitalSliceMatrix{LinearAlgebra.Diagonal{ComplexF64, Vector{ComplexF64}}}:
+2×2 OrbitalSliceMatrix{ComplexF64,LinearAlgebra.Diagonal{ComplexF64, Vector{ComplexF64}}}:
  5.61885e-12+1.38778e-17im           0.0+0.0im
          0.0+0.0im          -5.61882e-12+2.77556e-17im
 ```
@@ -1760,9 +1760,10 @@ diagonal
     sitepairs(s::HopSelector; kernel = missing)
 
 Create a selection of site pairs `s::SparseIndices` used to sparsely index into a
-`g::GreenFunction` or `g::GreenSolution`, as `g[s]`. Of the resulting matrix only the
-selected pairs of matrix elements will be computed, leaving the rest as zero (sparse
-matrix).
+`g::GreenFunction` or `g::GreenSolution`, as `g[s]`. Of the resulting `OrbitalSliceMatrix`
+only the selected pairs of matrix elements will be computed, leaving the rest as zero
+(sparse matrix). The sparse matrix spans the minimum number of complete unit cells to
+include all site pairs
 
 If `kernel = Q` (a matrix instead of `missing`), each of these site blocks `gᵢⱼ` will be
 replaced by `Tr(kernel * gᵢⱼ)`.
@@ -1780,10 +1781,10 @@ Equivalent to `sitepairs(hopselector(; hops...); kernel)`
 julia> g = HP.graphene(orbitals = 2, a0 = 1) |> attach(nothing, cells = (0,0)) |> greenfunction();
 
 julia> summary(g(1)[sitepairs(range = 1)])     # g(ω=1) site blocks between all sites in zero cell and all other sites at distance 1
-"28×4 OrbitalSliceMatrix{SparseArrays.SparseMatrixCSC{ComplexF64, Int64}}"
+"28×4 OrbitalSliceMatrix{ComplexF64,SparseArrays.SparseMatrixCSC{ComplexF64, Int64}}"
 
 julia> summary(g(1)[sitepairs(range = 1, kernel = SA[1 0; 0 -1])])    # σz spin density of the above
-"14×2 OrbitalSliceMatrix{SparseArrays.SparseMatrixCSC{ComplexF64, Int64}}"
+"14×2 OrbitalSliceMatrix{ComplexF64,SparseArrays.SparseMatrixCSC{ComplexF64, Int64}}"
 ```
 
 # See also
@@ -2090,7 +2091,7 @@ julia> ρ = densitymatrix(g[region = RP.circle(0.5)])
 DensityMatrix: density matrix on specified sites with solver of type DensityMatrixSpectrumSolver
 
 julia> ρ()  # with mu = kBT = 0 by default
-2×2 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+2×2 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
        0.5+0.0im  -0.262865+0.0im
  -0.262865+0.0im        0.5+0.0im
 ```
@@ -2256,7 +2257,7 @@ Note: `diagonal` indexing is currently not supported by `OrbitalSliceArray`.
 julia> g = LP.linear() |> hamiltonian(hopping(SA[0 1; 1 0]) + onsite(I), orbitals = 2) |> supercell(4) |> greenfunction;
 
 julia> mat = g(0.2)[region = r -> 2<=r[1]<=4]
-6×6 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+6×6 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
  -1.93554e-9-0.545545im          0.0-0.0im               0.0-0.0im              -0.5+0.218218im          0.4+0.37097im           0.0+0.0im
          0.0-0.0im       -1.93554e-9-0.545545im         -0.5+0.218218im          0.0-0.0im               0.0+0.0im               0.4+0.37097im
          0.0-0.0im              -0.5+0.218218im  -1.93554e-9-0.545545im          0.0-0.0im               0.0+0.0im              -0.5+0.218218im
@@ -2265,7 +2266,7 @@ julia> mat = g(0.2)[region = r -> 2<=r[1]<=4]
          0.0+0.0im               0.4+0.37097im          -0.5+0.218218im          0.0+0.0im               0.0-0.0im       -1.93554e-9-0.545545im
 
 julia> mat[(; cells = SA[1]), (; cells = SA[0])]
-2×4 OrbitalSliceMatrix{Matrix{ComplexF64}}:
+2×4 OrbitalSliceMatrix{ComplexF64,Matrix{ComplexF64}}:
  0.4+0.37097im  0.0+0.0im       0.0+0.0im       -0.5+0.218218im
  0.0+0.0im      0.4+0.37097im  -0.5+0.218218im   0.0+0.0im
 
@@ -2277,11 +2278,67 @@ julia> mat[sites(SA[1], 1)]
 
 # See also
     `siteselector`, `cellindices`, `orbaxes`
-
 """
 OrbitalSliceArray, OrbitalSliceVector, OrbitalSliceMatrix
 
 
+"""
+    orbaxes(A::OrbitalSliceArray)
+
+Return the orbital axes of `A`. This is a tuple of `OrbitalSliceGrouped` objects that can be
+used e.g. to index another `OrbitalSliceArray` or to inspect the indices of each site with
+`siteindexdict`.
+
+# Examples
+
+```jldoctest
+julia> g = HP.graphene(orbitals = 2) |> supercell((1,-1)) |> greenfunction;
+
+julia> d = ldos(g[cells = SA[0]])(2)
+8-element OrbitalSliceVector{Float64,Vector{Float64}}:
+ 0.1127904162648383
+ 0.11279041626483835
+ 0.06188774489933697
+ 0.06188774489933696
+ 0.061887744899337044
+ 0.06188774489933703
+ 0.11279041626483817
+ 0.11279041626483818
+
+julia> a = only(orbaxes(d))
+OrbitalSliceGrouped{Float64,2,1} : collection of subcells of orbitals (grouped by sites) for a 1D lattice in 2D space
+  Cells       : 1
+  Cell range  : ([0], [0])
+  Total sites : 4
+
+julia> siteindexdict(a)
+4-element Dictionaries.Dictionary{Quantica.CellIndices{1, Int64, Quantica.SiteLike}, UnitRange{Int64}}
+ CellSites{1,Int64} : 1 site in cell zero
+  Sites : 1 │ 1:2
+ CellSites{1,Int64} : 1 site in cell zero
+  Sites : 2 │ 3:4
+ CellSites{1,Int64} : 1 site in cell zero
+  Sites : 3 │ 5:6
+ CellSites{1,Int64} : 1 site in cell zero
+  Sites : 4 │ 7:8
+```
+
+# See also
+    `siteindexdict`
+"""
+orbaxes
+
+"""
+    siteindexdict(axis::OrbitalSliceGrouped)
+
+Return a dictionary of `CellSite`s representing single sites in an orbital axis of an
+`OrbitalSliceArray`, typically obtained with `orbaxes`. See `orbaxes` for an example.
+
+# See also
+    `orbaxes`, `OrbitalSliceArray`
+"""
+siteindexdict
+
 """
     serializer(T::Type, selectors...; parameter = :stream, encoder = identity, decoder = identity)
 

src/greenfunction.jl

@@ -227,7 +227,7 @@ function getindex_diagonal!(output, g, i::OrbitalSliceGrouped, ker; post = ident
 end
 
 function getindex_sparse!(output, g, hmask::Hamiltonian, ker; post = identity)
-    bs = blockstructure(hmask)
+    bs = blockstructure(g)  # used to find blocks of g for nonzeros in hmask
     oj = sites_to_orbs(sites(zerocell(g), :), g)
     for har in harmonics(hmask)
         oi = CellIndices(dcell(har), oj)

src/hamiltonian.jl

@@ -483,6 +483,9 @@ Base.getindex(h::Hamiltonian, i, j) = getindex(h, sites_to_orbs(i, h), sites_to_
 # we need AbstractHamiltonian here to avoid ambiguities with dn above
 Base.getindex(h::AbstractHamiltonian, i) = (i´ = sites_to_orbs(i, h); getindex(h, i´, i´))
 
+Base.getindex(h::AbstractHamiltonian, ::SparseIndices...) =
+    argerror("Sparse indexing not yet supported for AbstractHamiltonian")
+
 # wrapped matrix for end user consumption
 Base.getindex(h::Hamiltonian, i::OrbitalSliceGrouped, j::OrbitalSliceGrouped) =
     OrbitalSliceMatrix(

src/show.jl

@@ -57,7 +57,8 @@ displayname(s::Sublat) = sublatname(s) == Symbol(:_) ? "pending" : string(":", s
 
 function Base.show(io::IO, c::CellSites)
     i = get(io, :indent, "")
-    print(io, i, summary(c), "\n")
+    print(io, i, summary(c), "\n",
+"$i  Sites : $(siteindices(c))")
 end
 
 Base.summary(c::CellSites{L,V}) where {L,V} =
@@ -527,14 +528,14 @@ Base.summary(::CurrentDensitySlice{T}) where {T} =
 
 # For simplified printing of the array typename
 
-function Base.showarg(io::IO, ::OrbitalSliceMatrix{<:Any,M}, toplevel) where {M}
+function Base.showarg(io::IO, ::OrbitalSliceMatrix{T,M}, toplevel) where {T,M}
     toplevel || print(io, "::")
-    print(io,  "OrbitalSliceMatrix{$M}")
+    print(io,  "OrbitalSliceMatrix{$T,$M}")
 end
 
-function Base.showarg(io::IO, ::OrbitalSliceVector{<:Any,M}, toplevel) where {M}
+function Base.showarg(io::IO, ::OrbitalSliceVector{T,M}, toplevel) where {T,M}
     toplevel || print(io, "::")
-    print(io,  "OrbitalSliceVector{$M}")
+    print(io,  "OrbitalSliceVector{$T,$M}")
 end
 
 #endregion

src/specialmatrices.jl

@@ -550,6 +550,7 @@ function Base.view(a::OrbitalSliceMatrix, i::AnyCellSites, j::AnyCellSites = i)
     i´, j´ = apply(i, lattice(rowslice)), apply(j, lattice(colslice))
     rows = indexcollection(rowslice, i´)
     cols = j === i && rowslice === colslice ? rows : indexcollection(colslice, j´)
+    @show rows, cols
     return view(parent(a), rows, cols)
 end
 

src/tools.jl

@@ -137,8 +137,8 @@ lengths_to_offsets(v) = prepend!(cumsum(v), 0)
 lengths_to_offsets(f::Function, v) = prepend!(accumulate((i,j) -> i + f(j), v; init = 0), 0)
 
 
-# fast tr(A*B) without allocating an A*B product. Doesn't check dimensions or index offsets
-trace_prod(A::AbstractMatrix, B::AbstractMatrix) = sum(splat(*), zip(transpose(A), B))
+# fast tr(A*B)
+trace_prod(A::AbstractMatrix, B::AbstractMatrix) = (check_sizes(A,B); sum(splat(*), zip(transpose(A), B)))
 trace_prod(A::Number, B::AbstractMatrix) = A*tr(B)
 trace_prod(A::AbstractMatrix, B::Number) = trace_prod(B, A)
 trace_prod(A::UniformScaling, B::AbstractMatrix) = A.λ*tr(B)
@@ -151,6 +151,8 @@ trace_prod(A::Number, B::Diagonal) = trace_prod(B, A)
 trace_prod(A::Union{SMatrix,UniformScaling,Number}, B::Union{SMatrix,UniformScaling,Number}) =
      tr(A*B)
 
+check_sizes(A,B) = size(A,2) == size(B,1) || throw(DimensionMismatch("A has dimensions $(size(A)) but B has dimensions $(size(B))"))
+
 # Taken from Base julia, now deprecated there
 function permute!!(a, p::AbstractVector{<:Integer})
     Base.require_one_based_indexing(a, p)

test/test_greenfunction.jl

@@ -310,10 +310,18 @@ end
         @test g(0.2)[diagonal(i)] isa Union{OrbitalSliceMatrix, AbstractMatrix}
     end
     gg = g(ω)[sitepairs(range = 1/sqrt(3))]
-    @test gg isa OrbitalSliceMatrix{SparseMatrixCSC{ComplexF64, Int64}}
-    @test size(gg) == Quantica.flatsize(g) .* (length(Quantica.harmonics(hamiltonian(g))), 1)
-    gg = g(ω)[sitepairs(range = 1/sqrt(3), sublats = :A => :A, kernel = I)]
-    # gg = g(ω)[sitepairs(range = 1/sqrt(3), sublats = :A => :A, kernel = I)]
+    @test gg isa OrbitalSliceMatrix{ComplexF64,Quantica.SparseMatrixCSC{ComplexF64, Int64}}
+    @test size(gg) == Quantica.flatsize(g) .* (3, 1)
+    @test_throws DimensionMismatch g(ω)[sitepairs(range = 1, sublats = :B => :B, kernel = SA[1 0; 0 -1])]
+    gg = g(ω)[sitepairs(range = 1, sublats = :A => :A, kernel = I)]
+    @test size(gg) == size(g, 1) .* (3, 1)
+
+    g = HP.graphene(a0 = 1, t0 = 1, orbitals = (2,1)) |> supercell((1, -1)) |>
+           greenfunction(GS.Schur(boundary = -2))
+    ω = 0.6
+    gg = g(ω)[sitepairs(range = 1)]
+    gg´ = g(ω)[orbaxes(gg)...]
+    @test gg .* gg´ ≈ gg .* gg
 end
 
 @testset "densitymatrix diagonal slicing" begin