Adding DFTKMakieExt.jl & BandsUnfold.jl

Project.toml

@@ -6,6 +6,7 @@ authors = ["Michael F. Herbst <info@michael-herbst.com>", "Antoine Levitt <antoi
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+AtomBase = "ccf6c0ac-cd45-4eaf-bc3d-e8176fb87b3b"
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1"
 Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0"
@@ -52,6 +53,7 @@ GeometryOptimization = "673bf261-a53d-43b9-876f-d3c1fc8329c2"
 IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
@@ -65,6 +67,7 @@ DFTKGeometryOptimizationExt = "GeometryOptimization"
 DFTKIntervalArithmeticExt = "IntervalArithmetic"
 DFTKJLD2Ext = "JLD2"
 DFTKJSON3Ext = "JSON3"
+DFTKMakieExt = "Makie"
 DFTKPlotsExt = "Plots"
 DFTKWannier90Ext = "wannier90_jll"
 DFTKWannierExt = "Wannier"
@@ -76,6 +79,7 @@ ASEconvert = "0.1, 0.2"
 AbstractFFTs = "1"
 Aqua = "0.8.5"
 Artifacts = "1"
+AtomBase = "0.1.0"
 AtomsBase = "0.5.2"
 AtomsBuilder = "0.2"
 AtomsCalculators = "0.2.3"
@@ -108,6 +112,7 @@ LinearAlgebra = "1"
 LinearMaps = "3"
 Logging = "1"
 MPI = "0.20.22"
+Makie = "0.24.8"
 Markdown = "1"
 Optim = "1"
 PeriodicTable = "1"
@@ -156,6 +161,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
@@ -167,4 +173,4 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 
 [targets]
-test = ["Test", "TestItemRunner", "AMDGPU", "ASEconvert", "AtomsBuilder", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "GeometryOptimization", "IntervalArithmetic", "JLD2", "JSON3", "Logging", "Plots", "PythonCall", "QuadGK", "Random", "KrylovKit", "Wannier", "WriteVTK", "wannier90_jll"]
+test = ["Test", "TestItemRunner", "AMDGPU", "ASEconvert", "AtomsBuilder", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "GeometryOptimization", "IntervalArithmetic", "JLD2", "JSON3", "Logging", "Plots", "Makie", "PythonCall", "QuadGK", "Random", "KrylovKit", "Wannier", "WriteVTK", "wannier90_jll"]

ext/DFTKMakieExt.jl

@@ -0,0 +1,214 @@
+module DFTKMakieExt
+
+using DFTK
+using DFTK: is_metal, data_for_plotting, spin_components, default_band_εrange
+using Makie
+using Unitful
+using UnitfulAtomic
+using LinearAlgebra
+
+# Import DFTK stub functions so we can extend them natively
+import DFTK: plot_bandstructure, plot_dos, plot_bandstructure!, plot_dos!, plot_unfolded_bands!, plot_unfolded_bands, plot_folded_bands, plot_folded_bands!, plot_dos_rotated!
+
+@info "DFTKMakieExt successfully loaded!"
+
+function plot_bandstructure!(pos, band_data; title_text="Band Structure", y_limits=nothing)
+    data = DFTK.data_for_plotting(band_data)
+    eshift = something(band_data.εF, 0.0)
+    to_unit = ustrip(auconvert(u"eV", 1.0)) 
+
+    ax = Makie.Axis(pos, title=title_text, xlabel="Wave Vector", 
+                    ylabel=isnothing(band_data.εF) ? "Energy (eV)" : "Energy - εF (eV)")
+
+    for σ = 1:data.n_spin, iband = 1:data.n_bands, branch in data.kbranches
+        energies = (data.eigenvalues[:, iband, σ][branch] .- eshift) .* to_unit
+        Makie.lines!(ax, data.kdistances[branch], energies, 
+                     color = σ == 1 ? :blue : :red, linewidth=2)
+    end
+
+    for branch in data.kbranches[1:end-1]
+        Makie.vlines!(ax, [data.kdistances[last(branch)]], color=:black, linewidth=1)
+    end
+
+    ax.xticks = (data.ticks.distances, data.ticks.labels)
+    Makie.xlims!(ax, 0, data.kdistances[end])
+
+    if !isnothing(band_data.εF)
+        Makie.hlines!(ax, [0.0], color=:green, linewidth=2, linestyle=:dash)
+    end
+
+    # --- NEW: Apply custom Y-axis limits if provided ---
+    if !isnothing(y_limits)
+        Makie.ylims!(ax, y_limits[1], y_limits[2])
+    end
+
+    return ax
+end
+
+function plot_bandstructure(band_data; kwargs...)
+    fig = Figure(size=(800, 600), fontsize=18)
+    plot_bandstructure!(fig[1, 1], band_data; kwargs...)
+    return fig
+end
+
+function plot_dos!(pos, scfres; n_points=500, title_text="Density of States")
+    εF = scfres.εF
+    eshift = something(εF, 0.0)
+    to_unit = ustrip(auconvert(u"eV", 1.0))
+
+    εrange = DFTK.default_band_εrange(scfres.eigenvalues; εF=εF)
+    εs = range(εrange[1], εrange[2], length=n_points)
+
+    Dεs = DFTK.compute_dos.(εs, Ref(scfres.basis), Ref(scfres.eigenvalues); 
+                            smearing=scfres.basis.model.smearing, 
+                            temperature=scfres.basis.model.temperature)
+
+    energies = (εs .- eshift) .* to_unit
+    n_spin = scfres.basis.model.n_spin_components
+
+    ax = Axis(pos, title=title_text, 
+              xlabel=isnothing(εF) ? "Energy (eV)" : "Energy - εF (eV)", 
+              ylabel="Density of States")
+
+    for σ = 1:n_spin
+        D = [Dσ[σ] for Dσ in Dεs]
+        lines!(ax, energies, D, label="Spin $σ", color = σ == 1 ? :blue : :red, linewidth=2)
+    end
+
+    if !isnothing(εF)
+        vlines!(ax, [0.0], color=:green, linewidth=2, linestyle=:dash)
+    end
+    if n_spin > 1
+        axislegend(ax)
+    end
+    return ax
+end
+
+function plot_dos(scfres; kwargs...)
+    fig = Figure(size=(800, 600), fontsize=18)
+    plot_dos!(fig[1, 1], scfres; kwargs...)
+    return fig
+end
+
+function plot_dos_rotated!(pos, scfres; n_points=500, title_text="Total DOS")
+    εF = scfres.εF
+    eshift = something(εF, 0.0)
+    to_unit = ustrip(auconvert(u"eV", 1.0))
+
+    εrange = DFTK.default_band_εrange(scfres.eigenvalues; εF=εF)
+    εs = range(εrange[1], εrange[2], length=n_points)
+
+    Dεs = DFTK.compute_dos.(εs, Ref(scfres.basis), Ref(scfres.eigenvalues); 
+                            smearing=scfres.basis.model.smearing, 
+                            temperature=scfres.basis.model.temperature)
+
+    energies = (εs .- eshift) .* to_unit
+    n_spin = scfres.basis.model.n_spin_components
+
+    ax = Makie.Axis(pos, title=title_text, xlabel="DOS (states/eV)")
+
+    for σ = 1:n_spin
+        D = [Dσ[σ] for Dσ in Dεs]
+        # SWAP AXES: Plot DOS on X, Energy on Y
+        Makie.lines!(ax, D, energies, color = :blue, linewidth=2)
+    end
+
+    if !isnothing(εF)
+        Makie.hlines!(ax, [0.0], color=(:black, 0.4), linewidth=1, linestyle=:dash)
+    end
+
+    Makie.xlims!(ax, low=0.0) # Force DOS axis to start at 0
+    return ax
+end
+
+function plot_unfolded_bands!(pos, unfold_data; title_text="Unfolded Bands", colormap=:turbo)
+    eshift = something(unfold_data.εF, 0.0)
+    to_unit = ustrip(auconvert(u"eV", 1.0))
+
+    ax = Makie.Axis(pos, title=title_text, xlabel="Wavevector", 
+                    ylabel=isnothing(unfold_data.εF) ? "Energy (eV)" : "Energy (eV - Ef)")
+
+    sc = nothing # Initialize variable to hold the scatter object
+
+    for spin in 1:length(unfold_data.k_indices)
+        k_idxs = unfold_data.k_indices[spin]
+        x_vals = unfold_data.kdistances[k_idxs]
+        y_vals = (unfold_data.eigenvalues[spin] .- eshift) .* to_unit
+        weights = unfold_data.spectral_weights[spin]
+
+        # --- FIX: Sort data so high-weight points are drawn last (on top) ---
+        perm = sortperm(weights)
+        x_sorted = x_vals[perm]
+        y_sorted = y_vals[perm]
+        weights_sorted = weights[perm]
+
+        # Plot with sorted arrays
+        sc = Makie.scatter!(ax, x_sorted, y_sorted;
+                            color=weights_sorted, colormap=colormap, 
+                            colorrange=(0.0, 1.0), # Lock scale from 0 to 1
+                            markersize=4, strokewidth=0)
+    end
+
+    # Formatting X-axis with high-symmetry points
+    ax.xticks = (unfold_data.ticks.distances, unfold_data.ticks.labels)
+    for dist in unfold_data.ticks.distances[1:end-1]
+        Makie.vlines!(ax, [dist], color=(:black, 0.2), linewidth=1)
+    end
+
+    Makie.xlims!(ax, 0, unfold_data.kdistances[end])
+    if !isnothing(unfold_data.εF)
+        Makie.hlines!(ax, [0.0], color=(:black, 0.4), linewidth=1, linestyle=:dash)
+    end
+
+    # Return both the axis AND the scatter object so we can build a colorbar
+    return ax, sc 
+end
+
+function plot_unfolded_bands(unfold_data; kwargs...)
+    fig = Figure(size=(800, 600), fontsize=18)
+    plot_unfolded_bands!(fig[1, 1], unfold_data; kwargs...)
+    return fig
+end
+
+function plot_folded_bands!(pos, band_data_raw, unfold_data; title_text="Raw Supercell Bands (Folded)")
+    eshift = something(unfold_data.εF, 0.0)
+    to_unit = ustrip(auconvert(u"eV", 1.0))
+
+    ax = Makie.Axis(pos, title=title_text, xlabel="Wave Vector", 
+                    ylabel=isnothing(unfold_data.εF) ? "Energy (eV)" : "Energy (eV - Ef)")
+
+    n_kpoints = length(unfold_data.kdistances)
+    # Get the number of bands from the first k-point
+    n_bands = length(band_data_raw.eigenvalues[1]) 
+
+    # Trace a continuous line for each individual band
+    for ib in 1:n_bands
+        # Extract the energy of band `ib` across all k-points
+        band_energies = [(band_data_raw.eigenvalues[ik][ib] - eshift) * to_unit for ik in 1:n_kpoints]
+
+        # Plot as a solid line
+        Makie.lines!(ax, unfold_data.kdistances, band_energies, 
+                     color=(:black, 0.5), linewidth=1.5)
+    end
+
+    # Format X-axis with high-symmetry points
+    ax.xticks = (unfold_data.ticks.distances, unfold_data.ticks.labels)
+    for dist in unfold_data.ticks.distances[1:end-1]
+        Makie.vlines!(ax, [dist], color=(:black, 0.3), linewidth=1, linestyle=:dash)
+    end
+
+    Makie.xlims!(ax, 0, unfold_data.kdistances[end])
+    if !isnothing(unfold_data.εF)
+        Makie.hlines!(ax, [0.0], color=:green, linewidth=2, linestyle=:dash)
+    end
+
+    return ax
+end
+
+function plot_folded_bands(band_data_raw; kwargs...)
+    fig = Figure(size =(800, 600), fontsize=18)
+    plot_folded_bands!(fig[1,1], band_data_raw; kwargs...)
+    return fig
+end
+
+end # module

src/DFTK.jl

@@ -252,6 +252,22 @@ include("workarounds/forwarddiff_rules.jl")
 include("gpu/linalg.jl")
 include("gpu/gpu_arrays.jl")
 
+#\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+# Added for DFTKMakieExt.jl plot & subplot
+export plot_bandstructure!  
+export plot_dos!
+export plot_unfolded_bands
+export plot_unfolded_bands!
+export plot_folded_bands
+export plot_folded_bands!
+export plot_dos_rotated!
+
+# Added for postprocess of BandsUnfold.jl
+export unfold_bands
+export compute_folded_bands
+include("postprocess/BandsUnfold.jl")
+#\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
+
 # Precompilation block with a basic workflow
 
 function precompilation_workflow(lattice, atoms, positions, magnetic_moments;