diff --git a/Project.toml b/Project.toml index a924edc35a..af222b1da6 100644 --- a/Project.toml +++ b/Project.toml @@ -20,12 +20,10 @@ Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e" KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" Libxc = "66e17ffc-8502-11e9-23b5-c9248d0eb96d" -LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" -Optim = "429524aa-4258-5aef-a3af-852621145aeb" PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884" PkgVersion = "eebad327-c553-4316-9ea0-9fa01ccd7688" Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45" @@ -53,7 +51,10 @@ GeometryOptimization = "673bf261-a53d-43b9-876f-d3c1fc8329c2" IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" +Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" +Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9" WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9" @@ -66,6 +67,7 @@ DFTKGeometryOptimizationExt = "GeometryOptimization" DFTKIntervalArithmeticExt = "IntervalArithmetic" DFTKJLD2Ext = "JLD2" DFTKJSON3Ext = "JSON3" +DFTKManifoldsManoptRATExt = ["Manifolds", "Manopt", "RecursiveArrayTools"] DFTKPlotsExt = "Plots" DFTKWannier90Ext = "wannier90_jll" DFTKWannierExt = "Wannier" @@ -104,13 +106,13 @@ JLD2 = "0.4, 0.5, 0.6" JSON3 = "1" KrylovKit = "0.8.3, 0.9, 0.10" Libxc = "0.3.17" -LineSearches = "7" LinearAlgebra = "1" LinearMaps = "3" Logging = "1" MPI = "0.20.22" +Manifolds = "0.10.17" +Manopt = "0.5.23" Markdown = "1" -Optim = "1" PeriodicTable = "1" PkgVersion = "0.3" Plots = "1" @@ -123,6 +125,7 @@ PseudoPotentialIO = "0.1" PythonCall = "0.9" QuadGK = "2" Random = "1" +RecursiveArrayTools = "3.33.0" Roots = "2" SparseArrays = "1" SpecialFunctions = "2" @@ -158,10 +161,13 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" +Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" +Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9" @@ -169,4 +175,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", "Manifolds", "Manopt", "Plots", "PythonCall", "QuadGK", "Random", "RecursiveArrayTools", "KrylovKit", "Wannier", "WriteVTK", "wannier90_jll"] diff --git a/docs/Project.toml b/docs/Project.toml index 060f4a6a43..1af446ab04 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -11,6 +11,8 @@ Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" DFTK = "acf6eb54-70d9-11e9-0013-234b7a5f5337" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244" +DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GeometryOptimization = "673bf261-a53d-43b9-876f-d3c1fc8329c2" IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" @@ -21,12 +23,15 @@ LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3" LibGit2 = "76f85450-5226-5b5a-8eaa-529ad045b433" LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e" Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" +Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" +Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" PseudoPotentialData = "5751a51d-ac76-4487-a056-413ecf6fbe19" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" @@ -35,9 +40,15 @@ Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9" WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9" +[sources] +DFTK = {path = ".."} + [compat] AtomsBuilder = "0.2" AtomsCalculators = "0.2.3" Documenter = "~1.8" GeometryOptimization = "0.1" PseudoPotentialData = "0.2.0" + +[extras] +CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" diff --git a/docs/make.jl b/docs/make.jl index 3e708ac06f..858c3994db 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -93,6 +93,7 @@ PAGES = [ "developer/symmetries.md", "developer/gpu_computations.md", ], + "Extensions" => "extensions/index.md", "api.md", "publications.md", ] @@ -121,7 +122,6 @@ DFTKREPO = DFTKGH * ".git" # Setup julia dependencies for docs generation if not yet done Pkg.activate(@__DIR__) if !isfile(joinpath(@__DIR__, "Manifest.toml")) - Pkg.develop(Pkg.PackageSpec(; path=ROOTPATH)) Pkg.instantiate() end @@ -133,8 +133,10 @@ ENV["PLOTS_TEST"] = "true" # Import packages for docs generation using DFTK using Documenter +using DocumenterCitations, DocumenterInterLinks using Literate - +# For documentation in extensions: +using Manifolds, Manopt, RecursiveArrayTools # # Generate the docs # @@ -187,7 +189,7 @@ for file in literate_files badges = [ "[![](https://mybinder.org/badge_logo.svg)]" * "(@__BINDER_ROOT_URL__/$subfolder/@__NAME__.ipynb)", - "[![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)]" * + "[![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)]" * "(@__NBVIEWER_ROOT_URL__/$subfolder/@__NAME__.ipynb)", ] else @@ -216,8 +218,18 @@ mathengine = Documenter.MathJax3(Dict( ), )) +#Setup Citations and InterLinks +links = InterLinks( + "Manopt" => ("https://manoptjl.org/stable/"), + "ManifoldsBase" => ("https://juliamanifolds.github.io/ManifoldsBase.jl/stable/"), + "Manifolds" => ("https://juliamanifolds.github.io/Manifolds.jl/stable/"), +) +bib = CitationBibliography(joinpath(@__DIR__, "src", "references.bib"); style=:alpha) makedocs(; - modules=[DFTK], + modules=[ + DFTK, + Base.get_extension(DFTK, :DFTKManifoldsManoptRATExt), + ], format=Documenter.HTML(; # Use clean URLs, unless built as a "local" build prettyurls=CONTINUOUS_INTEGRATION, @@ -231,6 +243,7 @@ makedocs(; authors = "Michael F. Herbst, Antoine Levitt and contributors.", pages=transform_to_md(PAGES), checkdocs=:exports, + plugins=[bib, links], warnonly=DEBUG, remote_args..., ) diff --git a/docs/src/extensions/index.md b/docs/src/extensions/index.md new file mode 100644 index 0000000000..313485af55 --- /dev/null +++ b/docs/src/extensions/index.md @@ -0,0 +1,6 @@ +# Extensions + +## Manopt.jl + +```@docs +``` diff --git a/docs/src/references.bib b/docs/src/references.bib new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/anyons.jl b/examples/anyons.jl index 8c4fd33aea..8c702b6af0 100644 --- a/examples/anyons.jl +++ b/examples/anyons.jl @@ -5,6 +5,7 @@ using DFTK using StaticArrays using Plots +using Manifolds, Manopt, RecursiveArrayTools ## Unit cell. Having one of the lattice vectors as zero means a 2D system a = 14 diff --git a/examples/compare_solvers.jl b/examples/compare_solvers.jl index 775e1f00a3..7b42d2381c 100644 --- a/examples/compare_solvers.jl +++ b/examples/compare_solvers.jl @@ -8,6 +8,7 @@ using AtomsBuilder using DFTK using LinearAlgebra using PseudoPotentialData +using Manopt, Manifolds, RecursiveArrayTools pseudopotentials = PseudoFamily("dojo.nc.sr.pbesol.v0_4_1.standard.upf") model = model_DFT(bulk(:Si); functionals=PBEsol(), pseudopotentials) diff --git a/examples/gross_pitaevskii.jl b/examples/gross_pitaevskii.jl index 2372d13cba..358c8ffaeb 100644 --- a/examples/gross_pitaevskii.jl +++ b/examples/gross_pitaevskii.jl @@ -24,7 +24,7 @@ lattice = a .* [[1 0 0.]; [0 0 0]; [0 0 0]]; # which is special cased in DFTK to support 1D models. # # For the potential term `V` we pick a harmonic -# potential. We use the function `ExternalFromReal` which uses +# potential. We use the function `ExternalFromReal` which uses # Cartesian coordinates ( see [Lattices and lattice vectors](@ref conventions-lattice)). pot(x) = (x - a/2)^2; @@ -39,6 +39,7 @@ C = 1.0 # … and with this build the model using DFTK using LinearAlgebra +using Manopt, Manifolds, RecursiveArrayTools n_electrons = 1 # Increase this for fun terms = [Kinetic(), diff --git a/examples/gross_pitaevskii_2D.jl b/examples/gross_pitaevskii_2D.jl index 4e0a6f83a5..605bf8ce47 100644 --- a/examples/gross_pitaevskii_2D.jl +++ b/examples/gross_pitaevskii_2D.jl @@ -9,6 +9,7 @@ using DFTK using StaticArrays using Plots +using Manifolds, Manopt, RecursiveArrayTools ## Unit cell. Having one of the lattice vectors as zero means a 2D system a = 15 diff --git a/ext/DFTKManifoldsManoptRATExt.jl b/ext/DFTKManifoldsManoptRATExt.jl new file mode 100644 index 0000000000..95378ebb2f --- /dev/null +++ b/ext/DFTKManifoldsManoptRATExt.jl @@ -0,0 +1,535 @@ +module DFTKManifoldsManoptRATExt + using DFTK + using Manopt + using Manifolds + using RecursiveArrayTools + using LinearAlgebra + +""" + ManoptPreconditionersWrapper!!{T,S} + +A wrapper for a DFTK preconditioner to be used within Manopt, +which implements the `Manopt.Preconditioner` interface. + +This wrapper provides both an allocating `(M, p, X) -> Y` method to precondition +`X`, as well as the non-allocating `(M, Y, p, X) -> Y` that works in-place of `Y`. + +# Fields + +* `Nk::Int`: Number of k-points +* `Pks::Vector{T}`: (DFTK) Preconditioners, one for each k-point +* `kweights::Vector{S}`: weights for each k-point + +# Constructor + +mpw = ManoptPreconditionersWrapper!!(Nk, Pks, kweights) +""" +struct ManoptPreconditionersWrapper!!{T,S} + Nk::Int + Pks::Vector{T} + kweights::Vector{S} +end +function (mp::ManoptPreconditionersWrapper!!)( + M::ProductManifold,Y,p,X +) + # Update preconditioner + for ik = 1:mp.Nk + DFTK.precondprep!(mp.Pks[ik], p[M, ik]) + end + # Precondition the gradient in-place + for ik = 1:mp.Nk + DFTK.ldiv!(Y[M, ik], mp.Pks[ik], X[M, ik]) + DFTK.ldiv!(mp.kweights[ik], Y[M, ik]) # maybe remove local_Y + end + return Y +end +function (mp::ManoptPreconditionersWrapper!!)(M::ProductManifold, p, X) + Y = copy(M, p, X) + mp(M, Y, p, X) + return Y +end +""" + InsulatorEnergy{T,S,P,X,R,E,H} + +A functor to represent both the energy cost function from direct minimization and its gradient. +To call the cost function, use `cgf(M,p)`, the gradient can be evaluated in-place of a tangent vector `X` +with `cgf(M, X, p)`. + +For a derivation, see Sec. 7.2 and 7.5.2 of Cancès, Friesecke: [Density Functional Theory](https://doi.org/10.1007/978-3-031-22340-2), Springer. + +This functor is designed to be used within `Manopt.jl` and caches the last computed cost, +gradient and the interim values to spare calls to [`compute_density`](@ref) and [`energy_hamiltonian`](@ref). + +# Fields + +* `basis::PlaneWaveBasis{T}`: The plane wave basis used for the DFT calculation. +* `occupation::Vector{S}`: The occupation numbers for each k-point. +* `Nk::Int`: The number of k-points. +* `filled_occ::Int`: The number of filled occupation states. +* `ψ::P`: The last iterate, stored as a DFTK wave representation. +* `X::X`: The last gradient, stored as a `Manopt.jl` tangent vector, i.e. using an `ArrayPartition` +* `ρ::R`: The last density, computed from the wave functions. +* `energies::E`: The last vector of energies, computed from the wave functions. +* `ham::H`: The last Hamiltonian, computed from the wave functions. +""" +mutable struct InsulatorEnergy{T,S,P,X,R,E,H} + basis::PlaneWaveBasis{T} + occupation::Vector{S} + Nk::Int + filled_occ::Int + ψ::P # last iterate, While we usually have an iterate `ArrayPartition` `p` iterate in Manopt, we store it as a vector, i.e. in the ψ format for DFTK + X::X # last gradient + ρ::R # the last density + energies::E # the last vector of energies + ham::H # the last Hamiltonian + count::Int # Count hamiltonian calls. +end +# Function shared by both cost and gradient of cost: +function _compute_density_energy_hamiltonian!(cgf::InsulatorEnergy, M::ProductManifold, p) + # Can we improve this by copying elementwise? + # copyto!(cgf.ψ, (copy(x) for x in p.x)) # deepcopyto! + # Maybe like + for ik in eachindex(cgf.ψ) + copyto!(cgf.ψ[ik], p[M,ik]) + end + copyto!(cgf.ρ, compute_density(cgf.basis, cgf.ψ, cgf.occupation)) + # Below not inplace, but probably not that important. + cgf.energies, cgf.ham = energy_hamiltonian(cgf.basis, cgf.ψ, cgf.occupation; cgf.ρ) + cgf.count += 1 + return cgf +end +# The cost function: +function (cgf::InsulatorEnergy)(M::ProductManifold,p) + # Memoization check: Are we still at the same point? + if all(cgf.ψ[i] == p[M, i] for i in eachindex(cgf.ψ)) + _compute_density_energy_hamiltonian!(cgf, M, p) + end + return cgf.energies.total +end +# The gradient of cost function: +function (cgf::InsulatorEnergy)(M::ProductManifold, X, p) + # Memoization check: Is this X allready been computed? + if all(cgf.X[M, i] == X[M, i] for i in eachindex(cgf.ψ)) + # Are we still at the same point? + if all(cgf.ψ[i] == p[M, i] for i in eachindex(cgf.ψ)) + return X + end + end + # Memoization check: Are we still at the same point? + if all(cgf.ψ[i] == p[M, i] for i in eachindex(cgf.ψ)) + _compute_density_energy_hamiltonian!(cgf, M, p) + end + # Compute the Euclidean gradient in-place + for ik = 1:cgf.Nk + DFTK.mul!(X[M, ik], cgf.ham.blocks[ik], p[M, ik]) # mul! overload in DFTK + # Using get_component(), as "X[M, ik] .*=" is not yet supported in ManifoldsBase.jl + Manifolds.get_component(M, X, ik) .*= 2 * cgf.filled_occ * cgf.basis.kweights[ik] + end + riemannian_gradient!(M, X, p, X) # Convert to Riemannian gradient + copyto!(cgf.X, X) # Memoization + return X +end +# Access nested fields of the cost function +function Manopt.get_parameter(objective::Manopt.AbstractManifoldCostObjective, s::Symbol) + return Manopt.get_parameter(Manopt.get_cost_function(objective), s) +end +function Manopt.get_parameter(energy_costgrad::InsulatorEnergy, s::Symbol) + return Manopt.get_parameter(energy_costgrad, Val(s)) +end +Manopt.get_parameter(energy_costgrad::InsulatorEnergy, ::Val{:ρ}) = energy_costgrad.ρ +function Manopt.get_parameter(energy_costgrad::InsulatorEnergy, ::Val{:Hamiltonian}) + return energy_costgrad.ham +end +function Manopt.get_parameter(energy_costgrad::InsulatorEnergy, ::Val{:Energies}) + return energy_costgrad.energies +end +function Manopt.get_parameter( + energy_costgrad::InsulatorEnergy, ::Val{:HamiltonianEvaluations} + ) + return energy_costgrad.count +end + +# +# +# Debugs +# Δρ +mutable struct DebugDensityChange{F} <: DebugAction + io::IO + last_ρ::F + prefix::String +end +function DFTK.DebugDensityChange(ρ::T; prefix = "Δρ:", io::IO=stdout) where {T} + return DebugDensityChange{T}(io, ρ, prefix) +end +function (d::DebugDensityChange)( + problem::AbstractManoptProblem, ::AbstractManoptSolverState, k::Int +) + current_ρ = Manopt.get_parameter(Manopt.get_objective(problem), :ρ) + ch = norm(current_ρ - d.last_ρ) + d.last_ρ .= current_ρ + (k >= 0) && print(d.io, "$(d.prefix)$(ch)") +end +# ‖ρ‖ +mutable struct DebugDensityNorm <: DebugAction + io::IO + prefix::String +end +DFTK.DebugDensityNorm(; io::IO=stdout, prefix="‖ρ‖: ") = DebugDensityNorm(io, prefix) +function (::DebugDensityNorm)( + problem::AbstractManoptProblem, ::AbstractManoptSolverState, k::Int +) + current_ρ = Manopt.get_parameter(Manopt.get_objective(problem), :ρ) + return (k >= 0) && print(d.io, "$(d.prefix)$(norm(current_ρ))") +end +Base.show(io::IO, ::DebugDensityNorm) = print(io, "DebugDensityNorm()") +# +# +# Records +# ρ +mutable struct RecordDensity{F} <: RecordAction + recorded_values::Array{F,1} +end +DFTK.RecordDensity(ρ::F) where {F} = RecordDensity{F}(Array{F,1}()) +function (r::RecordDensity)( + problem::AbstractManoptProblem, ::AbstractManoptSolverState, k::Int +) + current_ρ = Manopt.get_parameter(Manopt.get_objective(problem), :ρ) + return Manopt.record_or_reset!(r, copy(current_ρ), k) +end +Base.show(io::IO, ::RecordDensity) = print(io, "RecordDensity()") +# Δρ +mutable struct RecordDensityChange{F,T} <: RecordAction + recorded_values::Array{F,1} + last_ρ::T +end +function DFTK.RecordDensityChange(ρ::T) where {T} + return RecordDensityChange{typeof(norm(ρ)),typeof(ρ)}(Array{typeof(norm(ρ)),1}(), ρ) +end +function (r::RecordDensityChange)( + problem::AbstractManoptProblem, s::AbstractManoptSolverState, k::Int +) + current_ρ = Manopt.get_parameter(Manopt.get_objective(problem), :ρ) + last_change = norm(current_ρ - r.last_ρ) + r.last_ρ .= current_ρ + return Manopt.record_or_reset!(r, last_change, k) +end +Base.show(io::IO, ::RecordDensityChange) = print(io, "RecordDensityChange()") +# +# norm ρ +mutable struct RecordDensityNorm{F} <: RecordAction + recorded_values::Array{F,1} +end +DFTK.RecordDensityNorm(ρ::F) where {F} = RecordDensityNorm(typeof(norm(ρ))) +DFTK.RecordDensityNorm(T::Type=Float64) = RecordDensityNorm{T}(Array{T,1}()) +function (r::RecordDensityNorm)( + problem::AbstractManoptProblem, ::AbstractManoptSolverState, k::Int +) + current_ρ = Manopt.get_parameter(Manopt.get_objective(problem), :ρ) + return Manopt.record_or_reset!(r, norm(current_ρ), k) +end +Base.show(io::IO, ::RecordDensityNorm) = print(io, "RecordDensityNorm()") + +# +# +# Stopping Criteria + +""" + StopWhenDensityChangeLess{T} + +A `Manopt.jl` stopping criterion that indicates to stop then the change in the density `ρ` +is less than a given tolerance `tol`. + +The stopping criterion assuemes that the density is either stored the objective, like the +`InsulatorEnergy` or is set as a parameter vie `get_parameter(objective, :ρ)` + +# Fields +* `tolerance::F`: The tolerance for the change in density. +* `at_iteration::Int`: The iteration at which the stopping criterion was met. +* `last_ρ::T`: The last value of the density. +* `last_change::F`: The last change in the density in order to generate the reason for stopping. + +# Constructor + +``` +StopWhenDensityChangeLess(tol::F, ρ::T) where {T,F<:Real} +``` + +Create the stopping criterion with a given tolerance `tol`. The provided density `ρ` +is only required to intialize the internal state. +""" +DFTK.StopWhenDensityChangeLess(tol::F, ρ::T) where {T,F<:Real} + +mutable struct StopWhenDensityChangeLess{T,F<:Real} <: Manopt.StoppingCriterion + tolerance::F + at_iteration::Int + last_ρ::T + last_change::F +end +function DFTK.StopWhenDensityChangeLess(tol::F, ρ::T) where {T,F<:Real} + StopWhenDensityChangeLess{T,F}(tol, -1, ρ, 2 * tol) +end +function (c::StopWhenDensityChangeLess)( + problem::P, state::S, k::Int + ) where {P<:Manopt.AbstractManoptProblem,S<:Manopt.AbstractManoptSolverState} + current_ρ = Manopt.get_parameter(Manopt.get_objective(problem), :ρ) + if k == 0 # reset on init + c.at_iteration = -1 + c.last_ρ .= 0 + c.last_change = 2 * c.tolerance + return false + end + c.last_change = norm(current_ρ - c.last_ρ) + c.last_ρ .= current_ρ + if c.last_change < c.tolerance + c.at_iteration = k + return true + end + return false +end +function Manopt.get_reason(c::StopWhenDensityChangeLess) + if c.at_iteration >= 0 + return "At iteration $(c.at_iteration) the algorithm performed a step with a Density change ($(c.last_change)) less than $(c.tolerance)." + end + return "" +end +function Manopt.status_summary(c::StopWhenDensityChangeLess) + has_stopped = (c.at_iteration >= 0) + s = has_stopped ? "reached" : "not reached" + return "|Δρ| = $(c.last_change) < $(c.tolerance):\t$s" +end +function Base.show(io::IO, c::StopWhenDensityChangeLess) + return print(io, "StopWhenDensityChangeLess with threshold $(c.tolerance).\n $(Manopt.status_summary(c))") +end + +# +# +# The direct minimization interface – the simple case +""" + direct_minimization(basis::DFTK.PlaneWaveBasis{T}; kwargs...) + +Compute a minimizer of the density energy function energy functional using the direct minimization method. + +# Argument + +* `basis::DFTK.PlaneWaveBasis{T}`: The plane wave basis to use for the DFT calculation. + +# Keyword Arguments + +* `maxiter=1000`: The maximum number of iterations for the optimization. If you set the `stopping_criterion=` directly, this keyword has no effect. +* `ψ=nothing`: The initial guess for the wave functions. If not provided, random orbitals will be generated. +* `ρ=guess_density(basis), # would be consistent with other scf solvers +* `tol=1e-6`: stop when the change in the density is less than this tolerance. If you set the `stopping_criterion=` directly, this keyword has no effect. +* `manifold=`[`Stiefel`](@extref `Manifolds.Stiefel`): the manifold the optimisation is running on. + The current default cost function allows to also use ∞`Grassmann`](@extref `Manifolds.Grassmann`), + both in their complex form. +* `record=[Manopt.RecordCost(), DFTK.RecordDensityChange(ρ), Manopt.RecordTime(:total)]`: + specify what to record during the Iterations. If present, these are included in the returned named tuple + +This uses several defaults from [`Manopt.jl`](@extref), for example it always uses +the [`quasi_newton`](@ref) solver. + +To change this solver, use [`direct_minimizartion`](@ref direct_minimization(::PlaneWaveBasis, ::Manopt.AbstractManoptSolverState))`(basis, solver_state)` +""" +function DFTK.direct_minimization(basis::PlaneWaveBasis; + ψ=nothing, + ρ=guess_density(basis), # would be consistent with other scf solvers + tol=1e-6, + maxiter=1_000, + manifold=Manifolds.Stiefel, + record=[Manopt.RecordCost() => :Etot, DFTK.RecordDensityChange(ρ) => :Δρ, Manopt.RecordTime(; mode=:total) => :time] +) + DFTK.direct_minimization( + basis, Manopt.QuasiNewtonState; + ψ=ψ, ρ=ρ, tol=tol, maxiter=maxiter, manifold=manifold, record=record + ) +end + +""" + direct_minimization(basis, state; kwargs...) + +Compute a minimizer of the Hartree-Fock energy functional using the direct minimization method. + +# Argument + +* `basis::DFTK.PlaneWaveBasis{T}`: The plane wave basis to use for the DFT calculation. +* `state_type::Type{<:Manopt.AbstractManoptSolverState}`: The type of the Manopt solver to use. + recommended: [`QuasiNewtonState`](@extref `Manopt.QuasiNewtonState`). + +# Keyword Arguments + +Similar to the simpler [`direct_minimizartion`](@ref direct_minimization(::PlaneWaveBasis))`(basis)` + +* `maxiter=1_000`: The maximum number of iterations for the optimization. If you set the `stopping_criterion=` directly, this keyword has no effect. + If you set the `stopping_criterion=` directly, this keyword has no effect. +* `ψ=nothing`: The initial guess for the wave functions. If not provided, random orbitals will be generated. +* `ρ=guess_density(basis), # would be consistent with other scf solvers +* `tol=1e-6`: stop when the change in the density is less than this tolerance. If you set the `stopping_criterion=` directly, this keyword has no effect. + If you set the `stopping_criterion=` directly, this keyword has no effect. +* `cost=nothing`: provide an individual cost function. You then also have to provide `gradient=` +* `gradient=nothing`: provide an individual gradient function. You then also have to provide `cost=` + If you provide both of these, for the access to `:Energies`, `:Hamiltonian`, `:HamiltonianEvaluation`, and `:ρ` + you also have to implement `Manopt.get_parameter` for these. +* `evaluation`=[`InplaceEvaluation`](@extref `Manopt.InplaceEvaluation`) whether the gradient of the objective is allocating or in-place., +* `manifold=`[`Stiefel`](@extref `Manifolds.Stiefel`): the manifold the optimisation is running on. + The default cost function allows to also use [`Grassmann`](@extref `Manifolds.Grassmann`). + If you set the `manifold_constructor=` directly, e.g. to switch to a real manifold, this keyword is ignored +* `preconditioner=`[`PreconditionerTPA`](@ref)` a preconditioner to use for the Newton equation or the gradient. +* `manifold_constructor=(n, k) -> manifold(n, k, ℂ)` the complete constructor for a single manifold within the + product manifold the optimization is defined on. +* `stopping_criterion=[`StopAfterIteration`](@extref `Manopt.StopAfterIteration`)` `[`|`](@extref Manopt.StopWhenAny)` `[`StopWhenDensityChangeLess`](@ref)`(tol,deepcopy(ρ))`: + a stopping criterion for the algorithm when to stop. Uses `maxiter=` and `tol=` as defaults, requires that `ρ` a density. +* `record=[[`RecordCost`](@extref `Manopt.RecordCost`)`() => :Etot, `[`RecordDensityChange`](@ref)`(ρ) => :Δρ`, `[`RecordTime`](@extref `Manopt.RecordTime`)`(; mode=:total) => :time]` + specify values to record diring the iterations, where in a `pair` the second determines how to access the values; the three default ones are included in the returned named tuple +* `retraction_method = `[`ProjectionRetraction`](@extref `Manifolds.ProjectionRetraction`) the retration to use on the manifold to more into a certain direction. +* `vector_transport_method=`[`ProjectionTransport`](@extref `Manifolds.ProjectionTransport`)`()` the vector transport to use to move tangent vectors between tangent spaces. +* `stepsize = `[`ArmijoLinesearch`](@extref `Manopt.ArmijoLinesearch`)`(; retraction_method=retraction_method)` + specify a step size rule. + +All other keyword arguments are passed to the solver state constructor as well as to +both a decorate for the objective and the state. +This allows both to set other setting for a solver but also to add `debug=` funtionality +or even a `cache=` to the objective, though the default objective already caches parts +that both a cost and a gradient at the same point would require. + +!!! note "Technical Note" + This interface is still work in progress and might change in the furute even with changes + that break compatibility. +""" +function DFTK.direct_minimization( + basis::PlaneWaveBasis{T}, state_type::Type{<:Manopt.AbstractManoptSolverState}; + ψ=nothing, + ρ=guess_density(basis), # would be consistent with other scf solvers + tol=1e-6, + maxiter=1_000, + cost=nothing, + gradient=nothing, + preconditioner=DFTK.PreconditionerTPA, + manifold=Manifolds.Stiefel, + manifold_constructor=(n, k) -> manifold(n, k, ℂ), + stopping_criterion = Manopt.StopAfterIteration(maxiter) | DFTK.StopWhenDensityChangeLess(tol,deepcopy(ρ)), + evaluation = Manopt.InplaceEvaluation(), + record=[Manopt.RecordCost() => :Etot, DFTK.RecordDensityChange(ρ) => :Δρ, Manopt.RecordTime(; mode=:total) => :time], + retraction_method = Manifolds.ProjectionRetraction(), + vector_transport_method=Manifolds.ProjectionTransport(), + stepsize=Manopt.ArmijoLinesearch(; retraction_method=retraction_method), + kwargs... +) where {T} + # Part 1: Get DFTK variables + # + # + model = basis.model + @assert iszero(model.temperature) # temperature is not yet supported + @assert isnothing(model.εF) # neither are computations with fixed + filled_occ = DFTK.filled_occupation(model) + Nk = length(basis.kpoints) + n_spin = model.n_spin_components # Int. 2 if :collinear, 1 otherwise + n_bands = div(model.n_electrons, n_spin * filled_occ, RoundUp) # Int + ψ = isnothing(ψ) ? [DFTK.random_orbitals(basis, kpt, n_bands) for kpt in basis.kpoints] : ψ + occupation = [filled_occ * ones(T, n_bands) for _ = 1:Nk] + energies, ham = energy_hamiltonian(basis, ψ, occupation; ρ=ρ) + # Part II: Setup solver + # + # + # Initialize the product manifold + dimensions = size.(ψ) # Vector of touples of ψ dimensions + manifold_array = map(dim -> manifold_constructor(dim[1], dim[2]), dimensions) + product_manifold = ProductManifold(manifold_array...) + Pks = [preconditioner(basis, kpt) for kpt in basis.kpoints] + Preconditioner = ManoptPreconditionersWrapper!!(Nk, Pks, basis.kweights) + # Repackage the ψ into a more efficient structure + recursive_ψ = ArrayPartition(ψ...) + if isnothing(cost) && isnothing(gradient) + cost_rgrad! = InsulatorEnergy(basis, + occupation, + Nk, + filled_occ, + zero.(ψ), #init different from ψ to avoid caching errors + rand(product_manifold; vector_at=recursive_ψ), # init different from X to avoid caching errors + deepcopy(ρ), + deepcopy(energies), + deepcopy(ham), + 0, # count of hamiltonian calls + ) + cost = cost_rgrad! # local cost function + if evaluation == InplaceEvaluation() + grad = (M,X,p) -> cost_rgrad!(M, X, p) + else + grad = (M,p) -> cost_rgrad!(M, zero_vector(M,p), p) + end + else + if isnothing(cost) | isnothing(gradient) + error("Providing a cost or gradient function directly also requires the other one to be provided") + end + end + # Build Objective & Problem + objective = Manopt.ManifoldGradientObjective(cost, grad; evaluation=evaluation) + deco_obj = Manopt.decorate_objective!(product_manifold, objective; kwargs...) + problem = Manopt.DefaultManoptProblem(product_manifold, deco_obj) + _stepsize = Manopt._produce_type(stepsize, product_manifold) + state = state_type( + product_manifold; + p=recursive_ψ, + stopping_criterion=stopping_criterion, + preconditioner=QuasiNewtonPreconditioner( + (M, Y, p, X) -> Preconditioner(M, Y, p, X); evaluation=InplaceEvaluation() + ), + direction=Manopt.PreconditionedDirectionRule(product_manifold, + (M, Y, p, X) -> Preconditioner(M, Y, p, X); + evaluation=InplaceEvaluation() + ), + # Set default, can still be overwritten by kwargs... + stepsize=_stepsize, + vector_transport_method=vector_transport_method, + retraction_method=retraction_method, + memory_size=10, + X=cost_rgrad!( # Initial gradient + product_manifold, + zero_vector(product_manifold, recursive_ψ), + recursive_ψ + ), + kwargs... + ) + deco_state = Manopt.decorate_state!(state; record=record, kwargs...) + Manopt.solve!(problem, deco_state) + # Parti III: Collect result in a struct and return that + # + # + # Bundle the variables in a NamedTuple for debugging: + p = Manopt.get_solver_result(deco_state) + recorded_values = Manopt.get_record_action(deco_state) + t = recorded_values[:time] + + # The NamedTuple that is returned, collecting all results + ( + info = "This object is summarizing variables for debugging purposes", + algorithm = "$(state)", + converged = Manopt.has_converged(Manopt.get_state(deco_state, true)), + product_manifold = product_manifold, + basis = basis, + history_Δρ = recorded_values[:Δρ], + history_Etot = recorded_values[:Etot], + runtime_ns = length(t) > 0 ? last(t) : zero(eltype(t)), # take the last recorded time if any recorded + ψ = collect(p.x), # reformat ArrayPartition back to a vector of matrices + model = model, + n_bands_converge = n_bands, + n_count = Manopt.get_count(Manopt.get_state(deco_state, true), :Iterations), + Nk = Nk, + occupation = occupation, + cost_value=get_cost(problem, p), + cost_grad_value=get_gradient(problem, p), + ρ = Manopt.get_parameter(Manopt.get_objective(problem), :ρ), + energies = Manopt.get_parameter(Manopt.get_objective(problem), :Energies), + ham = Manopt.get_parameter(Manopt.get_objective(problem), :Hamiltonian), + n_matvec = Manopt.get_parameter( + Manopt.get_objective(problem), + :HamiltonianEvaluations + ), + # The “pure” solver state without debug/records. + solver_state=Manopt.get_state(deco_state,true), + ) +end +#= +TODO +we could adapt to has_converged to report on convergence as well +=# +end \ No newline at end of file diff --git a/src/scf/direct_minimization.jl b/src/scf/direct_minimization.jl index da2403bb26..39cfc4fa63 100644 --- a/src/scf/direct_minimization.jl +++ b/src/scf/direct_minimization.jl @@ -1,197 +1,15 @@ # Direct minimization of the energy -using Optim -using LineSearches - -# This is all a bit annoying because our ψ is represented as ψ[k][G,n], and Optim accepts -# only dense arrays. We do a bit of back and forth using custom `pack` (ours -> optim's) and -# `unpack` (optim's -> ours) functions - -# Orbitals inside each kblock must be kept orthogonal: the -# project_tangent and retract work per kblock -struct DMManifold <: Optim.Manifold - Nk::Int - unpack::Function -end -function Optim.project_tangent!(m::DMManifold, g, x) - g_unpack = m.unpack(g) - x_unpack = m.unpack(x) - for ik = 1:m.Nk - Optim.project_tangent!(Optim.Stiefel(), g_unpack[ik], x_unpack[ik]) - end - g -end -function Optim.retract!(m::DMManifold, x) - x_unpack = m.unpack(x) - for ik = 1:m.Nk - Optim.retract!(Optim.Stiefel(), x_unpack[ik]) - end - x -end - -# Array of preconditioners -struct DMPreconditioner - Nk::Int - Pks::Vector # Pks[ik] is the preconditioner for k-point ik - unpack::Function - kweights::Vector -end -function LinearAlgebra.ldiv!(p, P::DMPreconditioner, d) - p_unpack = P.unpack(p) - d_unpack = P.unpack(d) - for ik = 1:P.Nk - ldiv!(p_unpack[ik], P.Pks[ik], d_unpack[ik]) - ldiv!(P.kweights[ik], p_unpack[ik]) - end - p -end -function LinearAlgebra.dot(x, P::DMPreconditioner, y) - x_unpack = P.unpack(x) - y_unpack = P.unpack(y) - sum(dot(P.kweights[ik]*x_unpack[ik], P.Pks[ik], y_unpack[ik]) - for ik = 1:P.Nk) -end -function precondprep!(P::DMPreconditioner, x) - x_unpack = P.unpack(x) - for ik = 1:P.Nk - precondprep!(P.Pks[ik], x_unpack[ik]) - end - P -end - - -""" -Computes the ground state by direct minimization. `kwargs...` are -passed to `Optim.Options()` and `optim_method` selects the optim approach -which is employed. """ -function direct_minimization(basis::PlaneWaveBasis{T}; - ψ=nothing, - tol=1e-6, - is_converged=ScfConvergenceDensity(tol), - maxiter=1_000, - prec_type=PreconditionerTPA, - callback=ScfDefaultCallback(), - optim_method=Optim.LBFGS, - alphaguess=LineSearches.InitialStatic(), - linesearch=LineSearches.BackTracking(), - kwargs...) where {T} - if mpi_nprocs() > 1 - # need synchronization in Optim - error("Direct minimization with MPI is not supported yet") - end - model = basis.model - @assert iszero(model.temperature) # temperature is not yet supported - @assert isnothing(model.εF) # neither are computations with fixed Fermi level - filled_occ = filled_occupation(model) - n_spin = model.n_spin_components - n_bands = div(model.n_electrons, n_spin * filled_occ, RoundUp) - Nk = length(basis.kpoints) - - if isnothing(ψ) - ψ = [random_orbitals(basis, kpt, n_bands) for kpt in basis.kpoints] - end - occupation = [filled_occ * ones(T, n_bands) for _ = 1:Nk] - - # we need to copy the reinterpret array here to not raise errors in Optim.jl - # TODO raise this issue in Optim.jl - pack(ψ) = copy(reinterpret_real(pack_ψ(ψ))) - unpack(x) = unpack_ψ(reinterpret_complex(x), size.(ψ)) - unsafe_unpack(x) = unsafe_unpack_ψ(reinterpret_complex(x), size.(ψ)) - - # This will get updated along the iterations - ρ = nothing - ham = nothing - info = nothing - energies = nothing - converged = false - start_ns = time_ns() - history_Etot = T[] - history_Δρ = T[] - - # Will be later overwritten by the Optim-internal state, which we need in the - # callback to access certain quantities for convergence control. - optim_state = nothing - - function compute_ρout(ψ, optim_state) - # This is the current preconditioned, but unscaled gradient, which implies that - # the next step would be ρout - ρ. We thus record convergence, but let Optim do - # one more step. - δψ = unsafe_unpack(optim_state.s) - # TODO This looks weird ... should there not be a retraction ? - ψ_next = [ortho_qr(ψ[ik] - δψ[ik]) for ik in 1:Nk] - compute_density(basis, ψ_next, occupation) - end + direct_minimization() - function optim_callback(ts) - ts.iteration < 1 && return false - converged && return true - ρout = compute_ρout(ψ, optim_state) - Δρ = ρout - ρ - push!(history_Δρ, norm(Δρ) * sqrt(basis.dvol)) - push!(history_Etot, energies.total) - - info = (; ham, basis, energies, occupation, ρout, ρin=ρ, ψ, - runtime_ns=time_ns() - start_ns, history_Δρ, history_Etot, - stage=:iterate, algorithm="DM", n_iter=ts.iteration, optim_state) - - converged = is_converged(info) - info = callback(info) - - false - end - - # computes energies and gradients - function fg!(::Any, G, x) - ψ = unpack(x) - ρ = compute_density(basis, ψ, occupation) - energies, ham = energy_hamiltonian(basis, ψ, occupation; ρ) - - # The energy has terms like occ * <ψ|H|ψ>, so the gradient is 2occ Hψ - if G !== nothing - G = unsafe_unpack(G) - for ik = 1:Nk - mul!(G[ik], ham.blocks[ik], ψ[ik]) - G[ik] .*= 2 * filled_occ * basis.kweights[ik] - end - end - energies.total - end - - manifold = DMManifold(Nk, unsafe_unpack) - Pks = [prec_type(basis, kpt) for kpt in basis.kpoints] - P = DMPreconditioner(Nk, Pks, unsafe_unpack, basis.kweights) - - optim_options = Optim.Options(; allow_f_increases=true, - callback=optim_callback, - # Disable convergence control by Optim - x_tol=-1, f_tol=-1, g_tol=-1, - iterations=maxiter, kwargs...) - optim_solver = optim_method(; P, precondprep=precondprep!, manifold, linesearch, alphaguess) - ψ_packed = pack(ψ) - objective = OnceDifferentiable(Optim.only_fg!(fg!), ψ_packed, zero(T); inplace=true) - optim_state = Optim.initial_state(optim_solver, optim_options, objective, ψ_packed) - res = Optim.optimize(objective, ψ_packed, optim_solver, optim_options, optim_state) - ψ = unpack(Optim.minimizer(res)) - - # Final Rayleigh-Ritz (not strictly necessary, but sometimes useful) - eigenvalues = Vector{T}[] - for ik = 1:Nk - Hψk = ham[ik] * ψ[ik] - F = eigen(Hermitian(ψ[ik]'Hψk)) - push!(eigenvalues, F.values) - ψ[ik] .= ψ[ik] * F.vectors - end - - εF = nothing # does not necessarily make sense here, as the - # Aufbau property might not even be true - - # We rely on the fact that the last point where fg! was called is the minimizer to - # avoid recomputing at ψ - info = (; ham, basis, energies, converged, ρ, eigenvalues, occupation, εF, - n_bands_converge=n_bands, n_iter=Optim.iterations(res), - runtime_ns=time_ns() - start_ns, history_Δρ, history_Etot, - ψ, stage=:finalize, algorithm="DM", optim_res=res) - callback(info) - info -end +Note. Requires now Manopt to work. +""" +function direct_minimization end + +function StopWhenDensityChangeLess end +function RecordDensityNorm end +function RecordDensityChange end +function RecordDensity end +function DebugDensityNorm end +function DebugDensityChange end diff --git a/src/terms/Hamiltonian.jl b/src/terms/Hamiltonian.jl index 869f64f4ab..70c4a0f584 100644 --- a/src/terms/Hamiltonian.jl +++ b/src/terms/Hamiltonian.jl @@ -185,10 +185,26 @@ end """ -Get energies and Hamiltonian -kwargs is additional info that might be useful for the energy terms to precompute -(eg the density ρ) + (; energies, hamiltonian) = energy_hamiltonian(basis::PlaneWaveBasis, ψ, occupation; kwargs...) + +Compute the energies and Hamiltonian. + +# Arguments +- `basis::PlaneWaveBasis`: The plane wave basis to use. +- `ψ`: The wavefunctions, as a vector of arrays, one for each k-point. +- `occupation`: The occupation numbers for each k-point and band. + +The keyword arguments `kwargs...` are passed to the internal function `ene_ops`. +These can be used to pass additional information to that function, e.g., the density ρ. + +# Returns + +The [`Energies`](@ref) object with the energy values for each term and the [`Hamiltonian`](@ref) object +in a dictionary with the keys `energies` and `ham`. + """ +energy_hamiltonian(basis::PlaneWaveBasis, ψ, occupation; kwargs...) + @timing function energy_hamiltonian(basis::PlaneWaveBasis, ψ, occupation; kwargs...) # it: index into terms, ik: index into kpoints @timing "ene_ops" ene_ops_arr = [ene_ops(term, basis, ψ, occupation; kwargs...) diff --git a/test/external_potential.jl b/test/external_potential.jl index 9c40620703..5ba4a2ecad 100644 --- a/test/external_potential.jl +++ b/test/external_potential.jl @@ -1,6 +1,7 @@ @testitem "External potential from Fourier coefficients" #= =# tags=[:core, :dont_test_mpi] begin using DFTK + using Manopt, Manifolds, RecursiveArrayTools lattice = [[10 0 0.]; [0 0 0]; [0 0 0]] diff --git a/test/scf_compare.jl b/test/scf_compare.jl index 022f8ecc7e..32ff99308c 100644 --- a/test/scf_compare.jl +++ b/test/scf_compare.jl @@ -1,6 +1,7 @@ @testitem "Compare different SCF algorithms (no spin, no temperature)" #= =# tags=[:core] setup=[TestCases] begin using DFTK + using Manopt, Manifolds, RecursiveArrayTools using DFTK: Applyχ0Model, select_occupied_orbitals silicon = TestCases.silicon tol = 1e-7