v0.2 (#4)

* initial commit

* remove type restriction on tolerances

* add AuxQuad extension

* update Project.toml

* getting ready for v0.2

* add Fourier & HDF5 tests

* migrate to IntegralCache

* fix the doc demo

* finalize interface

* finalize interface

* replace Integrals.jl with my own thing

* simplify

* make batchsolve behavior consistent

* update alloc_segbufs

* make parallel an argument to ptr and autoptr

* fix parameter numbering

* update deps for v0.2

Project.toml

@@ -1,38 +1,38 @@
 name = "AutoBZCore"
 uuid = "66bd3e16-1600-45cf-8f55-0b550710682b"
 authors = ["lxvm <lorenzo@vanmunoz.com>"]
-version = "0.1.6"
+version = "0.2.0"
 
 [deps]
 AutoSymPTR = "78a0c066-08f1-49a8-82f0-b29cd485e1d3"
 FourierSeriesEvaluators = "2a892dea-6eef-4bb5-9d1c-de966c9f6db5"
-Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
+HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 IteratedIntegration = "3ecdc4d6-ee34-4049-885a-a4e3631db98b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
+[weakdeps]
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+
+[extensions]
+HDF5Ext = "HDF5"
+
 [compat]
-AutoSymPTR = "0.1"
+AutoSymPTR = "0.2"
 FourierSeriesEvaluators = "0.1"
-Integrals = "3.1.1"
-IteratedIntegration = "0.1.5"
+HCubature = "1.4"
+HDF5 = "0.16.15"
+IteratedIntegration = "0.3"
 Reexport = "1"
-Requires = "1"
 StaticArrays = "1"
-julia = "1.6"
-
-[extensions]
-HDF5Ext = "HDF5"
+julia = "1.9"
 
 [extras]
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
-
-[weakdeps]
-HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+test = ["Test", "HDF5", "StaticArrays", "OffsetArrays"]

ext/HDF5Ext.jl

@@ -1,16 +1,11 @@
 module HDF5Ext
 
 using Printf: @sprintf
+using LinearAlgebra: norm
 
-if isdefined(Base, :get_extension)
-    using AutoBZCore: IntegralSolver, MixedParameters
-    using HDF5: h5open, write_dataset, read_dataset, Group, H5DataStore, create_dataset, create_group
-    import AutoBZCore: batchsolve
-else
-    using ..AutoBZCore: IntegralSolver, MixedParameters
-    using ..HDF5: h5open, write_dataset, read_dataset, Group, H5DataStore, create_dataset, create_group
-    import ..AutoBZCore: batchsolve
-end
+using AutoBZCore: IntegralSolver, MixedParameters, solver_type, AuxValue, IntegralSolution
+using HDF5: h5open, write_dataset, read_dataset, Group, H5DataStore, create_dataset, create_group, delete_object, name
+import AutoBZCore: batchsolve
 
 
 """
@@ -45,16 +40,40 @@ end
 # parallelization
 
 # returns (dset, ax) to allow for array-valued data types
-function autobz_create_dataset(parent, path, T::Type, dims_::Tuple{Vararg{Int}})
+function autobz_create_dataset(parent, path, T::Type, dims_)
     dims = ((ndims(T) == 0 ? () : size(T))..., dims_...)
     ax = ntuple(_-> :, Val(ndims(T)))
     return create_dataset(parent, path, eltype(T), dims), ax
 end
+function autobz_create_dataset(parent, path, ::Type{AuxValue{T}}, dims) where T
+    # split auxvalue into two groups for easier interoperability with other languages, since
+    # the HDF5 compound type could be a challenge
+    g = create_group(parent, "I")
+    gval, axval = autobz_create_dataset(g, "val", T, dims)
+    gaux, axaux = autobz_create_dataset(g, "aux", T, dims)
+    return (gval, gaux), (axval, axaux)
+end
 
-
-function param_group(parent, T, dims)
-    return create_dataset(parent, "p", T, dims)
+set_value(g, i::CartesianIndex, val) = g[i.I...] = val
+set_value(parent, ax::Tuple, i::CartesianIndex, sol) = parent[ax...,i.I...] = sol
+function set_value((gval, gaux)::Tuple, (axval, axaux)::Tuple, i::CartesianIndex, sol::AuxValue)
+    set_value(gval, axval, i, sol.val)
+    set_value(gaux, axaux, i, sol.aux)
+    return nothing
 end
+# special cases for 0-d arrays with scalar fields, for which rewrite is not supported
+function set_value(g, ::CartesianIndex{0}, val)
+    p = parent(g)
+    s = basename(name(g))
+    delete_object(p, s)
+    write_dataset(p, s, val)
+    return nothing
+end
+set_value(parent, ::Tuple{}, i::CartesianIndex{0}, sol) = set_value(parent, i, sol)
+
+
+
+param_group(parent, T, dims) = create_dataset(parent, "p", T, dims)
 function param_group(parent, ::Type{T}, dims) where {T<:Tuple}
     g = create_group(parent, "params")
     for (i, S) in enumerate(T.parameters)
@@ -74,51 +93,56 @@ function param_group(parent, ::Type{MixedParameters{T,NamedTuple{K,V}}}, dims) w
     return (g,q)
 end
 function param_record(group, p, i)
-    group[i] = p
+    set_value(group, i, p)
+    return nothing
 end
 function param_record(group, p::Tuple, i)
     for (j,e) in enumerate(p)
-        group[string(j)][i] = e
+        set_value(group[string(j)], i, e)
     end
+    return nothing
 end
 function param_record((g, q), p::MixedParameters, i)
     for (j,e) in enumerate(getfield(p, :args))
-        g[string(j)][i] = e
+        set_value(g[string(j)], i, e)
     end
     for (k,v) in pairs(getfield(p, :kwargs))
-        q[string(k)][i] = v
+        set_value(q[string(k)], i, v)
     end
+    return nothing
 end
 
 """
-    batchsolve(h5::H5DataStore, f::IntegralSolver, ps, T=Base.promote_op(f, eltype(ps)); nthreads=Threads.nthreads())
+    batchsolve(h5::H5DataStore, f::IntegralSolver, ps, [T]; verb=true, nthreads=Threads.nthreads())
 
 Batchsolver
 """
-function batchsolve(h5::H5DataStore, f::IntegralSolver, ps, T=Base.promote_op(f, eltype(ps)), verb=true; nthreads=Threads.nthreads())
-    dims = tuple(length(ps))
-
+function batchsolve(h5::H5DataStore, f::IntegralSolver, ps::AbstractArray, T=solver_type(f, ps[begin]); verb=true, nthreads=Threads.nthreads())
+    isconcretetype(T) || throw(ArgumentError("Result type of integrand is abstract or could not be inferred. Please provide the concrete return type to save to HDF5"))
+    len = length(ps)
+    dims = size(ps)
     gI, ax = autobz_create_dataset(h5, "I", T, dims)
     gE = create_dataset(h5, "E", Float64, dims)
     gt = create_dataset(h5, "t", Float64, dims)
     gr = create_dataset(h5, "retcode", Int32, dims)
     gp = param_group(h5, eltype(ps), dims)
 
-    function h5callback(f, i, p, sol, t)
-        verb && @info @sprintf "parameter %i finished in %e (s)" i t
-        gI[ax...,i] = sol
-        gE[i] = isnothing(sol.resid) ? NaN : convert(Float64, sol.resid)
-        gt[i] = t
-        gr[i] = Integer(sol.retcode)
+    function h5callback(_, i, n, p, sol, t)
+        verb && @info @sprintf "%5i / %i done in %e (s)" n len t
+        set_value(gI, ax, i, sol.u)
+        set_value(gE, i, isnothing(sol.resid) ? NaN : convert(Float64, T<:AuxValue ? sol.resid.val : sol.resid))
+        set_value(gt, i, t)
+        set_value(gr, i, Integer(sol.retcode))
         param_record(gp, p, i)
     end
 
-    @info "Started parameter sweep"
+    verb && @info "Started parameter sweep"
     t = time()
     sol = batchsolve(f, ps, T; callback=h5callback, nthreads=nthreads)
     t = time()-t
-    @info @sprintf "Finished parameter sweep in %e (s) CPU time and %e (s) wall clock time" sum(read(gt)) t
-    sol
+    verb && @info @sprintf "Finished parameter sweep in %e (s) CPU time and %e (s) wall clock time" sum(read(gt)) t
+
+    return sol
 end
 
-end
+end

src/AutoBZCore.jl

@@ -16,70 +16,60 @@ For example, computing the local Green's function can be done as follows:
     using FourierSeriesEvaluators
     using AutoBZCore
 
-    gloc_integrand(h_k; η, ω) = inv(complex(ω,η)*I-h_k)     # define integrand evaluator
-    h = FourierSeries([0.5, 0.0, 0.5]; offset=-2)           # construct cos(k) 1D integer lattice Hamiltonian
-    bz = FullBZ(2pi*I(1))                                   # construct BZ from lattice vectors A=2pi*I
-    integrand = FourierIntegrand(gloc_integrand, h, η=0.1)    # construct integrand with Fourier series h and parameter η=0.1
+    gloc_integrand(k, h; η, ω) = inv(complex(ω,η)*I-h(k))   # define integrand evaluator
+    h = FourierSeries([0.5, 0.0, 0.5]; period=1, offset=-2) # construct cos(2πk) 1D integer lattice Hamiltonian
+    integrand = Integrand(gloc_integrand, h, η=0.1)         # construct integrand with Fourier series h and parameter η=0.1
+    prob = IntegralProblem(integrand, 0, 1)                 # setup the integral problem
+    alg = QuadGKJL()                                        # choose integration algorithm (also AutoPTR() and PTR())
+    gloc = IntegralSolver(prob, alg; abstol=1e-3)           # construct a solver for gloc to within specified tolerance
+    gloc(ω=0.0)                                             # evaluate gloc at frequency ω=0.0
+
+
+    gloc_integrand(h_k::FourierValue; η, ω) = inv(complex(ω,η)*I-h_k.s)     # define integrand evaluator
+    h = FourierSeries([0.0; 0.5; 0.0;; 0.5; 0.0; 0.5;; 0.0; 0.5; 0.0]; period=1, offset=-2) # construct cos(2πk) 1D integer lattice Hamiltonian
+    bz = FullBZ(2pi*I(2))                                   # construct BZ from lattice vectors A=2pi*I
+    integrand = FourierIntegrand(gloc_integrand, h, η=0.1)   # construct integrand with Fourier series h and parameter η=0.1
+    prob = IntegralProblem(integrand, bz)                   # setup the integral problem
     alg = IAI()                                             # choose integration algorithm (also AutoPTR() and PTR())
-    gloc = IntegralSolver(integrand, bz, alg; abstol=1e-3)  # construct a solver for gloc to within specified tolerance
-    gloc(ω=0.0)                                               # evaluate gloc at frequency ω=0.0
+    gloc = IntegralSolver(prob, alg; abstol=1e-3)           # construct a solver for gloc to within specified tolerance
+    gloc(ω=0.0)                                             # evaluate gloc at frequency ω=0.0
 
 !!! note "Assumptions"
     `AutoBZCore` assumes that all calculations occur in the
     reciprocal lattice basis, since that is the basis in which Wannier
     interpolants are most efficiently described. See [`SymmetricBZ`](@ref) for
-    details.
+    details. We also assume that the integrands are cheap to evaluate, which is why we
+    provide adaptive methods in the first place, so that return types can be determined at
+    runtime (and mechanisms are in place for compile time as well)
 """
 module AutoBZCore
 
-!isdefined(Base, :get_extension) && using Requires
-@static if !isdefined(Base, :get_extension)
-    function __init__()
-        @require HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" include("../ext/HDF5Ext.jl")
-    end
-end
-
 using LinearAlgebra: I, norm, det, checksquare
 
-using StaticArrays: SMatrix
+using StaticArrays: SVector, SMatrix, pushfirst
 using Reexport
-@reexport using Integrals
-@reexport using FourierSeriesEvaluators
 @reexport using AutoSymPTR
+@reexport using FourierSeriesEvaluators
 @reexport using IteratedIntegration
 
-import Integrals: IntegralProblem, __solvebp_call, SciMLBase.NullParameters, ReCallVJP, ZygoteVJP
-import AutoSymPTR: autosymptr, symptr, symptr_rule!, symptr_rule, ptr, ptr_rule!, ptrindex, alloc_autobuffer
-import IteratedIntegration: iterated_integrand, iterated_pre_eval, alloc_segbufs
-
+using IteratedIntegration: alloc_segbufs, nextrule, nextdim, RuleQuad.GaussKronrod, NestedGaussKronrod
+using HCubature: hcubature
 
 export SymmetricBZ, FullBZ, nsyms
-export AbstractSymRep, SymRep, UnknownRep, TrivialRep, FaithfulRep, LatticeRep
-include("brillouin_zone.jl")
+export AbstractSymRep, SymRep, UnknownRep, TrivialRep
+include("domains.jl")
 
-export AbstractAutoBZAlgorithm, IAI, PTR, AutoPTR, PTR_IAI, AutoPTR_IAI, TAI
+export IntegralProblem, solve
+export IAI, PTR, AutoPTR, PTR_IAI, AutoPTR_IAI, TAI, QuadGKJL, HCubatureJL
 include("algorithms.jl")
 
+export MixedParameters, paramzip, paramproduct
 export IntegralSolver, batchsolve
-include("solver.jl")
-
-export MixedParameters, AbstractAutoBZIntegrand, paramzip, paramproduct
-include("parameter_interface.jl")
-
 export Integrand
-include("integrand.jl")
-
-"""
-    NTHREADS_KSUM = fill(Threads.nthreads())
+include("interfaces.jl")
 
-This represents the number of threads to parallelize over for k-sums in PTR. To
-change the number of threads to 1 to completely bypass multi-threading, write
-`AutoBZCore.NTHREADS_KSUM[] = 1`.
-"""
-const NTHREADS_KSUM = fill(Threads.nthreads())
-
-export FourierIntegrand
-include("fourier_integration.jl")
+export FourierIntegrand, FourierValue
+include("rules.jl")
 
 
-end
+end