Refactor bandstructure routines and json results (#895)

Project.toml

@@ -11,6 +11,7 @@ Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DftFunctionals = "6bd331d2-b28d-4fd3-880e-1a1c7f37947f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
@@ -34,7 +35,6 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
-ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 PseudoPotentialIO = "cb339c56-07fa-4cb2-923a-142469552264"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -64,6 +64,7 @@ Brillouin = "0.5.14"
 ChainRulesCore = "1.15"
 Dates = "1"
 DftFunctionals = "0.2"
+DocStringExtensions = "0.9"
 FFTW = "1.5"
 ForwardDiff = "0.10"
 GPUArraysCore = "0.1"
@@ -89,7 +90,6 @@ PrecompileTools = "1"
 Preferences = "1"
 Primes = "0.5"
 Printf = "1"
-ProgressMeter = "1"
 PseudoPotentialIO = "0.1"
 Random = "1"
 Requires = "1"

docs/src/developer/setup.md

@@ -74,6 +74,10 @@ TestEnv.activate()  # for tests in a temporary environment.
 using TestItemRunner
 @run_package_tests filter = ti -> :core ∈ ti.tags
 ```
+Or to only run the tests of a particular file `serialisation.jl` use
+```julia
+@run_package_tests filter = ti -> occursin("serialisation.jl", ti.filename)
+```
 
 If you need to write tests, note that you can create modules with `@testsetup`. To use
 a function `my_function` of a module `MySetup` in a `@testitem`, you can import it with

docs/src/guide/tutorial.jl

@@ -100,8 +100,15 @@ x = [r[1] for r in rvecs]                   # only keep the x coordinate
 plot(x, scfres.ρ[1, :, 1, 1], label="", xlabel="x", ylabel="ρ", marker=2)
 
 # We can also perform various postprocessing steps:
-# for instance compute a band structure
-plot_bandstructure(scfres; kline_density=10)
-# or get the cartesian forces (in Hartree / Bohr)
+# We can get the Cartesian forces (in Hartree / Bohr):
 compute_forces_cart(scfres)
 # As expected, they are numerically zero in this highly symmetric configuration.
+# We could also compute a band structure,
+plot_bandstructure(scfres; kline_density=10)
+# or plot a density of states, for which we increase the kgrid a bit
+# to get smoother plots:
+bands = compute_bands(scfres, MonkhorstPack(6, 6, 6))
+plot_dos(bands; temperature=1e-3, smearing=Smearing.FermiDirac())
+# Note that directly employing the `scfres` also works, but the results
+# are much cruder:
+plot_dos(scfres; temperature=1e-3, smearing=Smearing.FermiDirac())

docs/src/publications.md

@@ -22,8 +22,9 @@ Additionally the following publications describe DFTK or one of its algorithms:
 
 - E. Cancès, M. Hassan and L. Vidal.
   [*Modified-Operator Method for the Calculation of Band Diagrams of Crystalline Materials.*](https://arxiv.org/abs/2210.00442)
-  (Submitted).
+  Mathematics of Computation (2023).
   [ArXiv:2210.00442](https://arxiv.org/abs/2210.00442).
+  ([Supplementary material and computational scripts](https://github.com/LaurentVidal95/ModifiedOp).
 
 - E. Cancès, M. F. Herbst, G. Kemlin, A. Levitt and B. Stamm.
   [*Numerical stability and efficiency of response property calculations in density functional theory*](https://arxiv.org/abs/2210.04512)

examples/cohen_bergstresser.jl

@@ -29,5 +29,6 @@ eigres = diagonalize_all_kblocks(DFTK.lobpcg_hyper, ham, 6)
 using Plots
 using Unitful
 
-p = plot_bandstructure(basis; n_bands=8, εF, kline_density=10, unit=u"eV")
+bands = compute_bands(basis; n_bands=8, εF, kline_density=10)
+p = plot_bandstructure(bands; unit=u"eV")
 ylims!(p, (-5, 6))

examples/collinear_magnetism.jl

@@ -101,11 +101,14 @@ idown = iup + length(scfres.basis.kpoints) ÷ 2
 
 # We can observe the spin-polarization by looking at the density of states (DOS)
 # around the Fermi level, where the spin-up and spin-down DOS differ.
-
 using Plots
-plot_dos(scfres)
+bands_666 = compute_bands(scfres, MonkhorstPack(6, 6, 6))  # Increase kgrid to get nicer DOS.
+plot_dos(bands_666)
+# Note that if same k-grid as SCF should be employed, a simple `plot_dos(scfres)`
+# is sufficient.
 
 # Similarly the band structure shows clear differences between both spin components.
 using Unitful
 using UnitfulAtomic
-plot_bandstructure(scfres; kline_density=6)
+bands_kpath = compute_bands(scfres; kline_density=6)
+plot_bandstructure(bands_kpath)

examples/graphene.jl

@@ -34,6 +34,6 @@ basis = PlaneWaveBasis(model; Ecut, kgrid)
 scfres = self_consistent_field(basis)
 
 ## Construct 2D path through Brillouin zone
-sgnum = 13  # Graphene space group number
-kpath = irrfbz_path(model; dim=2, sgnum)
-plot_bandstructure(scfres, kpath; kline_density=20)
+kpath = irrfbz_path(model; dim=2, space_group_number=13)  # graphene space group number
+bands = compute_bands(scfres, kpath; kline_density=20)
+plot_bandstructure(bands)

examples/input_output.jl

@@ -78,11 +78,20 @@ end
 println(read("iron_afm_energies.json", String))
 
 # Once JSON3 is loaded, additionally a convenience function for saving
-# a parsable summary of `scfres` objects using `save_scfres` is available:
+# a summary of `scfres` objects using `save_scfres` is available:
 
 using JSON3
 save_scfres("iron_afm.json", scfres)
 
+# Similarly a summary of the band data (occupations, eigenvalues, εF, etc.)
+# for post-processing can be dumped using `save_bands`:
+save_bands("iron_afm_scfres.json", scfres)
+
+# Notably this function works both for the results obtained
+# by `self_consistent_field` as well as `compute_bands`:
+bands = compute_bands(scfres, kline_density=10)
+save_bands("iron_afm_bands.json", bands)
+
 # ## Writing and reading JLD2 files
 # The full state of a DFTK self-consistent field calculation can be
 # stored on disk in form of an [JLD2.jl](https://github.com/JuliaIO/JLD2.jl) file.
@@ -99,6 +108,6 @@ save_scfres("iron_afm.jld2", scfres);
 # See [Saving SCF results on disk and SCF checkpoints](@ref) for details.
 
 # (Cleanup files generated by this notebook.)
-rm("iron_afm.vts")
-rm("iron_afm.jld2")
-rm("iron_afm_energies.json")
+rm.(["iron_afm.vts", "iron_afm.jld2",
+     "iron_afm.json", "iron_afm_energies.json", "iron_afm_scfres.json",
+     "iron_afm_bands.json"])

examples/metallic_systems.jl

@@ -41,7 +41,7 @@ smearing = DFTK.Smearing.FermiDirac() # Smearing method
 
 model = model_DFT(lattice, atoms, positions, [:gga_x_pbe, :gga_c_pbe];
                   temperature, smearing)
-kgrid = kgrid_from_minimal_spacing(lattice, kspacing)
+kgrid = kgrid_from_maximal_spacing(lattice, kspacing)
 basis = PlaneWaveBasis(model; Ecut, kgrid);
 
 # Finally we run the SCF. Two magnesium atoms in
@@ -60,4 +60,7 @@ scfres.energies
 
 # The fact that magnesium is a metal is confirmed
 # by plotting the density of states around the Fermi level.
-plot_dos(scfres)
+# To get better plots, we decrease the k spacing a bit for this step
+kgrid_dos = kgrid_from_maximal_spacing(lattice, 0.7 / u"Å")
+bands = compute_bands(scfres, kgrid_dos)
+plot_dos(bands)

examples/pseudopotentials.jl

@@ -92,7 +92,7 @@ function run_bands(psp)
     basis = PlaneWaveBasis(model; Ecut=12, kgrid=(4, 4, 4))
 
     scfres   = self_consistent_field(basis; tol=1e-4)
-    bandplot = plot_bandstructure(scfres)
+    bandplot = plot_bandstructure(compute_bands(scfres))
     (; scfres, bandplot)
 end;
 

examples/publications/2022_cazalis.jl

@@ -60,8 +60,9 @@ basis = PlaneWaveBasis(model; Ecut, kgrid)
 scfres = self_consistent_field(basis, is_converged=DFTK.ScfConvergenceEnergy(1e-10))
 
 ## Plot bands
-sgnum = 13  # Graphene space group number
-p = plot_bandstructure(scfres, irrfbz_path(model; sgnum); n_bands=5)
+kpath = irrfbz_path(model; dim=2, space_group_number=13)
+bands = compute_bands(scfres, kpath; n_bands=5)
+p = plot_bandstructure(bands)
 Plots.hline!(p, [scfres.εF], label="", color="black")
 Plots.ylims!(p, (-Inf, Inf))
 p

examples/wannier90.jl

@@ -34,7 +34,8 @@ scfres = self_consistent_field(basis; nbandsalg, tol=1e-5);
 
 # Plot bandstructure of the system
 
-plot_bandstructure(scfres; kline_density=10)
+bands = compute_bands(scfres; kline_density=10)
+plot_bandstructure(bands)
 
 # Now we use the `run_wannier90` routine to generate all files needed by
 # wannier90 and to perform the wannierization procedure.

ext/DFTKJSON3Ext.jl

@@ -1,11 +1,14 @@
 module DFTKJSON3Ext
-using JSON3
 using DFTK
-using DFTK: todict
+using DFTK: todict, band_data_to_dict
+using JSON3
+using MPI
 
 function DFTK.save_scfres_master(filename::AbstractString, scfres::NamedTuple, ::Val{:json})
     # TODO Quick and dirty solution for now.
     #      The better approach is to integrate with StructTypes.jl
+    # Also should probably be merged with save_bands eventually
+    # once the treatment of MPI distributed data is uniform.
 
     data = Dict("energies" => todict(scfres.energies), "damping" => scfres.α)
     for key in (:converged, :occupation_threshold, :εF, :eigenvalues,
@@ -18,6 +21,18 @@ function DFTK.save_scfres_master(filename::AbstractString, scfres::NamedTuple, :
     end
 end
 
-#TODO introduce `todict` functions for all sorts of datastructures (basis, ...)
+function DFTK.save_bands(filename::AbstractString, band_data::NamedTuple, ::Val{:json};
+                    save_ψ=false)
+    save_ψ && @warn "save_ψ not supported with json files"
+
+    data = band_data_to_dict(band_data)
+    if mpi_master()
+        open(filename, "w") do io
+            JSON3.pretty(io, data)
+        end
+    end
+    MPI.Barrier(MPI.COMM_WORLD)
+    nothing
+end
 
 end

src/DFTK.jl

@@ -1,9 +1,6 @@
-"""
-DFTK --- The density-functional toolkit. Provides functionality for experimenting
-with plane-wave density-functional theory algorithms.
-"""
 module DFTK
 
+using DocStringExtensions
 using LinearAlgebra
 using Markdown
 using Printf
@@ -18,6 +15,13 @@ using Random
 using ChainRulesCore
 using PrecompileTools
 
+@template METHODS =
+"""
+$(TYPEDSIGNATURES)
+
+$(DOCSTRING)
+"""
+
 export Vec3
 export Mat3
 export mpi_nprocs
@@ -62,6 +66,7 @@ include("SymOp.jl")
 
 export Smearing
 export Model
+export MonkhorstPack, ExplicitKpoints
 export PlaneWaveBasis
 export compute_fft_size
 export G_vectors, G_vectors_cart, r_vectors, r_vectors_cart
@@ -72,15 +77,18 @@ export irfft
 export ifft!
 export fft
 export fft!
-export create_supercell
-export cell_to_supercell
+export kgrid_from_maximal_spacing, kgrid_from_minimal_n_kpoints
 include("Smearing.jl")
 include("Model.jl")
 include("structure.jl")
+include("bzmesh.jl")
 include("PlaneWaveBasis.jl")
 include("fft.jl")
 include("orbitals.jl")
-include("show.jl")
+include("input_output.jl")
+
+export create_supercell
+export cell_to_supercell
 include("supercell.jl")
 
 export Energies
@@ -159,11 +167,7 @@ include("scf/potential_mixing.jl")
 
 export symmetry_operations
 export standardize_atoms
-export bzmesh_uniform
-export bzmesh_ir_wedge
-export kgrid_from_minimal_spacing, kgrid_from_minimal_n_kpoints
 include("symmetry.jl")
-include("bzmesh.jl")
 
 export DensityConstructionMethod
 export AtomicDensity
@@ -191,6 +195,7 @@ include("external/stubs.jl")  # Function stubs for conditionally defined methods
 export compute_bands
 export plot_bandstructure
 export irrfbz_path
+export save_bands
 include("postprocess/band_structure.jl")
 
 export compute_forces

src/Model.jl

@@ -106,7 +106,7 @@ function Model(lattice::AbstractMatrix{T},
                terms=[Kinetic()],
                temperature=zero(T),
                smearing=temperature > 0 ? Smearing.FermiDirac() : Smearing.None(),
-               spin_polarization=default_spin_polarization(magnetic_moments),
+               spin_polarization=determine_spin_polarization(magnetic_moments),
                symmetries=default_symmetries(lattice, atoms, positions, magnetic_moments,
                                              spin_polarization, terms),
                ) where {T <: Real}
@@ -249,8 +249,6 @@ normalize_magnetic_moment(mm::AbstractVector)::Vec3{Float64} = mm
 """Number of valence electrons."""
 n_electrons_from_atoms(atoms) = sum(n_elec_valence, atoms; init=0)
 
-default_spin_polarization(magnetic_moments) = determine_spin_polarization(magnetic_moments)
-
 """
 Default logic to determine the symmetry operations to be used in the model.
 """
@@ -307,7 +305,6 @@ spin_components(model::Model) = spin_components(model.spin_polarization)
 import Base.Broadcast.broadcastable
 Base.Broadcast.broadcastable(model::Model) = Ref(model)
 
-
 #=
 There are two types of quantities, depending on how they transform under change of coordinates.