|
| 1 | +import Plots |
| 2 | + |
| 3 | +# This is needed to flag that the plots-dependent code has been loaded |
| 4 | +const PLOTS_LOADED = true |
| 5 | + |
| 6 | +""" |
| 7 | +Plot the trace of an SCF, i.e. the absolute error of the total energy at |
| 8 | +each iteration versus the converged energy in a semilog plot. By default |
| 9 | +a new plot canvas is generated, but an existing one can be passed and reused |
| 10 | +along with `kwargs` for the call to `plot!`. |
| 11 | +""" |
| 12 | +function ScfPlotTrace(plt=Plots.plot(yaxis=:log); kwargs...) |
| 13 | + energies = Float64[] |
| 14 | + function callback(info) |
| 15 | + if info.stage == :finalize |
| 16 | + minenergy = minimum(energies[max(1, end-5):end]) |
| 17 | + error = abs.(energies .- minenergy) |
| 18 | + error[error .== 0] .= NaN |
| 19 | + extra = ifelse(:mark in keys(kwargs), (), (mark=:x, )) |
| 20 | + Plots.plot!(plt, error; extra..., kwargs...) |
| 21 | + display(plt) |
| 22 | + else |
| 23 | + push!(energies, info.energies.total) |
| 24 | + end |
| 25 | + end |
| 26 | +end |
| 27 | + |
| 28 | + |
| 29 | +function plot_band_data(band_data; εF=nothing, |
| 30 | + klabels=Dict{String, Vector{Float64}}(), unit=:eV, kwargs...) |
| 31 | + eshift = isnothing(εF) ? 0.0 : εF |
| 32 | + data = prepare_band_data(band_data, klabels=klabels) |
| 33 | + |
| 34 | + # For each branch, plot all bands, spins and errors |
| 35 | + p = Plots.plot(xlabel="wave vector") |
| 36 | + for ibranch = 1:data.n_branches |
| 37 | + kdistances = data.kdistances[ibranch] |
| 38 | + for spin in data.spins, iband = 1:data.n_bands |
| 39 | + yerror = nothing |
| 40 | + if hasproperty(data, :λerror) |
| 41 | + yerror = data.λerror[ibranch][spin][iband, :] ./ unit_to_au(unit) |
| 42 | + end |
| 43 | + energies = (data.λ[ibranch][spin][iband, :] .- eshift) ./ unit_to_au(unit) |
| 44 | + |
| 45 | + color = (spin == :up) ? :blue : :red |
| 46 | + Plots.plot!(p, kdistances, energies; color=color, label="", yerror=yerror, |
| 47 | + kwargs...) |
| 48 | + end |
| 49 | + end |
| 50 | + |
| 51 | + # X-range: 0 to last kdistance value |
| 52 | + Plots.xlims!(p, (0, data.kdistances[end][end])) |
| 53 | + Plots.xticks!(p, data.ticks["distance"], |
| 54 | + [replace(l, raw"$\mid$" => " | ") for l in data.ticks["label"]]) |
| 55 | + |
| 56 | + ylims = [-4, 4] |
| 57 | + !isnothing(εF) && is_metal(band_data, εF) && (ylims = [-10, 10]) |
| 58 | + ylims = round.(ylims * units.eV ./ unit_to_au(unit), sigdigits=2) |
| 59 | + if isnothing(εF) |
| 60 | + Plots.ylabel!(p, "eigenvalues ($(string(unit))") |
| 61 | + else |
| 62 | + Plots.ylabel!(p, "eigenvalues - ε_f ($(string(unit)))") |
| 63 | + Plots.ylims!(p, ylims...) |
| 64 | + end |
| 65 | + |
| 66 | + p |
| 67 | +end |
| 68 | + |
| 69 | + |
| 70 | +function plot_dos(basis, eigenvalues; εF=nothing) |
| 71 | + n_spin = basis.model.n_spin_components |
| 72 | + εs = range(minimum(minimum(eigenvalues)) - .5, |
| 73 | + maximum(maximum(eigenvalues)) + .5, length=1000) |
| 74 | + |
| 75 | + p = Plots.plot() |
| 76 | + spinlabels = spin_components(basis.model) |
| 77 | + colors = [:blue, :red] |
| 78 | + for σ in 1:n_spin |
| 79 | + D = DOS.(εs, Ref(basis), Ref(eigenvalues), spins=(σ, )) |
| 80 | + label = n_spin > 1 ? "DOS $(spinlabels[σ]) spin" : "DOS" |
| 81 | + Plots.plot!(p, εs, D, label=label, color=colors[σ]) |
| 82 | + end |
| 83 | + if !isnothing(εF) |
| 84 | + Plots.vline!(p, [εF], label="εF", color=:green, lw=1.5) |
| 85 | + end |
| 86 | + p |
| 87 | +end |
| 88 | +plot_dos(scfres; kwargs...) = plot_dos(scfres.basis, scfres.eigenvalues; εF=scfres.εF, kwargs...) |
0 commit comments