Move all requires-based optional code to extension modules (#930)

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.14"
+version = "0.6.15"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -37,7 +37,6 @@ Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PseudoPotentialIO = "cb339c56-07fa-4cb2-923a-142469552264"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
@@ -49,45 +48,59 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a"
+IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9"
 wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [extensions]
+DFTKCUDAExt = "CUDA"
+DFTKGenericLinearAlgebraExt = "GenericLinearAlgebra"
+DFTKIntervalArithmeticExt = "IntervalArithmetic"
+DFTKJLD2Ext = "JLD2"
 DFTKJSON3Ext = "JSON3"
-DFTKWannierExt = "Wannier"
+DFTKPlotsExt = "Plots"
 DFTKWannier90Ext = "wannier90_jll"
-DFTKWriteVTK = "WriteVTK"
+DFTKWannierExt = "Wannier"
+DFTKWriteVTKExt = "WriteVTK"
 
 [compat]
 AbstractFFTs = "1"
 Artifacts = "1"
 AtomsBase = "0.3.1"
 Brillouin = "0.5.14"
 ChainRulesCore = "1.15"
+CUDA = "5"
 Dates = "1"
 DftFunctionals = "0.2"
 DocStringExtensions = "0.9"
 FFTW = "1.5"
 ForwardDiff = "0.10"
+GenericLinearAlgebra = "0.3"
 GPUArraysCore = "0.1"
 Interpolations = "0.14, 0.15"
 IntervalArithmetic = "0.20"
-IterTools = "1"
 IterativeSolvers = "0.9"
+IterTools = "1"
+JLD2 = "0.4"
 JSON3 = "1"
 LazyArtifacts = "1.3"
 Libxc = "0.3.17"
-LineSearches = "7"
 LinearAlgebra = "1"
 LinearMaps = "3"
+LineSearches = "7"
 LoopVectorization = "0.12"
 MPI = "0.20.13"
 Markdown = "1"
 Optim = "1"
 OrderedCollections = "1"
 PeriodicTable = "1"
+Plots = "1"
 PkgVersion = "0.3"
 Polynomials = "3, 4"
 PrecompileTools = "1"
@@ -96,7 +109,6 @@ Primes = "0.5"
 Printf = "1"
 PseudoPotentialIO = "0.1"
 Random = "1"
-Requires = "1"
 Roots = "2"
 SparseArrays = "1"
 SpecialFunctions = "2"

examples/arbitrary_floattype.jl

@@ -52,9 +52,10 @@ eltype(scfres.ρ)
 #
 # !!! note "Generic linear algebra routines"
 #     For more unusual floating-point types (like IntervalArithmetic or DoubleFloats),
-#     which are not directly supported in the standard `LinearAlgebra` library of Julia
-#     one additional step is required: One needs to explicitly enable the generic versions
-#     of standard linear-algebra operations like `cholesky` or `qr`, which are needed
-#     inside DFTK by loading the `GenericLinearAlgebra` package in the user script
+#     which are not directly supported in the standard `LinearAlgebra` and `FFTW`
+#     libraries one additional step is required: One needs to explicitly enable the generic
+#     versions of standard linear-algebra operations like `cholesky` or `qr` or standard
+#     `fft` operations, which DFTK requires. THis is done by loading the
+#     `GenericLinearAlgebra` package in the user script
 #     (i.e. just add ad `using GenericLinearAlgebra` next to your `using DFTK` call).
 #

src/workarounds/cuda_arrays.jl → ext/DFTKCUDAExt.jl

@@ -1,4 +1,11 @@
-synchronize_device(::GPU{<:CUDA.CuArray}) = CUDA.synchronize()
+module DFTKCUDAExt
+using CUDA
+import DFTK: GPU, DispatchFunctional
+using DftFunctionals
+using DFTK
+using Libxc
+
+DFTK.synchronize_device(::GPU{<:CUDA.CuArray}) = CUDA.synchronize()
 
 for fun in (:potential_terms, :kernel_terms)
     @eval function DftFunctionals.$fun(fun::DispatchFunctional,
@@ -7,3 +14,5 @@ for fun in (:potential_terms, :kernel_terms)
         $fun(fun.inner, ρ, args...)
     end
 end
+
+end

src/workarounds/fft_generic.jl → ext/DFTKGenericLinearAlgebraExt/DFTKGenericLinearAlgebraExt.jl

@@ -1,29 +1,35 @@
-include("FourierTransforms.jl/FourierTransforms.jl")
+module DFTKGenericLinearAlgebraExt
+using DFTK
+using DFTK: DummyInplace
+using LinearAlgebra
+using AbstractFFTs
+import AbstractFFTs: Plan, ScaledPlan,
+                     fft, ifft, bfft, fft!, ifft!, bfft!,
+                     plan_fft, plan_ifft, plan_bfft, plan_fft!, plan_ifft!, plan_bfft!,
+                     rfft, irfft, brfft, plan_rfft, plan_irfft, plan_brfft,
+                     fftshift, ifftshift,
+                     rfft_output_size, brfft_output_size,
+                     plan_inv, normalization
+import Base: show, summary, size, ndims, length, eltype, *, inv, \
+import LinearAlgebra: mul!
 
-# This is needed to flag that the fft_generic.jl file has already been loaded
-const GENERIC_FFT_LOADED = true
-
-if !isdefined(Main, :GenericLinearAlgebra)
-    @warn("Code paths for generic floating-point types activated in DFTK. Remember to " *
-          "add 'using GenericLinearAlgebra' to your user script. " *
-          "See https://docs.dftk.org/stable/examples/arbitrary_floattype/ for details.")
-end
+include("ctfft.jl")  # Main file of FourierTransforms.jl
 
 # Utility functions to setup FFTs for DFTK. Most functions in here
 # are needed to correct for the fact that FourierTransforms is not
 # yet fully compliant with the AbstractFFTs interface and has still
 # various bugs we work around.
 
-function next_working_fft_size(::Any, size::Integer)
+function DFTK.next_working_fft_size(::Any, size::Integer)
     # TODO FourierTransforms has a bug, which is triggered
     #      only in some factorizations, see
     #      https://github.com/JuliaComputing/FourierTransforms.jl/issues/10
     nextpow(2, size)  # We fall back to powers of two to be safe
 end
-default_primes(::Any) = (2, )
+DFTK.default_primes(::Any) = (2, )
 
 # Generic fallback function, Float32 and Float64 specialization in fft.jl
-function build_fft_plans!(tmp::AbstractArray{<:Complex})
+function DFTK.build_fft_plans!(tmp::AbstractArray{<:Complex})
     # Note: FourierTransforms has no support for in-place FFTs at the moment
     # ... also it's extension to multi-dimensional arrays is broken and
     #     the algo only works for some cases
@@ -40,14 +46,12 @@ function build_fft_plans!(tmp::AbstractArray{<:Complex})
     ipFFT, opFFT, ipBFFT, opBFFT
 end
 
-
-
 struct GenericPlan{T}
     subplans
     factor::T
 end
 
-function generic_apply(p::GenericPlan, X::AbstractArray)
+function Base.:*(p::GenericPlan, X::AbstractArray)
     pl1, pl2, pl3 = p.subplans
     ret = similar(X)
     for i = 1:size(X, 1), j = 1:size(X, 2)
@@ -65,21 +69,21 @@ end
 LinearAlgebra.mul!(Y, p::GenericPlan, X) = Y .= p * X
 LinearAlgebra.ldiv!(Y, p::GenericPlan, X) = Y .= p \ X
 
-import Base: *, \, inv, length
 length(p::GenericPlan) = prod(length, p.subplans)
-*(p::GenericPlan, X::AbstractArray) = generic_apply(p, X)
 *(p::GenericPlan{T}, fac::Number) where {T} = GenericPlan{T}(p.subplans, p.factor * T(fac))
 *(fac::Number, p::GenericPlan{T}) where {T} = p * fac
 \(p::GenericPlan, X) = inv(p) * X
 inv(p::GenericPlan{T}) where {T} = GenericPlan{T}(inv.(p.subplans), 1 / p.factor)
 
 function generic_plan_fft(data::AbstractArray{T, 3}) where {T}
-    GenericPlan{T}([FourierTransforms.plan_fft(data[:, 1, 1]),
-                    FourierTransforms.plan_fft(data[1, :, 1]),
-                    FourierTransforms.plan_fft(data[1, 1, :])], T(1))
+    GenericPlan{T}([plan_fft(data[:, 1, 1]),
+                    plan_fft(data[1, :, 1]),
+                    plan_fft(data[1, 1, :])], T(1))
 end
 function generic_plan_bfft(data::AbstractArray{T, 3}) where {T}
-    GenericPlan{T}([FourierTransforms.plan_bfft(data[:, 1, 1]),
-                    FourierTransforms.plan_bfft(data[1, :, 1]),
-                    FourierTransforms.plan_bfft(data[1, 1, :])], T(1))
+    GenericPlan{T}([plan_bfft(data[:, 1, 1]),
+                    plan_bfft(data[1, :, 1]),
+                    plan_bfft(data[1, 1, :])], T(1))
+end
+
 end

src/workarounds/FourierTransforms.jl/ctfft.jl → ext/DFTKGenericLinearAlgebraExt/ctfft.jl

@@ -2,6 +2,31 @@ using Primes
 
 # COV_EXCL_START
 
+# These files are copied (and very slightly modified) from [FourierTransforms.jl](https://github.com/JuliaComputing/FourierTransforms.jl)
+# commit [6a206bcfc8f49a129ca34beaf57d05cdc148dc8a](https://github.com/JuliaComputing/FourierTransforms.jl/tree/6a206bcfc8f49a129ca34beaf57d05cdc148dc8a).
+# For their license see [LICENCE.md](LICENCE.md).
+
+# The following licence applies:
+# Copyright (c) 2017-2019: Steven G. Johnson, Yingbo Ma, and Julia Computing.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
 # 1d Cooley-Tukey FFTs, using an FFTW-like (version 1) approach: automatic
 # generation of fixed-size FFT kernels (with and without twiddle factors)
 # which are combined to make arbitrary-size FFTs (plus generic base
@@ -301,7 +326,8 @@ macro nontwiddle(args...)
     end
     @assert isa(T, Type)
     quote
-        function FourierTransforms.applystep(ns::NontwiddleKernelStep{T,$n,$forward},
+        function DFTKGenericLinearAlgebraExt.applystep(
+                                  ns::NontwiddleKernelStep{T,$n,$forward},
                                   vn::Integer,
                                   X::AbstractArray{T},
                                   x0::Integer, xs::Integer, xvs::Integer,
@@ -340,7 +366,8 @@ macro twiddle(args...)
     end
     @assert isa(T, Type)
     quote
-        function FourierTransforms.applystep(ts::TwiddleKernelStep{T,$n,$forward},
+        function DFTKGenericLinearAlgebraExt.applystep(
+                                   ts::TwiddleKernelStep{T,$n,$forward},
                                    vn::Integer,
                                    X::AbstractArray{T},
                                    x0::Integer, xs::Integer, xvs::Integer,

src/workarounds/intervals.jl → ext/DFTKIntervalArithmeticExt.jl

@@ -1,20 +1,26 @@
+module DFTKIntervalArithmeticExt
+using DFTK
+using IntervalArithmetic
+using LinearAlgebra
+import DFTK: symmetry_operations, _is_well_conditioned, compute_Glims_fast
+import DFTK: local_potential_fourier
+import IntervalArithmetic: Interval
 import SpecialFunctions: erfc
-const Interval = IntervalArithmetic.Interval
 
 # Monkey-patch a few functions for Intervals
 # ... this is far from proper and a bit specific for our use case here
 # (that's why it's not contributed upstream).
 # should be done e.g. by changing  the rounding mode ...
 erfc(i::Interval) = Interval(prevfloat(erfc(i.lo)), nextfloat(erfc(i.hi)))
+Base.nextfloat(x::Interval) = Interval(nextfloat(x.lo), nextfloat(x.hi))
+Base.prevfloat(x::Interval) = Interval(prevfloat(x.lo), prevfloat(x.hi))
 
 # This is done to avoid using sincospi(x), called by cispi(x),
 # which has not been implemented in IntervalArithmetic
 # see issue #513 on IntervalArithmetic repository
-cis2pi(x::Interval) = exp(2 * (pi * (im * x)))
+DFTK.cis2pi(x::Interval) = exp(2 * (pi * (im * x)))
 
-Base.nextfloat(x::Interval) = Interval(nextfloat(x.lo), nextfloat(x.hi))
-Base.prevfloat(x::Interval) = Interval(prevfloat(x.lo), prevfloat(x.hi))
-value_type(::Type{<:Interval{T}}) where {T} = T
+DFTK.value_type(::Type{<:Interval{T}}) where {T} = T
 
 function compute_Glims_fast(lattice::AbstractMatrix{<:Interval}, args...; kwargs...)
     # This is done to avoid a call like ceil(Int, ::Interval)
@@ -25,7 +31,7 @@ function compute_Glims_fast(lattice::AbstractMatrix{<:Interval}, args...; kwargs
     # their midpoints should be good.
     compute_Glims_fast(IntervalArithmetic.mid.(lattice), args...; kwargs...)
 end
-function compute_Glims_precise(::AbstractMatrix{<:Interval}, args...; kwargs...)
+function DFTK.compute_Glims_precise(::AbstractMatrix{<:Interval}, args...; kwargs...)
     error("fft_size_algorithm :precise not supported with intervals")
 end
 
@@ -36,8 +42,8 @@ function _is_well_conditioned(A::AbstractArray{<:Interval}; kwargs...)
 end
 
 function symmetry_operations(lattice::AbstractMatrix{<:Interval}, atoms, positions,
-                             magnetic_moments=[];
-                             tol_symmetry=max(SYMMETRY_TOLERANCE, maximum(radius, lattice)))
+                                  magnetic_moments=[];
+                                  tol_symmetry=max(SYMMETRY_TOLERANCE, maximum(radius, lattice)))
     @assert tol_symmetry < 1e-2
     symmetry_operations(IntervalArithmetic.mid.(lattice), atoms, positions, magnetic_moments;
                         tol_symmetry)
@@ -50,10 +56,13 @@ function local_potential_fourier(el::ElementCohenBergstresser, q::T) where {T <:
     T(local_potential_fourier(el, IntervalArithmetic.mid(q)))
 end
 
-function estimate_integer_lattice_bounds(M::AbstractMatrix{<:Interval}, δ, shift=zeros(3))
+function DFTK.estimate_integer_lattice_bounds(M::AbstractMatrix{<:Interval}, δ,
+                                              shift=zeros(3))
     # As a general statement, with M a lattice matrix, then if ||Mx|| <= δ,
     # then xi = <ei, M^-1 Mx> = <M^-T ei, Mx> <= ||M^-T ei|| δ.
     # Below code does not support non-3D systems.
     xlims = [norm(inv(M')[:, i]) * δ + shift[i] for i = 1:3]
     map(x -> ceil(Int, x.hi), xlims)
 end
+
+end

src/external/jld2io.jl → ext/DFTKJLD2Ext.jl

@@ -1,4 +1,10 @@
-function save_scfres_master(file::AbstractString, scfres::NamedTuple, ::Val{:jld2})
+module DFTKJLD2Ext
+using DFTK
+using DFTK: energy_hamiltonian, AbstractArchitecture, AbstractKgrid
+using JLD2
+using MPI
+
+function DFTK.save_scfres_master(file::AbstractString, scfres::NamedTuple, ::Val{:jld2})
     !mpi_master() && error(
         "This function should only be called on MPI master after the k-point data has " *
         "been gathered with `gather_kpts`."
@@ -19,7 +25,7 @@ function save_scfres_master(file::AbstractString, scfres::NamedTuple, ::Val{:jld
 end
 
 
-function load_scfres(jld::JLD2.JLDFile)
+function DFTK.load_scfres(jld::JLD2.JLDFile)
     basis = jld["basis"]
     scfdict = Dict{Symbol, Any}(
         :ρ     => jld["ρ"],
@@ -45,7 +51,7 @@ function load_scfres(jld::JLD2.JLDFile)
     scfdict[:ham]      = ham
     (; (sym => scfdict[sym] for sym in jld["__propertynames"])...)
 end
-load_scfres(file::AbstractString) = JLD2.jldopen(load_scfres, file, "r")
+DFTK.load_scfres(file::AbstractString) = JLD2.jldopen(DFTK.load_scfres, file, "r")
 
 
 #
@@ -93,3 +99,5 @@ function Base.convert(::Type{PlaneWaveBasis{T,T,Arch,GT,RT,KGT}},
                    MPI.COMM_WORLD,
                    serial.architecture)
 end
+
+end