fixes to sparse indexing of Schur densitymatrix

pablosanjose · Oct 24, 2024 · 39a2457 · 39a2457
1 parent 6f28c90
commit 39a2457
Show file tree

Hide file tree

Showing 7 changed files with 96 additions and 74 deletions.
diff --git a/src/greenfunction.jl b/src/greenfunction.jl
@@ -67,9 +67,6 @@ call!(g::GreenSlice{T}, ω::T; kw...) where {T} =
 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...)
-
 real_or_complex_convert(::Type{T}, ω::Real) where {T<:Real} = convert(T, ω)
 real_or_complex_convert(::Type{T}, ω::Complex) where {T<:Real} = convert(Complex{T}, ω)
 
@@ -189,8 +186,10 @@ getindex!(output, g, i::OrbitalPairIndices, j::OrbitalPairIndices = i; kw...) =
 ############################################################################################
 # getindex_diagonal! and getindex_sparse! : GreenSolution diagonal and sparse indexing
 #   They rely on fill_diagonal! and fill_sparse! to fill the output matrix with mat elements
-#   We can swith `g` with another object that supports getindex_diag or getindex_sparse
-#   returning some `mat` object, that supports getindex and view
+#   We can swith `g` with another object `o` that supports getindex_diag or getindex_sparse
+#   returning `mat::AbstractMatrix`
+#   `mat` is not an OrbitalSliceGrouped for performance reasons. Instead the object `o`
+#   should also implement a `blockstructure` method so we know how to map sites to indices.
 #region
 
 getindex_diagonal!(output, g::GreenSolution, i::Union{Colon,Integer}, ker; post = identity) =
@@ -203,6 +202,7 @@ function getindex_diagonal!(output, g, o::CellOrbitalsGrouped, ker; post = ident
     return output
 end
 
+# g may be a GreenSolution or any other type with a minimal API (e.g. EigenProduct)
 function getindex_diagonal!(output, g, i::OrbitalSliceGrouped, ker; post = identity)
     offset = 0
     for o in cellsdict(i)
@@ -216,20 +216,20 @@ end
 
 function getindex_sparse!(output, g, hmask::Hamiltonian, ker; post = identity)
     bs = blockstructure(g)  # used to find blocks of g for nonzeros in hmask
-    oj = sites_to_orbs(sites(zerocell(g), :), g)
+    oj = sites_to_orbs(sites(zerocell(hmask), :), bs)  # full cell
     for har in harmonics(hmask)
         oi = CellIndices(dcell(har), oj)
-        mat = getindex_sparse(g, oi, oj, har)
-        fill_sparse!(har, mat, bs, ker; post)
+        mat_cell = getindex_sparse(g, oi, oj, har)
+        fill_sparse!(har, mat_cell, bs, ker; post)
     end
     fill_sparse!(parent(output), hmask)
     return output
 end
 
 # non-optimized fallback, can be specialized by solvers
-# optimized version would only compute site blocks in diagonal
+# optimized version could compute only site blocks in diagonal
 getindex_diag(g, oi) = g[oi, oi]
-# optimized version would only compute nonzero blocks from har::Harmonic
+# optimized version could compute only nonzero blocks from har::Harmonic
 getindex_sparse(g, oi, oj, har) = g[oi, oj]
 
 # input index ranges for each output index to fill
@@ -241,26 +241,26 @@ orbital_kernel_ranges(::Missing, o::CellOrbitalsGrouped) = eachindex(orbindices(
 orbital_kernel_ranges(kernel, o::CellOrbitalsGrouped) = orbranges(o)
 
 ## core drivers
-#   write to the diagonal of output, with a given offset, the elements Tr(kernel*mat[i,i])
+#   write to the diagonal of output, with a given offset, the elements Tr(kernel*mat_cell[i,i])
 #   for each orbital or site encoded in range i
-function fill_diagonal!(output, mat, blockrngs, kernel; post = identity, offset = 0)
+function fill_diagonal!(output, mat_cell, blockrngs, kernel; post = identity, offset = 0)
     for (n, rng) in enumerate(blockrngs)
         i = n + offset
-        output[i, i] = post(apply_kernel(kernel, mat, rng))
+        output[i, i] = post(apply_kernel(kernel, mat_cell, rng))
     end
     return output
 end
 
-#   overwrite nonzero h elements with Tr(kernel*mat[i,j]) for each site pair i,j
-function fill_sparse!(har::Harmonic, mat, bs, kernel; post = identity)
+#   overwrite nonzero h elements with Tr(kernel*mat_cell[irng,jrng]) for each site pair i,j
+function fill_sparse!(har::Harmonic, mat_cell, bs, kernel; post = identity)
     B = blocktype(har)
     hmat = unflat(har)
     rows = rowvals(hmat)
     nzs = nonzeros(hmat)
     for col in axes(hmat, 2), ptr in nzrange(hmat, col)
         row = rows[ptr]
         orows, ocols = flatrange(bs, row), flatrange(bs, col)
-        val = post(apply_kernel(kernel, mat, orows, ocols))
+        val = post(apply_kernel(kernel, mat_cell, orows, ocols))
         nzs[ptr] = mask_block(B, val)
     end
     needs_flat_sync!(matrix(har))
@@ -285,12 +285,12 @@ function fill_sparse!(output::SparseMatrixCSC, h::Hamiltonian)
     return output
 end
 
-apply_kernel(kernel, mat, rows, cols = rows) = apply_kernel(kernel, view_or_scalar(mat, rows, cols))
+apply_kernel(kernel, mat_cell, rows, cols = rows) = apply_kernel(kernel, view_or_scalar(mat_cell, rows, cols))
 apply_kernel(kernel::Missing, v) = v
 apply_kernel(kernel, v) = trace_prod(kernel, v)
 
-view_or_scalar(mat, rows::UnitRange, cols::UnitRange) = view(mat, rows, cols)
-view_or_scalar(mat, orb::Integer, orb´::Integer) = mat[orb, orb´]
+view_or_scalar(mat_cell, rows::UnitRange, cols::UnitRange) = view(mat_cell, rows, cols)
+view_or_scalar(mat_cell, orb::Integer, orb´::Integer) = mat_cell[orb, orb´]
 
 #endregion
 

diff --git a/src/observables.jl b/src/observables.jl
@@ -565,18 +565,10 @@ override_path!(override_path, ptsvec, pts) =
 
 (gf::DensityMatrixIntegrand)(ω; params...) = copy(call!(gf, ω; params...))
 
-# # extrinsic version - unused
-# function call!(output, gf::DensityMatrixIntegrand, ω; params...)
-#     call!(output, gf.gs, ω; gf.omegamap(ω)..., params...)
-#     output .*= fermi(ω - gf.mu, inv(gf.kBT))
-#     return output
-# end
-
-# intrinsic version
 function call!(gf::DensityMatrixIntegrand, ω; params...)
     output = call!(gf.gs, ω; gf.omegamap(ω)..., params...)
-    output .*= fermi(ω - gf.mu, inv(gf.kBT))
-    return output    # turns into a bare Array
+    serialize(output) .*= fermi(ω - gf.mu, inv(gf.kBT))
+    return output
 end
 
 function gf_to_rho!(x::AbstractMatrix)
@@ -711,10 +703,6 @@ numphaseshifts(phaseshifts) = length(phaseshifts)
 
 (J::JosephsonIntegrand)(ω; params...) = copy(call!(J, ω; params...))
 
-# # extrinsic version (fallback) - unused
-# call!(output, J::JosephsonIntegrand, ω; params...) = (output .= call!(J, ω; params...))
-
-# intrinsic version
 function call!(J::JosephsonIntegrand, ω; params...)
     gω = call!(J.g, ω; J.omegamap(ω)..., params...)
     f = fermi(ω, inv(J.kBT))
@@ -723,7 +711,7 @@ function call!(J::JosephsonIntegrand, ω; params...)
 end
 
 function josephson_traces(J, gω, f)
-    gr = gω[J.contactind, J.contactind]
+    gr = view(gω, J.contactind, J.contactind)
     Σi = selfenergy!(J.Σ, gω, J.contactind)
     return josephson_traces!(J, gr, Σi, f)
 end

diff --git a/src/show.jl b/src/show.jl
@@ -379,6 +379,23 @@ Base.summary(g::GreenSlice{T,E,L}) where {T,E,L} =
 
 #endregion
 
+############################################################################################
+# SparseIndices
+#region
+
+function Base.show(io::IO, s::SparseIndices)
+    i = get(io, :indent, "")
+    print(io, i, summary(s))
+end
+
+Base.summary(::PairIndices) =
+    "PairIndices: a form of SparseIndices used to sparsely index an object such as a GreenFunction on a selection of site pairs"
+
+Base.summary(::DiagIndices) =
+    "DiagIndices: a form of SparseIndices used to index the diagonal of an object such as a GreenFunction"
+
+#endregion
+
 ############################################################################################
 # Conductance
 #region

diff --git a/src/solvers/green/schur.jl b/src/solvers/green/schur.jl
@@ -706,13 +706,14 @@ end
 
 function fermi_h!(result, s, ϕ, µ, β = 0; params...)
     h = hamiltonian(s.gs)
+    bs = blockstructure(h)
     # Similar to spectrum(h, ϕ; params...), but less work (no sort! or sanitization)
     copy!(s.hmat, call!(h, ϕ; params...))  # sparse to dense
     ϵs, psis = eigen!(s.hmat)
     # special-casing β = Inf with views turns out to be slower
     fs = (@. ϵs = fermi(ϵs - µ, β))
     fpsis = (s.psis .= psis .* transpose(fs))
-    ρcell = EigenProduct(psis, fpsis, ϕ)
+    ρcell = EigenProduct(bs, psis, fpsis, ϕ)
     getindex!(result, ρcell, s.orbaxes...)
     return result
 end

diff --git a/src/solvers/green/spectrum.jl b/src/solvers/green/spectrum.jl
@@ -100,11 +100,12 @@ end
 ## call
 
 function (s::DensityMatrixSpectrumSolver)(µ, kBT; params...)
+    bs = blockstructure(s.gs)
     psis, fpsis, es, fs = s.psis, s.fpsis, s.es, s.fs
     β = inv(kBT)
     (@. fs = fermi(es - µ, β))
     fpsis .= psis .* transpose(fs)
-    ρcell = EigenProduct(psis, fpsis)
+    ρcell = EigenProduct(bs, psis, fpsis)
     result = call!_output(s.gs)
     getindex!(result, ρcell, s.orbaxes...)
     return result

diff --git a/src/specialmatrices.jl b/src/specialmatrices.jl
@@ -538,9 +538,11 @@ function Base.getindex(a::OrbitalSliceMatrix, i::SiteSelector, j::SiteSelector =
     return OrbitalSliceMatrix(m, (rowslice´, colslice´))
 end
 
-# CellSites: return an unwrapped Matrix or a view thereof (non-allocating)
+# CellSites: return an unwrapped Matrix or a scalar
 Base.getindex(a::OrbitalSliceMatrix, i::AnyCellSites, j::AnyCellSites = i) =
     copy(view(a, i, j))
+Base.getindex(a::OrbitalSliceMatrix, i::CellSite, j::CellSite = i) =
+    only(view(a, i, j))
 
 Base.getindex(a::OrbitalSliceMatrix, i::C, j::C = i) where {B,C<:CellSitePos{<:Any,<:Any,<:Any,B}} =
     sanitize_block(B, view(a, i, j))
@@ -573,18 +575,21 @@ LinearAlgebra.norm(a::OrbitalSliceMatrix) = norm(parent(a))
 ## broadcasting
 
 # following the manual: https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array
+# disabled for now, since it seems to cause issues with OrbitalSliceMatrix{<:SparseMatrixCSC}
 
-Broadcast.BroadcastStyle(::Type{<:OrbitalSliceArray}) = Broadcast.ArrayStyle{OrbitalSliceArray}()
+# Broadcast.BroadcastStyle(::Type{O}) where {O<:OrbitalSliceArray} = Broadcast.ArrayStyle{O}()
 
-Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{OrbitalSliceArray}}, ::Type{ElType}) where {ElType} =
-    OrbitalSliceArray(similar(Array{ElType}, axes(bc)), orbaxes(find_osa(bc)))
+# Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{O}}, ::Type{ElType}) where {ElType,O<:OrbitalSliceMatrix{<:Any,<:Array}} =
+#     OrbitalSliceArray(similar(Array{ElType}, axes(bc)), orbaxes(find_osa(bc)))
+# Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{O}}, ::Type{ElType}) where {ElType,O<:OrbitalSliceMatrix{<:Any,<:SparseMatrixCSC}} =
+#     OrbitalSliceArray(similar(parent(find_osa(bc)), ElType), orbaxes(find_osa(bc)))
 
-find_osa(bc::Base.Broadcast.Broadcasted) = find_osa(bc.args)
-find_osa(args::Tuple) = find_osa(find_osa(args[1]), Base.tail(args))
-find_osa(x) = x
-find_osa(::Tuple{}) = nothing
-find_osa(a::OrbitalSliceArray, rest) = a
-find_osa(::Any, rest) = find_osa(rest)
+# find_osa(bc::Base.Broadcast.Broadcasted) = find_osa(bc.args)
+# find_osa(args::Tuple) = find_osa(find_osa(args[1]), Base.tail(args))
+# find_osa(x) = x
+# find_osa(::Tuple{}) = nothing
+# find_osa(a::OrbitalSliceArray, rest) = a
+# find_osa(::Any, rest) = find_osa(rest)
 
 # taken from https://github.com/JuliaArrays/OffsetArrays.jl/blob/756e839563c88faa4ebe4ff971286747863aaff0/src/OffsetArrays.jl#L469
 
@@ -622,19 +627,28 @@ Broadcast.broadcast_unalias(dest::OrbitalSliceArray, src::OrbitalSliceArray) =
 #   the product. x is a scalar. Inspired by LowRankMatrices.jl
 #region
 
-struct EigenProduct{T,M<:AbstractMatrix{Complex{T}}} <: AbstractMatrix{T}
+struct EigenProduct{T,M<:AbstractMatrix{Complex{T}},O<:OrbitalBlockStructure} <: AbstractMatrix{T}
+    blockstruct::O
     U::M
     V::M
     phi::T
     x::Complex{T}
-    function EigenProduct{T,M}(U::M, V::M, phi::T, x::Complex{T}) where {T,M<:AbstractMatrix{Complex{T}}}
+    function EigenProduct{T,M,O}(blockstruct::O, U::M, V::M, phi::T, x::Complex{T}) where {T,M<:AbstractMatrix{Complex{T}},O<:OrbitalBlockStructure}
         axes(U, 2) != axes(V, 2) && argerror("U and V must have identical column axis")
-        return new(U, V, phi, x)
+        return new(blockstruct, U, V, phi, x)
     end
 end
 
-EigenProduct(U::M, V::M, phi = zero(T), x = one(Complex{T})) where {T,M<:AbstractMatrix{Complex{T}}} =
-    EigenProduct{T,M}(U, V, T(phi), Complex{T}(x))
+#region ## Constructors ##
+
+EigenProduct(blockstruct::O, U::M, V::M, phi = zero(T), x = one(Complex{T})) where {T,M<:AbstractMatrix{Complex{T}},O<:OrbitalBlockStructure} =
+    EigenProduct{T,M,O}(blockstruct, U, V, T(phi), Complex{T}(x))
+
+#endregion
+
+#region ## API ##
+
+blockstructure(P::EigenProduct) = P.blockstruct
 
 Base.@propagate_inbounds function Base.getindex(P::EigenProduct, i::Integer, j::Integer)
     ret = zero(eltype(P))
@@ -656,20 +670,22 @@ function Base.getindex(P::EigenProduct{T}, ci::AnyCellOrbitals, cj::AnyCellOrbit
     return view(phase * P, orbindices(ci), orbindices(cj))
 end
 
-Base.view(P::EigenProduct, is, js) = EigenProduct(view(P.U, is, :), view(P.V, js, :), P.phi, P.x)
+Base.view(P::EigenProduct, is, js) =
+    EigenProduct(P.blockstruct, view(P.U, is, :), view(P.V, js, :), P.phi, P.x)
 
 # AbstractArray interface
 Base.axes(P::EigenProduct) = axes(P.U, 1), axes(P.V, 1)
 Base.size(P::EigenProduct) = size(P.U, 1), size(P.V, 1)
 # Base.iterate(a::EigenProduct, i...) = iterate(parent(a), i...)
 Base.length(P::EigenProduct) = prod(size(P))
 Base.IndexStyle(::Type{T}) where {M,T<:EigenProduct{<:Any,M}} = IndexStyle(M)
-Base.similar(P::EigenProduct) = EigenProduct(similar(P.U), similar(P.V), P.phi, P.x)
-Base.similar(P::EigenProduct, t::Type) = EigenProduct(similar(P.U, t), similar(P.V, t), P.phi, P.x)
-Base.copy(P::EigenProduct) = EigenProduct(copy(P.U), copy(P.V), P.phi, P.x)
+Base.similar(P::EigenProduct) = EigenProduct(P.blockstruct, similar(P.U), similar(P.V), P.phi, P.x)
+Base.similar(P::EigenProduct, t::Type) = EigenProduct(P.blockstruct, similar(P.U, t), similar(P.V, t), P.phi, P.x)
+Base.copy(P::EigenProduct) = EigenProduct(P.blockstruct, copy(P.U), copy(P.V), P.phi, P.x)
 Base.eltype(::EigenProduct{T}) where {T} = Complex{T}
 
-Base.:*(x::Number, P::EigenProduct) = EigenProduct(P.U, P.V, P.phi, P.x * x)
+Base.:*(x::Number, P::EigenProduct) = EigenProduct(P.blockstruct, P.U, P.V, P.phi, P.x * x)
 Base.:*(P::EigenProduct, x::Number) = x*P
 
 #endregion
+#endregion
diff --git a/test/test_greenfunction.jl b/test/test_greenfunction.jl
@@ -296,6 +296,7 @@ end
     ω = 0.6
     @test g[diagonal(2)](ω) isa OrbitalSliceMatrix
     @test g[diagonal(2)](ω) == g(ω)[diagonal(2)]
+    @test g[diagonal(:)](ω) == g(ω)[diagonal(:)]
     @test size(g[diagonal(1)](ω)) == (12,12)
     @test size(g[diagonal(1, kernel = I)](ω)) == (8,8)
     @test size(g[diagonal(:)](ω),1) == size(g[diagonal(1)](ω),1) + size(g[diagonal(2)](ω),1) == 48
@@ -324,30 +325,26 @@ end
     @test gg .* gg´ ≈ gg .* gg
 end
 
-@testset "densitymatrix diagonal slicing" begin
+@testset "densitymatrix sparse slicing" begin
     g1 = LP.honeycomb() |> hamiltonian(hopping(0.1I, sublats = (:A,:B) .=> (:A,:B), range = 1) - plusadjoint(@hopping((; t=1) -> t*SA[0.4 1], sublats = :B => :A)), orbitals = (1,2)) |>
         supercell((1,-1), region = r -> -2<=r[2]<=2) |> attach(nothing, region = RP.circle(1), sublats = :B) |>
         attach(nothing, sublats = :A, cells = 0, region = r->r[2]<0) |> greenfunction(GS.Schur())
     g2 = LP.square() |> hopping(1) - onsite(1) |> supercell(3) |> supercell |> attach(nothing, region = RP.circle(2)) |> greenfunction(GS.Spectrum())
     g3 = LP.triangular() |> hamiltonian(hopping(-I) + onsite(1.8I), orbitals = 2) |> supercell(10) |> supercell |>
         attach(nothing, region = RP.circle(2)) |> attach(nothing, region = RP.circle(2, SA[1,2])) |> greenfunction(GS.KPM(bandrange=(-4,5)))
-    for g in (g1, g2, g3)
-        ρ0 = densitymatrix(g[diagonal(1)], 5)
-        ρ = densitymatrix(g[diagonal(1)])
-        @test g[diagonal(1)](0.2) == g(0.2)[diagonal(1)]
-        @test g[diagonal(:)](0.2) == g(0.2)[diagonal(:)]
-        @test g[diagonal(1)](0.2) isa OrbitalSliceMatrix
-        @test isapprox(ρ0(), ρ())
-        @test isapprox(ρ0(0.2), ρ(0.2))
-        if g !== g3 # KPM doesn't support finite temperatures yet
-            @test isapprox(ρ0(0.2, 0.3), ρ(0.2, 0.3))
-        end
-        if g === g1 || g === g3
-            ρ0 = densitymatrix(g[diagonal(1, kernel = SA[0 1; 1 0])], 5)
-            ρ = densitymatrix(g[diagonal(1, kernel = SA[0 1; 1 0])])
-            @test isapprox(ρ0(), ρ(); atol = 1e-8)
-            @test isapprox(ρ0(0.2), ρ(0.2); atol = 1e-8)
-        end
+    # KPM doesn't support finite temperatures or generic indexing yet, so g3 excluded from this loop
+    for g in (g1, g2), inds in (diagonal(1), diagonal(:), sitepairs(range = 1))
+        ρ0 = densitymatrix(g[inds], 5)
+        ρ = densitymatrix(g[inds])
+        @test isapprox(ρ0(), ρ(); atol = 1e-7)
+        @test isapprox(ρ0(0.2), ρ(0.2); atol = 1e-7)
+        @test isapprox(ρ0(0.2, 0.3), ρ(0.2, 0.3))
+    end
+    for g in (g1, g3)
+        ρ0 = densitymatrix(g[diagonal(1, kernel = SA[0 1; 1 0])], 5)
+        ρ = densitymatrix(g[diagonal(1, kernel = SA[0 1; 1 0])])
+        @test isapprox(ρ0(), ρ(); atol = 1e-8)
+        @test isapprox(ρ0(0.2), ρ(0.2); atol = 1e-8)
     end
 end
 
@@ -488,6 +485,8 @@ end
     g0 = LP.square() |> hamiltonian(hopping(1)) |> supercell(region = RP.square(10)) |> attach(glead, reverse = true; region = contact2) |> attach(glead; region = contact1) |> greenfunction;
     testcond(g0)
 
+    contact1 = r -> r[1] ≈ 5 && -1 <= r[2] <= 1
+    contact2 = r -> r[1] ≈ -5 && -1 <= r[2] <= 1
     glead = LP.square() |> hamiltonian(hopping(I) + onsite(SA[0 1; 1 0]), orbitals = 2) |> supercell((1,0), region = r -> -1 <= r[2] <= 1) |> greenfunction(GS.Schur(boundary = 0));
     g0 = LP.square() |> hamiltonian(hopping(I), orbitals = 2) |> supercell(region = RP.square(10)) |> attach(glead, reverse = true; region = contact2) |> attach(glead; region = contact1) |> greenfunction;
     testcond(g0; nambu = true)