Merge pull request #313 from pablosanjose/applyrange

Apply ranges to lattices before `combine`

src/apply.jl

@@ -20,8 +20,7 @@ function apply(s::SiteSelector, lat::Lattice{T,E,L}, cells...) where {T,E,L}
     unique!(sort!(cells))
     # isnull: to distinguish in a type-stable way between s.cells === missing and no-selected-cells
     # and the same for sublats
-    isnull = (s.cells !== missing && isempty(cells)) ||
-        (s.sublats !== missing && isempty(sublats))
+    isnull = (s.cells !== missing && isempty(cells)) || (s.sublats !== missing && isempty(sublats))
     return AppliedSiteSelector{T,E,L}(lat, region, sublats, cells, isnull)
 end
 
@@ -40,9 +39,10 @@ function apply(s::HopSelector, lat::Lattice{T,E,L}, cells...) where {T,E,L}
         sublats .= reverse.(sublats)
         dcells .*= -1
     end
-    isnull = (s.dcells !== missing && isempty(dcells)) ||
-        (s.sublats !== missing && isempty(sublats))
-    return AppliedHopSelector{T,E,L}(lat, region, sublats, dcells, (rmin, rmax), isnull)
+    includeonsite = s.includeonsite
+    # isnull: see above
+    isnull = (s.dcells !== missing && isempty(dcells)) || (s.sublats !== missing && isempty(sublats))
+    return AppliedHopSelector{T,E,L}(lat, region, sublats, dcells, (rmin, rmax), includeonsite, isnull)
 end
 
 sublatindex_or_zero(lat, ::Missing) = missing
@@ -66,9 +66,6 @@ sanitize_cells(::Missing, ::Val{L}) where {L} = missing
 sanitize_cells(f::Function, ::Val{L}) where {L} = f
 sanitize_cells(cells, ::Val{L}) where {L} = sanitize_SVector.(SVector{L,Int}, cells)
 
-applyrange(r::Neighbors, lat) = nrange(Int(r), lat)
-applyrange(r::Real, lat) = r
-
 padrange(r::Real, m) = isfinite(r) ? float(r) + m * sqrt(eps(float(r))) : float(r)
 
 applied_region(r, ::Missing) = true
@@ -128,6 +125,15 @@ end
 
 recursive_push!(v::Vector, f, cells) = recursive_push!(v, f)
 
+applyrange(ss::Tuple, h::AbstractHamiltonian) = applyrange.(ss, Ref(h))
+applyrange(s::Modifier, h::AbstractHamiltonian) = s
+applyrange(s::AppliedHoppingModifier, h::ParametricHamiltonian) =
+    AppliedHoppingModifier(s, applyrange(selector(s), lattice(h)))
+applyrange(s::HopSelector, lat::Lattice) = hopselector(s; range = sanitize_minmaxrange(hoprange(s), lat))
+
+applyrange(r::Neighbors, lat::Lattice) = nrange(Int(r), lat)
+applyrange(r::Real, lat::Lattice) = r
+
 #endregion
 
 ############################################################################################
@@ -248,7 +254,7 @@ function push_pointers!(ptrs, h, har, s::AppliedHopSelector, shifts = missing, b
         rrow = site(lat, row, dn)
         r, dr = rdr(rcol => rrow)
         r = apply_shift(shifts, r, col)
-        if (scol => srow, (r, dr), dn) in s
+        if (col => row, scol => srow, (r, dr), dn) in s
             ncol = norbs[scol]
             nrow = norbs[srow]
             sprow, spcol = CellSitePos(dn, row, rrow, B), CellSitePos(dn0, col, rcol, B)

src/builders.jl

@@ -197,8 +197,9 @@ function IJVBuilder(lat::Lattice{T}, hams::AbstractHamiltonian...) where {T}
     builder = IJVBuilderWithModifiers(lat, bs)
     push_ijvharmonics!(builder, hams...)
     mss = modifiers.(hams)
+    mss´ = applyrange.(mss, hams)
     bis = blockindices(hams)
-    unapplied_block_modifiers = ((ms, bi) -> intrablock.(parent.(ms), Ref(bi))).(mss, bis)
+    unapplied_block_modifiers = ((ms, bi) -> intrablock.(parent.(ms), Ref(bi))).(mss´, bis)
     push!(builder, tupleflatten(unapplied_block_modifiers...)...)
     return builder
 end

src/docstrings.jl

@@ -2221,12 +2221,15 @@ OrbitalSliceArray, OrbitalSliceVector, OrbitalSliceMatrix
 
 
 """
-    diagonal(args...)
+    diagonal(inds; kernel = missing)
 
-Wrapper over indices `args` used to obtain the diagonal of a `gω::GreenSolution`. If `d =
+Wrapper over indices `inds` used to obtain the diagonal of a `gω::GreenSolution`. If `d =
 diagonal(sel)`, then `gω[d] = diag(gω[sel, sel])`, although in most cases the computation
 is done more efficiently internally.
 
+If `kernel = Q` (a matrix) instead of `missing`, a vector with `Tr(gᵢᵢQ)` on
+each site `i` is returned (instead of of `vcat(diag(gᵢᵢ)...)`).
+
 """
 diagonal
 

src/greenfunction.jl

@@ -140,55 +140,64 @@ sanitize_cellorbs(c::CellOrbitalsGrouped) = CellOrbitals(cell(c), orbindices(c))
 
 ############################################################################################
 # GreenSolution diagonal indexing
-#   returns a vector of AbstractMatrices or scalars, one per site
+#   returns a vector of AbstractMatrices or scalars,
+#   one per site (if kernel is not missing) or one per orbital (if kernel is missing)
 #region
 
 Base.getindex(gω::GreenSolution{T}, i::DiagIndices, ::DiagIndices = i) where {T} =
-    append_diagonal!(Complex{T}[], gω, parent(i), kernel(i)) |>
+    append_diagonal!(Complex{T}[], gω, parent(i), kernel(i), gω) |>
         maybe_OrbitalSliceArray(sites_to_orbs(i, gω))
 
+#endregion
+
+############################################################################################
+# append_diagonal!(d, x, i, kernel, g; kw...)
+#   append to d the elements Tr(kernel*x[i,i]) for each site encoded in i, or each orbital
+#   if kernel is missing. g is the GreenFunction, used to convert i to CellOrbitals
+#region
+
 # If i::Union{SiteSlice,CellSites}, convert to orbitals
-append_diagonal!(d, gω, i, kernel; kw...) =
-    append_diagonal!(d, gω, sites_to_orbs(i, gω), kernel; kw...)
+append_diagonal!(d, x, i, kernel, g; kw...) =
+    append_diagonal!(d, x, sites_to_orbs(i, g), kernel; kw...)
 
-append_diagonal!(d, gω, s::OrbitalSlice, kernel; kw...) =
-    append_diagonal!(d, gω, cellsdict(s), kernel; kw...)    # no OrbitalSliceVector here
+# but not if i is a contact. Jumpt to core driver
+append_diagonal!(d, gω::GreenSolution, o::Union{Colon,Integer}, kernel, g; kw...) =
+    append_diagonal!(d, gω[o,o], contactranges(kernel, o, gω), kernel; kw...)
 
-append_diagonal!(d, gω, s::AnyOrbitalSlice, kernel; kw...) =
-    append_diagonal!(d, gω, cellsdict(s), kernel; kw...)
+# decompose any LatticeSlice into CellOrbitals
+append_diagonal!(d, x, s::AnyOrbitalSlice, kernel; kw...) =
+    append_diagonal!(d, x, cellsdict(s), kernel; kw...)
 
-function append_diagonal!(d, gω, s::AnyCellOrbitalsDict, kernel; kw...)
+function append_diagonal!(d, x, s::AnyCellOrbitalsDict, kernel; kw...)
     sizehint!(d, length(s))
     for sc in s
-        append_diagonal!(d, gω, sc, kernel; kw...)
+        append_diagonal!(d, x, sc, kernel; kw...)
     end
     return d
 end
 
-function append_diagonal!(d, gω, o::Union{AnyCellOrbitals,Integer,Colon}, kernel; post = identity)
-    gblock = diagonal_slice(gω, o)
-    rngs = orbranges_or_allorbs(kernel, o, gω)
-    for rng in rngs
-        push!(d, post(apply_kernel(kernel, view_or_scalar(gblock, rng))))
+append_diagonal!(d, gω::GreenSolution, o::AnyCellOrbitals, kernel; kw...) =
+    append_diagonal!(d, gω[o,o], orbindranges(kernel, o), kernel; kw...)
+
+# core driver
+function append_diagonal!(d, blockmat::AbstractMatrix, blockrngs, kernel; post = identity)
+    for rng in blockrngs
+        val = apply_kernel(kernel, blockmat, rng)
+        push!(d, post(val))
     end
     return d
 end
 
-# fallback, may be overloaded 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]
-
 # If no kernel is provided, we return the whole diagonal
-orbranges_or_allorbs(kernel::Missing, o::AnyCellOrbitals, gω) = eachindex(orbindices(o))
-orbranges_or_allorbs(kernel::Missing, o::Colon, gω) = 1:norbitals(contactorbitals(gω))
-orbranges_or_allorbs(kernel::Missing, i::Integer, gω) = 1:norbitals(contactorbitals(gω), i)
-orbranges_or_allorbs(kernel, o, gω) = orbranges_or_allorbs(o, gω)
-orbranges_or_allorbs(contact::Integer, gω) = orbranges(contactorbitals(gω), contact)
-orbranges_or_allorbs(::Colon, gω) = orbranges(contactorbitals(gω))
-orbranges_or_allorbs(o::CellOrbitalsGrouped, gω) = orbranges(o)
+contactranges(::Missing, o::Colon, gω) = 1:norbitals(contactorbitals(gω))
+contactranges(::Missing, i::Integer, gω) = 1:norbitals(contactorbitals(gω), i)
+contactranges(kernel, ::Colon, gω) = orbranges(contactorbitals(gω))
+contactranges(kernel, i::Integer, gω) = orbranges(contactorbitals(gω), i)
 
-view_or_scalar(gblock, rng::UnitRange) = view(gblock, rng, rng)
-view_or_scalar(gblock, i::Integer) = gblock[i, i]
+orbindranges(::Missing, o) = eachindex(orbindices(o))
+orbindranges(kernel, o) = orbranges(o)
+
+apply_kernel(kernel, gblock, rng) = apply_kernel(kernel, view_or_scalar(gblock, rng))
 
 apply_kernel(kernel::Missing, v::Number) = v
 apply_kernel(kernel::AbstractMatrix, v) = tr(kernel * v)
@@ -197,6 +206,9 @@ apply_kernel(kernel::Number, v) = kernel * tr(v)
 apply_kernel(kernel::Diagonal, v::AbstractMatrix) = sum(i -> kernel[i] * v[i, i], eachindex(kernel))
 apply_kernel(kernel::Diagonal, v::Number) = only(kernel) * v
 
+view_or_scalar(gblock, rng::UnitRange) = view(gblock, rng, rng)
+view_or_scalar(gblock, orb::Integer) = gblock[orb, orb]
+
 maybe_scalarize(s::OrbitalSliceGrouped, kernel::Missing) = s
 maybe_scalarize(s::OrbitalSliceGrouped, kernel) = scalarize(s)
 

src/meanfield.jl

@@ -0,0 +1,8 @@
+############################################################################################
+# hartree constructors
+#region
+
+# function hartree
+
+
+#endregion

src/observables.jl

@@ -218,7 +218,7 @@ Base.getindex(d::LocalSpectralDensitySolution, s::SiteSelector) =
     d[getindex(lattice(d.gω), s)]
 
 Base.getindex(d::LocalSpectralDensitySolution{T}, i) where {T} =
-    append_diagonal!(T[], d.gω, i, d.kernel, post = x -> -imag(x)/π) |>
+    append_diagonal!(T[], d.gω, i, d.kernel, d.gω; post = x -> -imag(x)/π) |>
         maybe_OrbitalSliceArray(sites_to_orbs(i, d.gω)) # see greenfunction.jl
 
 (d::LocalSpectralDensitySlice)(ω; params...) = copy(call!(d, ω; params...))
@@ -502,8 +502,8 @@ densitymatrix(s::AppliedGreenSolver, gs::GreenSlice; kw...) =
 densitymatrix(gs::GreenSlice, ωmax::Number; opts...) = densitymatrix(gs, (-ωmax, ωmax); opts...)
 
 function densitymatrix(gs::GreenSlice{T}, ωpoints; omegamap = Returns((;)), imshift = missing, atol = 1e-7, opts...) where {T}
-    check_nodiag_axes(gs)
-    result = similar_Matrix(gs)
+    # check_nodiag_axes(gs)
+    result = similar_Array(gs)
     opts´ = (; imshift, slope = 1, post = gf_to_rho!, atol, opts...)
     ωpoints_vec = collect(promote_type(T, typeof.(ωpoints)...), ωpoints)
     function ifunc(mu, kBT, override)
@@ -564,12 +564,19 @@ function call!(gf::DensityMatrixIntegrand, ω; params...)
     return gω
 end
 
-function gf_to_rho!(x)
+function gf_to_rho!(x::AbstractMatrix)
     x .= x .- x'
     x .*= -1/(2π*im)
     return x
 end
 
+# For diagonal indexing
+function gf_to_rho!(x::AbstractVector)
+    x .= x .- conj.(x)
+    x .*= -1/(2π*im)
+    return x
+end
+
 integrand(ρ::DensityMatrix{<:DensityMatrixIntegratorSolver}, mu = 0.0, kBT = 0.0) = integrand(ρ.solver(mu, kBT))
 
 path(ρ::DensityMatrix{<:DensityMatrixIntegratorSolver}, mu = 0.0, kBT = 0.0) = path(ρ.solver(mu, kBT))

src/selectors.jl

@@ -9,10 +9,10 @@ siteselector(s::SiteSelector; region = s.region, sublats = s.sublats, cells = s.
     SiteSelector(region, sublats, cells)
 
 hopselector(h::NamedTuple) = hopselector(; h...)
-hopselector(; region = missing, sublats = missing, dcells = missing, range = neighbors(1)) =
-    HopSelector(region, sublats, dcells, range)
-hopselector(s::HopSelector; region = s.region, sublats = s.sublats, dcells = s.dcells, range = s.range) =
-    HopSelector(region, sublats, dcells, range, s.adjoint)
+hopselector(; region = missing, sublats = missing, dcells = missing, range = neighbors(1), includeonsite = false) =
+    HopSelector(region, sublats, dcells, range, includeonsite)
+hopselector(s::HopSelector; region = s.region, sublats = s.sublats, dcells = s.dcells, range = s.range, includeonsite = s.includeonsite) =
+    HopSelector(region, sublats, dcells, range, includeonsite, s.adjoint)
 
 neighbors(n::Int) = Neighbors(n)
 neighbors(n::Int, lat::Lattice) = nrange(n, lat)
@@ -50,8 +50,8 @@ end
 #     return ((j, i), (r, dr), dcell) in sel
 # end
 
-function Base.in(((sj, si), (r, dr), dcell)::Tuple{Pair,Tuple,SVector}, sel::AppliedHopSelector)
-    return !isnull(sel) && !isonsite(dr) &&
+function Base.in(((j, i), (sj, si), (r, dr), dcell)::Tuple{Pair, Pair,Tuple,SVector}, sel::AppliedHopSelector)
+    return !isnull(sel) && (includeonsite(sel) || !isonsite((j, i), dcell)) &&
             indcells(dcell, sel) &&
             insublats(sj => si, sel) &&
             iswithinrange(dr, sel) &&
@@ -137,7 +137,7 @@ function foreach_hop(f, sel::AppliedHopSelector, kdtrees::Vector{<:KDTree}, ni::
         for j in js
             is = inrange_targets(site(lat, j, nj - ni), lat, si, rmax, kdtrees)
             for i in is
-                isonsite((j, i), ni - nj) && continue
+                !includeonsite(sel) && isonsite((j, i), ni - nj) && continue
                 r, dr = rdr(site(lat, j, nj) => site(lat, i, ni))
                 # Make sure we don't stop searching cells until we reach minimum range
                 isbelowrange(dr, sel) && (found = true)

src/solvers/green/bands.jl

@@ -849,6 +849,21 @@ function muladd_ψPψ⁺!(gmat, α, ψ, (rows, cols)::Tuple)
     return gmat
 end
 
+# fallback for a generic rngs collection (used e.g. by rngs = orbs in append_diagonal! below)
+function muladd_ψPψ⁺!(gmat, α, ψ, rngs)
+    for rng in rngs
+        if length(rng) == 1
+            i = only(rng)
+            gmat[i, i] += α * abs2(ψ[i])
+        else
+            gv = view(gmat, rng, rng)
+            ψv = view(ψ, rng, :)
+            mul!(gv, ψv, ψv', α, 1)
+        end
+    end
+    return gmat
+end
+
 view_or_copy(ψ, rows::Union{Colon,AbstractRange}, cols::Union{Colon,AbstractRange}) =
     view(ψ, rows, cols)
 view_or_copy(ψ, rows, cols) = ψ[rows, cols]
@@ -858,12 +873,12 @@ minimal_callsafe_copy(s::BandsGreenSlicer, parentham, parentcontacts) = s  # it
 #endregion
 
 ############################################################################################
-# diagonal_slice
+# append_diagonal!
 #   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 append_diagonal!(d, gω::GreenSolution{T,<:Any,<:Any,G}, o::AnyCellOrbitals, kernel; kw...) where {T,G<:BandsGreenSlicer{<:Any,Missing}}
     s = slicer(gω)   # BandsGreenSlicer
     norbs = flatsize(gω)
     rngs = orbranges(o)
@@ -877,21 +892,8 @@ function diagonal_slice(gω::GreenSolution{T,<:Any,<:Any,G}, o::CellOrbitals) wh
     simpinds = eachindex(simplices(subband))
     # to main driver with orbs = rngs
     inf_band_slice!(gmat, s.ω, dist, rngs, subband, subbandsimps, simpinds)
-    return gmat
-end
-
-function muladd_ψPψ⁺!(gmat, α, ψ, rngs::Vector{<:UnitRange})
-    for rng in rngs
-        if length(rng) == 1
-            i = only(rng)
-            gmat[i, i] += α * abs2(ψ[i])
-        else
-            gv = view(gmat, rng, rng)
-            ψv = view(ψ, rng, :)
-            mul!(gv, ψv, ψv', α, 1)
-        end
-    end
-    return gmat
+    orbs = orbindranges(kernel, o)
+    return append_diagonal!(d, gmat, orbs, kernel; kw...)
 end
 
 #endregion

src/solvers/green/kpm.jl

@@ -244,7 +244,6 @@ end
 ## Constructor
 
 function densitymatrix(s::AppliedKPMGreenSolver, gs::GreenSlice{T}) where {T}
-    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 +253,23 @@ 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(blockmat, gs) = _maybe_diagonal(blockmat, gs, rows(gs))
+_maybe_diagonal(blockmat, gs, is) = blockmat
+
+function _maybe_diagonal(blockmat::AbstractArray{C}, gs, is::DiagIndices) where {C}
+    blockrngs = contactranges(kernel(is), parent(is), gs)          # see greenfunction.jl
+    return append_diagonal!(C[], blockmat, blockrngs, kernel(is))  # parent(is) should be Colon or Integer
+end
+
 ## call
 
 function (d::DensityMatrixKPMSolver)(mu, kBT; params...)