add KPM diagonal support

src/greenfunction.jl

@@ -160,9 +160,6 @@ Base.getindex(gω::GreenSolution{T}, i::DiagIndices, ::DiagIndices = i) where {T
 append_diagonal!(d, x, i, kernel, g; kw...) =
     append_diagonal!(d, x, sites_to_orbs(i, g), kernel, g; kw...)
 
-# append_diagonal!(d, x, s::OrbitalSlice, kernel, g; kw...) =
-#     append_diagonal!(d, x, cellsdict(s), kernel, g; kw...)    # no OrbitalSliceVector here
-
 append_diagonal!(d, x, s::AnyOrbitalSlice, kernel, g; kw...) =
     append_diagonal!(d, x, cellsdict(s), kernel, g; kw...)
 
@@ -178,7 +175,7 @@ end
 function append_diagonal!(d, x, o::Union{AnyCellOrbitals,Colon,Integer}, kernel, g; post = identity)
     # Note that o can be a contact index (Colon or Integer), since sites_to_orbs doesn't
     # reduce these for performance (they are already precomputed if x::GreenSolution)
-    xblock = diagonal_slice(x, o)
+    xblock = maybe_diagonal_slice(x, o)
     rngs = orbranges_or_allorbs(kernel, o, g)
     for rng in rngs
         val = apply_kernel(kernel, xblock, rng)
@@ -187,9 +184,11 @@ function append_diagonal!(d, x, o::Union{AnyCellOrbitals,Colon,Integer}, kernel,
     return d
 end
 
-# fallback, may be overloaded for gω's that know how to do this more efficiently
+# generic fallback, may be specialized for gω's that know how to do this more efficiently
 # it should include all diagonal blocks for each site, not just the orbital diagonal
-diagonal_slice(gω, o) = gω[o, o]
+maybe_diagonal_slice(gω::GreenSolution, o) = gω[o, o]
+# no-op for objects different than GreenSolution. Useful for direct calls to append_diagonal!
+maybe_diagonal_slice(x, _) = x
 
 # If no kernel is provided, we return the whole diagonal
 orbranges_or_allorbs(kernel::Missing, o::AnyCellOrbitals, gω) = eachindex(orbindices(o))
@@ -210,7 +209,7 @@ apply_kernel(kernel::Diagonal, v::AbstractMatrix) = sum(i -> kernel[i] * v[i, i]
 apply_kernel(kernel::Diagonal, v::Number) = only(kernel) * v
 
 view_or_scalar(gblock, rng::UnitRange) = view(gblock, rng, rng)
-view_or_scalar(gblock, i::Integer) = gblock[i, i]
+view_or_scalar(gblock, orb::Integer) = gblock[orb, orb]
 
 maybe_scalarize(s::OrbitalSliceGrouped, kernel::Missing) = s
 maybe_scalarize(s::OrbitalSliceGrouped, kernel) = scalarize(s)

src/solvers/green/bands.jl

@@ -858,12 +858,12 @@ minimal_callsafe_copy(s::BandsGreenSlicer, parentham, parentcontacts) = s  # it
 #endregion
 
 ############################################################################################
-# diagonal_slice
+# maybe_diagonal_slice
 #   optimized block diagonal of g[::CellOrbitals] for BandsGreenSlicer without boundary
 #   otherwise we need to do it the normal way
 #region
 
-function diagonal_slice(gω::GreenSolution{T,<:Any,<:Any,G}, o::CellOrbitals) where {T,G<:BandsGreenSlicer{<:Any,Missing}}
+function maybe_diagonal_slice(gω::GreenSolution{T,<:Any,<:Any,G}, o::CellOrbitals) where {T,G<:BandsGreenSlicer{<:Any,Missing}}
     s = slicer(gω)   # BandsGreenSlicer
     norbs = flatsize(gω)
     rngs = orbranges(o)

src/solvers/green/kpm.jl

@@ -244,7 +244,7 @@ end
 ## Constructor
 
 function densitymatrix(s::AppliedKPMGreenSolver, gs::GreenSlice{T}) where {T}
-    check_nodiag_axes(gs)
+    # check_nodiag_axes(gs)
     has_selfenergy(gs) && argerror("The KPM densitymatrix solver currently support only `nothing` contacts")
     momenta = slice_momenta(s.momenta, gs)
     solver = DensityMatrixKPMSolver(momenta, s.bandCH)
@@ -254,14 +254,20 @@ end
 function slice_momenta(momenta, gs)
     r = _maybe_contactinds(gs, rows(gs))
     c = _maybe_contactinds(gs, cols(gs))
-    momenta´ = [view(m, r, c) for m in momenta]
+    momenta´ = [_maybe_diagonal(view(m, r, c), gs) for m in momenta]
     return momenta´
 end
 
 _maybe_contactinds(gs, ::Colon) = Colon()
 _maybe_contactinds(gs, i::Integer) = contactinds(gs, i)
+_maybe_contactinds(gs, i::DiagIndices) = _maybe_contactinds(gs, parent(i))
 _maybe_contactinds(gs, _) = throw(argerror("KPM doesn't support generic indexing"))
 
+_maybe_diagonal(v, gs) = _maybe_diagonal(v, gs, rows(gs))
+_maybe_diagonal(v, gs, is) = v
+_maybe_diagonal(v::AbstractArray{C}, gs, is::DiagIndices) where {C} =
+    append_diagonal!(C[], v, parent(is), kernel(is), parent(gs))  # parent(is) should be Colon or Integer
+
 ## call
 
 function (d::DensityMatrixKPMSolver)(mu, kBT; params...)

src/solvers/green/schur.jl

@@ -759,9 +759,6 @@ end
 fill_rho_blocks!(s::DensityMatrixSchurSolver, psis, fpsis, ϕ, di::DiagIndices, ::DiagIndices) =
     append_diagonal!(empty!(s.ρmat), FermiEigenstates(psis, fpsis), parent(di), kernel(di), s.g)
 
-# no-op overload
-diagonal_slice(fe::FermiEigenstates, _) = fe
-
 # specialized methods for function called by append_diagonal!
 # rng can be an Integer or a UnitRange
 function apply_kernel(kernel, fe::FermiEigenstates{T}, rng) where {T}

test/test_greenfunction.jl

@@ -418,9 +418,6 @@ end
     @test ρ0[sites(1), sites(SA[1], 1)] isa Matrix
     @test size(view(ρ0, sites(1), sites(SA[1], 1))) == (2, 2)
 
-    # Diagonal slicing not yet supported
-    @test_broken densitymatrix(gKPM[diagonal(1)])
-
     glead = LP.square() |> hamiltonian(hopping(1)) |> supercell((0,1), region = r -> -1 <= r[1] <= 1) |> attach(nothing; cells = SA[10]) |> greenfunction(GS.Schur(boundary = 0));
     contact1 = r -> r[1] ≈ 5 && -1 <= r[2] <= 1
     contact2 = r -> r[2] ≈ 5 && -1 <= r[1] <= 1