Proposal: Flat wrapper

[pablosanjose]

Context

Way back, we made in Quantica a pretty fundamental design choice to allow multi-orbital systems that pack several orbitals into single sites, with the only restriction that the number of orbitals of all sites in a sublattice is constant. This packing means that sparse Hamiltonian matrices are internally stored with SMatrix{N,M} eltype, encoding hoppings between two sites with M and N orbitals each. Since we cannot have, for performance reasons, an inhomogeneous eltype in a sparse matrix, we enlarge N and M to the highest number of orbitals in any sublattice, padding the extra entries with zeros.

This design is complicated, but it is both fast and general, with little overhead. It does require that we often need to convert sparse matrices of SMatrices to flattened sparse matrices of scalars, to e.g. call external linear algebra libraries like LAPACK. Other times we don't want to flatten, so we can take advantage of the dense and SIMD-favorable nature of inter-orbital SMatrices, e.g. in KPM routines. We also don't want to flatten when we want to write parametric Hamiltonians that have modifiers operating on orbital blocks, although that could perhaps be now made to work on views over sparse matrices (which are no longer allocating), at the expense of some serious performance traps. But in this issue I don't want to discuss the possible alternative designs (I still believe we made the right choice with the design). What I want to reconsider is how we deal with flattening.

Problem

We do some serious acrobatics to allow flattening of Bloch Hamiltonians. I would even say our current solution is clunky. We need to store flattened information in the OrbitalStructure metadata of Hamiltonians, alongside the non-flattened data. Slices of Hamiltonians and ParametricHamiltonians also need to keep track of flattened and unflattened slice indices. We need to specify the eltype when building the target matrix m=similarmatrix(h, flatten) or similarmatrix(h, Matrix{ComplexF64}) for bloch!, so that the flattening codepath for bloch!(m, h, ϕs) is triggered. Most of this is an internal burden, with little impact for the user, but I wonder if we could do something more elegant to encode flattened Hamiltonians.

Proposal

The idea would be to encode a Flat trait like one does Adjoint in Julia, by wrapping a Hamiltonian. So we could have flatten(h) or flat(h) be actually lazy, and produce an object of type

struct Flat{H<:Union{Hamiltonian, ParametricHamiltonian, Slice, Ket}, O<:OrbitalStructure}
    h::H
    flatorbstruct::O
end

with the actual, eager, flattening occurring only as needed (e.g. upon calling bloch or getindex). The unflat orbitalstructure(h) would be kept inside h, and would hold only the unflattened sublattice offsets, with flattened offsets in flatorbstruct.

Discussion

We could then specialize similarmatrix(h::Flat, A::Type{AbstractArray}) without needing to specify the eltype. Thus similarmatrix(flat(h), Matrix) would be the equivalent to the current similarmatrix(h, Matrix{Quantica.blockeltype(h)}) and similarmatrix(flat(h)) would be the equivalent to the current similarmatrix(h, flatten) which I never quite liked.

Similarly, the flatten/unflatten code would become trivial, as everything is lazy. Currently, we always need to say unflatten(x, o::OrbitalStructure) to unflatten x to a form compatible with o, because flattening loses the unflattened o. With this proposal, unflat would merely unwrap h, since nothing is lost as flat is lazy.

Another advantage is that the Flat functionality would become orthogonal to all the rest, and would reside in an independent flat.jl file

A possible shortcoming could be that repeated bloch of a lazy-flattened hamiltonian (e.g. for bandstructures) would be slower than a bloch of an eager-flattened hamiltonian. But experience now suggests that the flattening bloch mechanism is very fast, and is a negligible part of the runtime when computing bandstructures.

Another shortcoming is that we might not be able to use duck typing with Flat, unlike with Adjoint{<:AbstractArray} which is itself a subtype of AbstractArray, simply because the wrapped h is not always of the same abstract type. The workaround is to use an alias such as const MaybeFlat{H} = Union{H,Flat{H}} and use that in the signatures of functions