Skip to content

Commit

Permalink
bands boundary working. WIP optimize ldos
Browse files Browse the repository at this point in the history
optimized ldos for case without boundaries
  • Loading branch information
pablosanjose committed Oct 11, 2023
1 parent 80a3ca1 commit 7015f04
Show file tree
Hide file tree
Showing 9 changed files with 286 additions and 93 deletions.
13 changes: 13 additions & 0 deletions src/apply.jl
Original file line number Diff line number Diff line change
Expand Up @@ -281,4 +281,17 @@ end

#endregion


############################################################################################
# apply CellSites
#region

apply(c::CellSites{L,Vector{Int}}, ::Lattice{<:Any,<:Any,L}) where {L} = c

apply(c::CellSites{L,Colon}, l::Lattice{<:Any,<:Any,L}) where {L} =
CellSites(cell(c), collect(siterange(l)))

apply(c::CellSites{L}, l::Lattice{<:Any,<:Any,L}) where {L} =
CellSites(cell(c), [i for i in siteindices(c) if i in siterange(l)])

#endregion
62 changes: 35 additions & 27 deletions src/greenfunction.jl
Original file line number Diff line number Diff line change
Expand Up @@ -105,21 +105,21 @@ sites_to_orbs(cell::Union{SVector,Tuple}, g::GreenSolution{<:Any,<:Any,L}) where
sites_to_orbs(c::CellSites{<:Any,Colon}, g) =
CellOrbitals(cell(c), 1:flatsize(hamiltonian(g)))
# for sublat name or site range we use a UnitRange instead of a Vector
sites_to_orbs(c::CellSites{<:Any,<:Union{Symbol,AbstractUnitRange}}, g) =
sites_to_orbs(c::CellSites{<:Any,<:Union{Symbol,AbstractUnitRange,Integer}}, g) =
CellOrbitals(cell(c), flatrange(hamiltonian(g), siteindices(c)))
sites_to_orbs(c::CellOrbitals, g) = c
# fallback for cases where i and j are not *both* contact indices -> convert to OrbitalSlice
sites_to_orbs(::Colon, g) = orbslice(blockstructure(g)) # reads ContactBlockStructure
sites_to_orbs(i::Integer, g) = orbslice(blockstructure(g), i) # reads ContactBlockStructure
sites_to_orbs(::Colon, g) = orbslice(blockstructure(g)) # reads ContactBlockStructure
sites_to_orbs(i::Integer, g) = orbslice(blockstructure(g), i) # reads ContactBlockStructure

Base.getindex(g::GreenSolution, i::OrbitalSlice, j::OrbitalSlice) =
mortar([g[si, sj] for si in subcells(i), sj in subcells(j)])
mortar((g[si, sj] for si in subcells(i), sj in subcells(j)))

Base.getindex(g::GreenSolution, i::OrbitalSlice, j::CellOrbitals) =
mortar([g[si, sj] for si in subcells(i), sj in (j,)])
mortar((g[si, sj] for si in subcells(i), sj in (j,)))

Base.getindex(g::GreenSolution, i::CellOrbitals, j::OrbitalSlice) =
mortar([g[si, sj] for si in (i,), sj in subcells(j)])
mortar((g[si, sj] for si in (i,), sj in subcells(j)))

Base.getindex(g::GreenSolution, i::CellOrbitals, j::CellOrbitals) =
slicer(g)[sanitize_cellorbs(i), sanitize_cellorbs(j)]
Expand All @@ -139,25 +139,24 @@ Base.getindex(s::GreenSlicer, ::CellOrbitals, ::CellOrbitals) =
# returns a vector of AbstractMatrices or scalars, one per site
#region

Base.getindex(g::GreenSolution{T}, i::DiagIndices, ::DiagIndices = i) where {T} =
append_diagonal!(Complex{T}[], g, parent(i), kernel(i))
Base.getindex(::GreenSolution{T}, i::DiagIndices, ::DiagIndices = i) where {T} =
append_diagonal!(Complex{T}[], , parent(i), kernel(i))

# If i::Union{LatticeSlice,CellOrbitals}, convert to orbitals
append_diagonal!(d, g, i, kernel; kw...) =
append_diagonal!(d, g, sites_to_orbs(i, g), kernel; kw...)
# If i::Union{LatticeSlice,CellSites}, convert to orbitals
append_diagonal!(d, , i, kernel; kw...) =
append_diagonal!(d, , sites_to_orbs(i, ), kernel; kw...)

function append_diagonal!(d, g, s::OrbitalSlice, kernel; kw...)
function append_diagonal!(d, , s::OrbitalSlice, kernel; kw...)
sizehint!(d, length(s))
for sc in subcells(s)
append_diagonal!(d, g, sc, kernel; kw...)
append_diagonal!(d, , sc, kernel; kw...)
end
return d
end

function append_diagonal!(d, g, o::Union{CellOrbitals,Integer,Colon}, kernel; post = identity)
gblock = g[o, o]
norbs = size(gblock, 1)
rngs = orbranges_or_allorbs(kernel, norbs, o, g)
function append_diagonal!(d, gω, o::Union{CellOrbitals,Integer,Colon}, kernel; post = identity)
gblock = diagonal_slice(gω, o)
rngs = orbranges_or_allorbs(kernel, o, gω)
l = length(rngs)
l0 = length(d)
resize!(d, l0 + l)
Expand All @@ -167,9 +166,15 @@ function append_diagonal!(d, g, o::Union{CellOrbitals,Integer,Colon}, kernel; po
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, norbs, o, g) = 1:norbs
orbranges_or_allorbs(kernel, norbs, o, g) = orbranges(o, g)
orbranges_or_allorbs(kernel::Missing, ::Union{Integer,Colon,CellOrbitals}, gω) = 1:flatsize(gω)
orbranges_or_allorbs(kernel, contact::Integer, gω) = orbranges(blockstructure(gω), contact)
orbranges_or_allorbs(kernel, ::Colon, gω) = orbranges(blockstructure(gω))
orbranges_or_allorbs(kernel, o::CellOrbitals, gω) = orbranges(o)

view_or_scalar(gblock, rng::UnitRange) = view(gblock, rng, rng)
view_or_scalar(gblock, i::Integer) = gblock[i, i]
Expand Down Expand Up @@ -402,11 +407,14 @@ struct TMatrixSlicer{C,L,V<:AbstractArray{C},S} <: GreenSlicer{C}
blockstruct::ContactBlockStructure{L}
end

struct DummySlicer{C} <: GreenSlicer{C}
struct NothingSlicer{C} <: GreenSlicer{C}
end

#region ## Constructors ##

# if there are no Σblocks, return g0slicer
maybe_TMatrixSlicer(g0slicer::GreenSlicer, Σblocks::Tuple{}, blockstruct) = g0slicer

# Uses getindex(g0slicer) to construct g0contacts
function TMatrixSlicer(g0slicer::GreenSlicer{C}, Σblocks, blockstruct) where {C}
if isempty(Σblocks)
Expand All @@ -431,7 +439,7 @@ end
# Takes a precomputed g0contacts (for dummy g0slicer that doesn't implement indexing)
function TMatrixSlicer(g0contacts::AbstractMatrix{C}, Σblocks, blockstruct) where {C}
tmatrix, gcontacts = t_g_matrices(g0contacts, blockstruct, Σblocks...)
g0slicer = DummySlicer{C}()
g0slicer = NothingSlicer{C}()
return TMatrixSlicer(g0slicer, tmatrix, gcontacts, blockstruct)
end

Expand Down Expand Up @@ -490,21 +498,21 @@ function Base.getindex(s::TMatrixSlicer, i::CellOrbitals, j::CellOrbitals)
tkk´ = s.tmatrix
isempty(tkk´) && return g0ij
k = orbslice(s.blockstruct)
g0ik = mortar([g0[si, sk] for si in (i,), sk in subcells(k)])
g0k´j = mortar([g0[sk´, sj] for sk´ in subcells(k), sj in (j,)])
g0ik = mortar((g0[si, sk] for si in (i,), sk in subcells(k)))
g0k´j = mortar((g0[sk´, sj] for sk´ in subcells(k), sj in (j,)))
gij = mul!(g0ij, g0ik, tkk´ * g0k´j, 1, 1) # = g0ij + g0ik * tkk´ * g0k´j
return gij
end

minimal_callsafe_copy(s::TMatrixSlicer) = TMatrixSlicer(minimal_callsafe_copy(s.g0slicer),
s.tmatrix, s.gcontacts, s.blockstruct)

Base.view(::DummySlicer, i::Union{Integer,Colon}...) =
internalerror("view(::DummySlicer): unreachable reached")
Base.view(::NothingSlicer, i::Union{Integer,Colon}...) =
internalerror("view(::NothingSlicer): unreachable reached")

Base.getindex(::DummySlicer, i::CellOrbitals...) = argerror("Slicer does not support generic indexing")
Base.getindex(::NothingSlicer, i::CellOrbitals...) = argerror("Slicer does not support generic indexing")

minimal_callsafe_copy(s::DummySlicer) = s
minimal_callsafe_copy(s::NothingSlicer) = s

#endregion
#endregion
Expand Down
15 changes: 10 additions & 5 deletions src/slices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,14 @@ function Base.getindex(lat::Lattice, as::AppliedSiteSelector)
return latslice
end

Base.getindex(l::Lattice, c::CellSites{L}) where {L} =
LatticeSlice(l, [CellSites{L,Vector{Int}}(c)])
Base.getindex(l::Lattice, c::CellSites) = LatticeSlice(l, [apply(c, l)])

Base.getindex(ls::LatticeSlice; kw...) = getindex(ls, siteselector(; kw...))

Base.getindex(ls::LatticeSlice, ss::SiteSelector) = getindex(ls, apply(ss, parent(ls)))

function Base.getindex(l::LatticeSlice{<:Any,<:Any,L}, i::Integer) where {L}
# return cell, siteindex of the i-th site of LatticeSlice
function Base.getindex(l::LatticeSlice, i::Integer)
offset = 0
for scell in subcells(l)
ninds = length(siteindices(scell))
Expand All @@ -46,6 +46,7 @@ function Base.getindex(l::LatticeSlice{<:Any,<:Any,L}, i::Integer) where {L}
end
end
@boundscheck(boundserror(l, i))
return zerocell(lattice(l)), 0
end

# indexlist is populated with latslice indices of selected sites
Expand Down Expand Up @@ -312,23 +313,27 @@ missing_or_push!(v, n) = push!(v, n)
orbslice(x, g::GreenSolution) = orbslice(x, hamiltonian(g))
orbslice(x, h::AbstractHamiltonian) = orbslice(x, blockstructure(h))
orbslice(sc::CellSites, bs::OrbitalBlockStructure) = _orbslice((sc,), bs)
orbslice(scs::Vector{<:CellSites}, bs::OrbitalBlockStructure) = _orbslice(scs, bs)
orbslice(ls::LatticeSlice, bs::OrbitalBlockStructure) = _orbslice(subcells(ls), bs)

function _orbslice(subcells, bs::OrbitalBlockStructure)
subcells = [cellorbs(sc, bs) for sc in subcells]
return OrbitalSlice(subcells)
end

# simple version of the above: single cellorbs & no storage
cellorbs(sc::CellSites, g::GreenSolution) = cellorbs(sc, blockstructure(hamiltonian(g)))
cellorbs(sc::CellSites, h::AbstractHamiltonian) = cellorbs(sc, blockstructure(h))

function cellorbs(sc::CellSites, bs::OrbitalBlockStructure)
orbinds = Int[]
orbrngs = UnitRange{Int}[]
offset = 0
for i in siteindices(sc)
irng = flatrange(bs, i)
len = length(irng)
append!(orbinds, irng)
push!(orbrngs, irng)
push!(orbrngs, offset+1:offset+len)
offset += len
end
return CellOrbitals(cell(sc), orbinds, orbrngs)
end
Expand Down
4 changes: 1 addition & 3 deletions src/solvers/green.jl
Original file line number Diff line number Diff line change
Expand Up @@ -74,13 +74,11 @@ struct Bands{B<:Union{Missing,Pair},A,K} <: AbstractGreenSolver
end

Bands(bandargs...; boundary = missing, bandskw...) =
Bands(bandargs, NamedTuple(bandskw), boundary)
Bands(sort.(bandargs), NamedTuple(bandskw), boundary)

Bands(bandargs::Quantica.Mesh; kw...) =
argerror("Positional arguments of GS.Bands should be collections of Bloch phases or parameters")

maybesort!(l) = (issorted(l) || sort!(l); l)

end # module

const GS = GreenSolvers
Expand Down
Loading

0 comments on commit 7015f04

Please sign in to comment.