Initial support for serialisation to JSON (#875)

* Basic, untested json output

* Test stub

* move things around

* Serialisation test

* Add json example

* docfix

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.6.10
+current_version = 0.6.11
 commit = True
 tag = False
 

Project.toml

@@ -1,7 +1,7 @@
 name = "DFTK"
 uuid = "acf6eb54-70d9-11e9-0013-234b7a5f5337"
 authors = ["Michael F. Herbst <[email protected]>", "Antoine Levitt <[email protected]>"]
-version = "0.6.10"
+version = "0.6.11"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -100,6 +100,7 @@ FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a"
 IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -110,4 +111,4 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 
 [targets]
-test = ["Test", "ASEconvert", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "IntervalArithmetic", "JLD2", "Logging", "Plots", "QuadGK", "Random", "KrylovKit", "WriteVTK", "wannier90_jll"]
+test = ["Test", "ASEconvert", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "IntervalArithmetic", "JLD2", "JSON3", "Logging", "Plots", "QuadGK", "Random", "KrylovKit", "WriteVTK", "wannier90_jll"]

docs/Project.toml

@@ -10,6 +10,7 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 LibGit2 = "76f85450-5226-5b5a-8eaa-529ad045b433"

examples/input_output.jl

@@ -61,6 +61,28 @@ save_scfres("iron_afm.vts", scfres; save_ψ=true);
 # This will save the iron calculation above into the file `iron_afm.vts`,
 # using `save_ψ=true` to also include the KS orbitals.
 
+# ## Parsable data-export using json
+# Many structures in DFTK support the (unexported) `todict` function,
+# which returns a simplified dictionary representation of the data.
+
+DFTK.todict(scfres.energies)
+
+# This in turn can be easily written to disk using a JSON library.
+# Currently we integrate most closely with `JSON3`,
+# which is thus recommended.
+
+using JSON3
+open("iron_afm_energies.json", "w") do io
+    JSON3.pretty(io, DFTK.todict(scfres.energies))
+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:
+
+using JSON3
+save_scfres("iron_afm.json", scfres)
+
 # ## 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.
@@ -79,3 +101,4 @@ save_scfres("iron_afm.jld2", scfres);
 # (Cleanup files generated by this notebook.)
 rm("iron_afm.vts")
 rm("iron_afm.jld2")
+rm("iron_afm_energies.json")

src/DFTK.jl

@@ -230,8 +230,9 @@ function __init__()
     @require DoubleFloats="497a8b3b-efae-58df-a0af-a86822472b78" begin
         !isdefined(DFTK, :GENERIC_FFT_LOADED) && include("workarounds/fft_generic.jl")
     end
-    @require Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80" include("plotting.jl")
-    @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3819"  include("external/jld2io.jl")
+    @require Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80"    include("plotting.jl")
+    @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3819"     include("external/jld2io.jl")
+    @require JSON3="0f8b85d8-7281-11e9-16c2-39a750bddbf1"    include("external/json3.jl")
     @require WriteVTK="64499a7a-5c06-52f2-abe2-ccb03c286192" include("external/vtkio.jl")
     @require wannier90_jll="c5400fa0-8d08-52c2-913f-1e3f656c1ce9" begin
         include("external/wannier90.jl")

src/Energies.jl

@@ -33,12 +33,22 @@ function Energies(term_types::Vector, energies::Vector{T}) where {T}
 end
 
 function Base.propertynames(energies::Energies, private::Bool=false)
-    ret = keys(energies)
-    append!(ret, "total")
-    private && append!(ret, "energies")
+    ret = Symbol.(keys(energies))
+    push!(ret, :total)
+    private && push!(ret, :energies)
+    ret
 end
 function Base.getproperty(energies::Energies, x::Symbol)
-    x == :total && return sum(values(energies))
+    x == :total    && return sum(values(energies))
     x == :energies && return getfield(energies, x)
     energies.energies[string(x)]
 end
+
+"""
+Convert an `Energies` struct to a dictionary representation
+"""
+function todict(energies::Energies)
+    ret = convert(Dict, energies.energies)
+    ret["total"] = energies.total
+    ret
+end

src/external/json3.jl

@@ -0,0 +1,16 @@
+function 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
+
+    data = Dict("energies" => todict(scfres.energies), "damping" => scfres.α)
+    for key in (:converged, :occupation_threshold, :εF, :eigenvalues,
+                :occupation, :n_bands_converge, :n_iter, :algorithm, :norm_Δρ)
+        data[string(key)] = getproperty(scfres, key)
+    end
+
+    open(filename, "w") do io
+        JSON3.pretty(io, data)
+    end
+end
+
+#TODO introduce `todict` functions for all sorts of datastructures (basis, ...)

src/scf/scfres.jl

@@ -45,6 +45,9 @@ file extensions are recognized and supported:
       (may lead to larger files).
     * `extra_data`: `Dict{String,Array}` with additional data on the 3D real-space
       grid to store into the VTK file.
+- **json**: A JSON file with basic information about the SCF run. Stores for example
+   the number of iterations, occupations, norm of the most recent density change,
+   eigenvalues, Fermi level etc.
 
 !!! warning "No compatibility guarantees"
     No guarantees are made with respect to this function at this point.
@@ -56,7 +59,7 @@ function save_scfres(filename::AbstractString, scfres::NamedTuple; kwargs...)
     ext = Symbol(ext[2:end])
 
     # Whitelist valid extensions
-    !(ext in (:jld2, :vts)) && error("Extension '$ext' not supported by DFTK.")
+    !(ext in (:jld2, :vts, :json)) && error("Extension '$ext' not supported by DFTK.")
 
     # Gather scfres data on master MPI process
     scfres = gather_kpts_scfres(scfres)

test/serialisation.jl

@@ -1,5 +1,6 @@
 using DFTK
 using JLD2
+using JSON3
 using WriteVTK
 using MPI
 import DFTK: ScfDefaultCallback, ScfSaveCheckpoints
@@ -33,7 +34,7 @@ function test_scfres_agreement(tested, ref)
     @test tested.eigenvalues    == ref.eigenvalues
     @test tested.occupation     == ref.occupation
     @test tested.ψ              == ref.ψ
-    @test tested.ρ              ≈  ref.ρ     rtol=1e-14
+    @test tested.ρ              ≈  ref.ρ rtol=1e-14
 end
 
 
@@ -88,6 +89,26 @@ function test_serialisation(label;
             @test isfile(dumpfile)
         end
     end
+
+    @testset "JSON ($label)" begin
+        mktempdir() do tmpdir
+            dumpfile = joinpath(tmpdir, "scfres.json")
+            dumpfile = MPI.bcast(dumpfile, 0, MPI.COMM_WORLD)  # master -> everyone
+
+            save_scfres(dumpfile, scfres)
+            @test isfile(dumpfile)
+
+            data = open(JSON3.read, dumpfile)  # Get data back as dict
+            for key in (:converged, :occupation_threshold, :εF, :eigenvalues,
+                        :occupation, :n_bands_converge, :n_iter, :algorithm, :norm_Δρ)
+                @test data[key] == getproperty(scfres, key)
+            end
+            for key in keys(scfres.energies)
+                @test data["energies"][key] == scfres.energies[key]
+            end
+            @test data["energies"]["total"] == scfres.energies.total
+        end
+    end
 end
 
 @testset "Serialisation" begin