unflat of OrbitalSliceMatrix

src/docstrings.jl

@@ -1278,14 +1278,46 @@ scalars representing single orbitals. This is only relevant for multi-orbital Ha
 
 Equivalent to `unflat(())`
 
+    unflat(a::OrbitalSliceArray)
+
+Convert an `OrbitalSliceArray` into an array of views, where each view corresponds to a
+block for whole sites (see also `siteindexdict`).
+
+    unflat(f, a::OrbitalSliceArray)
+
+Like above, but apply `f` to each `view` block. Useful e.g. to convert views into StaticArrays with
+`f == SMatrix{N,N}`, see examples.
+
 # Examples
 
 ```
 julia> h = HP.graphene(orbitals = 2); h[unflat(0,0)]
 2×2 SparseArrays.SparseMatrixCSC{SMatrix{2, 2, ComplexF64, 4}, Int64} with 2 stored entries:
                      ⋅                       [2.7+0.0im 0.0+0.0im; 0.0+0.0im 2.7+0.0im]
  [2.7+0.0im 0.0+0.0im; 0.0+0.0im 2.7+0.0im]                      ⋅
+
+ julia> hmat = h[cells = (SA[0,0], SA[1,1])]
+8×8 OrbitalSliceMatrix{ComplexF64,SparseArrays.SparseMatrixCSC{ComplexF64, Int64}}:
+ 0.0+0.0im  0.0+0.0im  2.7+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
+ 0.0+0.0im  0.0+0.0im  0.0+0.0im  2.7+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
+ 2.7+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
+ 0.0+0.0im  2.7+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
+ 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  2.7+0.0im  0.0+0.0im
+ 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  2.7+0.0im
+ 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  2.7+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
+ 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im  2.7+0.0im  0.0+0.0im  0.0+0.0im
+
+julia> unflat(SMatrix{2,2}, hmat)
+4×4 Matrix{SMatrix{2, 2, ComplexF64, 4}}:
+ [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [2.7+0.0im 0.0+0.0im; 0.0+0.0im 2.7+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]
+ [2.7+0.0im 0.0+0.0im; 0.0+0.0im 2.7+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]
+ [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [2.7+0.0im 0.0+0.0im; 0.0+0.0im 2.7+0.0im]
+ [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]  [2.7+0.0im 0.0+0.0im; 0.0+0.0im 2.7+0.0im]  [0.0+0.0im 0.0+0.0im; 0.0+0.0im 0.0+0.0im]
 ```
+
+# See also
+    `flat`, `OrbitalSliceArray`
+
 """
 unflat
 
@@ -1749,7 +1781,7 @@ julia> summary(gω[ss])
 ```
 
 # See also
-    `GreenSolvers`, `diagonal`, `sitepairs`, `ldos`, `conductance`, `current`, `josephson`
+    `GreenSolvers`, `diagonal`, `sitepairs`, `ldos`, `conductance`, `current`, `josephson`, `OrbitalSliceArray`
 """
 greenfunction
 
@@ -1822,7 +1854,7 @@ julia> g(1)[diagonal(:, kernel = SA[1 0; 0 -1])]    # σz spin density of the ab
 ```
 
 # See also
-    `sitepairs`, `greenfunction`, `ldos`, `densitymatrix`
+    `sitepairs`, `greenfunction`, `ldos`, `densitymatrix`, `OrbitalSliceArray`
 """
 diagonal
 
@@ -1861,7 +1893,7 @@ julia> summary(g(1)[sitepairs(range = 1, kernel = SA[1 0; 0 -1])])    # σz spin
 ```
 
 # See also
-    `diagonal`, `hopselector`, `greenfunction`, `ldos`, `densitymatrix`
+    `diagonal`, `hopselector`, `greenfunction`, `ldos`, `densitymatrix`, `OrbitalSliceArray`
 """
 sitepairs
 
@@ -1999,7 +2031,7 @@ julia> J(0.2; B = 0.0)
 ```
 
 # See also
-    `greenfunction`, `ldos`, `conductance`, `josephson`, `transmission`
+    `greenfunction`, `ldos`, `conductance`, `josephson`, `densitymatrix`, `transmission`
 
 """
 current
@@ -2166,6 +2198,9 @@ julia> ρ()  # with mu = kBT = 0 by default
  -0.262865+0.0im        0.5+0.0im
 ```
 
+# See also
+    `greenfunction`, `josephson`, `ldos`, `current`, `conductance`, `transmission`, `OrbitalSliceArray`
+
 """
 densitymatrix
 
@@ -2242,7 +2277,7 @@ julia> J(0.0)
 ```
 
 # See also
-    `greenfunction`,`ldos`, `current`, `conductance`, `transmission`
+    `greenfunction`, `densitymatrix`, `ldos`, `current`, `conductance`, `transmission`, `OrbitalSliceArray`
 """
 josephson
 
@@ -2358,6 +2393,17 @@ Like the above, but returns a view instead of a copy of the indexed orbital matr
 
 Note: `diagonal` indexing is currently not supported by `OrbitalSliceArray`.
 
+# `unflat` conversion
+
+    unflat(mat)
+
+Convert a `mat::OrbitalSliceMatrix` (or an `OrbitalSliceArray` in general) to an array of
+site blocks, where each block is a view into the original `mat`.
+
+    unflat(SMatrix{N,N}, mat)
+
+Like the above but converting each block view into an NxN SMatrix, see `unflat` for details.
+
 # Examples
 
 ```
@@ -2384,7 +2430,7 @@ julia> mat[sites(SA[1], 1)]
 ```
 
 # See also
-    `siteselector`, `cellindices`, `orbaxes`
+    `siteselector`, `cellindices`, `orbaxes`, `unflat`
 """
 OrbitalSliceArray, OrbitalSliceVector, OrbitalSliceMatrix
 

src/specialmatrices.jl

@@ -549,8 +549,13 @@ Base.getindex(a::AbstractOrbitalMatrix, i::C, j::C = i) where {B,C<:CellSitePos{
 Base.checkbounds(::Type{Bool}, a::AbstractOrbitalMatrix, i, j) =
     checkbounds_axis_or_orbaxis(a, i, 1) && checkbounds_axis_or_orbaxis(a, j, 2)
 
-Base.view(a::AbstractOrbitalMatrix, i::AnyCellSites, j::AnyCellSites = i) =
-    view(parent(a), indexcollection(a, i, j)...)
+# we don't support a mix of index types (e.g. AnyCellSites and UnitRanges) because
+# AnyCellSites can be SparseIndices or DiagIndices
+Base.view(a::AbstractOrbitalArray, I::Vararg{AnyCellSites,M}) where {M} =
+    view(parent(a), indexcollection(a, I...)...)
+# but we have this fallback that works if I are all scalars or ranges
+Base.view(a::AbstractOrbitalArray, I::Vararg{Any,M}) where {M} =
+    view(parent(a), I...)
 
 function indexcollection(a::AbstractOrbitalMatrix, i::AnyCellSites, j::AnyCellSites = i)
     rowslice, colslice = orbaxes(a)
@@ -565,8 +570,6 @@ Base.getindex(a::OrbitalSliceVector, i::AnyCellSites) = copy(view(a, i))
 Base.getindex(a::OrbitalSliceVector, i::C) where {B,C<:CellSitePos{<:Any,<:Any,<:Any,B}} =
     sanitize_block(B, view(a, i))
 
-Base.view(a::OrbitalSliceVector, i::AnyCellSites) = view(parent(a), indexcollection(a, i))
-
 function indexcollection(a::OrbitalSliceVector, i::AnyCellSites)
     rowslice = only(orbaxes(a))
     rows = indexcollection(rowslice, apply(i, lattice(rowslice)))
@@ -630,6 +633,23 @@ function checkbounds_orbaxis(a::LatticeSlice, i::CellIndices)
     end
 end
 
+## unflat
+
+unflat(a::AbstractOrbitalArray) = unflat(identity, a)
+
+function unflat(type, a::AbstractOrbitalArray)
+    itr = Iterators.product(siteindexdict.(orbaxes(a))...)
+    matviews = (rngs -> type(view(a, rngs))).(itr)
+    return matviews
+end
+
+unflat(type, a::AbstractOrbitalMatrix) =
+    [type(view(a, is, js)) for is in siteindexdict(orbaxes(a, 1)), js in siteindexdict(orbaxes(a, 2))]
+
+unflat(type, a::AbstractOrbitalVector) =
+    [type(view(a, is)) for is in siteindexdict(orbaxes(a, 1))]
+
+
 ## broadcasting
 
 # following the manual: https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array

test/test_greenfunction.jl

@@ -616,6 +616,15 @@ end
     g = h |> attach(glead, cells = 1) |> greenfunction(GS.Schur(boundary = 0));
     @test sites(lattice(h)) == [SA[0.0]]
 
+    # test unflat of OrbitalSliceMatrix
+    lat = LP.linear()
+    g = combine(lat, translate(lat, SA[0.5])) |> hamiltonian(onsite(I) - hopping(I), orbitals = (1, 2)) |> greenfunction
+    gmat = g(0.0)[]
+    @test unflat(gmat) isa Matrix{<:SubArray}
+    g = combine(lat, translate(lat, SA[0.5])) |> hamiltonian(onsite(I) - hopping(I), orbitals = 2) |> greenfunction
+    gmat = g(0.0)[]
+    @test unflat(SMatrix{2,2}, gmat) isa Matrix{SMatrix{2,2,ComplexF64,4}}
+
     # check minimal_callsafe_copy for 2D Schur
     g = LP.square() |> supercell(1,3) |> @onsite((; mu = 0) -> mu) - hopping(1) |> greenfunction(GS.Schur());
     gs = g[cells = 1]