Affine decompositions with multi-indices and additional parameters

In this example, we want to explore the capabilities of the central AffineDecomposition type. To advance the previous examples where we covered the magnetization — a very simple observable that consists of only one affine term and a parameter-independent coefficient — we now turn to observables where the indices are multi-indices $r = (r_1, \dots, r_d)$ and the coefficients can depend on additional parameters $p$, aside from the $\bm{\mu}$ parameter points that are present in the Hamiltonian:

\[O(\bm{\mu}, p) = \sum_{q=1}^Q \alpha_q(\bm{\mu}, p)\, O_q\]

To stay within the realm of quantum magnetism, we will consider the so-called spin structure factor

\[\mathcal{S}(k) = \frac{1}{L} \sum_{r,r'=1}^L e^{-i (r - r') k} S^z_r S^z_{r'},\quad \alpha_{r,r'}(k) = \frac{e^{-i (r - r') k}}{L}, \quad O_{r,r'} = S^z_r S^z_{r'}\]

with a wave vector parameter $k$, to discuss the implementation of a more complicated observable.

To provide a physical setup, we again use the XXZ chain from The reduced basis workflow:

using LinearAlgebra
using SparseArrays
using ReducedBasis

σx = sparse([0.0 1.0; 1.0 0.0])
σy = sparse([0.0 -im; im 0.0])
σz = sparse([1.0 0.0; 0.0 -1.0])

function to_global(op::M, L::Int, i::Int) where {M<:AbstractMatrix}
    d = size(op, 1)
    if i == 1
        kron(op, M(I, d^(L - 1), d^(L - 1)))
    elseif i == L
        kron(M(I, d^(L - 1), d^(L - 1)), op)
    else
        kron(kron(M(I, d^(i - 1), d^(i - 1)), op), M(I, d^(L - i), d^(L - i)))
    end
end

function xxz_chain(L)
    H1 = 0.25 * sum(1:L-1) do i
        to_global(σx, L, i) * to_global(σx, L, i + 1) +
        to_global(σy, L, i) * to_global(σy, L, i + 1)
    end
    H2 = 0.25 * sum(1:L-1) do i
        to_global(σz, L, i) * to_global(σz, L, i + 1)
    end
    H3 = 0.5  * sum(1:L) do i
        to_global(σz, L, i)
    end
    AffineDecomposition([H1, H2, H3], μ -> [1.0, μ[1], -μ[2]])
end;

Construct the Hamiltonian for an XXZ chain with 6 sites ...

L = 6
H = xxz_chain(L);

... and generate a RB surrogate using an exact diagonalization solver.

Δ  = range(-1.0, 2.5; length=40)
hJ = range( 0.0, 3.5; length=40)
grid_train = RegularGrid(Δ, hJ)
greedy = Greedy(; estimator=Residual())
lobpcg = LOBPCG(; tol_degeneracy=1e-4)
qrcomp = QRCompress(; tol=1e-9)
rbres = assemble(H, grid_train, greedy, lobpcg, qrcomp);

n      max. err    ‖BᵀB-I‖     time      μ
------------------------------------------------------------
1      NaN         5.68e-15     199ms    [-1.0, 0.0]
2      2.33        5.7e-15      244ms    [2.231, 0.538]
3      0.946       5.8e-15     75.3ms    [0.795, 0.897]
4      0.508       5.91e-15    72.1ms    [1.154, 1.795]
5      0.472       6.21e-15    94.4ms    [2.5, 1.077]
6      0.385       6.51e-15    94.5ms    [0.346, 0.179]
7      0.0727      1.23e-14    74.4ms    [-0.551, 0.0]
8      0.0554      1.93e-14    83.8ms    [-0.462, 0.269]
9      0.0163      9.5e-14     92.0ms    [1.423, 0.538]
10     0.00522     1.61e-13     106ms    [-0.372, 0.449]
11     0.004       3.05e-13     116ms    [1.692, 1.346]
12     0.00151     4.61e-13     134ms    [-0.821, 0.0]
13     0.00039     1.08e-12     148ms    [-0.821, 0.09]
reached residual target

Now the task is to implement the double-sum in $\mathcal{S}$, as well as the $k$-dependency in the coefficients. The double-sum can be encoded by putting all $S^z_r S^z_{r'}$ combinations into an $L \times L$ matrix:

terms = map(idx -> to_global(σz, L, idx[1]) * to_global(σz, L, idx[2]),
            Iterators.product(1:L, 1:L));

Correspondingly, the coefficient function now has to map one $k$ value to a matrix of coefficients of the same size as the terms matrix:

coefficients = k -> map(idx -> cis(-(idx[1] - idx[2]) * k) / L,
                        Iterators.product(1:L, 1:L));

One feature of the structure factor that also shows up in many other affine decompositions with double-sums is that the term indices commute, i.e. $O_{r,r'} = O_{r',r}$. In that case, only the upper triangular matrix has to be computed since $B^\dagger O_{r,r'} B = B^\dagger O_{r',r} B$ are the same in the compressed affine decomposition. So let's create the AffineDecomposition and compress, exploiting this symmetry using the symmetric_terms keyword argument:

SFspin    = AffineDecomposition(terms, coefficients)
sfspin, _ = compress(SFspin, rbres.basis; symmetric_terms=true);

In the online evaluation of the structure factor, we then need to define some wave vector values and compute the structure factor at each of them. As usual, we first define a finer online grid of points and the matching online solver:

Δ_online    = range(first(Δ), last(Δ); length=100)
hJ_online   = range(first(hJ), last(hJ); length=100)
grid_online = RegularGrid(Δ_online, hJ_online)
fulldiag    = FullDiagonalization(lobpcg);

And then we map the grid points to the corresponding structure factor values for a set of different wave vectors:

using Statistics
wavevectors = [0.0, π/4, π/2, π]
sf = [zeros(size(grid_online)) for _ in 1:length(wavevectors)]
for (idx, μ) in pairs(grid_online)
    _, φ_rb = solve(rbres.h, rbres.basis.metric, μ, fulldiag)
    for (i, k) in enumerate(wavevectors)
        sf[i][idx] = mean(u -> real(dot(u, sfspin(k), u)), eachcol(φ_rb))
    end
end

Here we again see the convenience of measuring observables in the online stage; adding more wavevector values does not significantly increase the computational cost, since it corresponds to a mere reevaluation of the coefficient functions and small vector-matrix products. Finally, let us see how the structure factor behaves for the different wave vector values:

using Plots
kwargs = (; xlabel=raw"$\Delta$", ylabel=raw"$h/J$", colorbar=true, leg=false)
hms = []
for (i, q) in enumerate(wavevectors)
    push!(hms, heatmap(grid_online.ranges[1], grid_online.ranges[2], sf[i]';
                       title="\$k = $(round(q/π; digits=3))\\pi\$", kwargs...))
end
plot(hms...)

It can be nicely seen that the spin structure factor indicates the ferromagnetic phase at $k=0$ and then moves through the magnetization plateaus until it reaches the antiferromagnetic plateau at $k=\pi$.