1) Add plot and save utilities for ReciprocalPath, ReciprocalZone…

… and `BrillouinZone`; 2) Add a method `shrink` for `ReciprocalZone`.

.gitignore

@@ -5,6 +5,7 @@
 deps/deps.jl
 docs/build/
 test/*.png
+test/*.dat
 
 .ipynb_checkpoints/
 .vscode/

src/Frameworks.jl

@@ -12,6 +12,7 @@ using ..Spatials: AbstractLattice, Bond, Neighbors, bonds!, isintracell
 using ..Toolkit: atol, efficientoperations, rtol, decimaltostr
 
 import ..QuantumLattices: add!, expand, expand!, id, reset!, update, update!
+import ..Spatials: save
 import ..Toolkit: contentnames, getcontent
 
 export Action, Algorithm, AnalyticalExpression, Assignment, CompositeGenerator, Entry, Frontend, Image, OperatorGenerator, Parameters, RepresentationGenerator, initialize, prepare!, run!, save
@@ -751,24 +752,20 @@ Save the data of an assignment registered on an algorithm.
     delimited ? save(filename, assign.data) : serialize(filename, assign.data)
     return (alg, assign)
 end
-function save(filename::AbstractString, data::Tuple{AbstractVector{<:Number}, Union{AbstractVector{<:Number}, AbstractMatrix{<:Number}}})
+@inline save(filename::AbstractString, data::Tuple) = save(filename, data...)
+function save(filename::AbstractString, x::AbstractVector{<:Number}, y::Union{AbstractVector{<:Number}, AbstractMatrix{<:Number}})
+    @assert length(x)==size(y)[1] "save error: mismatched size of x and y."
     open(filename, "w") do f
-        writedlm(f, [data[1] data[2]])
+        writedlm(f, [x y])
     end
 end
-function save(filename::AbstractString, data::Tuple{AbstractVector{<:AbstractVector}, Union{AbstractVector{<:Number}, AbstractMatrix{<:Number}}})
+function save(filename::AbstractString, x::AbstractVector{<:Number}, y::AbstractVector{<:Number}, z::Union{AbstractMatrix{<:Number}, AbstractArray{<:Number, 3}})
+    @assert size(z)[1:2]==(length(y), length(x)) "save error: mismatched size of x, y and z."
     open(filename, "w") do f
-        writedlm(f, [reduce(hcat, data[1])' data[2]])
-    end
-end
-function save(filename::AbstractString, data::Tuple{AbstractVector{<:Number}, AbstractVector{<:Number}, Union{AbstractMatrix{<:Number}, AbstractArray{<:Number, 3}}})
-    open(filename, "w") do f
-        len = length(data[1])*length(data[2])
-        x = reshape(ones(length(data[2]))*data[1]', len)
-        y = reshape(data[2]*ones(length(data[1]))', len)
-        dim = ndims(data[3])==2 ? 1 : size(data[3])[3]
-        z = reshape(data[3], len, dim)
-        writedlm(f, [x y z])
+        new_x = kron(x, ones(length(y)))
+        new_y = kron(ones(length(x)), y)
+        new_z = reshape(z, length(x)*length(y), :)
+        writedlm(f, [new_x new_y new_z])
     end
 end
 
@@ -793,14 +790,22 @@ end
 Define the recipe for the visualization of an assignment of an algorithm.
 """
 @recipe function plot(pack::Tuple{Algorithm, Assignment})
-    title --> nameof(pack[1], pack[2])
-    titlefontsize --> 10
-    legend --> false
-    seriestype --> (isa(pack[2].data, Tuple{Any, Any, Any}) ? :heatmap : :path)
-    xminorticks --> 10
-    yminorticks --> 10
-    minorgrid --> true
+    title --> nameof(pack...)
+    attr = seriestype(pack)
+    isnothing(attr) || begin
+        seriestype --> attr
+        attr==:path && begin
+            legend --> false
+            minorgrid --> true
+            xminorticks --> 10
+            yminorticks --> 10
+        end
+    end
     pack[2].data
 end
+@inline seriestype(_...) = nothing
+@inline seriestype(pack::Tuple{Algorithm, Assignment}) = seriestype(pack[2].data...)
+@inline seriestype(::AbstractVector{<:Number}, ::Union{AbstractVector{<:Number}, AbstractMatrix{<:Number}}, _...) = :path
+@inline seriestype(::AbstractVector{<:Number}, ::AbstractVector{<:Number}, ::Union{AbstractMatrix{<:Number}, AbstractArray{<:Number, 3}}, _...) = :heatmap
 
 end  # module

src/QuantumLattices.jl

@@ -47,7 +47,7 @@ export particlenumbers, periods, spinfulparticles, spinzs, @abeliannumber
 include("Spatials.jl")
 using .Spatials
 export azimuth, azimuthd, distance, isintratriangle, isonline, isparallel, issubordinate, interlinks, minimumlengths, polar, polard, reciprocals, rotate, translate, tile, volume
-export AbstractLattice, Bond, BrillouinZone, Lattice, Neighbors, Point, ReciprocalSpace, ReciprocalZone, ReciprocalPath, bonds, bonds!, icoordinate, isintracell, rcoordinate, selectpath
+export AbstractLattice, Bond, BrillouinZone, Lattice, Neighbors, Point, ReciprocalSpace, ReciprocalZone, ReciprocalPath, bonds, bonds!, icoordinate, isintracell, rcoordinate, save, selectpath, shrink
 export @hexagon_str, @line_str, @rectangle_str
 
 # DegreesOfFreedom
@@ -67,6 +67,6 @@ export latexofphonons, Elastic, PID, Phonon, Kinetic, Hooke, PhononTerm
 # Frameworks
 include("Frameworks.jl")
 using .Frameworks
-export Action, Algorithm, AnalyticalExpression, Assignment, CompositeGenerator, Entry, Frontend, Image, OperatorGenerator, Parameters, RepresentationGenerator, initialize, prepare!, run!, save
+export Action, Algorithm, AnalyticalExpression, Assignment, CompositeGenerator, Entry, Frontend, Image, OperatorGenerator, Parameters, RepresentationGenerator, initialize, prepare!, run!
 
 end

src/Spatials.jl

@@ -1,10 +1,11 @@
 module Spatials
 
 using Base.Iterators: flatten, product
+using DelimitedFiles: writedlm
 using LinearAlgebra: cross, dot, norm
 using NearestNeighbors: KDTree, inrange, knn
 using Printf: @printf, @sprintf
-using RecipesBase: RecipesBase, @recipe, @series
+using RecipesBase: RecipesBase, @recipe, @series, @layout
 using StaticArrays: SVector
 using ..QuantumNumbers: Momenta, Momentum, periods
 using ..Toolkit: atol, rtol, efficientoperations, CompositeDict, Float, SimpleNamedVectorSpace, Segment, VectorSpaceCartesian, VectorSpaceDirectSummed, VectorSpaceStyle, getcontent
@@ -15,7 +16,7 @@ import ..QuantumNumbers: Momentum₁, Momentum₂, Momentum₃
 import ..Toolkit: contentnames, shape
 
 export azimuth, azimuthd, distance, isintratriangle, isonline, isparallel, issubordinate, interlinks, minimumlengths, polar, polard, reciprocals, rotate, translate, tile, volume
-export AbstractLattice, Bond, BrillouinZone, Lattice, Neighbors, Point, ReciprocalSpace, ReciprocalZone, ReciprocalPath, bonds, bonds!, icoordinate, isintracell, nneighbor, rcoordinate, selectpath, ticks
+export AbstractLattice, Bond, BrillouinZone, Lattice, Neighbors, Point, ReciprocalSpace, ReciprocalZone, ReciprocalPath, bonds, bonds!, icoordinate, isintracell, nneighbor, rcoordinate, save, selectpath, shrink, ticks, xaxis, yaxis, zaxis
 export hexagon120°map, hexagon60°map, linemap, rectanglemap, @hexagon_str, @line_str, @rectangle_str
 
 """
@@ -934,6 +935,9 @@ end
 @inline Base.keys(::BrillouinZone{K, P}) where {K, P<:Momentum} = Momenta(P)
 @inline Base.keytype(brillouinzone::BrillouinZone) = keytype(typeof(brillouinzone))
 @inline Base.keytype(::Type{<:BrillouinZone{K, P} where K}) where {P<:Momentum} = P
+@inline xaxis(brillouinzone::BrillouinZone) = (n=periods(keytype(brillouinzone))[1]; collect(Float64, 0:(n-1))/n)
+@inline yaxis(brillouinzone::BrillouinZone) = (n=periods(keytype(brillouinzone))[2]; collect(Float64, 0:(n-1))/n)
+@inline zaxis(brillouinzone::BrillouinZone) = (n=periods(keytype(brillouinzone))[3]; collect(Float64, 0:(n-1))/n)
 
 """
     BrillouinZone(reciprocals::AbstractVector{<:AbstractVector}, nk)
@@ -983,6 +987,19 @@ end
     end
     return result
 end
+@inline xaxis(reciprocalzone::ReciprocalZone) = collect(Float64, reciprocalzone.bounds[1])
+@inline yaxis(reciprocalzone::ReciprocalZone) = collect(Float64, reciprocalzone.bounds[2])
+@inline zaxis(reciprocalzone::ReciprocalZone) = collect(Float64, reciprocalzone.bounds[3])
+
+"""
+    shrink(reciprocalzone::ReciprocalZone{K}, ranges::Vararg{OrdinalRange{<:Integer}, N}) where {K, N} -> ReciprocalZone
+
+Shrink a reciprocal zone.
+"""
+function shrink(reciprocalzone::ReciprocalZone{K}, ranges::Vararg{OrdinalRange{<:Integer}, N}) where {K, N}
+    @assert length(ranges)==length(reciprocalzone.reciprocals) "shrink error: mismatched number of ranges and reciprocals."
+    return ReciprocalZone{K}(reciprocalzone.reciprocals, ntuple(i->reciprocalzone.bounds[i][ranges[i]], N|>Val))
+end
 
 """
     ReciprocalZone(
@@ -1333,4 +1350,105 @@ macro hexagon_str(str::String)
     return ntuple(i->map[points[i]]=>map[points[i+1]], length(points)-1), ntuple(i->points[i]=>points[i+1], length(points)-1)
 end
 
+# plot utilities
+block = quote
+    seriestype --> :path
+    legend --> false
+    minorgrid --> true
+    xminorticks --> 10
+    yminorticks --> 10
+    xticks --> ticks(path)
+    xlabel --> string(names(path)[1])
+    xlims --> (0, length(path)-1)
+    collect(0:length(path)-1), data
+end
+@eval @recipe plot(path::ReciprocalPath, data::AbstractVector) = $block
+@eval @recipe plot(path::ReciprocalPath, data::AbstractMatrix) = $block
+
+@recipe function plot(path::ReciprocalPath, y::AbstractVector, data::AbstractMatrix)
+    seriestype --> :heatmap
+    xticks --> ticks(path)
+    xlabel --> string(names(path)[1])
+    xlims --> (0, length(path)-1)
+    ylims --> (minimum(y), maximum(y))
+    clims --> extrema(data)
+    collect(0:length(path)-1), y, data
+end
+
+block = quote
+    @assert length(reciprocalspace.reciprocals)==2 "plot error: only two dimensional reciprocal spaces are supported."
+    x, y = xaxis(reciprocalspace), yaxis(reciprocalspace)
+    Δx, Δy= x[2]-x[1], y[2]-y[1]
+    seriestype --> :heatmap
+    aspect_ratio --> :equal
+    xlims --> (x[1]-Δx, x[end]+Δx)
+    ylims --> (y[1]-Δy, y[end]+Δy)
+    clims --> extrema(data)
+    xlabel --> string(names(reciprocalspace)[1], "₁")
+    ylabel --> string(names(reciprocalspace)[1], "₂")
+    x, y, data
+end
+@eval @recipe plot(reciprocalspace::BrillouinZone, data::AbstractMatrix) = $block
+@eval @recipe plot(reciprocalspace::ReciprocalZone, data::AbstractMatrix) = $block
+
+setup(expr::Expr) = quote
+    clims = extrema(data)
+    nr = round(Int, sqrt(size(data)[3]))
+    nc = ceil(Int, size(data)[3]/nr)
+    layout := @layout [(nr, nc); b{0.05h}]
+    colorbar := false
+    for i = 1:size(data)[3]
+        @series begin
+            isnothing(subtitles) || begin
+                title := subtitles[i]
+                isnothing(subtitlefontsize) || (titlefontsize := subtitlefontsize)
+            end
+            subplot := i
+            clims --> clims
+            $expr
+        end
+    end
+    subplot := nr*nc+1
+    seriestype := :heatmap
+    xlims := (minimum(clims), maximum(clims))
+    xlabel := ""
+    ylims := (0, 1)
+    yticks := (0:1, ("", ""))
+    ylabel := ""
+    LinRange(clims..., 100), [0, 1], [LinRange(clims..., 100)'; LinRange(clims..., 100)']
+end
+@eval @recipe plot(path::ReciprocalPath, y::AbstractVector, data::AbstractArray{<:Number, 3}; subtitles=nothing, subtitlefontsize=nothing) = $(setup(:(path, y, data[:, :, i])))
+@eval @recipe plot(reciprocalspace::BrillouinZone, data::AbstractArray{<:Number, 3}; subtitles=nothing, subtitlefontsize=nothing) = $(setup(:(reciprocalspace, data[:, :, i])))
+@eval @recipe plot(reciprocalspace::ReciprocalZone, data::AbstractArray{<:Number, 3}; subtitles=nothing, subtitlefontsize=nothing) = $(setup(:(reciprocalspace, data[:, :, i])))
+
+# save utilities
+function save(filename::AbstractString, path::ReciprocalPath, data::Union{AbstractVector{<:Number}, AbstractMatrix{<:Number}})
+    @assert length(path)==size(data)[1] "save error: mismatched size of path and data."
+    open(filename, "w") do f
+        writedlm(f, [matrix(path) data])
+    end
+end
+function save(filename::AbstractString, path::ReciprocalPath, y::AbstractVector, data::Union{AbstractMatrix{<:Number}, AbstractArray{<:Number, 3}})
+    @assert size(data)[1:2]==(length(y), length(path)) "save error: mismatched size of path, y and data."
+    open(filename, "w") do f
+        x = matrix(kron(path, ones(length(y))))
+        y = kron(ones(length(path)), y)
+        z = reshape(data, prod(size(data)[1:2]), :)
+        writedlm(f, [x y z])
+    end
+end
+function save(filename::AbstractString, reciprocalspace::Union{BrillouinZone, ReciprocalZone}, data::Union{AbstractMatrix{<:Number}, AbstractArray{<:Number, 3}})
+    @assert length(reciprocalspace)==prod(size(data)[1:2]) "save error: mismatched size of reciprocal space and data."
+    open(filename, "w") do f
+        writedlm(f, [matrix(reciprocalspace) reshape(data, prod(size(data)[1:2]), :)])
+    end
+end
+function matrix(vs::AbstractVector{<:AbstractVector})
+    result = zeros(eltype(eltype(vs)), length(vs), length(vs[1]))
+    for (i, v) in enumerate(vs)
+        result[i, :] = v
+    end
+    return result
+end
+
 end #module

test/Frameworks.jl

@@ -213,4 +213,7 @@ end
     @test dos.parameters == (t=1.0, U=6.0)
     @test dos.data == 28.0
     @test dos.action.μ == -3.0
+
+    save("path.dat", 1:100, rand(100, 2))
+    save("heatmap.dat", 1:100, 1:50, rand(50, 100))
 end

test/Spatials.jl

@@ -2,7 +2,7 @@ using Plots: plot, savefig, plot!
 using QuantumLattices.Spatials
 using QuantumLattices: decompose, dimension, dtype, expand
 using QuantumLattices.QuantumNumbers: AbelianNumbers, Momenta, Momentum₁, Momentum₂, Momentum₃
-using QuantumLattices.Toolkit: Float, contentnames, getcontent, shape
+using QuantumLattices.Toolkit: Float, Segment, contentnames, getcontent, shape
 using Random: seed!
 using StaticArrays: SVector
 
@@ -238,13 +238,19 @@ end
 
     recipls = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
     @test BrillouinZone(Momentum₃{10, 10, 10}, recipls) == BrillouinZone(recipls, 10)
+
+    bz = BrillouinZone(Momentum₃{10, 100, 1000}, recipls)
+    @test xaxis(bz) == collect(Float64, 0:9)/10
+    @test yaxis(bz) == collect(Float64, 0:99)/100
+    @test zaxis(bz) == collect(Float64, 0:999)/1000
 end
 
 @testset "ReciprocalZone" begin
     rz = ReciprocalZone([[1.0]], length=10)
     @test rz == ReciprocalZone([[1.0]], -1//2=>1//2; length=10)
     @test rz == ReciprocalZone([[1.0]], (-1//2=>1//2,); length=10)
     @test rz == ReciprocalZone([[1.0]], [-1//2=>1//2]; length=10)
+    @test all(shrink(rz, 1:5) .== ReciprocalZone([[1.0]], -0.5=>-0.1, length=5, ends=(true, true)))
 
     rz = ReciprocalZone{:q}([[1.0]], length=10)
     @test rz == ReciprocalZone{:q}([[1.0]], -1//2=>1//2; length=10)
@@ -256,6 +262,11 @@ end
     @test ReciprocalZone([b₁, b₂], -1=>1, -1=>1; length=10).volume == 4
     @test ReciprocalZone([b₁, b₂, b₃], -1=>1, -1=>1, -1=>1; length=10).volume == 8
 
+    rz = ReciprocalZone([b₁, b₂, b₃], -2=>2, -1=>1, -3=>3; length=10)
+    @test xaxis(rz) == collect(Segment(-2, 2, 10))
+    @test yaxis(rz) == collect(Segment(-1, 1, 10))
+    @test zaxis(rz) == collect(Segment(-3, 3, 10))
+
     bz = BrillouinZone{:q}(Momentum₂{8, 8}, [[1.0, 0.0], [0.0, 1.0]])
     rz = ReciprocalZone(bz)
     @test rz == ReciprocalZone{:q}([[1.0, 0.0], [0.0, 1.0]], 0=>1, 0=>1; length=8)
@@ -317,3 +328,51 @@ end
     plot!(plt, map(index->Tuple(bz[index]), indexes), seriestype=:scatter)
     savefig(plt, "PickPoint.png")
 end
+
+@testset "utilities" begin
+    path = ReciprocalPath([[2pi, 0], [0, 2pi]], rectangle"Γ-X-M-Γ")
+    band = map(k->-2cos(k[1])-2cos(k[2]), path)
+    savefig(plot(path, band), "SingleBand.png")
+    save("SingleBand.dat", path, band)
+    savefig(plot(path, [band -band]), "MultiBands.png")
+    save("MultiBands.dat", path, [band -band])
+
+    energies = LinRange(-6.0, 6.0, 401)
+    spectrum = [-imag(1/(energies[i]+0.1im-band[j])) for i=1:length(energies), j=1:length(band)]
+    savefig(plot(path, energies, spectrum), "SingleSpectrum.png")
+    save("SingleSpectrum.dat", path, energies, spectrum)
+
+    spectra = zeros(size(spectrum)..., 2)
+    spectra[:, :, 1] = spectrum
+    spectra[:, :, 2] = spectrum
+    savefig(plot(path, energies, spectra), "MultiSpectra.png")
+    save("MultiSpectra.dat", path, energies, spectra)
+
+    bz = BrillouinZone([[2pi, 0], [0, 2pi]], 200)
+    surface = zeros(Float64, (200, 200))
+    for (i, k) in enumerate(bz)
+        surface[i] = -imag(1/(0.1im+2cos(k[1])+2cos(k[2])))
+    end
+    savefig(plot(bz, surface), "SingleSurface.png")
+    save("SingleSurface.dat", bz, surface)
+
+    surfaces = zeros(size(surface)..., 2)
+    surfaces[:, :, 1] = surface
+    surfaces[:, :, 2] = surface
+    savefig(plot(bz, surfaces), "MultiSurfaces.png")
+    save("MultiSurfaces.dat", bz, surfaces)
+
+    rz = ReciprocalZone([[2pi, 0], [0, 2pi]], -2=>2, -1=>1, length=(400, 200))
+    surface = zeros(Float64, (200, 400))
+    for (i, k) in enumerate(rz)
+        surface[i] = -imag(1/(0.1im+2cos(k[1])+2cos(k[2])))
+    end
+    savefig(plot(rz, surface), "SingleExtendedSurface.png")
+    save("SingleExtendedSurface.dat", rz, surface)
+
+    surfaces = zeros(size(surface)..., 2)
+    surfaces[:, :, 1] = surface
+    surfaces[:, :, 2] = surface
+    savefig(plot(rz, surfaces), "MultiExtendedSurfaces.png")
+    save("MultiExtendedSurfaces.dat", rz, surfaces)
+end