Extend plot support for entanglementplot and transferplot to Makie.jl

Project.toml

@@ -17,17 +17,21 @@ OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorKitManifolds = "11fa318c-39cb-4a83-b1ed-cdc7ba1e3684"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 
 [extensions]
 MPSKitAdaptExt = "Adapt"
+MPSKitMakieExt = ["Makie", "LaTeXStrings"]
+MPSKitPlotsExt = "RecipesBase"
 
 [workspace]
 projects = ["test", "docs"]
@@ -40,16 +44,36 @@ Compat = "3.47, 4.10"
 DocStringExtensions = "0.9.3"
 HalfIntegers = "1.6.0"
 KrylovKit = "0.8.3, 0.9.2, 0.10"
+LaTeXStrings = "1"
 LinearAlgebra = "1.6"
 LoggingExtras = "~1.0"
+Makie = "0.24, 0.25"
 MatrixAlgebraKit = "0.6.5"
 OhMyThreads = "0.7, 0.8"
 OptimKit = "0.3.1, 0.4"
 Printf = "1"
 Random = "1"
-RecipesBase = "1.1"
+RecipesBase = "1"
 TensorKit = "0.16.5"
 TensorKitManifolds = "0.7, 0.8"
 TensorOperations = "5.5.1"
 VectorInterface = "0.2, 0.3, 0.4, 0.5"
 julia = "1.10"
+
+[extras]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+TensorKitTensors = "41b62e7d-e9d1-4e23-942c-79a97adf954b"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
+
+[targets]
+test = ["Aqua", "Adapt", "CUDA", "cuTENSOR", "Pkg", "Test", "TestExtras", "Plots", "Makie", "CairoMakie", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"]

ext/MPSKitMakieExt.jl

@@ -0,0 +1,192 @@
+module MPSKitMakieExt
+
+using Makie, LaTeXStrings
+using MPSKit, TensorKit
+
+#TODO?: add Colors.jl to access this, allows Plots extension to also use these colors
+const JLCOLORS = Makie.Colors.JULIA_LOGO_COLORS
+
+@recipe(EntanglementPlot, mps) do scene
+    Attributes(
+        site = 0,
+        expand_symmetry = false,
+        sortby = maximum,
+        sector_margin = 1 // 10,
+        sector_formatter = string,
+    )
+end
+
+function Makie.plot!(ep::EntanglementPlot)
+    #TODO: still want this style where sectors are separated?
+    mps = ep.mps[]
+    site = ep.site[]
+    margin = ep.sector_margin[]
+
+    (isa(mps, FiniteMPS) && (site == 0 || site > length(mps))) &&
+        throw(ArgumentError("Invalid site $site for the given mps."))
+
+    spectra = entanglement_spectrum(mps, site)
+
+    sectors = sectortype(mps)[]
+    spectrum = Vector{Vector{Float64}}()
+
+    for (c, b) in pairs(spectra)
+        if ep.expand_symmetry[]
+            b′ = repeat(b, dim(c))
+            sort!(b′; rev = true)
+            push!(spectrum, b′)
+        else
+            push!(spectrum, b)
+        end
+        push!(sectors, c)
+    end
+
+    # Sort sectors according to provided method
+    if length(spectrum) > 1
+        order = sortperm(spectrum; by = ep.sortby[], rev = true)
+        spectrum = spectrum[order]
+        sectors = sectors[order]
+    end
+
+    ax = Makie.current_axis()
+
+    # Axis styling
+    ax.title = L"\text{Entanglement Spectrum}"
+    ax.titlesize = 24
+
+    ax.xlabel = latexstring("\$\\chi\$ = $(round(Int, dim(left_virtualspace(mps, site))))") # still want this?
+    ax.xlabelsize = 24
+    ax.xticks = (1:length(sectors), ep.sector_formatter[].(sectors))
+    ax.xticklabelsize = 16
+    ax.xticklabelrotation = 45
+    ax.xticklabelalign = (:right, :top)
+    xlims!(ax, 1, length(sectors) + 1)
+
+    ax.ylabel = L"\log(\lambda)"
+    ax.ylabelsize = 24
+    bottom = floor(Int, log10(minimum(spectra)))
+    ax.yticks = (bottom:2:0, latexstring.(collect(bottom:2:0)))
+    ax.yticklabelsize = 16
+    ylims!(ax, bottom, 0 + 1.0e-1)
+
+    # Plot data
+    for (i, (partial_spectrum, sector)) in enumerate(zip(spectrum, sectors))
+        n_spectrum = length(partial_spectrum)
+        if n_spectrum == 1
+            x = [i + 0.5]
+        else
+            x = collect(range(i + float(margin), i + 1 - float(margin); length = n_spectrum))
+        end
+        scatter!(ep, x, log10.(partial_spectrum), color = JLCOLORS[mod1(i, length(JLCOLORS))])
+    end
+
+    return ep
+end
+
+function MPSKit.entanglementplot(args...; plotkwargs = (;), kwargs...)
+    p = entanglementplot(args...; kwargs...)
+    ax = p.axis
+
+    # overwrite user-provided axis attributes
+    for (k, v) in pairs(plotkwargs)
+        setproperty!(ax, k, v)
+    end
+    return p
+end
+
+#------------------------------------------------------------
+
+@recipe(TransferPlot, mps) do scene
+    Attributes(
+        below = nothing,
+        sectors = nothing,
+        transferkwargs = NamedTuple(),
+        thetaorigin = 0.0,
+        sector_formatter = string,
+    )
+end
+
+function Makie.plot!(tp::TransferPlot)
+    #TODO: consider radial plot
+    mps = tp.mps[]
+    below = tp.below[] === nothing ? mps : tp.below[]
+    sectors = tp.sectors[] === nothing ? [leftunit(mps)] : tp.sectors[]
+    transferkwargs = NamedTuple( # weird convert thing
+        k => (v isa Observable ? v[] : v) for (k, v) in pairs(tp.transferkwargs[])
+    )
+    thetaorigin = tp.thetaorigin[]
+    sector_formatter = tp.sector_formatter[]
+
+    ax = Makie.current_axis()
+    ax.title = L"\text{Transfer Spectrum}"
+    ax.titlesize = 24
+    ax.xlabel = L"\theta"
+    ax.xlabelsize = 24
+    ax.xticklabelsize = 16
+    ax.ylabel = L"r"
+    ax.ylabelsize = 24
+    ax.yticklabelsize = 16
+
+    ax.xticks = pitick(0, 2pi, 4; mode = :latex)
+    ax.yticks = (range(0, 1.0; length = 6), latexstring.(range(0, 1.0; length = 6)))
+    ax.xgridvisible = true
+    ax.ygridvisible = true
+
+    ax.leftspinevisible = true
+    ax.rightspinevisible = false
+    ax.bottomspinevisible = true
+    ax.topspinevisible = false
+
+    for (i, sector) in enumerate(sectors)
+        spectrum = transfer_spectrum(mps; below = below, sector = sector, transferkwargs...)
+        θ = mod2pi.(angle.(spectrum) .+ thetaorigin) .- thetaorigin
+        r = abs.(spectrum)
+        scatter!(tp, θ, r; label = sector_formatter(sector), color = JLCOLORS[mod1(i, length(JLCOLORS))])
+    end
+
+    xlims!(ax, thetaorigin - 0.1, thetaorigin + 2π + 0.1)
+    ylims!(ax, nothing, 1.05)
+    Legend(Makie.current_figure()[1, 1], tp.plots, [sector_formatter(s) for s in sectors]; tellwidth = false, halign = :center, valign = :top)
+    return tp
+end
+
+function MPSKit.transferplot(args...; plotkwargs = (;), kwargs...)
+    p = transferplot(args...; kwargs...)
+    ax = p.axis
+
+    # overwrite user-provided axis attributes
+    for (k, v) in pairs(plotkwargs)
+        setproperty!(ax, k, v)
+    end
+    return p
+end
+
+# utility for plotting
+
+function pitick(start, stop, denom; mode = :latex)
+    a = Int(cld(start, π / denom))
+    b = Int(fld(stop, π / denom))
+    tick = range(a * π / denom, b * π / denom; step = π / denom)
+    ticklabel = piticklabel.((a:b) .// denom, Val(mode))
+    return tick, ticklabel
+end
+
+function piticklabel(x::Rational, ::Val{:text})
+    iszero(x) && return "0"
+    S = x < 0 ? "-" : ""
+    n, d = abs(numerator(x)), denominator(x)
+    N = n == 1 ? "" : repr(n)
+    d == 1 && return S * N * "π"
+    return S * N * "π/" * repr(d)
+end
+
+function piticklabel(x::Rational, ::Val{:latex})
+    iszero(x) && return L"0"
+    S = x < 0 ? "-" : ""
+    n, d = abs(numerator(x)), denominator(x)
+    N = n == 1 ? "" : repr(n)
+    d == 1 && return L"%$S%$N\pi"
+    return L"%$S\frac{%$N\pi}{%$d}"
+end
+
+end

ext/MPSKitPlotsExt.jl

@@ -0,0 +1,117 @@
+module MPSKitPlotsExt
+
+using RecipesBase
+using MPSKit, TensorKit
+
+@userplot EntanglementPlot
+
+@recipe function f(
+        h::EntanglementPlot; site = 0, expand_symmetry = false, sortby = maximum,
+        sector_margin = 1 // 10, sector_formatter = string
+    )
+    mps = h.args[1]
+    (isa(mps, FiniteMPS) && (site == 0 || site > length(mps))) &&
+        throw(ArgumentError("Invalid site $site for the given mps."))
+
+    spectra = entanglement_spectrum(mps, site)
+    sectors = sectortype(mps)[]
+    spectrum = Vector{Vector{Float64}}()
+    for (c, b) in pairs(spectra)
+        if expand_symmetry # Duplicate entries according to the quantum dimension.
+            b′ = repeat(b, dim(c))
+            sort!(b′; rev = true)
+            push!(spectrum, b′)
+        else
+            push!(spectrum, b)
+        end
+        push!(sectors, c)
+    end
+
+    if length(spectrum) > 1
+        order = sortperm(spectrum; by = sortby, rev = true)
+        spectrum = spectrum[order]
+        sectors = sectors[order]
+    end
+
+    for (i, (partial_spectrum, sector)) in enumerate(zip(spectrum, sectors))
+        @series begin
+            title --> "Entanglement Spectrum"
+            legend --> false
+            grid --> :xy
+            widen --> true
+            bottom_margin -->(10, :mm)
+
+            xguide --> "χ = $(round(Int, dim(left_virtualspace(mps, site))))"
+            xticks --> (1:length(sectors), sector_formatter.(sectors))
+            xtickfonthalign --> :center
+            xtick_direction --> :out
+            xrotation --> 45
+            xlims --> (1, length(sectors) + 1)
+
+            ylims --> (-Inf, 1 + 1.0e-1)
+            yscale --> :log10
+            seriestype := :scatter
+            label := sector_formatter(sector)
+            n_spectrum = length(partial_spectrum)
+
+            # Put single dot in the middle, or a linear range with padding.
+            if n_spectrum == 1
+                x = [i + 1 // 2]
+            else
+                x = range(i + sector_margin, i + 1 - sector_margin; length = n_spectrum)
+            end
+            return x, partial_spectrum
+        end
+    end
+
+    return nothing
+end
+
+MPSKit.entanglementplot(args...; kwargs...) = entanglementplot(args...; kwargs...)
+
+#-----------------------------------------------------------------------------
+
+@userplot TransferPlot
+
+@recipe function f(
+        h::TransferPlot; sectors = nothing, transferkwargs = (;), thetaorigin = 0,
+        sector_formatter = string
+    )
+    if sectors === nothing
+        sectors = [leftunit(h.args[1])]
+    end
+
+    for sector in sectors
+        below = length(h.args) == 1 ? h.args[1] : h.args[2]
+        spectrum = transfer_spectrum(
+            h.args[1]; below = below, sector = sector,
+            transferkwargs...
+        )
+
+        @series begin
+            yguide --> "r"
+            ylims --> (-Inf, 1.05)
+
+            xguide --> "θ"
+            xlims --> (thetaorigin, thetaorigin + 2pi)
+            xticks --> range(0, 2pi; length = 7)
+            xformatter --> x -> "$(rationalize(x / π, tol = 0.05))π"
+            xwiden --> true
+            seriestype := :scatter
+            markershape --> :auto
+            label := sector_formatter(sector)
+            return mod2pi.(angle.(spectrum) .+ thetaorigin) .- thetaorigin, abs.(spectrum)
+        end
+    end
+
+    title --> "Transfer Spectrum"
+    legend --> false
+    grid --> :xy
+    framestyle --> :zerolines
+
+    return nothing
+end
+
+MPSKit.transferplot(args...; kwargs...) = transferplot(args...; kwargs...)
+
+end

src/MPSKit.jl

@@ -73,7 +73,6 @@ using KrylovKit: KrylovAlgorithm
 using OptimKit
 using Base.Threads
 using Base.Iterators
-using RecipesBase
 using VectorInterface
 using Accessors
 using HalfIntegers