fixing and adding tests

several fixes

sitepairs test

removed docstring test

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
@@ -1734,20 +1734,20 @@ Equivalent to `diagonal(siteselector(; sites...); kernel)`
     - `kernel`: if missing, all orbitals in the diagonal `g[i, i]` are returned when indexing `g[diagonal(i)]`. Otherwise, `Tr(g[site, site]*kernel)` for each site included in `i` is returned.
 
 # Example
-```jldoctest
+```julia
 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.38036e-12+0.0im          0.0+0.0im
-          0.0+0.0im  2.38036e-12+0.0im
+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
 ```
 
 # See also
@@ -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ᵢⱼ)`.
@@ -1776,14 +1777,14 @@ Equivalent to `sitepairs(hopselector(; hops...); kernel)`
     - `kernel`: if missing, all orbitals blocks `gᵢⱼ = g[i, j]` between selected sites pairs (i,j) are returned when indexing `g[sitepairs(...)]`. Otherwise, `gᵢⱼ` is replaced by `Tr(gᵢⱼ*kernel)`.
 
 # Example
-```jldoctest
+```julia
 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

@@ -64,8 +64,8 @@ end
 call!(g::GreenSlice{T}, ω::T; kw...) where {T} =
     call!(g, retarded_omega(ω, solver(parent(g))); kw...)
 
-call!(g::GreenSlice{T}, ω::Complex{T}; post = identity, params...) where {T} =
-    getindex!(call!_output(g), call!(greenfunction(g), ω; params...), orbinds_or_contactinds(g)...; post)
+call!(gs::GreenSlice{T}, ω::Complex{T}; post = identity, params...) where {T} =
+    getindex!(gs, call!(greenfunction(gs), ω; params...); post)
 
 ## extrinsic version - unused
 # call!(output, g::GreenSlice, ω; kw...) = call!(replace_output!(g, output), ω; kw...)
@@ -103,8 +103,7 @@ Base.view(g::GreenSolution, i::CT, j::CT = i) where {CT<:Union{Integer,Colon}} =
 # args could be anything: CellSites, CellOrbs, Colon, Integer, DiagIndices...
 function Base.getindex(g::GreenSolution, args...; post = identity, kw...)
     gs = getindex(parent(g), args...; kw...)  # get GreenSlice
-    output = call!_output(gs)
-    getindex!(output, g, orbinds_or_contactinds(gs)...; post)
+    output = getindex!(gs, g; post)
     return output
 end
 
@@ -126,14 +125,6 @@ getindex!(output, g::GreenSolution, i::Union{Integer,Colon}, j::Union{Integer,Co
 getindex!(output, g, i::AnyCellOrbitals, j::AnyCellOrbitals; post = identity) =
     (output .= post.(g[i, j]); output)
 
-function getindex!(output::OrbitalSliceMatrix, g, i::AnyCellOrbitals, j::AnyCellOrbitals)
-    oi, oj = orbaxes(output)
-    rows, cols = orbrange(oi, cell(i)), orbrange(oj, cell(j))
-    v = view(parent(output), rows, cols)
-    getindex!(v, g, i, j)
-    return output
-end
-
 # indexing over several cells
 getindex!(output, g, ci::AnyOrbitalSlice, cj::AnyOrbitalSlice; kw...) =
     getindex_cells!(output, g, cellsdict(ci), cellsdict(cj); kw...)
@@ -142,25 +133,34 @@ getindex!(output, g, ci::AnyOrbitalSlice, cj::AnyCellOrbitals; kw...) =
 getindex!(output, g, ci::AnyCellOrbitals, cj::AnyOrbitalSlice; kw...) =
     getindex_cells!(output, g, (ci,), cellsdict(cj); kw...)
 
-# we use the fact that the output AbstractMatrix is populated in the same site order of cis, cjs
 function getindex_cells!(output, g, cis, cjs; kw...)
-    ioffset = joffset = 0
-    for ci in cis
-        ilen = length(ci)
-        irng = ioffset+1:ioffset+ilen
-        joffset = 0
-        for cj in cjs
-            jlen = length(cj)
-            jrng = joffset+1:joffset+jlen
-            v = view(output, irng, jrng)
-            getindex!(v, g, ci, cj; kw...)
-            joffset += jlen
-        end
-        ioffset += ilen
+    for ci in cis, cj in cjs
+        getindex!(output, g, ci, cj; kw...) # will typically call the method below
     end
     return output
 end
 
+function getindex!(output::OrbitalSliceMatrix, g, i::AnyCellOrbitals, j::AnyCellOrbitals; kw...)
+    oi, oj = orbaxes(output)
+    rows, cols = orbrange(oi, cell(i)), orbrange(oj, cell(j))
+    v = view(parent(output), rows, cols)
+    getindex!(v, g, i, j; kw...)
+    return output
+end
+
+## common getindex! shortcut in terms of GreenSlice
+
+# index object g over a slice enconded in gs::GreenSlice, using its preallocated output
+getindex!(gs::GreenSlice, g; kw...) = getindex!(call!_output(gs), g, orbinds_or_contactinds(gs)...; kw...)
+
+# ifelse(rows && cols are contacts, (rows, cols), (orbrows, orbcols))
+# I.e: if rows, cols are contact indices retrieve them instead of orbslices.
+orbinds_or_contactinds(g) = orbinds_or_contactinds(rows(g), cols(g), orbrows(g), orbcols(g))
+orbinds_or_contactinds(
+    r::Union{Colon,Integer,DiagIndices{Colon},DiagIndices{<:Integer}},
+    c::Union{Colon,Integer,DiagIndices{Colon},DiagIndices{<:Integer}}, _, _) = (r, c)
+orbinds_or_contactinds(_, _, or, oc) = (or, oc)
+
 ## GreenSlicer -> Matrix, dispatches to each solver's implementation
 
 # fallback conversion to CellOrbitals
@@ -181,7 +181,7 @@ Base.view(g::GreenSlicer, args...) =
 getindex!(output, g, i::DiagIndices, j::DiagIndices = i; kw...) =
     getindex_diagonal!(output, g, parent(i), kernel(i); kw...)
 
-getindex!(output, g, i::OrbitalPairIndices, j::OrbitalPairIndices = i) =
+getindex!(output, g, i::OrbitalPairIndices, j::OrbitalPairIndices = i; kw...) =
     getindex_sparse!(output, g, parent(i), kernel(i); kw...)
 
 #endregion
@@ -215,7 +215,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/observables.jl

@@ -223,8 +223,7 @@ Base.getindex(d::LocalSpectralDensitySolution; selectors...) =
 function Base.getindex(d::LocalSpectralDensitySolution{T}, i) where {T}
     di = diagonal(i, kernel = d.kernel)
     gs = GreenSlice(parent(d.gω), di, di, T)  # get GreenSlice with real output
-    output = call!_output(gs)
-    getindex!(output, d.gω, orbinds_or_contactinds(gs)...; post = x -> -imag(x)/π)
+    output = getindex!(gs, d.gω; post = x -> -imag(x)/π)
     return diag(output)
 end
 

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