Merge pull request #280 from pablosanjose/hindexing

Generalize `AbstractHamiltonian` indexing

src/docstrings.jl

@@ -603,16 +603,24 @@ neighbors
 """
     hamiltonian(lat::Lattice, model; orbitals = 1)
 
-Create a `Hamiltonian` or `ParametricHamiltonian` by applying `model` to the lattice `lat`
-(see `onsite`, `@onsite`, `hopping` and `@hopping` for details on building tight-binding
-models).
+Create an `AbstractHamiltonian` (i.e. an `Hamiltonian` or `ParametricHamiltonian`) by
+applying `model` to the lattice `lat` (see `onsite`, `@onsite`, `hopping` and `@hopping` for
+details on building parametric and non-parametric tight-binding models).
 
     hamiltonian(lat::Lattice, model, modifiers...; orbitals = 1)
-    hamiltonian(h::AbstractHamiltonian, modifiers...; orbitals = 1)
 
 Create a `ParametricHamiltonian` where all onsite and hopping terms in `model` can be
-parametrically modified through the provided `modifiers` (see `@onsite!` and `@hopping!` for
-details on defining modifiers).
+parametrically modified through the provided parametric `modifiers` (see `@onsite!` and
+`@hopping!` for details on defining modifiers).
+
+    hamiltonian(h::AbstractHamiltonian, modifier, modifiers...)
+
+Add modifiers to an existing `AbstractHamiltonian`.
+
+    hamiltonian(h::ParametricHamiltonian)
+
+Return the base (non-parametric) `Hamiltonian` of `h`, with all modifiers and parametric
+model terms removed (see `@onsite`, `@hopping`, `@onsite!`, `@hopping!`).
 
 ## Keywords
 
@@ -650,9 +658,18 @@ Special syntax equivalent to `h[(0...)]`, which access the fundamental Bloch har
     h[i::CellSites, j::CellSites = i]
 
 With `i` and `j` of type `CellSites` and constructed with `sites([cell,] indices)`, return a
-sparse matrix block of `h` between the sites with the corresponding `indices` and in the
+`SparseMatrixCSC` block of `h` between the sites with the corresponding `indices` and in the
 given `cell`s.
 
+    h[srow::SiteSelector, scol::SiteSelector = srow]
+    h[kwrow::NamedTuple, kwcol::NamedTuple = kwrow]
+
+Return an `OrbitalSliceMatrix` of `h` between row and column sites selected by `srow` and
+`scol`, or by `siteselector(; kwrow...)` and `siteselector(; kwcol...)`
+
+Note: `CellSites` and `SiteSelector`s can be mixed when indexing, in which case the matrix
+block will be returned as a `SparseMatrixCSC`, instead of an `OrbitalSliceMatrix`.
+
 ## Call syntax
 
     ph(; params...)
@@ -690,10 +707,22 @@ julia> h((0,0))
      ⋅          ⋅      3.0+0.0im  0.0+0.0im
  0.0+0.0im  3.0+0.0im      ⋅          ⋅
  3.0+0.0im  0.0+0.0im      ⋅          ⋅
+
+julia> h[sites(1), sites(2)]
+2×2 SparseArrays.SparseMatrixCSC{ComplexF64, Int64} with 4 stored entries:
+ 0.0+0.0im  1.0+0.0im
+ 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}}:
+ 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
+ 3.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 ```
 
 # See also
-    `lattice`, `onsite`, `hopping`, `@onsite`, `@hopping`, `@onsite!`, `@hopping!`, `ishermitian`
+    `lattice`, `onsite`, `hopping`, `@onsite`, `@hopping`, `@onsite!`, `@hopping!`, `ishermitian`, `OrbitalSliceMatrix`
 """
 hamiltonian
 

src/greenfunction.jl

@@ -131,7 +131,6 @@ Base.getindex(g::GreenFunction; kw...) = g[siteselector(; kw...)]
 Base.getindex(g::GreenFunction, kw::NamedTuple) = g[siteselector(; kw...)]
 
 Base.getindex(g::GreenSolution; kw...) = g[getindex(lattice(g); kw...)]
-Base.getindex(g::GreenSolution, kw::NamedTuple) = g[getindex(lattice(g); kw...)]
 
 # g[::Integer, ::Integer] and g[:, :] - intra and inter contacts
 Base.view(g::GreenSolution, i::CT, j::CT = i) where {CT<:Union{Integer,Colon}} =

src/hamiltonian.jl

@@ -406,9 +406,11 @@ call!_output(p::ParametricHamiltonian) = call!_output(hamiltonian(p))
 #endregion
 
 ############################################################################################
-# indexing into AbstractHamiltonian (harmonic extraction) - see also slices.jl
+# indexing into AbstractHamiltonian - see also slices.jl
 #region
 
+# Extraction of Harmonics
+
 Base.getindex(h::AbstractHamiltonian, dn::Union{Tuple,Integer,SVector,AbstractVector}) =
     flat(h[hybrid(dn)])
 
@@ -441,10 +443,31 @@ function Base.isassigned(h::AbstractHamiltonian{<:Any,<:Any,L}, dn::SVector{L,In
     return false
 end
 
+# SiteSelector indexing - replicates GreenSolution indexing - see GreenFunctions.jl
+
+Base.getindex(h::ParametricHamiltonian; kw...) = getindex(call!(h); kw...)
+Base.getindex(h::Hamiltonian; kw...) = h[siteselector(; kw...)]
+
+# conversion down to CellOrbitals. See sites_to_orbs in slices.jl
+Base.getindex(h::ParametricHamiltonian, i, j) = getindex(call!(h), i, j)
+Base.getindex(h::Hamiltonian, i, j) = getindex(h, sites_to_orbs(i, h), sites_to_orbs(j, h))
+# 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´))
+
+# wrapped matrix for end user consumption
+Base.getindex(h::Hamiltonian, i::OrbitalSliceGrouped, j::OrbitalSliceGrouped) =
+    OrbitalSliceMatrix(
+        mortar((h[si, sj] for si in cellsdict(i), sj in cellsdict(j))),
+    (i, j))
+
+Base.getindex(h::Hamiltonian, i::AnyOrbitalSlice, j::AnyOrbitalSlice) =
+    mortar((h[si, sj] for si in cellsdict(i), sj in cellsdict(j)))
 
+Base.getindex(h::Hamiltonian, i::AnyOrbitalSlice, j::AnyCellOrbitals) =
+    mortar((h[si, sj] for si in cellsdict(i), sj in (j,)))
 
-Base.getindex(h::AbstractHamiltonian, i::AnyCellSites, j::AnyCellSites = i) =
-    hamiltonian(h)[sites_to_orbs(i, h), sites_to_orbs(j, h)]
+Base.getindex(h::Hamiltonian, i::AnyCellOrbitals, j::AnyOrbitalSlice) =
+    mortar((h[si, sj] for si in (i,), sj in cellsdict(j)))
 
 function Base.getindex(h::Hamiltonian{T}, i::AnyCellOrbitals, j::AnyCellOrbitals) where {T}
     dn = cell(i) - cell(j)

src/observables.jl

@@ -130,7 +130,7 @@ end
 ############################################################################################
 # ldos: local spectral density
 #   d = ldos(::GreenSolution; kernel = I)      -> d[sites...]::Vector
-#   d = ldos(::GreenSlice; kernel = I) -> d(ω; params...)::Vector
+#   d = ldos(::GreenSlice; kernel = I)         -> d(ω; params...)::Vector
 #   Here ldos is given as Tr(ρᵢᵢ * kernel) where ρᵢᵢ is the spectral function at site i
 #   Here is the generic fallback that uses G. Any more specialized methods need to be added
 #   to each GreenSolver
@@ -260,8 +260,9 @@ end
 
 function current_matrix(gω, ls, d)
     h = hamiltonian(parent(gω))
-    # see slices.jl for this form of getindex
-    current = h[ls, (hij, cij) -> maybe_project(apply_charge_current(hij, cij, d), d.direction)]
+    # see slices.jl for unflat_getindex
+    current = unflat_sparse_slice(h, ls, ls,
+        (hij, cij) -> maybe_project(apply_charge_current(hij, cij, d), d.direction))
     return current
 end
 

src/slices.jl

@@ -137,30 +137,28 @@ end
 #endregion
 
 ############################################################################################
-# slice of Hamiltonian h[latslice] - returns a SparseMatrix{B,Int}
-#   Elements::B can be transformed by `post(hij, (ci, cj))` using h[latslice; post]
-#   Here ci and cj are single-site CellSite for h
+# unflat_sparse_slice: slice of Hamiltonian h[latslice] - returns a SparseMatrix{B,Int}
+#   This is a more efficient slice builder than mortar, but does not flatten B
+#   Elements::B can be transformed by `post(hij, (ci, cj))` with `h[latslice, latslice, post]`
+#   Here ci and cj are single-site `CellSite` for h
 #   ParametricHamiltonian deliberately not supported, as the output is not updatable
 #region
 
-## disabled this method because h[] is too similar to h[()], and becomes confusing
-# Base.getindex(h::Hamiltonian; post = (hij, cij) -> hij, kw...) = h[getindex(lattice(h); kw...), post]
-
-function Base.getindex(h::Hamiltonian, ls::LatticeSlice, post = (hij, cij) -> hij)
+function unflat_sparse_slice(h::Hamiltonian, lsrows::LS, lscols::LS = lsrows, post = (hij, cij) -> hij) where {LS<:LatticeSlice}
     # @assert lattice(h) === lattice(ls)   # TODO: fails upon plotting a current density (see tutorial)
-    cszero = zerocellsites(h, 1)
+    cszero = zerocellsites(h, 1)           # like `sites(1)`, but with explicit cell
     B = typeof(post(zero(blocktype(h)), (cszero, cszero)))
-    ncols = nrows = length(ls)
+    nrows, ncols = length(lsrows), length(lscols)
     builder = CSC{B}(ncols)
     hars = harmonics(h)
-    for colcs in cellsites(ls)
+    for colcs in cellsites(lscols)
         colcell = cell(colcs)
         colsite = siteindices(colcs)
         for har in hars
             rowcell = colcell + dcell(har)
-            rowsubcell = findsubcell(rowcell, ls)
+            rowsubcell = findsubcell(rowcell, lsrows)
             rowsubcell === nothing && continue
-            rowoffset = offsets(ls, rowcell)
+            rowoffset = offsets(lsrows, rowcell)
             rowsubcellinds = siteindices(rowsubcell)
             # rowsubcellinds are the site indices in original unitcell for subcell = rowcell
             # rowoffset is the latslice site offset for the rowsubcellinds sites
@@ -185,7 +183,6 @@ function Base.getindex(h::Hamiltonian, ls::LatticeSlice, post = (hij, cij) -> hi
     return sparse(builder, nrows, ncols)
 end
 
-
 #endregion
 
 ############################################################################################

test/test_hamiltonian.jl

@@ -1,4 +1,4 @@
-using Quantica: Hamiltonian, ParametricHamiltonian, BarebonesOperator,
+using Quantica: Hamiltonian, ParametricHamiltonian, BarebonesOperator, OrbitalSliceMatrix, SparseMatrixCSC,
       sites, nsites, nonsites, nhoppings, coordination, flat, hybrid, transform!, nnz, nonzeros
 
 @testset "basic hamiltonians" begin
@@ -479,7 +479,15 @@ end
 end
 
 @testset "hamiltonian indexing" begin
-    h = HP.graphene(orbitals = (2, 1), a0 = 1) |> @hopping!((t; p = 0) -> p*t)
-    @test size(h[sites(1), sites(2)]) == (2, 1)
+    h = LP.honeycomb() |> hamiltonian(hopping(I), orbitals = (2, 1)) |> @hopping!((t; p = 1) -> p*t)
+    m = h[sites(1), sites(2)]
+    @test m isa SparseMatrixCSC
+    @test m == SA[1; 0;;]
     @test iszero(h[sites(1), sites(SA[2,3], 2)])
+    m = h[cells = 1]
+    @test m isa OrbitalSliceMatrix
+    @test m == SA[0 0 1; 0 0 0; 1 0 0]
+    m = h[sites(2), (; cells = 1)]
+    @test m isa SparseMatrixCSC
+    @test m == SA[1 0 0]
 end