Unflatten: restore orbital structure into elements

src/hamiltonian.jl

@@ -280,6 +280,53 @@ function flatoffsetorbs(i, lat, norbs, offsets´)
     return i´, N
 end
 
+## unflatten ##
+
+unflatten(v, h) = unflatten!(similar(v, orbitaltype(h), size(h, 2)), v, h)
+
+function unflatten!(v::AbstractVector{T}, vflat::AbstractArray, h::Hamiltonian) where {T}
+    norbs = length.(h.orbitals)
+    offsetsflat = flatoffsets(h.lattice.unitcell.offsets, norbs)
+    dimflat = last(offsetsflat)
+    check_unflatten_dst_dims(v, h)
+    check_unflatten_src_dims(vflat, dimflat)
+    check_unflatten_eltypes(v, h)
+    j = 0
+    for s in sublats(h.lattice)
+        N = norbs[s]
+        for i in offsetsflat[s]+1:N:offsetsflat[s+1]
+            j += 1
+            v[j] = padright(view(vflat, i:i+N-1), T)
+        end
+    end
+    return v
+end
+
+check_unflatten_dst_dims(v, h) =
+    size(v, 1) == size(h, 2) ||
+        throw(ArgumentError("Dimension of destination array is inconsistent with Hamiltonian"))
+
+check_unflatten_src_dims(vflat, dimflat) =
+    size(vflat, 1) == dimflat ||
+        throw(ArgumentError("Dimension of source array is inconsistent with Hamiltonian"))
+
+check_unflatten_eltypes(v::AbstractVector{T}, h) where {T} =
+    T === orbitaltype(h) ||
+        throw(ArgumentError("Eltype of desination array is inconsistent with Hamiltonian"))
+
+## maybe_unflatten: call unflatten but only if we cannot do it without copying
+maybe_unflatten(vflat, h) = _maybe_unflatten(vflat, h, orbitaltype(h), h.orbitals)
+# source is already of the correct orbitaltype(h)
+function _maybe_unflatten(v::AbstractVector{T}, h, ::Type{T}, orbs) where {T}
+    check_unflatten_dst_dims(v, h)
+    return v
+end
+# source can be reinterpreted, because the number of orbitals is the same M for all N sublattices
+_maybe_unflatten(v::AbstractVector{T}, h, ::Type{S}, ::NTuple{N,NTuple{M}}) where {N,M,T<:Number,S<:SVector{M}} =
+    reinterpret(SVector{M,T}, v)
+# otherwise call unflatten
+_maybe_unflatten(v, h, S, orbs) = unflatten(v, h)
+
 #######################################################################
 # similarmatrix
 #######################################################################
@@ -568,7 +615,7 @@ blockdim(::Type{S}) where {N,S<:SMatrix{N,N}} = N
 blockdim(::Type{T}) where {T<:Number} = 1
 
 # find SVector type that can hold all orbital amplitudes in any lattice sites
-orbitaltype(orbs, type::Type{Tv} = Complex{T}) where {T,Tv} =
+orbitaltype(orbs, type::Type{Tv}) where {Tv} =
     _orbitaltype(SVector{1,Tv}, orbs...)
 _orbitaltype(::Type{S}, ::NTuple{D,NameType}, os...) where {N,Tv,D,S<:SVector{N,Tv}} =
     (M = max(N,D); _orbitaltype(SVector{M,Tv}, os...))

src/tools.jl

@@ -86,12 +86,16 @@ end
 
 padright(sv::StaticVector{E,T}, x::T, ::Val{E}) where {E,T} = sv
 padright(sv::StaticVector{E1,T1}, x::T2, ::Val{E2}) where {E1,T1,E2,T2} =
-    (T = promote_type(T1,T2); SVector{E2, T}(ntuple(i -> i > E1 ? x : convert(T, sv[i]), Val(E2))))
+    (T = promote_type(T1,T2); SVector{E2,T}(ntuple(i -> i > E1 ? x : convert(T, sv[i]), Val(E2))))
 padright(sv::StaticVector{E,T}, ::Val{E2}) where {E,T,E2} = padright(sv, zero(T), Val(E2))
 padright(sv::StaticVector{E,T}, ::Val{E}) where {E,T} = sv
 padright(t::NTuple{N´,Any}, x, ::Val{N}) where {N´,N} = ntuple(i -> i > N´ ? x : t[i], Val(N))
 padright(t::NTuple{N´,Any}, ::Val{N}) where {N´,N} = ntuple(i -> i > N´ ? 0 : t[i], Val(N))
 
+padright(v, ::Type{S}) where {E,T,S<:SVector{E,T}} = padright(v, zero(T), S)
+padright(v, x::T, ::Type{S}) where {E,T,S<:SVector{E,T}} =
+    SVector{E,T}(ntuple(i -> i > length(v) ? x : convert(T, v[i]), Val(E)))
+
 # Pad element type to a "larger" type
 @inline padtotype(s::SMatrix{E,L}, ::Type{S}) where {E,L,E2,L2,S<:SMatrix{E2,L2}} =
     S(SMatrix{E2,E}(I) * s * SMatrix{L,L2}(I))

test/test_bandstructure.jl

@@ -76,4 +76,29 @@ end
     b = bandstructure(ph, mesh2D, mapping = (k, φ) -> ((;k = k, p = SA[1, φ]),))
     @test length(bands(b)) == 4
     @test_throws UndefKeywordError bandstructure(ph, mesh2D, mapping = (k, φ) -> ((;p = SA[1, φ]),))
+end
+
+@testset "unflatten" begin
+    h = LatticePresets.honeycomb() |> hamiltonian(onsite(2I) + hopping(I, range = 1), orbitals = (Val(2), Val(1))) |> unitcell(2) |> unitcell
+    sp = states(spectrum(h))[:,1]
+    sp´ = Quantica.maybe_unflatten(sp, h)
+    l = size(h, 1)
+    @test length(sp) == 1.5 * l
+    @test length(sp´) == l
+    @test all(x -> iszero(last(x)), sp´[l+1:end])
+    @test sp´ isa Vector
+    @test sp´ !== sp
+
+    h = LatticePresets.honeycomb() |> hamiltonian(onsite(2I) + hopping(I, range = 1), orbitals = Val(2)) |> unitcell(2) |> unitcell
+    sp = states(spectrum(h))[:,1]
+    sp´ = Quantica.maybe_unflatten(sp, h)
+    l = size(h, 1)
+    @test length(sp) == 2 * l
+    @test length(sp´) == l
+    @test sp´ isa Base.ReinterpretArray
+
+    h = LatticePresets.honeycomb() |> hamiltonian(onsite(2I) + hopping(I, range = 1), orbitals = Val(2)) |> unitcell(2) |> unitcell
+    sp = states(spectrum(h))[:,1]
+    sp´ = Quantica.maybe_unflatten(sp, h)
+    @test sp === sp
 end