From f7057412fa453166a46b82e9c0365b93923df8af Mon Sep 17 00:00:00 2001 From: lxvm Date: Tue, 4 Jun 2024 14:27:49 -0400 Subject: [PATCH 01/17] initial commit --- src/AutoBZCore.jl | 57 ++- src/algorithms.jl | 685 ++++++----------------------------- src/algorithms_autosymptr.jl | 105 ++++++ src/algorithms_iterated.jl | 136 +++++++ src/algorithms_meta.jl | 298 +++++++++++++++ src/batch.jl | 77 ---- src/brillouin.jl | 64 ++++ src/domains.jl | 19 + src/dos_interfaces.jl | 48 ++- src/fourier.jl | 148 ++++---- src/inplace.jl | 15 - src/interfaces.jl | 310 ++++++++-------- src/parameters.jl | 6 - test/interface_tests.jl | 86 +++-- 14 files changed, 1053 insertions(+), 1001 deletions(-) create mode 100644 src/algorithms_autosymptr.jl create mode 100644 src/algorithms_iterated.jl create mode 100644 src/algorithms_meta.jl delete mode 100644 src/batch.jl delete mode 100644 src/inplace.jl diff --git a/src/AutoBZCore.jl b/src/AutoBZCore.jl index 45c167c..3025ff0 100644 --- a/src/AutoBZCore.jl +++ b/src/AutoBZCore.jl @@ -1,7 +1,7 @@ """ A package providing a common interface to integration algorithms intended for applications including Brillouin-zone integration and Wannier interpolation. Its design is influenced by -high-level libraries like Integrals.jl, and it makes use of Julia's multiple dispatch to +high-level libraries like Integrals.jl to implement the CommonSolve.jl interface, and it makes use of Julia's multiple dispatch to provide the same interface for integrands with optimized inplace, batched, and Fourier series evaluation. @@ -11,23 +11,24 @@ As a first example, we integrate sine over [0,1] as a function of its period. ``` julia> using AutoBZCore -julia> f = IntegralSolver((x,p) -> sin(p*x), 0, 1, QuadGKJL()); +julia> prob = IntegralProblem((x,p) -> sin(p*x), (0, 1), 0.3); -julia> f(0.3) # solves the integral of sin(p*x) over [0,1] with p=0.3 +julia> solve(prob, QuadGKJL()).value # solves the integral of sin(p*x) over [0,1] with p=0.3 0.14887836958131329 ``` -Notice that we construct an [`IntegralSolver`](@ref) object that we can evaluate at -different parameters with a function-like interface. For more examples, see the +Notice that we construct an [`IntegralProblem`](@ref) object that we can [`solve`](@ref) at +with a choice of algorithm. For more examples, see the documentation. ### Features Special integrand interfaces -- [`ParameterIntegrand`](@ref): allows user integrand to use keyword arguments -- [`InplaceIntegrand`](@ref): allows an integrand to write its result inplace to an array -- [`BatchIntegrand`](@ref): allows user-side parallelization on e.g. shared memory, +- [`IntegralFunction`](@ref): generic user integrand of the form `f(x, p)` +- [`InplaceIntegralFunction`](@ref): allows an integrand to write its result inplace to an array +- [`InplaceBatchIntegralFunction`](@ref): allows user-side parallelization on e.g. shared memory, distributed memory, or the gpu -- [`FourierIntegrand`](@ref): efficient evaluation of Fourier series for cubatures with +- [`CommonSolveIntegralFunction`](@ref): define an integrand that also solves a problem +- [`FourierIntegralFunction`](@ref): efficient evaluation of Fourier series for cubatures with hierachical grids Quadrature algorithms: @@ -52,45 +53,41 @@ using LinearAlgebra: I, norm, det, checksquare, isdiag, Diagonal, tr, diag, eige using StaticArrays: SVector, SMatrix, pushfirst, sacollect using FunctionWrappers: FunctionWrapper using ChunkSplitters: chunks, getchunk -using Reexport -@reexport using AutoSymPTR -@reexport using FourierSeriesEvaluators -@reexport using IteratedIntegration -@reexport using QuadGK -@reexport using HCubature - +using AutoSymPTR +using FourierSeriesEvaluators +using IteratedIntegration +using QuadGK: quadgk, quadgk!, BatchIntegrand +using HCubature: hcubature using FourierSeriesEvaluators: workspace_allocate, workspace_contract!, workspace_evaluate!, workspace_evaluate, period using IteratedIntegration: limit_iterate, interior_point -using HCubature: hcubature +using HCubature: hcubature, hquadrature +using CommonSolve: solve +import CommonSolve: init, solve! +export init, solve!, solve -export PuncturedInterval, HyperCube include("domains.jl") -export InplaceIntegrand -include("inplace.jl") - -export BatchIntegrand, NestedBatchIntegrand -include("batch.jl") - -# export IntegralProblem, solve, init, solve! # we don't export the SciML interface -export IntegralSolver, batchsolve +export IntegralFunction, InplaceIntegralFunction, InplaceBatchIntegralFunction +export CommonSolveIntegralFunction +export IntegralProblem include("interfaces.jl") export QuadGKJL, HCubatureJL, QuadratureFunction +include("algorithms.jl") export AuxQuadGKJL, ContQuadGKJL, MeroQuadGKJL +# include("algorithms_iterated.jl") export MonkhorstPack, AutoSymPTRJL +# include("algorithms_autosymptr.jl") export NestedQuad, AbsoluteEstimate, EvalCounter -include("algorithms.jl") +include("algorithms_meta.jl") export SymmetricBZ, nsyms export load_bz, FBZ, IBZ, InversionSymIBZ, CubicSymIBZ export AbstractSymRep, SymRep, UnknownRep, TrivialRep +export AutoBZProblem export IAI, PTR, AutoPTR, TAI, PTR_IAI, AutoPTR_IAI include("brillouin.jl") -export ParameterIntegrand, paramzip, paramproduct -include("parameters.jl") - export FourierIntegrand, FourierValue include("fourier.jl") diff --git a/src/algorithms.jl b/src/algorithms.jl index 720ee9a..ecc66af 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -1,9 +1,6 @@ # Methods an algorithm must define # - init_cacheval -# - do_solve - - -# Here we replicate the algorithms provided by SciML +# - do_integral """ QuadGKJL(; order = 7, norm = norm) @@ -37,64 +34,77 @@ function init_midpoint_scale(a::T, b::T) where {T} end end init_midpoint_scale(dom::PuncturedInterval) = init_midpoint_scale(endpoints(dom)...) -function init_segbuf(f, dom, p, norm) - x, s = init_midpoint_scale(dom) - u = x/oneunit(x) - TX = typeof(u) - fx_s = f(x, p) * s/oneunit(s) - TI = typeof(fx_s) - TE = typeof(norm(fx_s)) - return IteratedIntegration.alloc_segbuf(TX, TI, TE) -end -function init_segbuf(f::InplaceIntegrand, dom, p, norm) - x, s = init_midpoint_scale(dom) - u = x/oneunit(x) - TX = typeof(u) - TI = typeof(f.I *s/oneunit(s)) - fill!(f.I, zero(eltype(f.I))) - TE = typeof(norm(f.I)*s/oneunit(s)) - return IteratedIntegration.alloc_segbuf(TX, TI, TE) -end -function init_segbuf(f::BatchIntegrand, dom, p, norm) - x, s = init_midpoint_scale(dom) +function init_segbuf(prototype, segs, alg) + x, s = init_midpoint_scale(segs) u = x/oneunit(x) TX = typeof(u) - fx_s = zero(eltype(f.y)) * s/oneunit(s) # TODO BatchIntegrand(InplaceIntegrand) should depend on size of result + fx_s = prototype * s/oneunit(s) TI = typeof(fx_s) - TE = typeof(norm(fx_s)) + TE = typeof(alg.norm(fx_s)) return IteratedIntegration.alloc_segbuf(TX, TI, TE) end -function init_cacheval(f, dom, p, alg::QuadGKJL) - f isa NestedBatchIntegrand && throw(ArgumentError("QuadGK doesn't support nested batching")) - f isa BatchIntegrand && throw(ArgumentError("QuadGK doesn't support batched integrands")) - return init_segbuf(f, dom, p, alg.norm) -end -function do_solve(f::F, dom, p, alg::QuadGKJL, cacheval; - reltol = nothing, abstol = nothing, maxiters = typemax(Int)) where {F} - segs = segments(dom) +function init_cacheval(f::IntegralFunction, dom, p, alg::QuadGKJL; kws...) + segs = PuncturedInterval(dom) + prototype = get_prototype(f, get_prototype(segs), p) + return init_segbuf(prototype, segs, alg) +end +function init_cacheval(f::InplaceIntegralFunction, dom, p, alg::QuadGKJL; kws...) + segs = PuncturedInterval(dom) + prototype = get_prototype(f, get_prototype(segs), p) + return init_segbuf(prototype, segs, alg), similar(prototype) +end +function init_cacheval(f::InplaceBatchIntegralFunction, dom, p, alg::QuadGKJL; kws...) + segs = PuncturedInterval(dom) + pt = get_prototype(segs) + prototype = get_prototype(f, pt, p) + prototype isa AbstractVector || throw(ArgumentError("QuadGKJL only supports batch integrands with vector outputs")) + pts = zeros(typeof(pt), 2*alg.order+1) + upts = pts / pt + return init_segbuf(first(prototype), segs, alg), similar(prototype), pts, upts +end +function init_cacheval(f::CommonSolveIntegralFunction, dom, p, alg::QuadGKJL; kws...) + segs = PuncturedInterval(dom) + x = get_prototype(segs) + cache, integrand, prototype = _init_commonsolvefunction(f, dom, p; x) + return init_segbuf(prototype, segs, alg), cache, integrand +end +function do_integral(f, dom, p, alg::QuadGKJL, cacheval; + reltol = nothing, abstol = nothing, maxiters = typemax(Int)) # we need to strip units from the limits since infinity transformations change the units # of the limits, which can break the segbuf u = oneunit(eltype(dom)) - usegs = map(x -> x/u, segs) - if f isa InplaceIntegrand - g! = (y, x) -> f.f!(y, u*x, p) - result = f.I / u - val, err = quadgk!(g!, result, usegs..., maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(f.I .= u .* val, u*err, true, -1) - else - g = x -> f(u*x, p) - val, err = quadgk(g, usegs..., maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(u*val, u*err, true, -1) - end + usegs = map(x -> x/u, dom) + atol = isnothing(abstol) ? abstol : abstol/u + val, err = call_quadgk(f, p, u, usegs, cacheval; maxevals = maxiters, rtol = reltol, atol, order = alg.order, norm = alg.norm) + value = u*val + retcode = err < max(something(atol, zero(err)), alg.norm(val)*something(reltol, isnothing(atol) ? sqrt(eps(one(eltype(usegs)))) : 0)) ? Success : Failure + stats = (; error=u*err) + return IntegralSolution(value, retcode, stats) +end +function call_quadgk(f::IntegralFunction, p, u, usegs, cacheval; kws...) + quadgk(x -> f.f(u*x, p), usegs...; kws..., segbuf=cacheval) +end +function call_quadgk(f::InplaceIntegralFunction, p, u, usegs, cacheval; kws...) + # TODO allocate everything in the QuadGK.InplaceIntegrand in the cacheval + quadgk!((y, x) -> f.f!(y, u*x, p), cacheval[2], usegs...; kws..., segbuf=cacheval[1]) +end +function call_quadgk(f::InplaceBatchIntegralFunction, p, u, usegs, cacheval; kws...) + pts = cacheval[3] + g = BatchIntegrand((y, x) -> f.f!(y, resize!(pts, length(x)) .= u .* x, p), cacheval[2], cacheval[4]; max_batch=f.max_batch) + quadgk(g, usegs...; kws..., segbuf=cacheval[1]) +end +function call_quadgk(f::CommonSolveIntegralFunction, p, u, usegs, cacheval; kws...) + # cache = cacheval[2] could call do_solve!(cache, f, x, p) to fully specialize + integrand = cacheval[3] + quadgk(x -> integrand(x, p), usegs...; kws..., segbuf=cacheval[1]) end + """ HCubatureJL(; norm=norm, initdiv=1) -A copy of `HCubatureJL` from Integrals.jl. +Multi-dimensional h-adaptive cubature from HCubature.jl. """ struct HCubatureJL{N} <: IntegralAlgorithm norm::N @@ -102,25 +112,30 @@ struct HCubatureJL{N} <: IntegralAlgorithm end HCubatureJL(; norm=norm, initdiv=1) = HCubatureJL(norm, initdiv) -function init_cacheval(f, dom, p, ::HCubatureJL) - f isa NestedBatchIntegrand && throw(ArgumentError("HCubatureJL doesn't support nested batching")) - f isa BatchIntegrand && throw(ArgumentError("HCubatureJL doesn't support batching")) - return Some(nothing) +function init_cacheval(f::IntegralFunction, dom, p, ::HCubatureJL; kws...) + # TODO utilize hcubature_buffer + return +end +function init_cacheval(f::CommonSolveIntegralFunction, dom, p, ::HCubatureJL; kws...) + cache, integrand, = _init_commonsolvefunction(f, dom, p) + return cache, integrand end -function do_solve(f, dom, p, alg::HCubatureJL, cacheval; - reltol = 0, abstol = 0, maxiters = typemax(Int)) - - g = if f isa InplaceIntegrand - fx = f.I / oneunit(eltype(dom)) - x -> (f.f!(fx, x, p); fx*one(eltype(dom))) - else - x -> f(x, p) - end +function do_integral(f, dom, p, alg::HCubatureJL, cacheval; reltol = 0, abstol = 0, maxiters = typemax(Int)) a, b = endpoints(dom) + g = hcubature_integrand(f, p, a, b, cacheval) routine = a isa Number ? hquadrature : hcubature - val, err = routine(g, a, b; norm = alg.norm, initdiv = alg.initdiv, atol=abstol, rtol=reltol, maxevals=maxiters) - return IntegralSolution(val, err, true, -1) + value, error = routine(g, a, b; norm = alg.norm, initdiv = alg.initdiv, atol=abstol, rtol=reltol, maxevals=maxiters) + retcode = error < max(something(abstol, zero(error)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(eltype(a))) : abstol)) ? Success : Failure + stats = (; error) + return IntegralSolution(value, retcode, stats) +end +function hcubature_integrand(f::IntegralFunction, p, a, b, cacheval) + x -> f.f(x, p) +end +function hcubature_integrand(f::CommonSolveIntegralFunction, p, a, b, cacheval) + integrand = cacheval[2] + return x -> integrand(x, p) end """ @@ -159,27 +174,36 @@ struct QuadratureFunction{F} <: IntegralAlgorithm nthreads::Int end QuadratureFunction(; fun=trapz, npt=50, nthreads=1) = QuadratureFunction(fun, npt, nthreads) -init_buffer(f, len) = nothing -init_buffer(f::BatchIntegrand, len) = Vector{eltype(f.y)}(undef, len) -function init_cacheval(f, dom::PuncturedInterval, p, alg::QuadratureFunction) - f isa NestedBatchIntegrand && throw(ArgumentError("QuadratureFunction doesn't support nested batching")) +function init_cacheval(f::IntegralFunction, dom, p, alg::QuadratureFunction; kws...) + x, w = alg.fun(alg.npt) + return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing) +end +function init_cacheval(f::InplaceIntegralFunction, dom, p, alg::QuadratureFunction; kws...) + x, w = alg.fun(alg.npt) + proto = get_prototype(f, first(x), p) + y = similar(proto) + ytmp = similar(proto) + I = y * first(w) + Itmp = similar(I) + return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing, I, Itmp, y, ytmp) +end +function init_cacheval(f::InplaceBatchIntegralFunction, dom, p, alg::QuadratureFunction; kws...) buf = init_buffer(f, alg.nthreads) x, w = alg.fun(alg.npt) - return (rule=[(w,x) for (w,x) in zip(w,x)], buffer=buf) + proto=get_prototype(f, first(x), p) + return (rule=[(w,x) for (w,x) in zip(w,x)], buffer=similar(proto, len), y=similar(proto, len), x=Vector{float(eltype(dom))}(undef, len)) +end +function init_cacheval(f::CommonSolveIntegralFunction, dom, p, alg::QuadratureFunction; kws...) + x, w = alg.fun(alg.npt) + cache, integrand, = _init_commonsolvefunction(f, dom, p) + return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing, cache, integrand) end -function do_solve(f, dom, p, alg::QuadratureFunction, cacheval; +function do_integral(f, dom, p, alg::QuadratureFunction, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) rule = cacheval.rule; buffer=cacheval.buffer segs = segments(dom) - g = if f isa BatchIntegrand - xx = eltype(f.x) === Nothing ? typeof((segs[1]+segs[end])/2)[] : f.x - AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), f.y, xx, max_batch=f.max_batch) - elseif f isa InplaceIntegrand - AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), f.I) - else - x -> f(x, p) - end + g = autosymptr_integrand(f, p, segs, cacheval) A = sum(1:length(segs)-1) do i a, b = segs[i], segs[i+1] s = (b-a)/2 @@ -187,505 +211,18 @@ function do_solve(f, dom, p, alg::QuadratureFunction, cacheval; return AutoSymPTR.quadsum(arule, g, s, buffer) end - return IntegralSolution(A, nothing, true, -1) -end - -# Here we put the quadrature algorithms from IteratedIntegration - -""" - AuxQuadGKJL(; order = 7, norm = norm) - -Generalization of the QuadGKJL provided by Integrals.jl that allows for `AuxValue`d -integrands for auxiliary integration and multi-threaded evaluation with the `batch` argument -to `IntegralProblem` -""" -struct AuxQuadGKJL{F} <: IntegralAlgorithm - order::Int - norm::F -end -function AuxQuadGKJL(; order = 7, norm = norm) - return AuxQuadGKJL(order, norm) -end - -function init_cacheval(f, dom, p, alg::AuxQuadGKJL) - f isa NestedBatchIntegrand && throw(ArgumentError("AuxQuadGKJL doesn't support nested batching")) - return init_segbuf(f, dom, p, alg.norm) -end - -function do_solve(f, dom, p, alg::AuxQuadGKJL, cacheval; - reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - - segs = segments(dom) - u = oneunit(eltype(dom)) - usegs = map(x -> x/u, segs) - if f isa InplaceIntegrand - g! = (y, x) -> f.f!(y, u*x, p) - result = f.I / u - val, err = auxquadgk!(g!, result, usegs, maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(f.I .= u .* val, u*err, true, -1) - elseif f isa BatchIntegrand - xx = eltype(f.x) === Nothing ? typeof((segs[1]+segs[end])/2)[] : f.x - g_ = (y, x) -> (resize!(xx, length(x)); f.f!(y, xx .= u .* x, p)) - g = IteratedIntegration.AuxQuadGK.BatchIntegrand(g_, f.y, xx/u, max_batch=f.max_batch) - val, err = auxquadgk(g, usegs, maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(u*val, u*err, true, -1) - else - g = x -> f(u*x, p) - val, err = auxquadgk(g, usegs, maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(u*val, u*err, true, -1) - end -end - -""" - ContQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.ContQuadGK.NewtonDeflation()) - -A 1d contour deformation quadrature scheme for scalar, complex-valued integrands. It -defaults to regular `quadgk` behavior on the real axis, but if it finds a root of 1/f -nearby, in the sense of Bernstein ellipse for the standard segment `[-1,1]` with semiaxes -`cosh(rho)` and `sinh(rho)`, on either the upper/lower half planes, then it dents the -contour away from the presumable pole. -""" -struct ContQuadGKJL{F,M} <: IntegralAlgorithm - order::Int - norm::F - rho::Float64 - rootmeth::M -end -function ContQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.ContQuadGK.NewtonDeflation()) - return ContQuadGKJL(order, norm, rho, rootmeth) -end - -function init_cacheval(f, dom, p, alg::ContQuadGKJL) - f isa NestedBatchIntegrand && throw(ArgumentError("ContQuadGK doesn't support nested batching")) - f isa BatchIntegrand && throw(ArgumentError("ContQuadGK doesn't support batching")) - f isa InplaceIntegrand && throw(ArgumentError("ContQuadGK doesn't support inplace integrands")) - - a, b = endpoints(dom) - x, s = (a+b)/2, (b-a)/2 - TX = typeof(x) - fx_s = one(ComplexF64) * s # currently the integrand is forcibly written to a ComplexF64 buffer - TI = typeof(fx_s) - TE = typeof(alg.norm(fx_s)) - r_segbuf = IteratedIntegration.ContQuadGK.PoleSegment{TX,TI,TE}[] - fc_s = f(complex(x), p) * complex(s) # the regular evalrule is used on complex segments - TCX = typeof(complex(x)) - TCI = typeof(fc_s) - TCE = typeof(alg.norm(fc_s)) - c_segbuf = IteratedIntegration.ContQuadGK.Segment{TCX,TCI,TCE}[] - return (r=r_segbuf, c=c_segbuf) -end - -function do_solve(f, dom, p, alg::ContQuadGKJL, cacheval; - reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - - segs = segments(dom) - g = x -> f(x, p) - val, err = contquadgk(g, segs, maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, - rtol = reltol, atol = abstol, order = alg.order, norm = alg.norm, r_segbuf=cacheval.r, c_segbuf=cacheval.c) - return IntegralSolution(val, err, true, -1) -end - -""" - MeroQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.MeroQuadGK.NewtonDeflation()) - -A 1d pole subtraction quadrature scheme for scalar, complex-valued integrands that are -meromorphic. It defaults to regular `quadgk` behavior on the real axis, but if it finds -nearby roots of 1/f, in the sense of Bernstein ellipse for the standard segment `[-1,1]` -with semiaxes `cosh(rho)` and `sinh(rho)`, it attempts pole subtraction on that segment. -""" -struct MeroQuadGKJL{F,M} <: IntegralAlgorithm - order::Int - norm::F - rho::Float64 - rootmeth::M -end -function MeroQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.MeroQuadGK.NewtonDeflation()) - return MeroQuadGKJL(order, norm, rho, rootmeth) -end - -function init_cacheval(f, dom, p, alg::MeroQuadGKJL) - f isa NestedBatchIntegrand && throw(ArgumentError("MeroQuadGK doesn't support nested batching")) - f isa BatchIntegrand && throw(ArgumentError("MeroQuadGK doesn't support batching")) - f isa InplaceIntegrand && throw(ArgumentError("MeroQuadGK doesn't support inplace integrands")) - a, b = endpoints(dom) - x, s = (a + b)/2, (b-a)/2 - fx_s = one(ComplexF64) * s # ignore the actual integrand since it is written to CF64 array - err = alg.norm(fx_s) - return IteratedIntegration.alloc_segbuf(typeof(x), typeof(fx_s), typeof(err)) -end - -function do_solve(f, dom, p, alg::MeroQuadGKJL, cacheval; - reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - - segs = segments(dom) - g = x -> f(x, p) - val, err = meroquadgk(g, segs, maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, - rtol = reltol, atol = abstol, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(val, err, true, -1) -end - -# Algorithms from AutoSymPTR.jl - -""" - MonkhorstPack(; npt=50, syms=nothing, nthreads=1) - -Periodic trapezoidal rule with a fixed number of k-points per dimension, `npt`, -using the `PTR` rule from [AutoSymPTR.jl](https://github.com/lxvm/AutoSymPTR.jl). -`nthreads` sets the numbers of threads used to parallelize the quadrature only when the -integrand is a [`BatchIntegrand`](@ref), in which case the user must parallelize the -integrand evaluations. For no threading set `nthreads=1`. -**The caller should check that the integral is converged w.r.t. `npt`**. -""" -struct MonkhorstPack{S} <: IntegralAlgorithm - npt::Int - syms::S - nthreads::Int -end -MonkhorstPack(; npt=50, syms=nothing, nthreads=1) = MonkhorstPack(npt, syms, nthreads) -function init_rule(dom::Basis, alg::MonkhorstPack) - # rule = AutoSymPTR.MonkhorstPackRule(alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) - # return rule(eltype(dom), Val(ndims(dom))) - if alg.syms === nothing - return AutoSymPTR.PTR(eltype(dom), Val(ndims(dom)), alg.npt) - else - return AutoSymPTR.MonkhorstPack(eltype(dom), Val(ndims(dom)), alg.npt, alg.syms) - end -end - -rule_type(::AutoSymPTR.PTR{N,T}) where {N,T} = SVector{N,T} -rule_type(::AutoSymPTR.MonkhorstPack{N,T}) where {N,T} = SVector{N,T} - -function init_cacheval(f, dom::Basis, p, alg::MonkhorstPack) - f isa NestedBatchIntegrand && throw(ArgumentError("MonkhorstPack doesn't support nested batching")) - rule = init_rule(dom, alg) - buf = init_buffer(f, alg.nthreads) - return (rule=rule, buffer=buf) -end - -function do_solve(f, dom, p, alg::MonkhorstPack, cacheval; - reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - g = if f isa BatchIntegrand - xx = eltype(f.x) === Nothing ? typeof(dom*zero(rule_type(cacheval.rule)))[] : f.x - AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), f.y, xx, max_batch=f.max_batch) - elseif f isa InplaceIntegrand - AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), f.I) - else - x -> f(x, p) - end - I = cacheval.rule(g, dom, cacheval.buffer) - return IntegralSolution(I, nothing, true, -1) -end - -""" - AutoSymPTRJL(; norm=norm, a=1.0, nmin=50, nmax=1000, n₀=6, Δn=log(10), keepmost=2, nthreads=1) - -Periodic trapezoidal rule with automatic convergence to tolerances passed to the -solver with respect to `norm` using the routine `autosymptr` from -[AutoSymPTR.jl](https://github.com/lxvm/AutoSymPTR.jl). -`nthreads` sets the numbers of threads used to parallelize the quadrature only when the -integrand is a [`BatchIntegrand`](@ref), in which case the user must parallelize the -integrand evaluations. For no threading set `nthreads=1`. -**This algorithm is the most efficient for smooth integrands**. -""" -struct AutoSymPTRJL{F,S} <: IntegralAlgorithm - norm::F - a::Float64 - nmin::Int - nmax::Int - n₀::Float64 - Δn::Float64 - keepmost::Int - syms::S - nthreads::Int -end -function AutoSymPTRJL(; norm=norm, a=1.0, nmin=50, nmax=1000, n₀=6.0, Δn=log(10), keepmost=2, syms=nothing, nthreads=1) - return AutoSymPTRJL(norm, a, nmin, nmax, n₀, Δn, keepmost, syms, nthreads) -end -function init_rule(dom::Basis, alg::AutoSymPTRJL) - return AutoSymPTR.MonkhorstPackRule(alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) -end -function init_cacheval(f, dom::Basis, p, alg::AutoSymPTRJL) - f isa NestedBatchIntegrand && throw(ArgumentError("AutoSymPTRJL doesn't support nested batching")) - rule = init_rule(dom, alg) - cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) - buffer = init_buffer(f, alg.nthreads) - return (rule=rule, cache=cache, buffer=buffer) -end - -function do_solve(f, dom, p, alg::AutoSymPTRJL, cacheval; - reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - - g = if f isa BatchIntegrand - xx = eltype(f.x) === Nothing ? typeof(dom*zero(rule_type(cacheval.cache[1])))[] : f.x - AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), f.y, xx, max_batch=f.max_batch) - elseif f isa InplaceIntegrand - AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), f.I) - else - x -> f(x, p) - end - val, err = autosymptr(g, dom; syms = alg.syms, rule = cacheval.rule, cache = cacheval.cache, keepmost = alg.keepmost, - abstol = abstol, reltol = reltol, maxevals = maxiters, norm=alg.norm, buffer=cacheval.buffer) - return IntegralSolution(val, err, true, -1) -end - -# Meta-algorithms - -""" - NestedQuad(alg::IntegralAlgorithm) - NestedQuad(algs::IntegralAlgorithm...) - -Nested integration by repeating one quadrature algorithm or composing a list of algorithms. -The domain of integration must be an `AbstractIteratedLimits` from the -IteratedIntegration.jl package. Analogous to `nested_quad` from IteratedIntegration.jl. -The integrand should expect `SVector` inputs. Do not use this for very high-dimensional -integrals, since the compilation time scales very poorly with respect to dimensionality. -In order to improve the compilation time, FunctionWrappers.jl is used to enforce type -stability of the integrand, so you should always pick the widest integration limit type so -that inference works properly. For example, if [`ContQuadGKJL`](@ref) is used as an -algorithm in the nested scheme, then the limits of integration should be made complex. -""" -struct NestedQuad{T} <: IntegralAlgorithm - algs::T - NestedQuad(alg::IntegralAlgorithm) = new{typeof(alg)}(alg) - NestedQuad(algs::Tuple{Vararg{IntegralAlgorithm}}) = new{typeof(algs)}(algs) -end -NestedQuad(algs::IntegralAlgorithm...) = NestedQuad(algs) - -# this function helps create a tree of the cachevals used by each quadrature -function nested_cacheval(f::F, p::P, algs, segs, lims, state, x, xs...) where {F,P} - dom = PuncturedInterval(segs) - a, b = segs[1], segs[2] - dim = ndims(lims) - alg = algs[dim] - mid = (a+b)/2 # sample point that should be safe to evaluate - next = limit_iterate(lims, state, mid) # see what the next limit gives - if xs isa Tuple{} # the next integral takes us to the inner integral - # base case test integrand of the inner integral - # we need to pass dummy integrands to all the outer integrals so that they can build - # their caches with the right types - if f isa BatchIntegrand || f isa NestedBatchIntegrand - # Batch integrate the inner integral only - cacheval = init_cacheval(BatchIntegrand(nothing, f.y, f.x, max_batch=f.max_batch), dom, p, alg) - return (nothing, cacheval, oneunit(eltype(f.y))*mid) - elseif f isa InplaceIntegrand - # Inplace integrate through the whole nest structure - fxi = f.I*mid/oneunit(prod(next)) - cacheval = init_cacheval(InplaceIntegrand(nothing, fxi), dom, p, alg) - return (nothing, cacheval, fxi) - else - fx = f(next,p) - cacheval = init_cacheval((x, p) -> fx, dom, p, alg) - return (nothing, cacheval, fx*mid) - end - elseif f isa NestedBatchIntegrand - algs_ = algs[1:dim-1] - # numbered names to avoid type instabilities (we are not using dispatch, but the - # compiler's optimization for the recursive function's argument types) - nest0 = nested_cacheval(f.f[1], p, algs_, next..., x, xs[1:dim-2]...) - cacheval = init_cacheval(BatchIntegrand(nothing, f.y, f.x, max_batch=f.max_batch), dom, p, alg) - return (ntuple(n -> n == 1 ? nest0 : deepcopy(nest0), Val(length(f.f))), cacheval, nest0[3]*mid) - else - algs_ = algs[1:dim-1] - nest1 = nested_cacheval(f, p, algs_, next..., x, xs[1:dim-2]...) - h = nest1[3] - hx = h*mid - # units may change for outer integral - if f isa InplaceIntegrand - cacheval = init_cacheval(InplaceIntegrand(nothing, hx), dom, p, alg) - return (nest1, cacheval, hx) - else - cacheval = init_cacheval((x, p) -> h, dom, p, alg) - return (nest1, cacheval, hx) - end - end -end -function init_cacheval(f, dom::AbstractIteratedLimits, p, alg::NestedQuad) - algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(dom))) : alg.algs - return nested_cacheval(f, p, algs, limit_iterate(dom)..., interior_point(dom)...) -end - -function init_nest(f::F, fxx, dom, p,lims, state, algs, cacheval; kws_...) where {F} - kws = NamedTuple(kws_) - xx = float(oneunit(eltype(dom))) - FX = typeof(fxx/xx) - TX = typeof(xx) - TP = Tuple{typeof(p),typeof(lims),typeof(state)} - if algs isa Tuple{} # inner integral - if f isa BatchIntegrand - return f - elseif f isa NestedBatchIntegrand - nchunk = length(f.f) - return BatchIntegrand(FunctionWrapper{Nothing,Tuple{typeof(f.y),typeof(f.x),TP}}() do y, x, (p, lims, state) - Threads.@threads for ichunk in 1:min(nchunk, length(x)) - for (i, j) in zip(getchunk(x, ichunk, nchunk, :scatter), getchunk(y, ichunk, nchunk, :scatter)) - xi = x[i] - y[j] = f.f[ichunk](limit_iterate(lims, state, xi), p) - end - end - return nothing - end, f.y, f.x, max_batch=f.max_batch) - elseif f isa InplaceIntegrand - return InplaceIntegrand(FunctionWrapper{Nothing,Tuple{FX,TX,TP}}() do y, x, (p, lims, state) - f.f!(y, limit_iterate(lims, state, x), p) - return nothing - end, f.I) - else - return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) - return f(limit_iterate(lims, state, x), p) - end - end - else - if f isa InplaceIntegrand - return InplaceIntegrand(FunctionWrapper{Nothing,Tuple{FX,TX,TP}}() do y, x, (p, lims, state) - segs, lims_, state_ = limit_iterate(lims, state, x) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - do_solve(InplaceIntegrand(f.f!, y), lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval; kwargs...) - return nothing - end, f.I) - elseif f isa NestedBatchIntegrand - nchunks = length(f.f) - return BatchIntegrand(FunctionWrapper{Nothing,Tuple{typeof(f.y),typeof(f.x),TP}}() do y, x, (p, lims, state) - Threads.@threads for ichunk in 1:min(nchunks, length(x)) - for (i, j) in zip(getchunk(x, ichunk, nchunks, :scatter), getchunk(y, ichunk, nchunks, :scatter)) - xi = x[i] - segs, lims_, state_ = limit_iterate(lims, state, xi) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - y[j] = do_solve(f.f[ichunk], lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval[ichunk]; kwargs...).u - end - end - return nothing - end, f.y, f.x, max_batch=f.max_batch) - else - return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) - segs, lims_, state_ = limit_iterate(lims, state, x) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - sol = do_solve(f, lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval; kwargs...) - return sol.u - end - end - end -end - -struct NestState{P,G,S} - p::P - segs::G - state::S -end - -function do_solve(f::F, lims::AbstractIteratedLimits, p_, alg::NestedQuad, cacheval; kws...) where {F} - g, p, segs, state = if p_ isa NestState - gg = if f isa NestedBatchIntegrand - fx = eltype(f.x) === Nothing ? float(eltype(p_.segs))[] : f.x - NestedBatchIntegrand(f.f, f.y, fx, max_batch=f.max_batch) - else - f - end - gg, p_.p, p_.segs, p_.state - else - seg, lim, sta = limit_iterate(lims) - gg = if f isa BatchIntegrand - fx = eltype(f.x) === Nothing ? typeof(interior_point(lims))[] : f.x - BatchIntegrand(f.y, similar(f.x, eltype(eltype(fx))), max_batch=f.max_batch) do y, xs, (p, lims, state) - resize!(fx, length(xs)) - f.f!(y, map!(x -> limit_iterate(lims, state, x), fx, xs), p) - end - elseif f isa NestedBatchIntegrand - # this should be done recursively at the outermost level, but it is lazy. - fx = eltype(f.x) === Nothing ? float(eltype(seg))[] : f.x - NestedBatchIntegrand(f.f, f.y, fx, max_batch=f.max_batch) - else - f - end - gg, p_, seg, sta - end - dom = PuncturedInterval(segs) - dim = ndims(lims) # constant propagation :) - algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(dim)) : alg.algs - nest = init_nest(g, cacheval[3], dom, p, lims, state, algs[1:dim-1], cacheval[1]; kws...) - return do_solve(nest, dom, (p, lims, state), algs[dim], cacheval[2]; kws...) -end - -""" - AbsoluteEstimate(est_alg, abs_alg; kws...) - -Most algorithms are efficient when using absolute error tolerances, but how do you know the -size of the integral? One option is to estimate it using second algorithm. - -A multi-algorithm to estimate an integral using an `est_alg` to generate a rough estimate of -the integral that is combined with a user's relative tolerance to re-calculate the integral -to higher accuracy using the `abs_alg`. The keywords passed to the algorithm may include -`reltol`, `abstol` and `maxiters` and are given to the `est_alg` solver. They should limit -the amount of work of `est_alg` so as to only generate an order-of-magnitude estimate of the -integral. The tolerances passed to `abs_alg` are `abstol=max(abstol,reltol*norm(I))` and -`reltol=0`. -""" -struct AbsoluteEstimate{E<:IntegralAlgorithm,A<:IntegralAlgorithm,F,K<:NamedTuple} <: IntegralAlgorithm - est_alg::E - abs_alg::A - norm::F - kws::K + return IntegralSolution(A, Success, (; numevals = length(cacheval.rule)*(length(segs)-1))) end -function AbsoluteEstimate(est_alg, abs_alg; norm=norm, kwargs...) - kws = NamedTuple(kwargs) - checkkwargs(kws) - return AbsoluteEstimate(est_alg, abs_alg, norm, kws) +function autosymptr_integrand(f::IntegralFunction, p, segs, cacheval) + x -> f.f(x, p) end - -function init_cacheval(f, dom, p, alg::AbsoluteEstimate) - return (est=init_cacheval(f, dom, p, alg.est_alg), - abs=init_cacheval(f, dom, p, alg.abs_alg)) +function autosymptr_integrand(f::InplaceIntegralFunction, p, segs, cacheval) + AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), cacheval.I, cacheval.Itmp, cacheval.y, cacheval.ytmp) end - -function do_solve(f, dom, p, alg::AbsoluteEstimate, cacheval; - abstol=nothing, reltol=nothing, maxiters=typemax(Int)) - sol = do_solve(f, dom, p, alg.est_alg, cacheval.est; alg.kws...) - val = alg.norm(sol.u) # has same units as sol - rtol = reltol === nothing ? sqrt(eps(one(val))) : reltol # use the precision of the solution to set the default relative tolerance - atol = max(abstol === nothing ? zero(val) : abstol, rtol*val) - return do_solve(f, dom, p, alg.abs_alg, cacheval.abs; - abstol=atol, reltol=zero(rtol), maxiters=maxiters) +function autosymptr_integrand(f::InplaceBatchIntegralFunction, p, segs, cacheval) + AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), cacheval.y, cacheval.x, max_batch=f.max_batch) end - - -""" - EvalCounter(::IntegralAlgorithm) - -An algorithm which counts the evaluations used by another algorithm. -The count is stored in the `sol.numevals` field. -""" -struct EvalCounter{T<:IntegralAlgorithm} <: IntegralAlgorithm - alg::T -end - -function init_cacheval(f, dom, p, alg::EvalCounter) - return init_cacheval(f, dom, p, alg.alg) -end - -function do_solve(f, dom, p, alg::EvalCounter, cacheval; kws...) - if f isa InplaceIntegrand - ni::Int = 0 - gi = (y, x, p) -> (ni += 1; f.f!(y, x, p)) - soli = do_solve(InplaceIntegrand(gi, f.I), dom, p, alg.alg, cacheval; kws...) - return IntegralSolution(soli.u, soli.resid, soli.retcode, ni) - elseif f isa BatchIntegrand - nb::Int = 0 - gb = (y, x, p) -> (nb += length(x); f.f!(y, x, p)) - solb = do_solve(BatchIntegrand(gb, f.y, f.x, max_batch=f.max_batch), dom, p, alg.alg, cacheval; kws...) - return IntegralSolution(solb.u, solb.resid, solb.retcode, nb) - elseif f isa NestedBatchIntegrand - # TODO allocate a bunch of accumulators associated with the leaves of the nested - # integrand or rewrap the algorithms in NestedQuad - error("NestedBatchIntegrand not yet supported with EvalCounter") - else - n::Int = 0 - g = (x, p) -> (n += 1; f(x, p)) # we need let to prevent Core.Box around the captured variable - sol = do_solve(g, dom, p, alg.alg, cacheval; kws...) - return IntegralSolution(sol.u, sol.resid, sol.retcode, n) - end +function autosymptr_integrand(f::CommonSolveIntegralFunction, p, segs, cacheval) + integrand = cacheval.integrand + return x -> integrand(x, p) end diff --git a/src/algorithms_autosymptr.jl b/src/algorithms_autosymptr.jl new file mode 100644 index 0000000..59fe2f9 --- /dev/null +++ b/src/algorithms_autosymptr.jl @@ -0,0 +1,105 @@ +# We could move these into an extension, although QuadratureFunction also uses AutoSymPTR.jl +# for evaluation + + +""" + MonkhorstPack(; npt=50, syms=nothing, nthreads=1) + +Periodic trapezoidal rule with a fixed number of k-points per dimension, `npt`, +using the `PTR` rule from [AutoSymPTR.jl](https://github.com/lxvm/AutoSymPTR.jl). +`nthreads` sets the numbers of threads used to parallelize the quadrature only when the +integrand is a [`BatchIntegrand`](@ref), in which case the user must parallelize the +integrand evaluations. For no threading set `nthreads=1`. +**The caller should check that the integral is converged w.r.t. `npt`**. +""" +struct MonkhorstPack{S} <: IntegralAlgorithm + npt::Int + syms::S + nthreads::Int +end +MonkhorstPack(; npt=50, syms=nothing, nthreads=1) = MonkhorstPack(npt, syms, nthreads) +function init_rule(dom::Basis, alg::MonkhorstPack) + # rule = AutoSymPTR.MonkhorstPackRule(alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) + # return rule(eltype(dom), Val(ndims(dom))) + if alg.syms === nothing + return AutoSymPTR.PTR(eltype(dom), Val(ndims(dom)), alg.npt) + else + return AutoSymPTR.MonkhorstPack(eltype(dom), Val(ndims(dom)), alg.npt, alg.syms) + end +end + +rule_type(::AutoSymPTR.PTR{N,T}) where {N,T} = SVector{N,T} +rule_type(::AutoSymPTR.MonkhorstPack{N,T}) where {N,T} = SVector{N,T} + +function init_cacheval(f, dom::Basis, p, alg::MonkhorstPack) + f isa NestedBatchIntegrand && throw(ArgumentError("MonkhorstPack doesn't support nested batching")) + rule = init_rule(dom, alg) + buf = init_buffer(f, alg.nthreads) + return (rule=rule, buffer=buf) +end + +function do_solve(f, dom, p, alg::MonkhorstPack, cacheval; + reltol = nothing, abstol = nothing, maxiters = typemax(Int)) + g = if f isa BatchIntegrand + xx = eltype(f.x) === Nothing ? typeof(dom*zero(rule_type(cacheval.rule)))[] : f.x + AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), f.y, xx, max_batch=f.max_batch) + elseif f isa InplaceIntegrand + AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), f.I) + else + x -> f(x, p) + end + I = cacheval.rule(g, dom, cacheval.buffer) + return IntegralSolution(I, nothing, true, -1) +end + +""" + AutoSymPTRJL(; norm=norm, a=1.0, nmin=50, nmax=1000, n₀=6, Δn=log(10), keepmost=2, nthreads=1) + +Periodic trapezoidal rule with automatic convergence to tolerances passed to the +solver with respect to `norm` using the routine `autosymptr` from +[AutoSymPTR.jl](https://github.com/lxvm/AutoSymPTR.jl). +`nthreads` sets the numbers of threads used to parallelize the quadrature only when the +integrand is a [`BatchIntegrand`](@ref), in which case the user must parallelize the +integrand evaluations. For no threading set `nthreads=1`. +**This algorithm is the most efficient for smooth integrands**. +""" +struct AutoSymPTRJL{F,S} <: IntegralAlgorithm + norm::F + a::Float64 + nmin::Int + nmax::Int + n₀::Float64 + Δn::Float64 + keepmost::Int + syms::S + nthreads::Int +end +function AutoSymPTRJL(; norm=norm, a=1.0, nmin=50, nmax=1000, n₀=6.0, Δn=log(10), keepmost=2, syms=nothing, nthreads=1) + return AutoSymPTRJL(norm, a, nmin, nmax, n₀, Δn, keepmost, syms, nthreads) +end +function init_rule(dom::Basis, alg::AutoSymPTRJL) + return AutoSymPTR.MonkhorstPackRule(alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) +end +function init_cacheval(f, dom::Basis, p, alg::AutoSymPTRJL) + f isa NestedBatchIntegrand && throw(ArgumentError("AutoSymPTRJL doesn't support nested batching")) + rule = init_rule(dom, alg) + cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) + buffer = init_buffer(f, alg.nthreads) + return (rule=rule, cache=cache, buffer=buffer) +end + +function do_solve(f, dom, p, alg::AutoSymPTRJL, cacheval; + reltol = nothing, abstol = nothing, maxiters = typemax(Int)) + + g = if f isa BatchIntegrand + xx = eltype(f.x) === Nothing ? typeof(dom*zero(rule_type(cacheval.cache[1])))[] : f.x + AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), f.y, xx, max_batch=f.max_batch) + elseif f isa InplaceIntegrand + AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), f.I) + else + x -> f(x, p) + end + val, err = autosymptr(g, dom; syms = alg.syms, rule = cacheval.rule, cache = cacheval.cache, keepmost = alg.keepmost, + abstol = abstol, reltol = reltol, maxevals = maxiters, norm=alg.norm, buffer=cacheval.buffer) + return IntegralSolution(val, err, true, -1) +end diff --git a/src/algorithms_iterated.jl b/src/algorithms_iterated.jl new file mode 100644 index 0000000..7ca6d5e --- /dev/null +++ b/src/algorithms_iterated.jl @@ -0,0 +1,136 @@ +# TODO move these to an extension if SymmetricBZ can be simplified + +""" + AuxQuadGKJL(; order = 7, norm = norm) + +Generalization of the QuadGKJL provided by Integrals.jl that allows for `AuxValue`d +integrands for auxiliary integration and multi-threaded evaluation with the `batch` argument +to `IntegralProblem` +""" +struct AuxQuadGKJL{F} <: IntegralAlgorithm + order::Int + norm::F +end +function AuxQuadGKJL(; order = 7, norm = norm) + return AuxQuadGKJL(order, norm) +end + +function init_cacheval(f, dom, p, alg::AuxQuadGKJL) + f isa NestedBatchIntegrand && throw(ArgumentError("AuxQuadGKJL doesn't support nested batching")) + return init_segbuf(f, dom, p, alg.norm) +end + +function do_solve(f, dom, p, alg::AuxQuadGKJL, cacheval; + reltol = nothing, abstol = nothing, maxiters = typemax(Int)) + + segs = segments(dom) + u = oneunit(eltype(dom)) + usegs = map(x -> x/u, segs) + if f isa InplaceIntegrand + g! = (y, x) -> f.f!(y, u*x, p) + result = f.I / u + val, err = auxquadgk!(g!, result, usegs, maxevals = maxiters, + rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) + return IntegralSolution(f.I .= u .* val, u*err, true, -1) + elseif f isa BatchIntegrand + xx = eltype(f.x) === Nothing ? typeof((segs[1]+segs[end])/2)[] : f.x + g_ = (y, x) -> (resize!(xx, length(x)); f.f!(y, xx .= u .* x, p)) + g = IteratedIntegration.AuxQuadGK.BatchIntegrand(g_, f.y, xx/u, max_batch=f.max_batch) + val, err = auxquadgk(g, usegs, maxevals = maxiters, + rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) + return IntegralSolution(u*val, u*err, true, -1) + else + g = x -> f(u*x, p) + val, err = auxquadgk(g, usegs, maxevals = maxiters, + rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) + return IntegralSolution(u*val, u*err, true, -1) + end +end + +""" + ContQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.ContQuadGK.NewtonDeflation()) + +A 1d contour deformation quadrature scheme for scalar, complex-valued integrands. It +defaults to regular `quadgk` behavior on the real axis, but if it finds a root of 1/f +nearby, in the sense of Bernstein ellipse for the standard segment `[-1,1]` with semiaxes +`cosh(rho)` and `sinh(rho)`, on either the upper/lower half planes, then it dents the +contour away from the presumable pole. +""" +struct ContQuadGKJL{F,M} <: IntegralAlgorithm + order::Int + norm::F + rho::Float64 + rootmeth::M +end +function ContQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.ContQuadGK.NewtonDeflation()) + return ContQuadGKJL(order, norm, rho, rootmeth) +end + +function init_cacheval(f, dom, p, alg::ContQuadGKJL) + f isa NestedBatchIntegrand && throw(ArgumentError("ContQuadGK doesn't support nested batching")) + f isa BatchIntegrand && throw(ArgumentError("ContQuadGK doesn't support batching")) + f isa InplaceIntegrand && throw(ArgumentError("ContQuadGK doesn't support inplace integrands")) + + a, b = endpoints(dom) + x, s = (a+b)/2, (b-a)/2 + TX = typeof(x) + fx_s = one(ComplexF64) * s # currently the integrand is forcibly written to a ComplexF64 buffer + TI = typeof(fx_s) + TE = typeof(alg.norm(fx_s)) + r_segbuf = IteratedIntegration.ContQuadGK.PoleSegment{TX,TI,TE}[] + fc_s = f(complex(x), p) * complex(s) # the regular evalrule is used on complex segments + TCX = typeof(complex(x)) + TCI = typeof(fc_s) + TCE = typeof(alg.norm(fc_s)) + c_segbuf = IteratedIntegration.ContQuadGK.Segment{TCX,TCI,TCE}[] + return (r=r_segbuf, c=c_segbuf) +end + +function do_solve(f, dom, p, alg::ContQuadGKJL, cacheval; + reltol = nothing, abstol = nothing, maxiters = typemax(Int)) + + segs = segments(dom) + g = x -> f(x, p) + val, err = contquadgk(g, segs, maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, + rtol = reltol, atol = abstol, order = alg.order, norm = alg.norm, r_segbuf=cacheval.r, c_segbuf=cacheval.c) + return IntegralSolution(val, err, true, -1) +end + +""" + MeroQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.MeroQuadGK.NewtonDeflation()) + +A 1d pole subtraction quadrature scheme for scalar, complex-valued integrands that are +meromorphic. It defaults to regular `quadgk` behavior on the real axis, but if it finds +nearby roots of 1/f, in the sense of Bernstein ellipse for the standard segment `[-1,1]` +with semiaxes `cosh(rho)` and `sinh(rho)`, it attempts pole subtraction on that segment. +""" +struct MeroQuadGKJL{F,M} <: IntegralAlgorithm + order::Int + norm::F + rho::Float64 + rootmeth::M +end +function MeroQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIntegration.MeroQuadGK.NewtonDeflation()) + return MeroQuadGKJL(order, norm, rho, rootmeth) +end + +function init_cacheval(f, dom, p, alg::MeroQuadGKJL) + f isa NestedBatchIntegrand && throw(ArgumentError("MeroQuadGK doesn't support nested batching")) + f isa BatchIntegrand && throw(ArgumentError("MeroQuadGK doesn't support batching")) + f isa InplaceIntegrand && throw(ArgumentError("MeroQuadGK doesn't support inplace integrands")) + a, b = endpoints(dom) + x, s = (a + b)/2, (b-a)/2 + fx_s = one(ComplexF64) * s # ignore the actual integrand since it is written to CF64 array + err = alg.norm(fx_s) + return IteratedIntegration.alloc_segbuf(typeof(x), typeof(fx_s), typeof(err)) +end + +function do_solve(f, dom, p, alg::MeroQuadGKJL, cacheval; + reltol = nothing, abstol = nothing, maxiters = typemax(Int)) + + segs = segments(dom) + g = x -> f(x, p) + val, err = meroquadgk(g, segs, maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, + rtol = reltol, atol = abstol, order = alg.order, norm = alg.norm, segbuf=cacheval) + return IntegralSolution(val, err, true, -1) +end diff --git a/src/algorithms_meta.jl b/src/algorithms_meta.jl new file mode 100644 index 0000000..4d4daa9 --- /dev/null +++ b/src/algorithms_meta.jl @@ -0,0 +1,298 @@ +""" + NestedQuad(alg::IntegralAlgorithm) + NestedQuad(algs::IntegralAlgorithm...) + +Nested integration by repeating one quadrature algorithm or composing a list of algorithms. +The domain of integration must be an `AbstractIteratedLimits` from the +IteratedIntegration.jl package. Analogous to `nested_quad` from IteratedIntegration.jl. +The integrand should expect `SVector` inputs. Do not use this for very high-dimensional +integrals, since the compilation time scales very poorly with respect to dimensionality. +In order to improve the compilation time, FunctionWrappers.jl is used to enforce type +stability of the integrand, so you should always pick the widest integration limit type so +that inference works properly. For example, if [`ContQuadGKJL`](@ref) is used as an +algorithm in the nested scheme, then the limits of integration should be made complex. +""" +struct NestedQuad{T} <: IntegralAlgorithm + algs::T + NestedQuad(alg::IntegralAlgorithm) = new{typeof(alg)}(alg) + NestedQuad(algs::Tuple{Vararg{IntegralAlgorithm}}) = new{typeof(algs)}(algs) +end +NestedQuad(algs::IntegralAlgorithm...) = NestedQuad(algs) + +# this function helps create a tree of the cachevals used by each quadrature +function nested_cacheval(f::F, p::P, algs, segs, lims, state, x, xs...) where {F,P} + dom = PuncturedInterval(segs) + a, b = segs[1], segs[2] + dim = ndims(lims) + alg = algs[dim] + mid = (a+b)/2 # sample point that should be safe to evaluate + next = limit_iterate(lims, state, mid) # see what the next limit gives + if xs isa Tuple{} # the next integral takes us to the inner integral + # base case test integrand of the inner integral + # we need to pass dummy integrands to all the outer integrals so that they can build + # their caches with the right types + if f isa BatchIntegrand || f isa NestedBatchIntegrand + # Batch integrate the inner integral only + cacheval = init_cacheval(BatchIntegrand(nothing, f.y, f.x, max_batch=f.max_batch), dom, p, alg) + return (nothing, cacheval, oneunit(eltype(f.y))*mid) + elseif f isa InplaceIntegrand + # Inplace integrate through the whole nest structure + fxi = f.I*mid/oneunit(prod(next)) + cacheval = init_cacheval(InplaceIntegrand(nothing, fxi), dom, p, alg) + return (nothing, cacheval, fxi) + else + fx = f(next,p) + cacheval = init_cacheval((x, p) -> fx, dom, p, alg) + return (nothing, cacheval, fx*mid) + end + elseif f isa NestedBatchIntegrand + algs_ = algs[1:dim-1] + # numbered names to avoid type instabilities (we are not using dispatch, but the + # compiler's optimization for the recursive function's argument types) + nest0 = nested_cacheval(f.f[1], p, algs_, next..., x, xs[1:dim-2]...) + cacheval = init_cacheval(BatchIntegrand(nothing, f.y, f.x, max_batch=f.max_batch), dom, p, alg) + return (ntuple(n -> n == 1 ? nest0 : deepcopy(nest0), Val(length(f.f))), cacheval, nest0[3]*mid) + else + algs_ = algs[1:dim-1] + nest1 = nested_cacheval(f, p, algs_, next..., x, xs[1:dim-2]...) + h = nest1[3] + hx = h*mid + # units may change for outer integral + if f isa InplaceIntegrand + cacheval = init_cacheval(InplaceIntegrand(nothing, hx), dom, p, alg) + return (nest1, cacheval, hx) + else + cacheval = init_cacheval((x, p) -> h, dom, p, alg) + return (nest1, cacheval, hx) + end + end +end +#= +function init_cacheval(f, dom::AbstractIteratedLimits, p, alg::NestedQuad) + algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(dom))) : alg.algs + return nested_cacheval(f, p, algs, limit_iterate(dom)..., interior_point(dom)...) +end +=# + +function _update!(cache, x, lims_state) + segs, lims, state = limit_iterate(lims_state..., x) + len = segs[end] - segs[begin] + kws = cache.kwargs + cache.cacheval.dom = segs + cache.cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws + cache.cacheval.p = (lims, state) + return +end +_postsolve(sol, x, p) = sol.value +function init_cacheval(f::IntegralFunction, nextdom, p, alg::NestedQuad; kws...) + x0, (segs, lims, state) = if nextdom isa AbstractIteratedLimits + interior_point(nextdom), limit_iterate(nextdom) + else + nothing, nextdom + end + algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(lims))) : alg.algs + proto = get_prototype(f, x0, p) + func = if ndims(lims) == 1 + IntegralFunction(proto) do x, lims_state + f.f(limit_iterate(lims_state..., x), p) + end + else + len = segs[end] - segs[begin] + a, b, = segs + x = (a+b)/2 + next = limit_iterate(lims, state, x) + kws = NamedTuple(kws) + kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws + subprob = IntegralProblem(IntegralFunction(f.f, proto), next, p; kwargs...) + CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1]), _update!, _postsolve, proto*x^(ndims(lims)-1), FunctionWrapperSpecialize()) + end + prob = IntegralProblem(func, segs, (lims, state); kws...) + return init(prob, algs[ndims(lims)]) +end + +function do_integral(f, dom, p, alg::NestedQuad, cacheval; kws...) + return solve!(cacheval) +end +#= +function init_nest(f::F, fxx, dom, p,lims, state, algs, cacheval; kws_...) where {F} + kws = NamedTuple(kws_) + xx = float(oneunit(eltype(dom))) + FX = typeof(fxx/xx) + TX = typeof(xx) + TP = Tuple{typeof(p),typeof(lims),typeof(state)} + if algs isa Tuple{} # inner integral + if f isa BatchIntegrand + return f + elseif f isa NestedBatchIntegrand + nchunk = length(f.f) + return BatchIntegrand(FunctionWrapper{Nothing,Tuple{typeof(f.y),typeof(f.x),TP}}() do y, x, (p, lims, state) + Threads.@threads for ichunk in 1:min(nchunk, length(x)) + for (i, j) in zip(getchunk(x, ichunk, nchunk, :scatter), getchunk(y, ichunk, nchunk, :scatter)) + xi = x[i] + y[j] = f.f[ichunk](limit_iterate(lims, state, xi), p) + end + end + return nothing + end, f.y, f.x, max_batch=f.max_batch) + elseif f isa InplaceIntegrand + return InplaceIntegrand(FunctionWrapper{Nothing,Tuple{FX,TX,TP}}() do y, x, (p, lims, state) + f.f!(y, limit_iterate(lims, state, x), p) + return nothing + end, f.I) + else + return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) + return f(limit_iterate(lims, state, x), p) + end + end + else + if f isa InplaceIntegrand + return InplaceIntegrand(FunctionWrapper{Nothing,Tuple{FX,TX,TP}}() do y, x, (p, lims, state) + segs, lims_, state_ = limit_iterate(lims, state, x) + len = segs[end] - segs[1] + kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws + do_solve(InplaceIntegrand(f.f!, y), lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval; kwargs...) + return nothing + end, f.I) + elseif f isa NestedBatchIntegrand + nchunks = length(f.f) + return BatchIntegrand(FunctionWrapper{Nothing,Tuple{typeof(f.y),typeof(f.x),TP}}() do y, x, (p, lims, state) + Threads.@threads for ichunk in 1:min(nchunks, length(x)) + for (i, j) in zip(getchunk(x, ichunk, nchunks, :scatter), getchunk(y, ichunk, nchunks, :scatter)) + xi = x[i] + segs, lims_, state_ = limit_iterate(lims, state, xi) + len = segs[end] - segs[1] + kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws + y[j] = do_solve(f.f[ichunk], lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval[ichunk]; kwargs...).u + end + end + return nothing + end, f.y, f.x, max_batch=f.max_batch) + else + return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) + segs, lims_, state_ = limit_iterate(lims, state, x) + len = segs[end] - segs[1] + kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws + sol = do_solve(f, lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval; kwargs...) + return sol.u + end + end + end +end + +struct NestState{P,G,S} + p::P + segs::G + state::S +end + +function do_solve(f::F, lims::AbstractIteratedLimits, p_, alg::NestedQuad, cacheval; kws...) where {F} + g, p, segs, state = if p_ isa NestState + gg = if f isa NestedBatchIntegrand + fx = eltype(f.x) === Nothing ? float(eltype(p_.segs))[] : f.x + NestedBatchIntegrand(f.f, f.y, fx, max_batch=f.max_batch) + else + f + end + gg, p_.p, p_.segs, p_.state + else + seg, lim, sta = limit_iterate(lims) + gg = if f isa BatchIntegrand + fx = eltype(f.x) === Nothing ? typeof(interior_point(lims))[] : f.x + BatchIntegrand(f.y, similar(f.x, eltype(eltype(fx))), max_batch=f.max_batch) do y, xs, (p, lims, state) + resize!(fx, length(xs)) + f.f!(y, map!(x -> limit_iterate(lims, state, x), fx, xs), p) + end + elseif f isa NestedBatchIntegrand + # this should be done recursively at the outermost level, but it is lazy. + fx = eltype(f.x) === Nothing ? float(eltype(seg))[] : f.x + NestedBatchIntegrand(f.f, f.y, fx, max_batch=f.max_batch) + else + f + end + gg, p_, seg, sta + end + dom = PuncturedInterval(segs) + dim = ndims(lims) # constant propagation :) + algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(dim)) : alg.algs + nest = init_nest(g, cacheval[3], dom, p, lims, state, algs[1:dim-1], cacheval[1]; kws...) + return do_solve(nest, dom, (p, lims, state), algs[dim], cacheval[2]; kws...) +end +=# +""" + AbsoluteEstimate(est_alg, abs_alg; kws...) + +Most algorithms are efficient when using absolute error tolerances, but how do you know the +size of the integral? One option is to estimate it using second algorithm. + +A multi-algorithm to estimate an integral using an `est_alg` to generate a rough estimate of +the integral that is combined with a user's relative tolerance to re-calculate the integral +to higher accuracy using the `abs_alg`. The keywords passed to the algorithm may include +`reltol`, `abstol` and `maxiters` and are given to the `est_alg` solver. They should limit +the amount of work of `est_alg` so as to only generate an order-of-magnitude estimate of the +integral. The tolerances passed to `abs_alg` are `abstol=max(abstol,reltol*norm(I))` and +`reltol=0`. +""" +struct AbsoluteEstimate{E<:IntegralAlgorithm,A<:IntegralAlgorithm,F,K<:NamedTuple} <: IntegralAlgorithm + est_alg::E + abs_alg::A + norm::F + kws::K +end +function AbsoluteEstimate(est_alg, abs_alg; norm=norm, kwargs...) + kws = NamedTuple(kwargs) + checkkwargs(kws) + return AbsoluteEstimate(est_alg, abs_alg, norm, kws) +end + +function init_cacheval(f, dom, p, alg::AbsoluteEstimate) + return (est=init_cacheval(f, dom, p, alg.est_alg), + abs=init_cacheval(f, dom, p, alg.abs_alg)) +end + +function do_solve(f, dom, p, alg::AbsoluteEstimate, cacheval; + abstol=nothing, reltol=nothing, maxiters=typemax(Int)) + sol = do_solve(f, dom, p, alg.est_alg, cacheval.est; alg.kws...) + val = alg.norm(sol.u) # has same units as sol + rtol = reltol === nothing ? sqrt(eps(one(val))) : reltol # use the precision of the solution to set the default relative tolerance + atol = max(abstol === nothing ? zero(val) : abstol, rtol*val) + return do_solve(f, dom, p, alg.abs_alg, cacheval.abs; + abstol=atol, reltol=zero(rtol), maxiters=maxiters) +end + + +""" + EvalCounter(::IntegralAlgorithm) + +An algorithm which counts the evaluations used by another algorithm. +The count is stored in the `sol.numevals` field. +""" +struct EvalCounter{T<:IntegralAlgorithm} <: IntegralAlgorithm + alg::T +end + +function init_cacheval(f, dom, p, alg::EvalCounter) + return init_cacheval(f, dom, p, alg.alg) +end + +function do_solve(f, dom, p, alg::EvalCounter, cacheval; kws...) + if f isa InplaceIntegrand + ni::Int = 0 + gi = (y, x, p) -> (ni += 1; f.f!(y, x, p)) + soli = do_solve(InplaceIntegrand(gi, f.I), dom, p, alg.alg, cacheval; kws...) + return IntegralSolution(soli.u, soli.resid, soli.retcode, ni) + elseif f isa BatchIntegrand + nb::Int = 0 + gb = (y, x, p) -> (nb += length(x); f.f!(y, x, p)) + solb = do_solve(BatchIntegrand(gb, f.y, f.x, max_batch=f.max_batch), dom, p, alg.alg, cacheval; kws...) + return IntegralSolution(solb.u, solb.resid, solb.retcode, nb) + elseif f isa NestedBatchIntegrand + # TODO allocate a bunch of accumulators associated with the leaves of the nested + # integrand or rewrap the algorithms in NestedQuad + error("NestedBatchIntegrand not yet supported with EvalCounter") + else + n::Int = 0 + g = (x, p) -> (n += 1; f(x, p)) # we need let to prevent Core.Box around the captured variable + sol = do_solve(g, dom, p, alg.alg, cacheval; kws...) + return IntegralSolution(sol.u, sol.resid, sol.retcode, n) + end +end diff --git a/src/batch.jl b/src/batch.jl deleted file mode 100644 index 1642c82..0000000 --- a/src/batch.jl +++ /dev/null @@ -1,77 +0,0 @@ -""" - BatchIntegrand(f!, y::AbstractArray, x::AbstractVector, max_batch=typemax(Int)) - -Constructor for a `BatchIntegrand` accepting an integrand of the form `f!(y,x,p) = y .= f!.(x, Ref(p))` -that can evaluate the integrand at multiple quadrature nodes using, for example, threads, -the GPU, or distributed-memory. The `max_batch` keyword is a soft limit on the number of -nodes passed to the integrand. The buffers `y,x` must both be `resize!`-able since the -number of evaluation points may vary between calls to `f!`. -""" -struct BatchIntegrand{F,Y,X} - # in-place function f!(y, x, p) that takes an array of x values and outputs an array of results in-place - f!::F - y::Y - x::X - max_batch::Int # maximum number of x to supply in parallel - function BatchIntegrand(f!, y::AbstractArray, x::AbstractVector, max_batch::Integer=typemax(Int)) - max_batch > 0 || throw(ArgumentError("maximum batch size must be positive")) - return new{typeof(f!),typeof(y),typeof(x)}(f!, y, x, max_batch) - end -end - - -""" - BatchIntegrand(f!, y, x; max_batch=typemax(Int)) - -Constructor for a `BatchIntegrand` with pre-allocated buffers. -""" -BatchIntegrand(f!, y, x; max_batch::Integer=typemax(Int)) = - BatchIntegrand(f!, y, x, max_batch) - -""" - BatchIntegrand(f!, y::Type, x::Type=Nothing; max_batch=typemax(Int)) - -Constructor for a `BatchIntegrand` whose range type is known. The domain type is optional. -Array buffers for those types are allocated internally. -""" -BatchIntegrand(f!, Y::Type, X::Type=Nothing; max_batch::Integer=typemax(Int)) = - BatchIntegrand(f!, Y[], X[], max_batch) - - -""" - NestedBatchIntegrand(f::Tuple, y::AbstractVector, x::AbstractVector, max_batch::Integer) - -An integrand type intended for multi-threaded evaluation of [`NestedQuad`](@ref). The caller -provides a tuple `f` of worker functions that can evaluate the same integrand on different -threads, so as to avoid race conditions. These workers can also be `NestedBatchIntegrand`s -depending on if the user wants to parallelize the integration at multiple levels of nesting. -The other arguments are the same as for [`BatchIntegrand`](@ref). -""" -struct NestedBatchIntegrand{F,T,Y<:AbstractVector,X<:AbstractVector} - f::T - y::Y - x::X - max_batch::Int - function NestedBatchIntegrand(f::NTuple, y::Y, x::X, max_batch::Integer) where {Y,X} - if eltype(f) <: NestedBatchIntegrand - return new{_nesttype(eltype(f)),typeof(f),Y,X}(f, y, x, max_batch) - else - return new{eltype(f),typeof(f),Y,X}(f, y, x, max_batch) - end - end - function NestedBatchIntegrand(f::AbstractArray{F}, y::Y, x::X, max_batch::Integer) where {F,Y,X} - return new{F,typeof(f),Y,X}(f, y, x, max_batch) - end - function NestedBatchIntegrand(f::AbstractArray{T}, y::Y, x::X, max_batch::Integer) where {F,T<:NestedBatchIntegrand{F},Y,X} - return new{F,typeof(f),Y,X}(f, y, x, max_batch) - end -end - -_nesttype(::Type{<:NestedBatchIntegrand{F}}) where {F} = F -function NestedBatchIntegrand(f, y, x; max_batch::Integer=typemax(Int)) - return NestedBatchIntegrand(f, y, x, max_batch) -end - -function NestedBatchIntegrand(f, ::Type{Y}, ::Type{X}=Nothing; kws...) where {Y,X} - return NestedBatchIntegrand(f, Y[], X[]; kws...) -end diff --git a/src/brillouin.jl b/src/brillouin.jl index 84b1b71..cd4abdb 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -320,6 +320,69 @@ These algorithms also use the symmetries of the Brillouin zone and the integrand """ abstract type AutoBZAlgorithm <: IntegralAlgorithm end +""" + AutoBZProblem([rep], f, bz, [p]; kwargs...) + +Construct a BZ integration problem. + +## Arguments +- `rep::AbstractSymRep`: The symmetry representation of `f` (default: `UnknownRep()`) +- `f::AbstractIntegralFunction`: The integrand +- `bz::SymmetricBZ`: The Brillouin zone to integrate over +- `p`: parameters for the integrand (default: `NullParameters()`) + +## Keywords +Additional keywords are passed directly to the solver +""" +struct AutoBZProblem{R<:AbstractSymRep,F<:AbstractIntegralFunction,BZ<:SymmetricBZ,P,K<:NamedTuple} + rep::R + f::F + bz::BZ + p::P + kwargs::K +end +function AutoBZProblem(rep::AbstractSymRep, f::AbstractIntegralFunction, bz::SymmetricBZ, p=NullParameters(); kws...) + return AutoBZProblem(rep, f, bz, p, NamedTuple(kws)) +end +function AutoBZProblem(rep::AbstractSymRep, f, bz::SymmetricBZ, p=NullParameters(); kws...) + return AutoBZProblem(IntegralFunction(f), bz, p; kws...) +end +function AutoBZProblem(f, bz::SymmetricBZ, p=NullParameters(); kws...) + return AutoBZProblem(UnknownRep(), f, bz, p; kws...) +end + +mutable struct AutoBZCache{R,F,BZ,P,A,C,K} + rep::R + f::F + bz::BZ + p::P + alg::A + cacheval::C + kwargs::K +end + +function init(prob::AutoBZProblem, alg::AutoBZAlgorithm; kwargs...) + rep = prob.rep; f = prob.f; bz = prob.bz; p = prob.p + kws = (; prob.kwargs..., kwargs...) + checkkwargs(kws) + cacheval = init_cacheval(rep, f, bz, p, alg) + return AutoBZCache(rep, f, bz, p, alg, cacheval, kws) +end + +""" + solve(::AutoBZProblem, ::AutoBZAlgorithm; kws...)::IntegralSolution +""" +solve(prob::AutoBZProblem, alg::AutoBZAlgorithm; kwargs...) + +""" + solve!(::IntegralCache)::IntegralSolution + +Compute the solution to an [`IntegralProblem`](@ref) constructed from [`init`](@ref). +""" +function solve!(c::AutoBZCache) + return do_solve_autobz(c.rep, c.f, c.dom, c.p, c.alg, c.cacheval; c.kwargs...) +end +#= function init_cacheval(f, bz::SymmetricBZ, p, bzalg::AutoBZAlgorithm) _, dom, alg = bz_to_standard(bz, bzalg) return init_cacheval(f, dom, p, alg) @@ -497,3 +560,4 @@ end function do_solve(f, bz::SymmetricBZ, p, alg::EvalCounter{<:AutoBZAlgorithm}, cacheval; kws...) return do_solve_autobz(count_bz_to_standard, f, bz, p, alg.alg, cacheval; kws...) end +=# diff --git a/src/domains.jl b/src/domains.jl index c58abef..190870a 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -1,3 +1,11 @@ +endpoints(dom) = (first(dom), last(dom)) +breakpoints(dom) = dom[begin+1:end-1] # or Iterators.drop(Iterators.take(dom, length(dom)-1), 1) +segments(dom) = dom +function get_prototype(dom) + a, b, = dom + return (a+b)/2 +end + """ PuncturedInterval(s) @@ -10,9 +18,14 @@ struct PuncturedInterval{T,S} PuncturedInterval(s::S) where {N,S<:NTuple{N}} = new{eltype(s),S}(s) PuncturedInterval(s::S) where {T,S<:AbstractVector{T}} = new{T,S}(s) end +PuncturedInterval(s::PuncturedInterval) = s Base.eltype(::Type{PuncturedInterval{T,S}}) where {T,S} = T segments(p::PuncturedInterval) = p.s endpoints(p::PuncturedInterval) = (p.s[begin], p.s[end]) +function get_prototype(p::PuncturedInterval) + a, b, = segments(p) + return (a + b)/2 +end """ HyperCube(a, b) @@ -31,3 +44,9 @@ HyperCube(a, b) = HyperCube(promote(a...), promote(b...)) Base.eltype(::Type{HyperCube{d,T}}) where {d,T} = T endpoints(c::HyperCube) = (c.a, c.b) +function get_prototype(p::HyperCube) + a, b = endpoints(p) + return (a + b)/2 +end + +get_prototype(l::AbstractIteratedLimits) = interior_point(l) diff --git a/src/dos_interfaces.jl b/src/dos_interfaces.jl index e6b0f63..56be1f5 100644 --- a/src/dos_interfaces.jl +++ b/src/dos_interfaces.jl @@ -30,12 +30,15 @@ where ``E \\in \\text{domain}`` and ``\\delta`` is the Dirac Delta distribution. of freedom) or continuous (e.g. for `H` a Hamiltonian parameterized by crystal momentum). """ -struct DOSProblem{H,D,P} +struct DOSProblem{H,D,P,K<:NamedTuple} H::H domain::D p::P + kwargs::K +end +function DOSProblem(H, domain, p=NullParameters(); kws...) + return DOSProblem(H, domain, p, NamedTuple(kws)) end -DOSProblem(H, domain) = DOSProblem(H, domain, NullParameters()) struct DOSSolution{U,E} u::U @@ -47,13 +50,13 @@ end # store the data in a mutable cache so that the user can update the cache and # compute the result again without setting up new problems. mutable struct DOSCache{H,D,P,A,C,K} - H::H - domain::D - p::P - alg::A - cacheval::C - isfresh::Bool # true if H has been replaced/modified, otherwise false - kwargs::K + H::H + domain::D + p::P + alg::A + cacheval::C + isfresh::Bool # true if H has been replaced/modified, otherwise false + kwargs::K end function Base.setproperty!(cache::DOSCache, name::Symbol, item) @@ -66,11 +69,6 @@ end # by default, algorithms won't have anything in the cache init_cacheval(h, dom, p, ::DOSAlgorithm) = nothing -function make_cache(h, dom, p, alg::DOSAlgorithm; kwargs...) - cacheval = init_cacheval(h, dom, p, alg) - return DOSCache(h, dom, p, alg, cacheval, false, NamedTuple(kwargs)) -end - # check same keywords as for integral problems: abstol, reltol, maxiters checkkwargs_dos(kws) = checkkwargs(kws) @@ -81,19 +79,10 @@ Create a cache of the data used by an algorithm to solve the given problem. """ function init(prob::DOSProblem, alg::DOSAlgorithm; kwargs...) h = prob.H; dom = prob.domain; p = prob.p - checkkwargs_dos(NamedTuple(kwargs)) - return make_cache(h, dom, p, alg; kwargs...) -end - -""" - solve(::DOSProblem, ::DOSAlgorithm; kwargs...)::DOSSolution - -Compute the solution to a [`DOSProblem`](@ref) using the given algorithm. -The keyword arguments to the solver can be `abstol`, `reltol`, and `maxiters`. -""" -function solve(prob::DOSProblem, alg::DOSAlgorithm; kwargs...) - cache = init(prob, alg; kwargs...) - return solve!(cache) + kws = (; prob.kwargs..., kwargs...) + checkkwargs_dos(kws) + cacheval = init_cacheval(h, dom, p, alg) + return DOSCache(h, dom, p, alg, cacheval, false, kws) end """ @@ -112,3 +101,8 @@ function solve!(c::DOSCache) end function dos_solve end + +""" + solve(::DOSProblem, ::DOSAlgorithm; kws...)::DOSSolution +""" +solve(prob::DOSProblem, alg::DOSAlgorithm; kwargs...) diff --git a/src/fourier.jl b/src/fourier.jl index 98a0457..b37d979 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -18,35 +18,25 @@ # We use the pattern of allowing the user to pass a container with the integrand, Fourier # series and workspace, and use dispatch to enable the optimizations -# the nested batched integrand is optional, but when included it allows for thread-safe parallelization -struct FourierIntegrand{F,P,W,N} - f::ParameterIntegrand{F,P} - w::W - nest::N - function FourierIntegrand(f::ParameterIntegrand{F,P}, w::FourierWorkspace) where {F,P} - return new{F,P,typeof(w),Nothing}(f, w, nothing) - end - function FourierIntegrand(f::ParameterIntegrand{F,P}, w::FourierWorkspace, nest::NestedBatchIntegrand{<:ParameterIntegrand{F}}) where {F,P} - return new{F,P,typeof(w),typeof(nest)}(f, w, nest) - end - function FourierIntegrand(f::ParameterIntegrand{F,P}, w::FourierWorkspace, nest::ParameterIntegrand{F}) where {F,P} - return new{F,P,typeof(w),typeof(nest)}(f, w, nest) - end -end +# the nested batched integrand is optional, but when included it allows for thread-safe +# parallelization """ - FourierIntegrand(f, w::FourierWorkspace, args...; kws...) + FourierIntegralFunction(f, s) -Constructs an integrand of the form `f(FourierValue(x,w(x)), args...; kws...)` where the -Fourier series in `w` is evaluated efficiently, i.e. one dimension at a time, with -compatible algorithms. `f` should accept parameters as arguments and keywords, similar to a -[`ParameterIntegrand`](@ref) although the first argument to `f` will always be a -[`FourierValue`](@ref). +## Arguments +- `f::AbstractIntegralFunction`: The integrand, accepting inputs `f(x, s(x), p)` +- `s::AbstractFourierSeries`: The Fourier series to evaluate """ -function FourierIntegrand(f, w::FourierWorkspace, args...; kws...) - return FourierIntegrand(ParameterIntegrand(f, args...; kws...), w) +struct FourierIntegralFunction{F<:AbstractIntegralFunction,S<:AbstractFourierSeries} <: AbstractIntegralFunction + f::F + s::S end +function get_prototype(f::FourierIntegralFunction, x, p) + return get_prototype(f.f, x, p) +end +#= """ FourierIntegrand(f, s::AbstractFourierSeries, args...; kws...) @@ -85,20 +75,6 @@ function workspace_allocate_vec(s::AbstractFourierSeries{N}, x::NTuple{N,Any}, l return FourierWorkspace(s, ws) end - -function (s::IntegralSolver{<:FourierIntegrand})(args...; kwargs...) - p = MixedParameters(args...; kwargs...) - sol = solve_p(s, p) - return sol.u -end - -function remake_cache(f::FourierIntegrand, dom, p, alg, cacheval, kwargs) - # TODO decide what to do with the nest, since it may have allocated - fp = ParameterIntegrand(f.f.f) - new = f.nest === nothing ? FourierIntegrand(fp, f.w) : FourierIntegrand(fp, f.w, f.nest) - return remake_integrand_cache(new, dom, merge(f.f.p, p), alg, cacheval, kwargs) -end - # FourierIntegrands should expect a FourierValue as input """ @@ -117,9 +93,9 @@ Base.convert(::Type{T}, f::FourierValue) where {T<:FourierValue} = T(f.x,f.s) Base.zero(::Type{FourierValue{X,S}}) where {X,S} = FourierValue(zero(X), zero(S)) Base.:*(B::AbstractMatrix, f::FourierValue) = FourierValue(B*f.x, f.s) -(f::FourierIntegrand)(x::FourierValue, p) = f.f(x, p) +# (f::FourierIntegrand)(x::FourierValue, p) = f.f(x, p) # fallback evaluator of FourierIntegrand for algorithms without specialized rules -(f::FourierIntegrand)(x, p) = f(FourierValue(x, f.w(x)), p) +# (f::FourierIntegrand)(x, p) = f(FourierValue(x, f.w(x)), p) # PTR rules @@ -322,9 +298,9 @@ end # dispatch on PTR algorithms -function init_buffer(f::FourierIntegrand, len) - return f.nest isa NestedBatchIntegrand ? Vector{eltype(f.nest.y)}(undef, len) : nothing -end +# function init_buffer(f::FourierIntegrand, len) +# return f.nest isa NestedBatchIntegrand ? Vector{eltype(f.nest.y)}(undef, len) : nothing +# end rule_type(::FourierPTR{N,T,S}) where {N,T,S} = FourierValue{SVector{N,T},S} function init_fourier_rule(w::FourierWorkspace, dom::Basis, alg::MonkhorstPack) @@ -335,62 +311,62 @@ function init_fourier_rule(w::FourierWorkspace, dom::Basis, alg::MonkhorstPack) return FourierMonkhorstPack(w, eltype(dom), Val(ndims(dom)), alg.npt, alg.syms) end end -function init_cacheval(f::FourierIntegrand, dom::Basis , p, alg::MonkhorstPack) - rule = init_fourier_rule(f.w, dom, alg) - buf = init_buffer(f, alg.nthreads) - return (rule=rule, buffer=buf) -end +# function init_cacheval(f::FourierIntegrand, dom::Basis , p, alg::MonkhorstPack) +# rule = init_fourier_rule(f.w, dom, alg) +# buf = init_buffer(f, alg.nthreads) +# return (rule=rule, buffer=buf) +# end function init_fourier_rule(w::FourierWorkspace, dom::Basis, alg::AutoSymPTRJL) @assert ndims(w.series) == ndims(dom) return FourierMonkhorstPackRule(w, alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) end -function init_cacheval(f::FourierIntegrand, dom::Basis, p, alg::AutoSymPTRJL) - rule = init_fourier_rule(f.w, dom, alg) - cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) - buffer = init_buffer(f, alg.nthreads) - return (rule=rule, cache=cache, buffer=buffer) -end -function init_cacheval(f::FourierIntegrand, bz::SymmetricBZ, p, bzalg::AutoPTR) - bz_, dom, alg = bz_to_standard(bz, bzalg) - rule = SymmetricRuleDef(init_fourier_rule(f.w, dom, alg), SymRep(f), bz_) - cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) - buffer = init_buffer(f, alg.nthreads) - return (rule=rule, cache=cache, buffer=buffer) -end +# function init_cacheval(f::FourierIntegrand, dom::Basis, p, alg::AutoSymPTRJL) +# rule = init_fourier_rule(f.w, dom, alg) +# cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) +# buffer = init_buffer(f, alg.nthreads) +# return (rule=rule, cache=cache, buffer=buffer) +# end +# function init_cacheval(f::FourierIntegrand, bz::SymmetricBZ, p, bzalg::AutoPTR) +# bz_, dom, alg = bz_to_standard(bz, bzalg) +# rule = SymmetricRuleDef(init_fourier_rule(f.w, dom, alg), SymRep(f), bz_) +# cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) +# buffer = init_buffer(f, alg.nthreads) +# return (rule=rule, cache=cache, buffer=buffer) +# end function init_fourier_rule(s::AbstractFourierSeries, bz::SymmetricBZ, alg::PTR) dom = Basis(bz.B) return FourierMonkhorstPack(s, eltype(dom), Val(ndims(dom)), alg.npt, bz.syms) end -function nested_to_batched(f::FourierIntegrand, dom::Basis, rule) - ys = f.nest.y/prod(ntuple(n -> oneunit(eltype(dom)), Val(ndims(dom)))) - xs = rule_type(rule)[] - return BatchIntegrand(ys, xs, max_batch=f.nest.max_batch) do y, x, p - # would be better to fully unwrap the nested structure, but this is one level - nchunk = length(f.nest.f) - Threads.@threads for ichunk in 1:min(nchunk, length(x)) - for (i, j) in zip(getchunk(x, ichunk, nchunk, :batch), getchunk(y, ichunk, nchunk, :batch)) - y[j] = FourierIntegrand(f.nest.f[ichunk], FourierWorkspace(nothing,nothing))(x[i], p) - end - end - return nothing - end -end - -function do_solve(f::FourierIntegrand, dom, p, alg::MonkhorstPack, cacheval; kws...) - g = f.nest isa NestedBatchIntegrand ? nested_to_batched(f, dom, cacheval.rule) : f.f - return do_solve(g, dom, p, alg, cacheval; kws...) -end - -function do_solve(f::FourierIntegrand, dom, p, alg::AutoSymPTRJL, cacheval; kws...) - g = f.nest isa NestedBatchIntegrand ? nested_to_batched(f, dom, cacheval.cache[1]) : f.f - return do_solve(g, dom, p, alg, cacheval; kws...) -end +# function nested_to_batched(f::FourierIntegrand, dom::Basis, rule) +# ys = f.nest.y/prod(ntuple(n -> oneunit(eltype(dom)), Val(ndims(dom)))) +# xs = rule_type(rule)[] +# return BatchIntegrand(ys, xs, max_batch=f.nest.max_batch) do y, x, p +# # would be better to fully unwrap the nested structure, but this is one level +# nchunk = length(f.nest.f) +# Threads.@threads for ichunk in 1:min(nchunk, length(x)) +# for (i, j) in zip(getchunk(x, ichunk, nchunk, :batch), getchunk(y, ichunk, nchunk, :batch)) +# y[j] = FourierIntegrand(f.nest.f[ichunk], FourierWorkspace(nothing,nothing))(x[i], p) +# end +# end +# return nothing +# end +# end + +# function do_solve(f::FourierIntegrand, dom, p, alg::MonkhorstPack, cacheval; kws...) +# g = f.nest isa NestedBatchIntegrand ? nested_to_batched(f, dom, cacheval.rule) : f.f +# return do_solve(g, dom, p, alg, cacheval; kws...) +# end + +# function do_solve(f::FourierIntegrand, dom, p, alg::AutoSymPTRJL, cacheval; kws...) +# g = f.nest isa NestedBatchIntegrand ? nested_to_batched(f, dom, cacheval.cache[1]) : f.f +# return do_solve(g, dom, p, alg, cacheval; kws...) +# end # dispatch on IAI algorithms - +#= function nested_cacheval(f::FourierIntegrand, p::P, algs, segs, lims, state, x, xs...) where {P} dom = PuncturedInterval(segs) a, b = segs[1], segs[2] @@ -528,3 +504,5 @@ end function do_solve(f::FourierIntegrand, bz::SymmetricBZ, p, alg::EvalCounter{<:AutoBZAlgorithm}, cacheval; kws...) return do_solve_autobz(count_bz_to_standard, f, bz, p, alg.alg, cacheval; kws...) end +=# +=# diff --git a/src/inplace.jl b/src/inplace.jl deleted file mode 100644 index 6e5c375..0000000 --- a/src/inplace.jl +++ /dev/null @@ -1,15 +0,0 @@ -""" - InplaceIntegrand(f!, result::AbstractArray) - -Constructor for a `InplaceIntegrand` accepting an integrand of the form `f!(y,x,p)`. The -caller also provides an output array needed to store the result of the quadrature. -Intermediate `y` arrays are allocated during the calculation, and the final result is -may or may not be written to `result`, so use the IntegralSolution immediately after the -calculation to read the result, and don't expect it to persist if the same integrand is used -for another calculation. -""" -struct InplaceIntegrand{F,T<:AbstractArray} - # in-place function f!(y, x, p) that takes one x value and outputs an array of results in-place - f!::F - I::T -end diff --git a/src/interfaces.jl b/src/interfaces.jl index f5cb728..8372b97 100644 --- a/src/interfaces.jl +++ b/src/interfaces.jl @@ -1,8 +1,128 @@ -# we recreate a lot of the SciML Integrals.jl functionality, but only for our algorithms -# the features we omit are: inplace integrands, infinite limit transformations, nout and -# batch keywords. Otherwise, there is a correspondence between. -# solve -> solve! -# init -> init +abstract type AbstractIntegralFunction end +# should have at least two fields: +# - f +# - integrand_prototype + +""" + IntegralFunction(f, [prototype=nothing]) + +Constructor for an out-of-place integrand of the form `f(x, p)`. +Optionally, a `prototype` can be provided for the output of the function. +""" +struct IntegralFunction{F,P} <: AbstractIntegralFunction + f::F + prototype::P +end +IntegralFunction(f) = IntegralFunction(f, nothing) + +function get_prototype(f::IntegralFunction, x, p) + f.prototype === nothing ? f.f(x, p) : f.prototype +end + +""" + InplaceIntegralFunction(f!, prototype::AbstractArray) + +Constructor for an inplace integrand of the form `f!(y, x, p)`. +A `prototype` array is required to store the same type and size as the result, `y`. +""" +struct InplaceIntegralFunction{F,P<:AbstractArray} <: AbstractIntegralFunction + # in-place function f!(y, x, p) that takes one x value and outputs an array of results in-place + f!::F + prototype::P +end + +function get_prototype(f::InplaceIntegralFunction, x, p) + # iip is required to have a prototype array + f.prototype +end + +""" + InplaceBatchIntegralFunction(f!, prototype; max_batch::Integer=typemax(Int)) + +Constructor for an inplace, batched integrand of the form `f!(y, x, p)` that accepts an +array `x` containing a batch of evaluation points stored along the last axis of the array. +A `prototype` array is required to store the same type and size as the result, `y`, however +the last axis, which is reserved for batching, which should contain at least one element. +The `max_batch` keyword sets a soft limit on the number of points batched simultaneously. +""" +struct InplaceBatchIntegralFunction{F,P<:AbstractArray} <: AbstractIntegralFunction + f!::F + prototype::P + max_batch::Int +end + +function InplaceBatchIntegralFunction(f!, p::AbstractArray; max_batch::Integer=typemax(Int)) + return InplaceBatchIntegralFunction(f!, p, max_batch) +end + +function get_prototype(f::InplaceBatchIntegralFunction, x, p) + # iip is required to have a prototype array + f.prototype +end + +abstract type AbstractSpecialization end +struct NoSpecialize <: AbstractSpecialization end +struct FunctionWrapperSpecialize <: AbstractSpecialization end +struct FullSpecialize <: AbstractSpecialization end + +""" + CommonSolveIntegralFunction(prob, alg, update!, postsolve, [prototype, specialize]; kws...) + +Constructor for an integrand that solves a problem defined with the CommonSolve.jl +interface, `prob`, which is instantiated using `init(prob, alg; kws...)`. Helper functions +include: `update!(cache, x, p)` is called before +`solve!(cache)`, followed by `postsolve(sol, x, p)`, which should return the value of the solution. +The `prototype` argument can help control how much to `specialize` on the type of the +problem, which defaults to `FullSpecialize()` so that run times are improved. However +`FunctionWrapperSpecialize()` may help reduce compile times. +""" +struct CommonSolveIntegralFunction{P,A,K,U,S,T,M<:AbstractSpecialization} <: AbstractIntegralFunction + prob::P + alg::A + kwargs::K + update!::U + postsolve::S + prototype::T + specialize::M +end +function CommonSolveIntegralFunction(prob, alg, update!, postsolve, prototype=nothing, specialize=FullSpecialize(); kws...) + return CommonSolveIntegralFunction(prob, alg, NamedTuple(kws), update!, postsolve, prototype, specialize) +end + +function do_solve!(cache, f, x, p) + f.update!(cache, x, p) + sol = solve!(cache) + return f.postsolve(sol, x, p) +end +function get_prototype(f::CommonSolveIntegralFunction, x, p) + if isnothing(f.prototype) + cache = init(f.prob, f.alg; f.kwargs...) + do_solve!(cache, f, x, p) + else + f.prototype + end +end +function init_specialized_integrand(cache, f, dom, p; x=get_prototype(dom), prototype=f.prototype) + proto = prototype === nothing ? do_solve!(cache, f, x, p) : prototype + func = (x, p) -> do_solve!(cache, f, x, p) + integrand = if f.specialize isa FullSpecialize + func + elseif f.specialize isa FunctionWrapperSpecialize + FunctionWrapper{typeof(prototype), typeof((x, p))}(func) + else + throw(ArgumentError("$(f.specialize) is not implemented")) + end + return integrand, proto +end +function _init_commonsolvefunction(f, dom, p; kws...) + cache = init(f.prob, f.alg; f.kwargs...) + integrand, prototype = init_specialized_integrand(cache, f, dom, p; kws...) + return cache, integrand, prototype +end + +# TODO reimplement CommonSolveIntegralFunction as a CommonSolveIntegralProblem OR +# (why? so that FunctionWrappers can be used more systematically) +# TODO add ThreadedCommonSolveIntegralFunction and DistributedCommonSolveIntegralFunction """ IntegralAlgorithm @@ -11,10 +131,6 @@ Abstract supertype for integration algorithms. """ abstract type IntegralAlgorithm end -# Methods an algorithm must define -# - init_cacheval -# - solve! - """ NullParameters() @@ -23,28 +139,27 @@ A singleton type representing absent parameters struct NullParameters end """ - IntegralProblem(f, domain, [p=NullParameters]) - IntegralProblem(f, a::T, b::T, [p=NullParameters]) where {T} + IntegralProblem(f, domain, [p=NullParameters]; kwargs...) -Collects the data need to define an integral of a function `f(x, p)` over a `domain` -containing the points, `x`, and set with parameters `p` (default: [`NullParameters`](@ref)). -If the domain is an interval or hypercube, it can also be specified by its endpoints `a, b`, -and it gets converted to a [`PuncturedInterval`](@ref) or [`HyperCube`](@ref). +## Arguments +- `f::AbstractIntegralFunction`: The function to integrate +- `domain`: The domain to integrate over, e.g. `(lb, ub)` +- `p`: Parameters to pass to the integrand + +## Keywords +Additional keywords are passed directly to the solver """ -struct IntegralProblem{F,D,P} +struct IntegralProblem{F<:AbstractIntegralFunction,D,P,K<:NamedTuple} f::F dom::D p::P - function IntegralProblem{F,D,P}(f::F, dom::D, p::P) where {F,D,P} - return new{F,D,P}(f, dom, p) - end + kwargs::K end -function IntegralProblem(f::F, dom::D, p::P=NullParameters()) where {F,D,P} - return IntegralProblem{F,D,P}(f, dom, p) +function IntegralProblem(f::AbstractIntegralFunction, dom, p=NullParameters(); kws...) + return IntegralProblem(f, dom, p, NamedTuple(kws)) end -function IntegralProblem(f::F, a::T, b::T, p::P=NullParameters()) where {F,T,P} - dom = T <: Number ? PuncturedInterval((a, b)) : HyperCube(a, b) - return IntegralProblem{F,typeof(dom),P}(f, dom, p) +function IntegralProblem(f, dom, p=NullParameters(); kws...) + return IntegralProblem(IntegralFunction(f), dom, p; kws...) end mutable struct IntegralCache{F,D,P,A,C,K} @@ -56,11 +171,6 @@ mutable struct IntegralCache{F,D,P,A,C,K} kwargs::K end -function make_cache(f, dom, p, alg::IntegralAlgorithm; kwargs...) - cacheval = init_cacheval(f, dom, p, alg) - return IntegralCache(f, dom, p, alg, cacheval, NamedTuple(kwargs)) -end - function checkkwargs(kwargs) for key in keys(kwargs) key in (:abstol, :reltol, :maxiters) || throw(ArgumentError("keyword $key unrecognized")) @@ -76,9 +186,11 @@ keyword arguments to the solver (i.e. `abstol`, `reltol`, or `maxiters`) that ca for solving the problem for multiple different parameters of the same type. """ function init(prob::IntegralProblem, alg::IntegralAlgorithm; kwargs...) - checkkwargs(NamedTuple(kwargs)) f = prob.f; dom = prob.dom; p = prob.p - return make_cache(f, dom, p, alg; kwargs...) + kws = (; prob.kwargs..., kwargs...) + checkkwargs(kws) + cacheval = init_cacheval(f, dom, p, alg; kws...) + return IntegralCache(f, dom, p, alg, cacheval, kws) end """ @@ -103,10 +215,7 @@ tolerance can be used in combination with a smaller-than necessary absolute tole that the solution is resolved up to the requested significant digits, unless the integral is smaller than the absolute tolerance. """ -function solve(prob::IntegralProblem, alg::IntegralAlgorithm; kwargs...) - cache = init(prob, alg; kwargs...) - return solve!(cache) -end +solve(prob::IntegralProblem, alg::IntegralAlgorithm; kwargs...) """ solve!(::IntegralCache)::IntegralSolution @@ -114,130 +223,17 @@ end Compute the solution to an [`IntegralProblem`](@ref) constructed from [`init`](@ref). """ function solve!(c::IntegralCache) - return do_solve(c.f, c.dom, c.p, c.alg, c.cacheval; c.kwargs...) -end - -struct IntegralSolution{T,E} - u::T - resid::E - retcode::Bool - numevals::Int -end -# a value of numevals < 0 means it was undefined - - -""" - IntegralSolver(f, dom, alg; [abstol, reltol, maxiters]) - IntegralSolver(f, lb, ub, alg::AbstractIntegralAlgorithm; [abstol, reltol, maxiters]) - -Returns a functor, `fun`, that accepts input parameters `p` and solves the corresponding -integral `fun(p) -> solve(IntegralProblem(f, lb, ub, p), alg).u`. See [`solve`](@ref) for -details on the keywords. - -If `f` is a [`ParameterIntegrand`](@ref) or [`FourierIntegrand`](@ref), then the functor -interface is modified to accept parameters as function arguments, and the following is done: -`fun(args...; kwargs...) -> solve(IntegralProblem(f, lb, ub, merge(f.p, -MixedParameters(args...; kwargs...))), alg).u` where `f.p` are the preset parameters of `f`. -""" -struct IntegralSolver{F,D,A,T,K} <: Function - f::F - dom::D - alg::A - cacheval::T - kwargs::K -end - -# some algorithms can already allocate a rule based on `dom`, but we leave extending this -# method as an optimization to the caller -init_solver_cacheval(f, dom, alg) = nothing - -function IntegralSolver(f, dom, alg::IntegralAlgorithm; kwargs...) - checkkwargs(NamedTuple(kwargs)) - cacheval = init_solver_cacheval(f, dom, alg) - return IntegralSolver(f, dom, alg, cacheval, NamedTuple(kwargs)) -end - -function IntegralSolver(f, a::T, b::T, alg::IntegralAlgorithm; kwargs...) where {T} - dom = T <: Number ? PuncturedInterval((a, b)) : HyperCube(a, b) - return IntegralSolver(f, dom, alg; kwargs...) + return do_integral(c.f, c.dom, c.p, c.alg, c.cacheval; c.kwargs...) end -function IntegralSolver(prob::IntegralProblem, alg::IntegralAlgorithm; kwargs...) - return IntegralSolver(prob.f, prob.dom, alg; kwargs...) +@enum ReturnCode begin + Success + Failure + MaxIters end - -# layer to intercept the problem parameters & algorithm and transform them -remake_cache(args...) = IntegralCache(args...) -remake_cache(c, p) = remake_cache(c.f, c.dom, p, c.alg, c.cacheval, c.kwargs) - -function solve_p(s::IntegralSolver, p) - c = if s.cacheval===nothing - prob = IntegralProblem(s.f, s.dom, p) - init(prob, s.alg; s.kwargs...) - else - remake_cache(s, p) - end - return solve!(c) -end - -function (s::IntegralSolver)(p) - sol = solve_p(s, p) - return sol.u -end - -# parallelization - -""" - batchparam(ps, nthreads) - -If the cost of a calculation smoothly varies with the parameters `ps`, then -batch `ps` into `nthreads` groups where the `i`th element of group `j` is -`ps[j+(i-1)*nthreads]` along the longest axis of `ps`. We assume that multidimensional -arrays of parameters have smoothest cost along their longest axis -""" -function batchparam(xs::AbstractArray{T,N}, nthreads) where {T,N} - (s = size(xs)) === () && return (((CartesianIndex(()), only(xs)),),) - @assert nthreads >= 1 - len, dim = findmax(s) - batches = [Tuple{CartesianIndex{N},T}[] for _ in 1:min(nthreads, len)] - for i in CartesianIndices(xs) - push!(batches[mod(i[dim]-1, nthreads)+1], (i, xs[i])) - end - return batches -end - -function batchsolve!(out, f, ps, nthreads, callback) - n = Threads.Atomic{Int}(0) - Threads.@threads for batch in batchparam(ps, nthreads) - f_ = Threads.threadid() == 1 ? f : deepcopy(f) # avoid data races for in place integrators - for (i, p) in batch - t = time() - sol = solve_p(f_, p) - callback(f, i, Threads.atomic_add!(n, 1) + 1, p, sol, time() - t) - out[i] = sol.u - end - end - return out -end - -solver_type(::F, ::P) where {F,P} = Base.promote_op((f, p) -> solve_p(f, p).u, F, P) - -""" - batchsolve(f::IntegralSolver, ps::AbstractArray, [T]; nthreads=Threads.nthreads()) - -Evaluate the [`IntegralSolver`](@ref) `f` at each of the parameters `ps` in -parallel. Returns an array similar to `ps` containing the evaluated integrals `I`. This is -a form of multithreaded broadcasting. Providing the return type `f(eltype(ps))::T` is -optional, but will help in case inference of that type fails. -""" -function batchsolve(f::IntegralSolver, ps::AbstractArray, T=solver_type(f, ps[begin]); nthreads=Threads.nthreads(), callback=(x...)->nothing) - solver = if f.cacheval === nothing - prob = IntegralProblem(f.f, f.dom, ps[begin]) - cache = init(prob, f.alg; f.kwargs...) - IntegralSolver(f.f, f.dom, f.alg, cache.cacheval, f.kwargs) - else - f - end - return batchsolve!(similar(ps, T), solver, ps, nthreads, callback) +struct IntegralSolution{T,S} + value::T + retcode::ReturnCode + stats::S end diff --git a/src/parameters.jl b/src/parameters.jl index fa19f0d..c9d52ce 100644 --- a/src/parameters.jl +++ b/src/parameters.jl @@ -104,12 +104,6 @@ function remake_cache(f::ParameterIntegrand, dom, p, alg, cacheval, kwargs) return remake_integrand_cache(new, dom, merge(f.p, p), alg, cacheval, kwargs) end -function (s::IntegralSolver{<:ParameterIntegrand})(args...; kwargs...) - p = MixedParameters(args...; kwargs...) - sol = solve_p(s, p) - return sol.u -end - function do_solve(f::ParameterIntegrand, dom, p, alg::EvalCounter, cacheval; kws...) n::Int = 0 function g(x, args...; kwargs...) diff --git a/test/interface_tests.jl b/test/interface_tests.jl index cef2e26..e2cf745 100644 --- a/test/interface_tests.jl +++ b/test/interface_tests.jl @@ -1,8 +1,8 @@ using Test using LinearAlgebra using AutoBZCore -using AutoBZCore: IntegralProblem, solve, MixedParameters using AutoBZCore: PuncturedInterval, HyperCube, segments, endpoints +using AutoBZCore: CubicLimits @testset "domains" begin # PuncturedInterval @@ -25,23 +25,43 @@ using AutoBZCore: PuncturedInterval, HyperCube, segments, endpoints end @testset "quadrature" begin - # QuadratureFunction QuadGKJL AuxQuadGKJL ContQuadGKJL MeroQuadGKJL a = 0.0 b = 2pi abstol=1e-5 - p = 3.0 + p=3.0 + # QuadratureFunction QuadGKJL AuxQuadGKJL ContQuadGKJL MeroQuadGKJL for (f, ref) in ( ((x,p) -> p*sin(x), 0.0), ((x,p) -> p*one(x), p*(b-a)), ((x,p) -> inv(p-cos(x)), (b-a)/sqrt(p^2-1)), ) - prob = IntegralProblem(f, a, b, p) - for alg in (QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) - @test ref ≈ solve(prob, alg, abstol=abstol).u atol=abstol + prob = IntegralProblem(f, (a, b), p; abstol) + for alg in (QuadratureFunction(), QuadGKJL(),)#(QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) + sol = solve(prob, alg) + @test ref ≈ sol.value atol=abstol end end end +@testset "commonproblem" begin + a = 1.0 + b = 2pi + abstol=1e-5 + p=3.0 + # QuadratureFunction QuadGKJL AuxQuadGKJL ContQuadGKJL MeroQuadGKJL + update! = (cache, x, p) -> cache.p = x + postsolve = (sol, x, p) -> sol.value + f = (x, p) -> p + x + ref = (b-a)*(b^2-a^2) + subprob = IntegralProblem(f, (a, b), (a+b)/2; abstol) + integrand = CommonSolveIntegralFunction(subprob, QuadGKJL(), update!, postsolve) + prob = IntegralProblem(integrand, (a, b), p; abstol) + for alg in (QuadratureFunction(), QuadGKJL(), HCubatureJL()) #(QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) + sol = solve(prob, alg) + @test ref ≈ sol.value atol=abstol + end +end + @testset "cubature" begin # HCubatureJL MonkhorstPack AutoSymPTRJL NestedQuad a = 0.0 @@ -53,14 +73,16 @@ end ((x,p) -> p*one(eltype(x)), p*(b-a)^dim), ((x,p) -> prod(y -> inv(p-cos(y)), x), ((b-a)/sqrt(p^2-1))^dim), ) - prob = IntegralProblem(f, fill(a, dim), fill(b, dim), p) + prob = IntegralProblem(f, (fill(a, dim), fill(b, dim)), p; abstol) for alg in (HCubatureJL(),) - @test ref ≈ solve(prob, alg, abstol=abstol).u atol=abstol + @test ref ≈ solve(prob, alg).value atol=abstol end + #= prob = IntegralProblem(f, Basis(b*I(dim)), p) for alg in (MonkhorstPack(), AutoSymPTRJL(),) - @test ref ≈ solve(prob, alg, abstol=abstol).u atol=abstol + @test ref ≈ solve(prob, alg, abstol=abstol).value atol=abstol end + =# end end @@ -75,18 +97,21 @@ end ((y,x,p) -> y .= p*one(only(x)), [p*(b-a)]), ((y,x,p) -> y .= inv(p-cos(only(x))), [(b-a)/sqrt(p^2-1)]), ) - integrand = InplaceIntegrand(f, [0.0]) - inplaceprob = IntegralProblem(integrand, a, b, p) - for alg in (QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), HCubatureJL(),) - @test ref ≈ solve(inplaceprob, alg, abstol=abstol).u atol=abstol + integrand = InplaceIntegralFunction(f, [0.0]) + inplaceprob = IntegralProblem(integrand, (a, b), p; abstol) + for alg in (QuadGKJL(),)#(QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), HCubatureJL(),) + @test ref ≈ solve(inplaceprob, alg).value atol=abstol end + #= inplaceprob = IntegralProblem(integrand, Basis([b;;]), p) for alg in (MonkhorstPack(), AutoSymPTRJL()) - @test ref ≈ solve(inplaceprob, alg, abstol=abstol).u atol=abstol + @test ref ≈ solve(inplaceprob, alg, abstol=abstol).value atol=abstol end + =# end end + @testset "batch" begin # QuadratureFunction AuxQuadGKJL MonkhorstPack AutoSymPTRJL a = 0.0 @@ -98,15 +123,17 @@ end ((y,x,p) -> y .= p .* one.(only.(x)), p*(b-a)), ((y,x,p) -> y .= inv.(p .- cos.(only.(x))), (b-a)/sqrt(p^2-1)), ) - integrand = BatchIntegrand(f, Float64) - batchprob = IntegralProblem(integrand, a, b, p) - for alg in (QuadratureFunction(), AuxQuadGKJL()) - @test ref ≈ solve(batchprob, alg, abstol=abstol).u atol=abstol + integrand = InplaceBatchIntegralFunction(f, zeros(1)) + batchprob = IntegralProblem(integrand, (a, b), p) + for alg in (QuadGKJL(),)#(QuadratureFunction(), AuxQuadGKJL()) + @test ref ≈ solve(batchprob, alg, abstol=abstol).value atol=abstol end + #= batchprob = IntegralProblem(integrand, Basis([b;;]), p) for alg in (MonkhorstPack(), AutoSymPTRJL()) - @test ref ≈ solve(batchprob, alg, abstol=abstol).u atol=abstol + @test ref ≈ solve(batchprob, alg, abstol=abstol).value atol=abstol end + =# end end @@ -115,20 +142,18 @@ end f(x, p) = 1.0 + p*sum(cos, x) p = 7.0 abstol=1e-3 - for dim in 1:3, alg in (QuadratureFunction(), AuxQuadGKJL()) + for dim in 1:3, alg in (QuadratureFunction(), QuadGKJL())# AuxQuadGKJL()) ref = (2pi)^dim dom = CubicLimits(zeros(dim), 2pi*ones(dim)) - prob = IntegralProblem(f, dom, p) + prob = IntegralProblem(f, dom, p; abstol) ndalg = NestedQuad(alg) - @test ref ≈ solve(prob, ndalg, abstol=abstol).u atol=abstol - inplaceprob = IntegralProblem(InplaceIntegrand((y,x,p) -> y .= f(x,p), [0.0]), dom, p) - @test [ref] ≈ solve(inplaceprob, ndalg, abstol=abstol).u atol=abstol - batchprob = IntegralProblem(BatchIntegrand((y,x,p) -> y .= f.(x,Ref(p)), Float64), dom, p) - @test ref ≈ solve(batchprob, ndalg, abstol=abstol).u atol=abstol - nestedbatchprob = IntegralProblem(NestedBatchIntegrand(ntuple(n -> f, 3), Float64), dom, p) - @test ref ≈ solve(nestedbatchprob, ndalg, abstol=abstol).u atol=abstol + @test ref ≈ solve(prob, ndalg).value atol=abstol + # inplaceprob = IntegralProblem(InplaceIntegrand((y,x,p) -> y .= f(x,p), [0.0]), dom, p) + # @test [ref] ≈ solve(inplaceprob, ndalg, abstol=abstol).value atol=abstol + # batchprob = IntegralProblem(BatchIntegrand((y,x,p) -> y .= f.(x,Ref(p)), Float64), dom, p) + # @test ref ≈ solve(batchprob, ndalg, abstol=abstol).value atol=abstol end - + #= # AbsoluteEstimate est_alg = QuadratureFunction() abs_alg = QuadGKJL() @@ -137,7 +162,7 @@ end f2(x, p) = inv(complex(p...) - cos(x)) prob = IntegralProblem(f2, 0.0, 2pi, (0.5, 1e-3)) abstol = 1e-5; reltol=1e-5 - @test solve(prob, alg, reltol=reltol).u ≈ solve(prob, ref_alg, abstol=abstol).u atol=abstol + @test solve(prob, alg, reltol=reltol).value ≈ solve(prob, ref_alg, abstol=abstol).value atol=abstol # EvalCounter for prob in ( @@ -156,4 +181,5 @@ end @test solve(prob, EvalCounter(alg)).numevals == numevals end end + =# end From d2d526e3ac3eae171d545b02d1ffe2c1c245ab22 Mon Sep 17 00:00:00 2001 From: lxvm Date: Wed, 5 Jun 2024 16:47:17 -0400 Subject: [PATCH 02/17] progress on aps_example, but small eta problem? --- aps_example/aps_example.jl | 79 ++++++++++++++++++++++++++-------- src/AutoBZCore.jl | 4 +- src/algorithms_meta.jl | 45 ++++++++++++------- src/brillouin.jl | 88 +++++++++++++++++++------------------- src/fourier.jl | 87 ++++++++++++++++++++++++++++++++++--- test/fourier.jl | 56 +++++++++++++++++++++++- test/interface_tests.jl | 35 +++++++++------ 7 files changed, 296 insertions(+), 98 deletions(-) diff --git a/aps_example/aps_example.jl b/aps_example/aps_example.jl index b18a316..a5088f7 100644 --- a/aps_example/aps_example.jl +++ b/aps_example/aps_example.jl @@ -1,6 +1,6 @@ -import Pkg -Pkg.activate(".") # reproducible environment included -Pkg.instantiate() # install dependencies +# import Pkg +# Pkg.activate(".") # reproducible environment included +# Pkg.instantiate() # install dependencies using WannierIO @@ -20,36 +20,81 @@ for (i, h, n) in zip(hrdat.Rvectors, hrdat.H, hrdat.Rdegens) H_R[CartesianIndex(Tuple(i))] = h / n end -using AutoBZCore, LinearAlgebra +using FourierSeriesEvaluators, LinearAlgebra -bz = load_bz(CubicSymIBZ(), "svo.wout") -# bz = load_bz(IBZ(), "svo.wout") # works with SymmetryReduceBZ.jl installed h = FourierSeries(H_R, period=1.0) η = 1e-2 # 10 meV (scattering amplitude) -dos_integrand(h_k::FourierValue, η, ω) = -imag(tr(inv((ω+im*η)*I - h_k.s)))/pi -integrand = FourierIntegrand(dos_integrand, h, η) +ω_min = 10 +ω_max = 15 +p0 = (; η, ω=(ω_min + ω_max)/2) # initial parameters +# BUG cannot redefine this function without breaking functionwrappers +# https://github.com/JuliaLang/julia/issues/52635#issuecomment-2150808569 +greens_function(k, h_k, (; η, ω)) = tr(inv((ω+im*η)*I - h_k)) +prototype = let k = FourierSeriesEvaluators.period(h) + greens_function(k, h(k), p0) +end -dos_solver_iai = IntegralSolver(integrand, bz, IAI(); abstol=1e-3) -dos_solver_ptr = IntegralSolver(integrand, bz, PTR(npt=100); abstol=1e-3) +using AutoBZCore +bz = load_bz(CubicSymIBZ(), "svo.wout") +# bz = load_bz(IBZ(), "svo.wout") # works with SymmetryReduceBZ.jl installed -using HChebInterp +integrand = FourierIntegralFunction(greens_function, h, prototype) +prob_dos = AutoBZProblem(integrand, bz, p0; abstol=1e-3) -dos_iai = hchebinterp(dos_solver_iai, 10, 15; atol=1e-2) -dos_ptr = hchebinterp(dos_solver_ptr, 10, 15; atol=1e-2) +using HChebInterp +cheb_order = 15 +#= +batch_iai = let prob = prob_dos, alg = IAI(QuadGKJL()), nthreads=min(cheb_order+1, Threads.nthreads()) + cache = init(prob, alg) + ω -> begin + cache.p = (; cache.p..., ω) + solve!(cache).value + end + #= + caches = [init(prob, alg) for _ in 1:nthreads] + BatchFunction() do ωs + out = Vector{typeof(prototype)}(undef, length(ωs)) + Threads.@threads for i in 1:nthreads + cache = caches[i] + for j in i:nthreads:length(ωs) + ω = ωs[j] + cache.p = (; cache.p..., ω) + out[j] = solve!(cache).value + end + end + return out + end + =# +end +=# +batch_iai = let cache = init(prob_dos, IAI((QuadGKJL(), QuadGKJL(), QuadGKJL()), (AutoBZCore.FullSpecialize(), AutoBZCore.FunctionWrapperSpecialize(), AutoBZCore.FullSpecialize()))) + ω -> begin + cache.p = (; cache.p..., ω) + solve!(cache).value + end +end +@time greens_iai = hchebinterp(batch_iai, ω_min, ω_max; atol=1e-2) +#= +batch_ptr = let prob = prob_dos, nthreads=cheb_order+1 + ω -> batchsolve(prob, PTR(; nthreads), ω; nthreads=1) +end +greens_ptr = hchebinterp(dos_solver_ptr, 10, 15; atol=1e-2) +=# using CairoMakie set_theme!(fontsize=24, linewidth=4) fig1 = Figure() -ax1 = Axis(fig1[1,1], limits=((10,15), (0,det(bz.B)*6)), xlabel="ω (eV)", ylabel="SVO DOS (eV⁻¹ Å⁻³)") -p1 = lines!(ax1, 10:η/100:15, dos_iai; label="IAI(), η=$η") +ax1 = Axis(fig1[1,1], limits=((10,15), (0,6)), xlabel="ω (eV)", ylabel="SVO DOS (eV⁻¹)") +p1 = lines!(ax1, 10:η/100:15, ω -> -imag(greens_iai(ω))/pi/det(bz.B); label="IAI, η=$η") axislegend(ax1) save("iai_svo_dos.pdf", fig1) - +#= fig2 = Figure() ax2 = Axis(fig2[1,1], limits=((10,15), (0,det(bz.B)*6)), xlabel="ω (eV)", ylabel="SVO DOS (eV⁻¹ Å⁻³)") -p2 = lines!(ax2, 10:η/100:15, dos_ptr; label="PTR(), η=$η") +p2 = lines!(ax2, 10:η/100:15, ω -> -imag(greens_ptr(ω))/pi; label="PTR, η=$η") axislegend(ax2) save("ptr_svo_dos.pdf", fig2) +=# diff --git a/src/AutoBZCore.jl b/src/AutoBZCore.jl index 3025ff0..8f4bcea 100644 --- a/src/AutoBZCore.jl +++ b/src/AutoBZCore.jl @@ -83,12 +83,12 @@ include("algorithms_meta.jl") export SymmetricBZ, nsyms export load_bz, FBZ, IBZ, InversionSymIBZ, CubicSymIBZ -export AbstractSymRep, SymRep, UnknownRep, TrivialRep +export AbstractSymRep, UnknownRep, TrivialRep export AutoBZProblem export IAI, PTR, AutoPTR, TAI, PTR_IAI, AutoPTR_IAI include("brillouin.jl") -export FourierIntegrand, FourierValue +export FourierIntegralFunction include("fourier.jl") export DOSProblem diff --git a/src/algorithms_meta.jl b/src/algorithms_meta.jl index 4d4daa9..0dceb72 100644 --- a/src/algorithms_meta.jl +++ b/src/algorithms_meta.jl @@ -12,13 +12,15 @@ stability of the integrand, so you should always pick the widest integration lim that inference works properly. For example, if [`ContQuadGKJL`](@ref) is used as an algorithm in the nested scheme, then the limits of integration should be made complex. """ -struct NestedQuad{T} <: IntegralAlgorithm +struct NestedQuad{T,S} <: IntegralAlgorithm algs::T - NestedQuad(alg::IntegralAlgorithm) = new{typeof(alg)}(alg) - NestedQuad(algs::Tuple{Vararg{IntegralAlgorithm}}) = new{typeof(algs)}(algs) + specialize::S + NestedQuad(alg::IntegralAlgorithm, specialize::AbstractSpecialization=FunctionWrapperSpecialize()) = new{typeof(alg),typeof(specialize)}(alg, specialize) + NestedQuad(algs::Tuple{Vararg{IntegralAlgorithm}}, specialize::Tuple{Vararg{AbstractSpecialization}}=ntuple(_->FunctionWrapperSpecialize(), length(algs))) = new{typeof(algs),typeof(specialize)}(algs, specialize) end NestedQuad(algs::IntegralAlgorithm...) = NestedQuad(algs) +#= # this function helps create a tree of the cachevals used by each quadrature function nested_cacheval(f::F, p::P, algs, segs, lims, state, x, xs...) where {F,P} dom = PuncturedInterval(segs) @@ -67,20 +69,20 @@ function nested_cacheval(f::F, p::P, algs, segs, lims, state, x, xs...) where {F end end end -#= function init_cacheval(f, dom::AbstractIteratedLimits, p, alg::NestedQuad) algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(dom))) : alg.algs return nested_cacheval(f, p, algs, limit_iterate(dom)..., interior_point(dom)...) end =# -function _update!(cache, x, lims_state) +function _update!(cache, x, (; p, lims_state)) segs, lims, state = limit_iterate(lims_state..., x) len = segs[end] - segs[begin] kws = cache.kwargs + cache.p = p cache.cacheval.dom = segs cache.cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - cache.cacheval.p = (lims, state) + cache.cacheval.p = (; cache.cacheval.p..., lims_state=(lims, state)) return end _postsolve(sol, x, p) = sol.value @@ -91,11 +93,10 @@ function init_cacheval(f::IntegralFunction, nextdom, p, alg::NestedQuad; kws...) nothing, nextdom end algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(lims))) : alg.algs + spec = alg.specialize isa AbstractSpecialization ? ntuple(i -> alg.specialize, Val(ndims(lims))) : alg.specialize proto = get_prototype(f, x0, p) func = if ndims(lims) == 1 - IntegralFunction(proto) do x, lims_state - f.f(limit_iterate(lims_state..., x), p) - end + inner_integralfunction(f, proto) else len = segs[end] - segs[begin] a, b, = segs @@ -103,16 +104,29 @@ function init_cacheval(f::IntegralFunction, nextdom, p, alg::NestedQuad; kws...) next = limit_iterate(lims, state, x) kws = NamedTuple(kws) kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - subprob = IntegralProblem(IntegralFunction(f.f, proto), next, p; kwargs...) - CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1]), _update!, _postsolve, proto*x^(ndims(lims)-1), FunctionWrapperSpecialize()) + integrand = outer_integralfunction(f, proto) + subprob = IntegralProblem(integrand, next, p; kwargs...) + CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1], spec[1:ndims(lims)-1]), _update!, _postsolve, proto*x^(ndims(lims)-1), spec[ndims(lims)]) end - prob = IntegralProblem(func, segs, (lims, state); kws...) + prob = IntegralProblem(func, segs, (; p, lims_state=(lims, state)); kws...) return init(prob, algs[ndims(lims)]) + # the order of updates is somewhat tricky. I think some could be simplified if instead + # we use an IntegralProblem modified to contain lims_state, instead of passing the + # parameter as well end function do_integral(f, dom, p, alg::NestedQuad, cacheval; kws...) + cacheval.p = (; cacheval.p..., p) return solve!(cacheval) end +function inner_integralfunction(f::IntegralFunction, proto) + IntegralFunction(proto) do x, (; p, lims_state) + f.f(limit_iterate(lims_state..., x), p) + end +end +function outer_integralfunction(f::IntegralFunction, proto) + return IntegralFunction(f.f, proto) +end #= function init_nest(f::F, fxx, dom, p,lims, state, algs, cacheval; kws_...) where {F} kws = NamedTuple(kws_) @@ -243,7 +257,7 @@ function AbsoluteEstimate(est_alg, abs_alg; norm=norm, kwargs...) checkkwargs(kws) return AbsoluteEstimate(est_alg, abs_alg, norm, kws) end - +#= function init_cacheval(f, dom, p, alg::AbsoluteEstimate) return (est=init_cacheval(f, dom, p, alg.est_alg), abs=init_cacheval(f, dom, p, alg.abs_alg)) @@ -258,7 +272,7 @@ function do_solve(f, dom, p, alg::AbsoluteEstimate, cacheval; return do_solve(f, dom, p, alg.abs_alg, cacheval.abs; abstol=atol, reltol=zero(rtol), maxiters=maxiters) end - +=# """ EvalCounter(::IntegralAlgorithm) @@ -269,7 +283,7 @@ The count is stored in the `sol.numevals` field. struct EvalCounter{T<:IntegralAlgorithm} <: IntegralAlgorithm alg::T end - +#= function init_cacheval(f, dom, p, alg::EvalCounter) return init_cacheval(f, dom, p, alg.alg) end @@ -296,3 +310,4 @@ function do_solve(f, dom, p, alg::EvalCounter, cacheval; kws...) return IntegralSolution(sol.u, sol.resid, sol.retcode, n) end end +=# diff --git a/src/brillouin.jl b/src/brillouin.jl index cd4abdb..a8c8e84 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -71,19 +71,6 @@ Symmetry representation of objects with trivial transformation under the group. """ struct TrivialRep <: AbstractSymRep end - -""" - SymRep(f) - -`SymRep` specifies the symmetry representation of the integral of the function -`f`. When you define a new integrand, you can choose to implement this trait to -specify how the integral is transformed under the symmetries of the lattice in -order to map the integral of `f` on the IBZ to the result for the FBZ. - -New types for `SymRep` should also extend a corresponding method for -[`AutoBZCore.symmetrize_`](@ref). -""" -SymRep(::Any) = UnknownRep() const TrivialRepType = Union{Number,AbstractArray{<:Any,0}} """ @@ -106,6 +93,7 @@ the result of an integral on the FBZ from `x`, which was calculated on the IBZ. """ symmetrize_(::TrivialRep, bz::SymmetricBZ, x) = nsyms(bz)*x symmetrize_(::UnknownRep, ::SymmetricBZ, x) = x +symmetrize_(::UnknownRep, bz::SymmetricBZ, x::TrivialRepType) = symmetrize_(TrivialRep(), bz, x) symmetrize(_, ::FullBZ, x) = x symmetrize(_, ::FullBZ, x::TrivialRepType) = x @@ -341,8 +329,23 @@ struct AutoBZProblem{R<:AbstractSymRep,F<:AbstractIntegralFunction,BZ<:Symmetric p::P kwargs::K end + +const WARN_UNKNOWN_SYMMETRY = """ +A symmetric BZ was used with an integrand whose symmetry representation is unknown. +For correctness, the calculation will proceed on the full BZ. +However, it is better either to integrate without symmetries or to use symmetries by +implementing an AbstractSymRep for your type. +""" + function AutoBZProblem(rep::AbstractSymRep, f::AbstractIntegralFunction, bz::SymmetricBZ, p=NullParameters(); kws...) - return AutoBZProblem(rep, f, bz, p, NamedTuple(kws)) + proto = get_prototype(f, interior_point(bz.lims), p) + if rep isa UnknownRep && !(bz isa FullBZ) && !(proto isa TrivialRepType) + @warn WARN_UNKNOWN_SYMMETRY + fbz = SymmetricBZ(bz.A, bz.B, lattice_bz_limits(bz.B), nothing) + return AutoBZProblem(rep, f, fbz, p, NamedTuple(kws)) + else + return AutoBZProblem(rep, f, bz, p, NamedTuple(kws)) + end end function AutoBZProblem(rep::AbstractSymRep, f, bz::SymmetricBZ, p=NullParameters(); kws...) return AutoBZProblem(IntegralFunction(f), bz, p; kws...) @@ -365,7 +368,7 @@ function init(prob::AutoBZProblem, alg::AutoBZAlgorithm; kwargs...) rep = prob.rep; f = prob.f; bz = prob.bz; p = prob.p kws = (; prob.kwargs..., kwargs...) checkkwargs(kws) - cacheval = init_cacheval(rep, f, bz, p, alg) + cacheval = init_cacheval(rep, f, bz, p, alg; kws...) return AutoBZCache(rep, f, bz, p, alg, cacheval, kws) end @@ -380,41 +383,37 @@ solve(prob::AutoBZProblem, alg::AutoBZAlgorithm; kwargs...) Compute the solution to an [`IntegralProblem`](@ref) constructed from [`init`](@ref). """ function solve!(c::AutoBZCache) - return do_solve_autobz(c.rep, c.f, c.dom, c.p, c.alg, c.cacheval; c.kwargs...) + return do_solve_autobz(c.rep, c.f, c.bz, c.p, c.alg, c.cacheval; c.kwargs...) end -#= -function init_cacheval(f, bz::SymmetricBZ, p, bzalg::AutoBZAlgorithm) - _, dom, alg = bz_to_standard(bz, bzalg) - return init_cacheval(f, dom, p, alg) + +function init_cacheval(rep, f, bz::SymmetricBZ, p, bzalg::AutoBZAlgorithm; kws...) + prob, alg = bz_to_standard(f, bz, p, bzalg; kws...) + return init(prob, alg) end function do_solve(f, bz::SymmetricBZ, p, bzalg::AutoBZAlgorithm, cacheval; kws...) do_solve_autobz(bz_to_standard, f, bz, p, bzalg, cacheval; kws...) end -const WARN_UNKNOWN_SYMMETRY = """ -A symmetric BZ was used with an integrand whose symmetry representation is unknown. -For correctness, the calculation will be repeated on the full BZ. -However, it is better either to integrate without symmetries or to use symmetries by extending SymRep for your type. -""" -function do_solve_autobz(bz_to_standard, f, bz, p, bzalg::AutoBZAlgorithm, cacheval; _kws...) - bz_, dom, alg = bz_to_standard(bz, bzalg) - - j = abs(det(bz_.B)) # rescale tolerance to (I)BZ coordinate and get the right number of digits +function do_solve_autobz(rep, f, bz, p, bzalg::AutoBZAlgorithm, cacheval; _kws...) + j = abs(det(bz.B)) # rescale tolerance to (I)BZ coordinate and get the right number of digits kws = NamedTuple(_kws) - kws_ = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol / (j * nsyms(bz_)),)) : kws + cacheval.f = f + cacheval.p = p + cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol / (j * nsyms(bz)),)) : kws - sol = do_solve(f, dom, p, alg, cacheval; kws_...) + sol = solve!(cacheval) # TODO find a way to throw a warning when constructing the problem instead of after a solve - SymRep(f) isa UnknownRep && !(bz_ isa FullBZ) && !(sol.u isa TrivialRepType) && begin + rep isa UnknownRep && !(bz isa FullBZ) && !(sol.value isa TrivialRepType) && begin @warn WARN_UNKNOWN_SYMMETRY - fbz = SymmetricBZ(bz_.A, bz_.B, lattice_bz_limits(bz_.B), nothing) - _cacheval = init_cacheval(f, fbz, p, bzalg) - return do_solve(f, fbz, p, bzalg, _cacheval; _kws...) + error("not implemented") + # fbz = SymmetricBZ(bz_.A, bz_.B, lattice_bz_limits(bz_.B), nothing) + # _cacheval = init_cacheval(f, fbz, p, bzalg) + # return do_solve(f, fbz, p, bzalg, _cacheval; _kws...) end - val = j*symmetrize(f, bz_, sol.u) - err = sol.resid === nothing ? nothing : j*symmetrize(f, bz_, sol.resid) - return IntegralSolution(val, err, sol.retcode, sol.numevals) + value = j*symmetrize_(rep, bz, sol.value) + # err = sol.resid === nothing ? nothing : j*symmetrize(f, bz_, sol.resid) + return IntegralSolution(value, sol.retcode, sol.stats) end # AutoBZAlgorithms must implement: @@ -428,17 +427,18 @@ Iterated-adaptive integration using `nested_quad` from [IteratedIntegration.jl](https://github.com/lxvm/IteratedIntegration.jl). **This algorithm is the most efficient for localized integrands**. """ -struct IAI{T} <: AutoBZAlgorithm +struct IAI{T,S} <: AutoBZAlgorithm algs::T - IAI(alg::IntegralAlgorithm=AuxQuadGKJL()) = new{typeof(alg)}(alg) - IAI(algs::Tuple{Vararg{IntegralAlgorithm}}) = new{typeof(algs)}(algs) + specialize::S + IAI(alg::IntegralAlgorithm=AuxQuadGKJL(), specialize::AbstractSpecialization=FunctionWrapperSpecialize()) = new{typeof(alg),typeof(specialize)}(alg, specialize) + IAI(algs::Tuple{Vararg{IntegralAlgorithm}}, specialize::Tuple{Vararg{AbstractSpecialization}}=ntuple(_->FunctionWrapperSpecialize(),length(algs))) = new{typeof(algs),typeof(specialize)}(algs, specialize) end IAI(algs::IntegralAlgorithm...) = IAI(algs) -function bz_to_standard(bz::SymmetricBZ, alg::IAI) - return bz, bz.lims, NestedQuad(alg.algs) +function bz_to_standard(f, bz, p, bzalg::IAI; kws...)#(bz::SymmetricBZ, alg::IAI) + return IntegralProblem(f, bz.lims, p; kws...), NestedQuad(bzalg.algs, bzalg.specialize) end - +#= """ PTR(; npt=50, nthreads=1) diff --git a/src/fourier.jl b/src/fourier.jl index b37d979..31beec8 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -22,20 +22,97 @@ # parallelization """ - FourierIntegralFunction(f, s) + FourierIntegralFunction(f, s, [prototype]) ## Arguments -- `f::AbstractIntegralFunction`: The integrand, accepting inputs `f(x, s(x), p)` +- `f`: The integrand, accepting inputs `f(x, s(x), p)` - `s::AbstractFourierSeries`: The Fourier series to evaluate """ -struct FourierIntegralFunction{F<:AbstractIntegralFunction,S<:AbstractFourierSeries} <: AbstractIntegralFunction +struct FourierIntegralFunction{F,S,P} <: AbstractIntegralFunction f::F s::S + prototype::P end +FourierIntegralFunction(f, s::AbstractFourierSeries) = FourierIntegralFunction(f, s, nothing) -function get_prototype(f::FourierIntegralFunction, x, p) - return get_prototype(f.f, x, p) + +function get_prototype(f::FourierIntegralFunction, x, ws, p) + f.prototype === nothing ? f.f(x, ws(x), p) : f.prototype +end +get_prototype(f::FourierIntegralFunction, x, p) = get_prototype(f, x, f.s, p) + +function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadGKJL; kws...) + segs = PuncturedInterval(dom) + ws = FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + prototype = get_prototype(f, get_prototype(segs), ws, p) + return init_segbuf(prototype, segs, alg), ws +end +function call_quadgk(f::FourierIntegralFunction, p, u, usegs, cacheval; kws...) + segbuf, ws = cacheval + quadgk(x -> (ux = u*x; f.f(ux, ws(ux), p)), usegs...; kws..., segbuf) +end + +function init_cacheval(f::FourierIntegralFunction, dom, p, ::HCubatureJL; kws...) + # TODO utilize hcubature_buffer + ws = FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + return ws +end +function hcubature_integrand(f::FourierIntegralFunction, p, a, b, ws) + x -> f.f(x, ws(x), p) +end + +function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadratureFunction; kws...) + x, w = alg.fun(alg.npt) + ws = FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + # TODO: pre-allocate, later evaluate the series on the + return (rule=[(w, x) for (w,x) in zip(w,x)], buffer=nothing, ws) +end +function autosymptr_integrand(f::FourierIntegralFunction, p, segs, cacheval) + ws = cacheval.ws + x -> f.f(x, ws(x), p) end + +function _fourier_update!(cache, x, (; p, ws, lims_state)) + segs, lims, state = limit_iterate(lims_state..., x) + s = workspace_contract!(ws, x) + len = segs[end] - segs[begin] + kws = cache.kwargs + cache.p = p + cache.cacheval.dom = segs + cache.cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws + cache.cacheval.p = (; cache.cacheval.p..., ws=s, lims_state=(lims, state)) + return +end +function init_cacheval(f::FourierIntegralFunction, nextdom, p, alg::NestedQuad; kws...) + x0, (segs, lims, state) = if nextdom isa AbstractIteratedLimits + interior_point(nextdom), limit_iterate(nextdom) + else + nothing, nextdom + end + algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(lims))) : alg.algs + spec = alg.specialize isa AbstractSpecialization ? ntuple(i -> alg.specialize, Val(ndims(lims))) : alg.specialize + ws = f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + proto = get_prototype(f, x0, ws, p) + func = if ndims(lims) == 1 + IntegralFunction(proto) do x, (; p, ws, lims_state) + f.f(limit_iterate(lims_state..., x), workspace_evaluate!(ws, x), p) + end + else + len = segs[end] - segs[begin] + a, b, = segs + x = (a+b)/2 + next = limit_iterate(lims, state, x) + s = workspace_contract!(ws, x) + kws = NamedTuple(kws) + kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws + integrand = FourierIntegralFunction(f.f, s, proto) + subprob = IntegralProblem(integrand, next, p; kwargs...) + CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1], spec[1:ndims(lims)-1]), _fourier_update!, _postsolve, proto*x^(ndims(lims)-1), spec[ndims(lims)]) + end + prob = IntegralProblem(func, segs, (; p, ws, lims_state=(lims, state)); kws...) + return init(prob, algs[ndims(lims)]) +end + #= """ FourierIntegrand(f, s::AbstractFourierSeries, args...; kws...) diff --git a/test/fourier.jl b/test/fourier.jl index 150fcc8..9f237a4 100644 --- a/test/fourier.jl +++ b/test/fourier.jl @@ -1,11 +1,64 @@ using Test using LinearAlgebra using StaticArrays +using FourierSeriesEvaluators using AutoBZCore -using AutoBZCore: IntegralProblem, solve, MixedParameters +using AutoBZCore: CubicLimits using AutoBZCore: PuncturedInterval, HyperCube, segments, endpoints +@testset "FourierIntegralFunction" begin + @testset "quadrature" begin + a = 0 + b = 1 + p = 0.0 + t = 1.0 + s = FourierSeries([1, 0, 1]/2; period=t, offset=-2) + int(x, s, p) = x * s + p + ref = (b-a)*p + t*(b*sin(b/t*2pi) + t*cos(b/t*2pi) - (a*sin(a/t*2pi) + t*cos(a/t*2pi))) + abstol = 1e-5 + prob = IntegralProblem(FourierIntegralFunction(int, s), (a, b), p; abstol) + for alg in (QuadGKJL(), QuadratureFunction()) + @test solve(prob, alg).value ≈ ref atol=abstol + end + end + @testset "cubature" for dim in 2:3 + a = zeros(dim) + b = ones(dim) + p = 0.0 + t = 1.0 + s = FourierSeries([prod(x) for x in Iterators.product([(0.1, 0.5, 0.3) for i in 1:dim]...)]; period=t, offset=-2) + int(x, s, p) = prod(x) * s + p + abstol = 1e-4 + prob = IntegralProblem(FourierIntegralFunction(int, s), (a, b), p; abstol) + refprob = IntegralProblem(IntegralFunction((x, p) -> int(x, s(x), p)), (a, b), p; abstol) + for alg in (HCubatureJL(),) + @test solve(prob, alg).value ≈ solve(refprob, alg).value atol=abstol + end + end + @testset "meta-algorithms" for dim in 2:3 + # NestedQuad + a = zeros(dim) + b = ones(dim) + p0 = 0.0 + t = 1.0 + s = FourierSeries([prod(x) for x in Iterators.product([(0.1, 0.5, 0.3) for i in 1:dim]...)]; period=t, offset=-2) + int(x, s, p) = prod(x) * s + p + abstol = 1e-4 + prob = IntegralProblem(FourierIntegralFunction(int, s), CubicLimits(a, b), p0; abstol) + refprob = IntegralProblem(FourierIntegralFunction(int, s), (a, b), p0; abstol) + for alg in (QuadGKJL(),) + cache = init(prob, NestedQuad(alg)) + refcache = init(refprob, HCubatureJL()) + for p in [5.0, 6.0] + cache.p = p + refcache.p = p + @test solve!(cache).value ≈ solve!(refcache).value atol=abstol + end + end + end +end +#= @testset "FourierIntegrand" begin for dims in 1:3 s = FourierSeries(integer_lattice(dims), period=1) @@ -54,3 +107,4 @@ end end end end +=# diff --git a/test/interface_tests.jl b/test/interface_tests.jl index e2cf745..77e2c31 100644 --- a/test/interface_tests.jl +++ b/test/interface_tests.jl @@ -47,18 +47,22 @@ end a = 1.0 b = 2pi abstol=1e-5 - p=3.0 + p0=3.0 # QuadratureFunction QuadGKJL AuxQuadGKJL ContQuadGKJL MeroQuadGKJL - update! = (cache, x, p) -> cache.p = x + update! = (cache, x, p) -> cache.p = (x, p) postsolve = (sol, x, p) -> sol.value - f = (x, p) -> p + x - ref = (b-a)*(b^2-a^2) - subprob = IntegralProblem(f, (a, b), (a+b)/2; abstol) + f = (x, (y, p)) -> p*(y + x) + subprob = IntegralProblem(f, (a, b), ((a+b)/2, p0); abstol) integrand = CommonSolveIntegralFunction(subprob, QuadGKJL(), update!, postsolve) - prob = IntegralProblem(integrand, (a, b), p; abstol) + prob = IntegralProblem(integrand, (a, b), p0; abstol) for alg in (QuadratureFunction(), QuadGKJL(), HCubatureJL()) #(QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) - sol = solve(prob, alg) - @test ref ≈ sol.value atol=abstol + cache = init(prob, alg) + for p in [3.0, 4.0] + ref = p*(b-a)*(b^2-a^2) + cache.p = p + sol = solve!(cache) + @test ref ≈ sol.value atol=abstol + end end end @@ -136,18 +140,21 @@ end =# end end - @testset "multi-algorithms" begin # NestedQuad - f(x, p) = 1.0 + p*sum(cos, x) - p = 7.0 + f(x, p) = 1.0 + p*sum(abs2 ∘ cos, x) abstol=1e-3 + p0 = 0.0 for dim in 1:3, alg in (QuadratureFunction(), QuadGKJL())# AuxQuadGKJL()) - ref = (2pi)^dim dom = CubicLimits(zeros(dim), 2pi*ones(dim)) - prob = IntegralProblem(f, dom, p; abstol) + prob = IntegralProblem(f, dom, p0; abstol) ndalg = NestedQuad(alg) - @test ref ≈ solve(prob, ndalg).value atol=abstol + cache = init(prob, ndalg) + for p in [5.0, 7.0] + cache.p = p + ref = (2pi)^dim + dim*p*pi*(2pi)^(dim-1) + @test ref ≈ solve!(cache).value atol=abstol + end # inplaceprob = IntegralProblem(InplaceIntegrand((y,x,p) -> y .= f(x,p), [0.0]), dom, p) # @test [ref] ≈ solve(inplaceprob, ndalg, abstol=abstol).value atol=abstol # batchprob = IntegralProblem(BatchIntegrand((y,x,p) -> y .= f.(x,Ref(p)), Float64), dom, p) From 74d4e575c00b3631583d8e7cfb19e35b76750011 Mon Sep 17 00:00:00 2001 From: lxvm Date: Thu, 13 Jun 2024 10:27:26 -0400 Subject: [PATCH 03/17] ptr+tai working --- aps_example/aps_example.jl | 33 ++-- src/AutoBZCore.jl | 2 +- src/algorithms.jl | 2 +- src/algorithms_autosymptr.jl | 43 ++--- src/algorithms_meta.jl | 2 + src/brillouin.jl | 54 +++--- src/fourier.jl | 318 +++++++---------------------------- 7 files changed, 113 insertions(+), 341 deletions(-) diff --git a/aps_example/aps_example.jl b/aps_example/aps_example.jl index a5088f7..43f3b2f 100644 --- a/aps_example/aps_example.jl +++ b/aps_example/aps_example.jl @@ -24,11 +24,11 @@ using FourierSeriesEvaluators, LinearAlgebra h = FourierSeries(H_R, period=1.0) -η = 1e-2 # 10 meV (scattering amplitude) +η = 5e-2 # 10 meV (scattering amplitude) ω_min = 10 ω_max = 15 p0 = (; η, ω=(ω_min + ω_max)/2) # initial parameters -# BUG cannot redefine this function without breaking functionwrappers +# BUG cannot redefine this function without breaking functionwrappers in v1.10 # https://github.com/JuliaLang/julia/issues/52635#issuecomment-2150808569 greens_function(k, h_k, (; η, ω)) = tr(inv((ω+im*η)*I - h_k)) prototype = let k = FourierSeriesEvaluators.period(h) @@ -45,14 +45,15 @@ prob_dos = AutoBZProblem(integrand, bz, p0; abstol=1e-3) using HChebInterp cheb_order = 15 -#= -batch_iai = let prob = prob_dos, alg = IAI(QuadGKJL()), nthreads=min(cheb_order+1, Threads.nthreads()) + +function dos_solver(prob, alg) cache = init(prob, alg) ω -> begin cache.p = (; cache.p..., ω) solve!(cache).value end - #= +end +function threaded_dos_solver(prob, alg; nthreads=min(cheb_order, Threads.nthreads())) caches = [init(prob, alg) for _ in 1:nthreads] BatchFunction() do ωs out = Vector{typeof(prototype)}(undef, length(ωs)) @@ -66,22 +67,14 @@ batch_iai = let prob = prob_dos, alg = IAI(QuadGKJL()), nthreads=min(cheb_order+ end return out end - =# end -=# -batch_iai = let cache = init(prob_dos, IAI((QuadGKJL(), QuadGKJL(), QuadGKJL()), (AutoBZCore.FullSpecialize(), AutoBZCore.FunctionWrapperSpecialize(), AutoBZCore.FullSpecialize()))) - ω -> begin - cache.p = (; cache.p..., ω) - solve!(cache).value - end -end -@time greens_iai = hchebinterp(batch_iai, ω_min, ω_max; atol=1e-2) -#= -batch_ptr = let prob = prob_dos, nthreads=cheb_order+1 - ω -> batchsolve(prob, PTR(; nthreads), ω; nthreads=1) -end -greens_ptr = hchebinterp(dos_solver_ptr, 10, 15; atol=1e-2) -=# + +dos_solver_iai = dos_solver(prob_dos, IAI(QuadGKJL())) +@time greens_iai = hchebinterp(dos_solver_iai, ω_min, ω_max; atol=1e-2, order=cheb_order) + +# dos_solver_ptr = dos_solver(prob_dos, PTR(; npt=100)) +# @time greens_ptr = hchebinterp(dos_solver_ptr, ω_min, ω_max; atol=1e-2, order=cheb_order) + using CairoMakie set_theme!(fontsize=24, linewidth=4) diff --git a/src/AutoBZCore.jl b/src/AutoBZCore.jl index 8f4bcea..03ba301 100644 --- a/src/AutoBZCore.jl +++ b/src/AutoBZCore.jl @@ -77,7 +77,7 @@ include("algorithms.jl") export AuxQuadGKJL, ContQuadGKJL, MeroQuadGKJL # include("algorithms_iterated.jl") export MonkhorstPack, AutoSymPTRJL -# include("algorithms_autosymptr.jl") +include("algorithms_autosymptr.jl") export NestedQuad, AbsoluteEstimate, EvalCounter include("algorithms_meta.jl") diff --git a/src/algorithms.jl b/src/algorithms.jl index ecc66af..252e5de 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -174,6 +174,7 @@ struct QuadratureFunction{F} <: IntegralAlgorithm nthreads::Int end QuadratureFunction(; fun=trapz, npt=50, nthreads=1) = QuadratureFunction(fun, npt, nthreads) + function init_cacheval(f::IntegralFunction, dom, p, alg::QuadratureFunction; kws...) x, w = alg.fun(alg.npt) return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing) @@ -188,7 +189,6 @@ function init_cacheval(f::InplaceIntegralFunction, dom, p, alg::QuadratureFuncti return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing, I, Itmp, y, ytmp) end function init_cacheval(f::InplaceBatchIntegralFunction, dom, p, alg::QuadratureFunction; kws...) - buf = init_buffer(f, alg.nthreads) x, w = alg.fun(alg.npt) proto=get_prototype(f, first(x), p) return (rule=[(w,x) for (w,x) in zip(w,x)], buffer=similar(proto, len), y=similar(proto, len), x=Vector{float(eltype(dom))}(undef, len)) diff --git a/src/algorithms_autosymptr.jl b/src/algorithms_autosymptr.jl index 59fe2f9..f9c8347 100644 --- a/src/algorithms_autosymptr.jl +++ b/src/algorithms_autosymptr.jl @@ -18,7 +18,7 @@ struct MonkhorstPack{S} <: IntegralAlgorithm nthreads::Int end MonkhorstPack(; npt=50, syms=nothing, nthreads=1) = MonkhorstPack(npt, syms, nthreads) -function init_rule(dom::Basis, alg::MonkhorstPack) +function init_rule(dom, alg::MonkhorstPack) # rule = AutoSymPTR.MonkhorstPackRule(alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) # return rule(eltype(dom), Val(ndims(dom))) if alg.syms === nothing @@ -31,25 +31,20 @@ end rule_type(::AutoSymPTR.PTR{N,T}) where {N,T} = SVector{N,T} rule_type(::AutoSymPTR.MonkhorstPack{N,T}) where {N,T} = SVector{N,T} -function init_cacheval(f, dom::Basis, p, alg::MonkhorstPack) - f isa NestedBatchIntegrand && throw(ArgumentError("MonkhorstPack doesn't support nested batching")) +function init_cacheval(f::IntegralFunction, dom, p, alg::MonkhorstPack; kws...) rule = init_rule(dom, alg) buf = init_buffer(f, alg.nthreads) return (rule=rule, buffer=buf) end -function do_solve(f, dom, p, alg::MonkhorstPack, cacheval; +function do_integral(f, dom, p, alg::MonkhorstPack, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - g = if f isa BatchIntegrand - xx = eltype(f.x) === Nothing ? typeof(dom*zero(rule_type(cacheval.rule)))[] : f.x - AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), f.y, xx, max_batch=f.max_batch) - elseif f isa InplaceIntegrand - AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), f.I) - else - x -> f(x, p) - end - I = cacheval.rule(g, dom, cacheval.buffer) - return IntegralSolution(I, nothing, true, -1) + g = autosymptr_integrand(f, p, dom, cacheval) + # TODO convert the domain to a Basis + value = cacheval.rule(g, dom, cacheval.buffer) + retcode = Success + stats = (; numevals=length(cacheval.rule)) + return IntegralSolution(value, retcode, stats) end """ @@ -77,11 +72,10 @@ end function AutoSymPTRJL(; norm=norm, a=1.0, nmin=50, nmax=1000, n₀=6.0, Δn=log(10), keepmost=2, syms=nothing, nthreads=1) return AutoSymPTRJL(norm, a, nmin, nmax, n₀, Δn, keepmost, syms, nthreads) end -function init_rule(dom::Basis, alg::AutoSymPTRJL) +function init_rule(dom, alg::AutoSymPTRJL) return AutoSymPTR.MonkhorstPackRule(alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) end -function init_cacheval(f, dom::Basis, p, alg::AutoSymPTRJL) - f isa NestedBatchIntegrand && throw(ArgumentError("AutoSymPTRJL doesn't support nested batching")) +function init_cacheval(f::IntegralFunction, dom, p, alg::AutoSymPTRJL; kws...) rule = init_rule(dom, alg) cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) buffer = init_buffer(f, alg.nthreads) @@ -91,15 +85,10 @@ end function do_solve(f, dom, p, alg::AutoSymPTRJL, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - g = if f isa BatchIntegrand - xx = eltype(f.x) === Nothing ? typeof(dom*zero(rule_type(cacheval.cache[1])))[] : f.x - AutoSymPTR.BatchIntegrand((y,x) -> f.f!(y,x,p), f.y, xx, max_batch=f.max_batch) - elseif f isa InplaceIntegrand - AutoSymPTR.InplaceIntegrand((y,x) -> f.f!(y,x,p), f.I) - else - x -> f(x, p) - end - val, err = autosymptr(g, dom; syms = alg.syms, rule = cacheval.rule, cache = cacheval.cache, keepmost = alg.keepmost, + g = autosymptr_integrand(f, p, dom, cacheval) + value, error = autosymptr(g, dom; syms = alg.syms, rule = cacheval.rule, cache = cacheval.cache, keepmost = alg.keepmost, abstol = abstol, reltol = reltol, maxevals = maxiters, norm=alg.norm, buffer=cacheval.buffer) - return IntegralSolution(val, err, true, -1) + retcode = error < max(something(abstol, zero(error)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(eltype(a))) : abstol)) ? Success : Failure + stats = (; error) + return IntegralSolution(value, retcode, stats) end diff --git a/src/algorithms_meta.jl b/src/algorithms_meta.jl index 0dceb72..a69b818 100644 --- a/src/algorithms_meta.jl +++ b/src/algorithms_meta.jl @@ -115,8 +115,10 @@ function init_cacheval(f::IntegralFunction, nextdom, p, alg::NestedQuad; kws...) # parameter as well end + function do_integral(f, dom, p, alg::NestedQuad, cacheval; kws...) cacheval.p = (; cacheval.p..., p) + cacheval.kwargs = (; cacheval.kwargs..., kws...) return solve!(cacheval) end function inner_integralfunction(f::IntegralFunction, proto) diff --git a/src/brillouin.jl b/src/brillouin.jl index a8c8e84..e1fcdb2 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -403,17 +403,10 @@ function do_solve_autobz(rep, f, bz, p, bzalg::AutoBZAlgorithm, cacheval; _kws.. cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol / (j * nsyms(bz)),)) : kws sol = solve!(cacheval) - # TODO find a way to throw a warning when constructing the problem instead of after a solve - rep isa UnknownRep && !(bz isa FullBZ) && !(sol.value isa TrivialRepType) && begin - @warn WARN_UNKNOWN_SYMMETRY - error("not implemented") - # fbz = SymmetricBZ(bz_.A, bz_.B, lattice_bz_limits(bz_.B), nothing) - # _cacheval = init_cacheval(f, fbz, p, bzalg) - # return do_solve(f, fbz, p, bzalg, _cacheval; _kws...) - end value = j*symmetrize_(rep, bz, sol.value) + stats = (; sol.stats...) # err = sol.resid === nothing ? nothing : j*symmetrize(f, bz_, sol.resid) - return IntegralSolution(value, sol.retcode, sol.stats) + return IntegralSolution(value, sol.retcode, stats) end # AutoBZAlgorithms must implement: @@ -435,10 +428,10 @@ struct IAI{T,S} <: AutoBZAlgorithm end IAI(algs::IntegralAlgorithm...) = IAI(algs) -function bz_to_standard(f, bz, p, bzalg::IAI; kws...)#(bz::SymmetricBZ, alg::IAI) +function bz_to_standard(f, bz, p, bzalg::IAI; kws...) return IntegralProblem(f, bz.lims, p; kws...), NestedQuad(bzalg.algs, bzalg.specialize) end -#= + """ PTR(; npt=50, nthreads=1) @@ -452,8 +445,8 @@ struct PTR <: AutoBZAlgorithm end PTR(; npt=50, nthreads=1) = PTR(npt, nthreads) -function bz_to_standard(bz::SymmetricBZ, alg::PTR) - return bz, canonical_ptr_basis(bz.B), MonkhorstPack(npt=alg.npt, syms=bz.syms, nthreads=alg.nthreads) +function bz_to_standard(f, bz, p, alg::PTR; kws...) + return IntegralProblem(f, canonical_ptr_basis(bz.B), p; kws...), MonkhorstPack(npt=alg.npt, syms=bz.syms, nthreads=alg.nthreads) end @@ -478,9 +471,11 @@ end function AutoPTR(; norm=norm, a=1.0, nmin=50, nmax=1000, n₀=6.0, Δn=log(10), keepmost=2, nthreads=1) return AutoPTR(norm, a, nmin, nmax, n₀, Δn, keepmost, nthreads) end -function bz_to_standard(bz::SymmetricBZ, alg::AutoPTR) - return bz, canonical_ptr_basis(bz.B), AutoSymPTRJL(norm=alg.norm, a=alg.a, nmin=alg.nmin, nmax=alg.nmax, n₀=alg.n₀, Δn=alg.Δn, keepmost=alg.keepmost, syms=bz.syms, nthreads=alg.nthreads) + +function bz_to_standard(f, bz, p, alg::AutoPTR; kws...) + return IntegralProblem(f, canonical_ptr_basis(bz.B), p; kws...), AutoSymPTRJL(norm=alg.norm, a=alg.a, nmin=alg.nmin, nmax=alg.nmax, n₀=alg.n₀, Δn=alg.Δn, keepmost=alg.keepmost, syms=bz.syms, nthreads=alg.nthreads) end + function init_cacheval(f, bz::SymmetricBZ, p, bzalg::AutoPTR) bz_, dom, alg = bz_to_standard(bz, bzalg) f isa NestedBatchIntegrand && throw(ArgumentError("AutoSymPTRJL doesn't support nested batching")) @@ -490,20 +485,17 @@ function init_cacheval(f, bz::SymmetricBZ, p, bzalg::AutoPTR) return (rule=rule, cache=cache, buffer=buffer) end function do_solve_autobz(bz_to_standard, f, bz, p, bzalg::AutoPTR, cacheval; _kws...) - bz_, dom, alg = bz_to_standard(bz, bzalg) - j = abs(det(bz_.B)) # rescale tolerance to (I)BZ coordinate and get the right number of digits + j = abs(det(bz.B)) # rescale tolerance to (I)BZ coordinate and get the right number of digits kws = NamedTuple(_kws) - kws_ = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol / j,)) : kws + cacheval.f = f + cacheval.p = p + cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol / j,)) : kws - sol = do_solve(f, dom, p, alg, cacheval; kws_...) - # TODO find a way to throw a warning when constructing the problem instead of after a solve - SymRep(f) isa UnknownRep && !(bz_ isa FullBZ) && !(sol.u isa TrivialRepType) && begin - @warn WARN_UNKNOWN_SYMMETRY - fbz = SymmetricBZ(bz_.A, bz_.B, lattice_bz_limits(bz_.B), nothing) - _cacheval = init_cacheval(f, fbz, p, bzalg) - return do_solve(f, fbz, p, bzalg, _cacheval; _kws...) - end - return IntegralSolution(sol.u * j, sol.resid * j, sol.retcode, sol.numevals) + sol = solve!(cacheval) + value = j*sol.value + error = j*sol.stats.error + stats = (; sol.stats..., error) + return IntegralSolution(value, sol.retcode, stats) end """ @@ -519,13 +511,13 @@ struct TAI{N} <: AutoBZAlgorithm end TAI(; norm=norm, initdiv=1) = TAI(norm, initdiv) -function bz_to_standard(bz::SymmetricBZ, alg::TAI) - bz_ = bz.lims isa CubicLimits ? bz : SymmetricBZ(bz.A, bz.B, lattice_bz_limits(bz.B), nothing) - l = bz_.lims - return bz_, HyperCube(l.a, l.b), HCubatureJL(norm=alg.norm, initdiv = alg.initdiv) +function bz_to_standard(f, bz, p, alg::TAI; kws...) + @assert bz.lims isa CubicLimits "TAI can only integrate rectangular regions" + return IntegralProblem(f, HyperCube(bz.lims.a, bz.lims.b), p; kws...), HCubatureJL(norm=alg.norm, initdiv = alg.initdiv) end +#= """ PTR_IAI(; ptr=PTR(), iai=IAI()) diff --git a/src/fourier.jl b/src/fourier.jl index 31beec8..e695908 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -22,28 +22,64 @@ # parallelization """ - FourierIntegralFunction(f, s, [prototype]) + FourierIntegralFunction(f, s, [prototype=nothing]; alias=false) ## Arguments - `f`: The integrand, accepting inputs `f(x, s(x), p)` - `s::AbstractFourierSeries`: The Fourier series to evaluate +- `prototype`: +- `alias::Bool`: whether to `deepcopy` the series (false) or use the series as-is (true) """ struct FourierIntegralFunction{F,S,P} <: AbstractIntegralFunction f::F s::S prototype::P + alias::Bool end -FourierIntegralFunction(f, s::AbstractFourierSeries) = FourierIntegralFunction(f, s, nothing) - +FourierIntegralFunction(f, s, p=nothing; alias=false) = FourierIntegralFunction(f, s, p, alias) function get_prototype(f::FourierIntegralFunction, x, ws, p) f.prototype === nothing ? f.f(x, ws(x), p) : f.prototype end get_prototype(f::FourierIntegralFunction, x, p) = get_prototype(f, x, f.s, p) +# similar to workspace_allocate, but more type-stable because of loop unrolling and vector types +function workspace_allocate_vec(s::AbstractFourierSeries{N}, x::NTuple{N,Any}, len::NTuple{N,Integer}=ntuple(one,Val(N))) where {N} + # Only the top-level workspace has an AbstractFourierSeries in the series field + # In the lower level workspaces the series field has a cache that can be contract!-ed + # into a series + dim = Val(N) + if N == 1 + c = FourierSeriesEvaluators.allocate(s, x[N], dim) + ws = Vector{typeof(c)}(undef, len[N]) + ws[1] = c + for n in 2:len[N] + ws[n] = FourierSeriesEvaluators.allocate(s, x[N], dim) + end + else + c = FourierSeriesEvaluators.allocate(s, x[N], dim) + t = FourierSeriesEvaluators.contract!(c, s, x[N], dim) + c_ = FourierWorkspace(c, FourierSeriesEvaluators.workspace_allocate(t, x[1:N-1], len[1:N-1]).cache) + ws = Vector{typeof(c_)}(undef, len[N]) + ws[1] = c_ + for n in 2:len[N] + _c = FourierSeriesEvaluators.allocate(s, x[N], dim) + _t = FourierSeriesEvaluators.contract!(_c, s, x[N], dim) + ws[n] = FourierWorkspace(_c, FourierSeriesEvaluators.workspace_allocate(_t, x[1:N-1], len[1:N-1]).cache) + end + end + return FourierWorkspace(s, ws) +end + +struct FourierValue{X,S} + x::X + s::S +end +@inline AutoSymPTR.mymul(w, x::FourierValue) = FourierValue(AutoSymPTR.mymul(w, x.x), x.s) + function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) prototype = get_prototype(f, get_prototype(segs), ws, p) return init_segbuf(prototype, segs, alg), ws end @@ -54,7 +90,7 @@ end function init_cacheval(f::FourierIntegralFunction, dom, p, ::HCubatureJL; kws...) # TODO utilize hcubature_buffer - ws = FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) return ws end function hcubature_integrand(f::FourierIntegralFunction, p, a, b, ws) @@ -63,13 +99,13 @@ end function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadratureFunction; kws...) x, w = alg.fun(alg.npt) - ws = FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) # TODO: pre-allocate, later evaluate the series on the return (rule=[(w, x) for (w,x) in zip(w,x)], buffer=nothing, ws) end function autosymptr_integrand(f::FourierIntegralFunction, p, segs, cacheval) ws = cacheval.ws - x -> f.f(x, ws(x), p) + x -> x isa FourierValue ? f.f(x.x, x.s, p) : f.f(x, ws(x), p) end function _fourier_update!(cache, x, (; p, ws, lims_state)) @@ -91,7 +127,7 @@ function init_cacheval(f::FourierIntegralFunction, nextdom, p, alg::NestedQuad; end algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(lims))) : alg.algs spec = alg.specialize isa AbstractSpecialization ? ntuple(i -> alg.specialize, Val(ndims(lims))) : alg.specialize - ws = f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.s, FourierSeriesEvaluators.period(f.s)) + ws = f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) proto = get_prototype(f, x0, ws, p) func = if ndims(lims) == 1 IntegralFunction(proto) do x, (; p, ws, lims_state) @@ -105,7 +141,7 @@ function init_cacheval(f::FourierIntegralFunction, nextdom, p, alg::NestedQuad; s = workspace_contract!(ws, x) kws = NamedTuple(kws) kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - integrand = FourierIntegralFunction(f.f, s, proto) + integrand = FourierIntegralFunction(f.f, s, proto; alias=true) subprob = IntegralProblem(integrand, next, p; kwargs...) CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1], spec[1:ndims(lims)-1]), _fourier_update!, _postsolve, proto*x^(ndims(lims)-1), spec[ndims(lims)]) end @@ -113,67 +149,6 @@ function init_cacheval(f::FourierIntegralFunction, nextdom, p, alg::NestedQuad; return init(prob, algs[ndims(lims)]) end -#= -""" - FourierIntegrand(f, s::AbstractFourierSeries, args...; kws...) - -Outer constructor for `FourierIntegrand` that wraps the Fourier series `s` into a -single-threaded `FourierWorkspace`. -""" -function FourierIntegrand(f, s::AbstractFourierSeries, args...; kws...) - return FourierIntegrand(f, workspace_allocate_vec(s, period(s)), args...; kws...) -end - -# similar to workspace_allocate, but more type-stable because of loop unrolling and vector types -function workspace_allocate_vec(s::AbstractFourierSeries{N}, x::NTuple{N,Any}, len::NTuple{N,Integer}=ntuple(one,Val(N))) where {N} - # Only the top-level workspace has an AbstractFourierSeries in the series field - # In the lower level workspaces the series field has a cache that can be contract!-ed - # into a series - dim = Val(N) - if N == 1 - c = FourierSeriesEvaluators.allocate(s, x[N], dim) - ws = Vector{typeof(c)}(undef, len[N]) - ws[1] = c - for n in 2:len[N] - ws[n] = FourierSeriesEvaluators.allocate(s, x[N], dim) - end - else - c = FourierSeriesEvaluators.allocate(s, x[N], dim) - t = FourierSeriesEvaluators.contract!(c, s, x[N], dim) - c_ = FourierWorkspace(c, workspace_allocate_vec(t, x[1:N-1], len[1:N-1]).cache) - ws = Vector{typeof(c_)}(undef, len[N]) - ws[1] = c_ - for n in 2:len[N] - _c = FourierSeriesEvaluators.allocate(s, x[N], dim) - _t = FourierSeriesEvaluators.contract!(_c, s, x[N], dim) - ws[n] = FourierWorkspace(_c, workspace_allocate_vec(_t, x[1:N-1], len[1:N-1]).cache) - end - end - return FourierWorkspace(s, ws) -end - -# FourierIntegrands should expect a FourierValue as input - -""" - FourierValue(x, s) - -A container used by [`FourierIntegrand`](@ref) to pass a point, `x`, and the value of a -Fourier series evaluated at the point, `s`, to integrands. The properties `x` and `s` of a -`FourierValue` store the point and evaluated series, respectively. -""" -struct FourierValue{X,S} - x::X - s::S -end -Base.convert(::Type{T}, f::FourierValue) where {T<:FourierValue} = T(f.x,f.s) - -Base.zero(::Type{FourierValue{X,S}}) where {X,S} = FourierValue(zero(X), zero(S)) -Base.:*(B::AbstractMatrix, f::FourierValue) = FourierValue(B*f.x, f.s) - -# (f::FourierIntegrand)(x::FourierValue, p) = f.f(x, p) -# fallback evaluator of FourierIntegrand for algorithms without specialized rules -# (f::FourierIntegrand)(x, p) = f(FourierValue(x, f.w(x)), p) - # PTR rules # no symmetries @@ -280,7 +255,7 @@ function _fourier_symptr!(vals::AbstractVector, w::FourierWorkspace, x::Abstract @inbounds wi = wsym[i, idx...] iszero(wi) && continue @inbounds xi = x[i] - vals[o+(n+=1)] = (wi, FourierValue((xi, coord...), workspace_evaluate!(w, t*xi))) + vals[o+(n+=1)] = (wi, FourierValue(SVector(xi, coord...), workspace_evaluate!(w, t*xi))) end return vals end @@ -380,7 +355,7 @@ end # end rule_type(::FourierPTR{N,T,S}) where {N,T,S} = FourierValue{SVector{N,T},S} -function init_fourier_rule(w::FourierWorkspace, dom::Basis, alg::MonkhorstPack) +function init_fourier_rule(w::FourierWorkspace, dom, alg::MonkhorstPack) @assert ndims(w.series) == ndims(dom) if alg.syms === nothing return FourierPTR(w, eltype(dom), Val(ndims(dom)), alg.npt) @@ -388,198 +363,19 @@ function init_fourier_rule(w::FourierWorkspace, dom::Basis, alg::MonkhorstPack) return FourierMonkhorstPack(w, eltype(dom), Val(ndims(dom)), alg.npt, alg.syms) end end -# function init_cacheval(f::FourierIntegrand, dom::Basis , p, alg::MonkhorstPack) -# rule = init_fourier_rule(f.w, dom, alg) -# buf = init_buffer(f, alg.nthreads) -# return (rule=rule, buffer=buf) -# end +function init_cacheval(f::FourierIntegralFunction, dom , p, alg::MonkhorstPack; kws...) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + rule = init_fourier_rule(ws, dom, alg) + return (; rule, buffer=nothing, ws) +end function init_fourier_rule(w::FourierWorkspace, dom::Basis, alg::AutoSymPTRJL) @assert ndims(w.series) == ndims(dom) return FourierMonkhorstPackRule(w, alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) end -# function init_cacheval(f::FourierIntegrand, dom::Basis, p, alg::AutoSymPTRJL) -# rule = init_fourier_rule(f.w, dom, alg) -# cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) -# buffer = init_buffer(f, alg.nthreads) -# return (rule=rule, cache=cache, buffer=buffer) -# end -# function init_cacheval(f::FourierIntegrand, bz::SymmetricBZ, p, bzalg::AutoPTR) -# bz_, dom, alg = bz_to_standard(bz, bzalg) -# rule = SymmetricRuleDef(init_fourier_rule(f.w, dom, alg), SymRep(f), bz_) -# cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) -# buffer = init_buffer(f, alg.nthreads) -# return (rule=rule, cache=cache, buffer=buffer) -# end -function init_fourier_rule(s::AbstractFourierSeries, bz::SymmetricBZ, alg::PTR) - dom = Basis(bz.B) - return FourierMonkhorstPack(s, eltype(dom), Val(ndims(dom)), alg.npt, bz.syms) -end - -# function nested_to_batched(f::FourierIntegrand, dom::Basis, rule) -# ys = f.nest.y/prod(ntuple(n -> oneunit(eltype(dom)), Val(ndims(dom)))) -# xs = rule_type(rule)[] -# return BatchIntegrand(ys, xs, max_batch=f.nest.max_batch) do y, x, p -# # would be better to fully unwrap the nested structure, but this is one level -# nchunk = length(f.nest.f) -# Threads.@threads for ichunk in 1:min(nchunk, length(x)) -# for (i, j) in zip(getchunk(x, ichunk, nchunk, :batch), getchunk(y, ichunk, nchunk, :batch)) -# y[j] = FourierIntegrand(f.nest.f[ichunk], FourierWorkspace(nothing,nothing))(x[i], p) -# end -# end -# return nothing -# end -# end - -# function do_solve(f::FourierIntegrand, dom, p, alg::MonkhorstPack, cacheval; kws...) -# g = f.nest isa NestedBatchIntegrand ? nested_to_batched(f, dom, cacheval.rule) : f.f -# return do_solve(g, dom, p, alg, cacheval; kws...) -# end - -# function do_solve(f::FourierIntegrand, dom, p, alg::AutoSymPTRJL, cacheval; kws...) -# g = f.nest isa NestedBatchIntegrand ? nested_to_batched(f, dom, cacheval.cache[1]) : f.f -# return do_solve(g, dom, p, alg, cacheval; kws...) -# end - - -# dispatch on IAI algorithms -#= -function nested_cacheval(f::FourierIntegrand, p::P, algs, segs, lims, state, x, xs...) where {P} - dom = PuncturedInterval(segs) - a, b = segs[1], segs[2] - dim = ndims(lims) - alg = algs[dim] - mid = (a+b)/2 # sample point that should be safe to evaluate - next = limit_iterate(lims, state, mid) # see what the next limit gives - if xs isa Tuple{} # the next integral takes us to the inner integral - # base case test integrand of the inner integral - # we need to pass dummy integrands to all the outer integrals so that they can build - # their caches with the right types - if f.nest isa NestedBatchIntegrand - # Batch integrate the inner integral only - cacheval = init_cacheval(BatchIntegrand(nothing, f.nest.y, f.nest.x, max_batch=f.nest.max_batch), dom, p, alg) - return (nothing, cacheval, zero(eltype(f.nest.y))*mid) - else - fx = f(next,p) - cacheval = init_cacheval((x, p) -> fx, dom, p, alg) - return (nothing, cacheval, fx*mid) - end - elseif f.nest isa NestedBatchIntegrand - algs_ = algs[1:dim-1] - # numbered names to avoid type instabilities (we are not using dispatch, but the - # compiler's optimization for the recursive function's argument types) - nest0 = nested_cacheval(FourierIntegrand(f.f, f.w, f.nest.f[1]), p, algs_, next..., x, xs[1:dim-2]...) - cacheval = init_cacheval(BatchIntegrand(nothing, f.nest.y, f.nest.x, max_batch=f.nest.max_batch), dom, p, alg) - return (ntuple(n -> n == 1 ? nest0 : deepcopy(nest0), Val(length(f.nest.f))), cacheval, nest0[3]*mid) - else - algs_ = algs[1:dim-1] - nest1 = nested_cacheval(f, p, algs_, next..., x, xs[1:dim-2]...) - h = nest1[3] - hx = h*mid - # units may change for outer integral - cacheval = init_cacheval((x, p) -> h, dom, p, alg) - return (nest1, cacheval, hx) - end -end - -function init_nest(f::FourierIntegrand, fxx, dom, p,lims, state, algs, cacheval; kws_...) - kws = NamedTuple(kws_) - xx = float(oneunit(eltype(dom))) - FX = typeof(fxx/xx) - TX = typeof(xx) - TP = Tuple{typeof(p),typeof(lims),typeof(state)} - if algs isa Tuple{} # inner integral - if f.nest isa NestedBatchIntegrand - nchunk = min(length(f.nest.f), length(f.w.cache)) - return BatchIntegrand(f.nest.y, f.nest.x, max_batch=f.nest.max_batch) do y, x, (p, lims, state) - Threads.@threads for ichunk in 1:min(nchunk, length(x)) - for (i, j) in zip(getchunk(x, ichunk, nchunk, :scatter), getchunk(y, ichunk, nchunk, :scatter)) - xi = x[i] - v = FourierValue(limit_iterate(lims, state, xi), workspace_evaluate!(f.w, xi, ichunk)) - y[j] = f.nest.f[ichunk](v, p) - end - end - return nothing - end - else - return (x, (p, lims, state)) -> begin - # return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) - v = FourierValue(limit_iterate(lims, state, x), workspace_evaluate!(f.w, x)) - return f.f(v, p) - end - end - else - if f.nest isa NestedBatchIntegrand - nchunks = min(length(f.nest.f), length(f.w.cache)) - return BatchIntegrand(FunctionWrapper{Nothing,Tuple{typeof(f.nest.y),typeof(f.nest.x),TP}}() do y, x, (p, lims, state) - Threads.@threads for ichunk in 1:min(nchunks, length(x)) - for (i, j) in zip(getchunk(x, ichunk, nchunks, :scatter), getchunk(y, ichunk, nchunks, :scatter)) - xi = x[i] - segs, lims_, state_ = limit_iterate(lims, state, xi) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - fx = FourierIntegrand(f.f, workspace_contract!(f.w, xi, ichunk), f.nest.f[ichunk]) - y[j] = do_solve(fx, lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval[ichunk]; kwargs...).u - end - end - return nothing - end, f.nest.y, f.nest.x, max_batch=f.nest.max_batch) - else - g = f.nest === nothing ? f.f : f.nest - return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) - segs, lims_, state_ = limit_iterate(lims, state, x) - fx = FourierIntegrand(g, workspace_contract!(f.w, x)) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - sol = do_solve(fx, lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval; kwargs...) - return sol.u - end - end - end -end - -function init_cacheval(f::FourierIntegrand, dom::AbstractIteratedLimits, p, alg::NestedQuad) - algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(dom))) : alg.algs - return nested_cacheval(f, p, algs, limit_iterate(dom)..., interior_point(dom)...) -end - -function do_solve(f::FourierIntegrand, lims::AbstractIteratedLimits, p_, alg::NestedQuad, cacheval; kws...) - p, segs, state = if p_ isa NestState - p_.p, p_.segs, p_.state - else - seg, lim, sta = limit_iterate(lims) - p_, seg, sta - end - dom = PuncturedInterval(segs) - g = if f.nest isa NestedBatchIntegrand && eltype(f.nest.x) === Nothing - FourierIntegrand(f.f, f.w, NestedBatchIntegrand(f.nest.f, f.nest.y, float(eltype(dom))[], max_batch=f.nest.max_batch)) - else - f - end - (dim = ndims(lims)) == ndims(f.w.series) || throw(ArgumentError("variables in Fourier series don't match domain")) - algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(dim)) : alg.algs - nest = init_nest(g, cacheval[3], dom, p, lims, state, algs[1:dim-1], cacheval[1]; kws...) - return do_solve(nest, dom, (p, lims, state), algs[dim], cacheval[2]; kws...) -end - -function do_solve(f::FourierIntegrand, dom, p, alg::EvalCounter, cacheval; kws...) - if f.nest isa NestedBatchIntegrand - error("NestedBatchIntegrand not yet supported with EvalCounter") - else - n::Int = 0 - function g(args...; kwargs...) - n += 1 - return f.f.f(args...; kwargs...) - end - h = FourierIntegrand(ParameterIntegrand{typeof(g)}(g, f.f.p), f.w) - sol = do_solve(h, dom, p, alg.alg, cacheval; kws...) - return IntegralSolution(sol.u, sol.resid, true, n) - end -end - -# method needed to resolve ambiguities -function do_solve(f::FourierIntegrand, bz::SymmetricBZ, p, alg::EvalCounter{<:AutoBZAlgorithm}, cacheval; kws...) - return do_solve_autobz(count_bz_to_standard, f, bz, p, alg.alg, cacheval; kws...) +function init_cacheval(f::FourierIntegralFunction, dom, p, alg::AutoSymPTRJL; kws...) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + rule = init_fourier_rule(ws, dom, alg) + cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) + return (; rule, cache, buffer=nothing, ws) end -=# -=# From c78e9e0b2402b76ba9b60e06c5b2f0cf27b5c6ba Mon Sep 17 00:00:00 2001 From: lxvm Date: Thu, 13 Jun 2024 21:53:06 -0400 Subject: [PATCH 04/17] nestedquad simplifications --- src/algorithms_meta.jl | 195 +++++------------------------------------ src/fourier.jl | 110 ++++++++++++++++------- src/interfaces.jl | 2 +- 3 files changed, 99 insertions(+), 208 deletions(-) diff --git a/src/algorithms_meta.jl b/src/algorithms_meta.jl index a69b818..aacd3ba 100644 --- a/src/algorithms_meta.jl +++ b/src/algorithms_meta.jl @@ -19,61 +19,7 @@ struct NestedQuad{T,S} <: IntegralAlgorithm NestedQuad(algs::Tuple{Vararg{IntegralAlgorithm}}, specialize::Tuple{Vararg{AbstractSpecialization}}=ntuple(_->FunctionWrapperSpecialize(), length(algs))) = new{typeof(algs),typeof(specialize)}(algs, specialize) end NestedQuad(algs::IntegralAlgorithm...) = NestedQuad(algs) - -#= -# this function helps create a tree of the cachevals used by each quadrature -function nested_cacheval(f::F, p::P, algs, segs, lims, state, x, xs...) where {F,P} - dom = PuncturedInterval(segs) - a, b = segs[1], segs[2] - dim = ndims(lims) - alg = algs[dim] - mid = (a+b)/2 # sample point that should be safe to evaluate - next = limit_iterate(lims, state, mid) # see what the next limit gives - if xs isa Tuple{} # the next integral takes us to the inner integral - # base case test integrand of the inner integral - # we need to pass dummy integrands to all the outer integrals so that they can build - # their caches with the right types - if f isa BatchIntegrand || f isa NestedBatchIntegrand - # Batch integrate the inner integral only - cacheval = init_cacheval(BatchIntegrand(nothing, f.y, f.x, max_batch=f.max_batch), dom, p, alg) - return (nothing, cacheval, oneunit(eltype(f.y))*mid) - elseif f isa InplaceIntegrand - # Inplace integrate through the whole nest structure - fxi = f.I*mid/oneunit(prod(next)) - cacheval = init_cacheval(InplaceIntegrand(nothing, fxi), dom, p, alg) - return (nothing, cacheval, fxi) - else - fx = f(next,p) - cacheval = init_cacheval((x, p) -> fx, dom, p, alg) - return (nothing, cacheval, fx*mid) - end - elseif f isa NestedBatchIntegrand - algs_ = algs[1:dim-1] - # numbered names to avoid type instabilities (we are not using dispatch, but the - # compiler's optimization for the recursive function's argument types) - nest0 = nested_cacheval(f.f[1], p, algs_, next..., x, xs[1:dim-2]...) - cacheval = init_cacheval(BatchIntegrand(nothing, f.y, f.x, max_batch=f.max_batch), dom, p, alg) - return (ntuple(n -> n == 1 ? nest0 : deepcopy(nest0), Val(length(f.f))), cacheval, nest0[3]*mid) - else - algs_ = algs[1:dim-1] - nest1 = nested_cacheval(f, p, algs_, next..., x, xs[1:dim-2]...) - h = nest1[3] - hx = h*mid - # units may change for outer integral - if f isa InplaceIntegrand - cacheval = init_cacheval(InplaceIntegrand(nothing, hx), dom, p, alg) - return (nest1, cacheval, hx) - else - cacheval = init_cacheval((x, p) -> h, dom, p, alg) - return (nest1, cacheval, hx) - end - end -end -function init_cacheval(f, dom::AbstractIteratedLimits, p, alg::NestedQuad) - algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(dom))) : alg.algs - return nested_cacheval(f, p, algs, limit_iterate(dom)..., interior_point(dom)...) -end -=# +# TODO add a parallelization option for use when it is safe to do so function _update!(cache, x, (; p, lims_state)) segs, lims, state = limit_iterate(lims_state..., x) @@ -86,29 +32,29 @@ function _update!(cache, x, (; p, lims_state)) return end _postsolve(sol, x, p) = sol.value -function init_cacheval(f::IntegralFunction, nextdom, p, alg::NestedQuad; kws...) +function init_cacheval(f, nextdom, p, alg::NestedQuad; kws...) x0, (segs, lims, state) = if nextdom isa AbstractIteratedLimits interior_point(nextdom), limit_iterate(nextdom) else - nothing, nextdom + nextdom end algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(lims))) : alg.algs spec = alg.specialize isa AbstractSpecialization ? ntuple(i -> alg.specialize, Val(ndims(lims))) : alg.specialize - proto = get_prototype(f, x0, p) - func = if ndims(lims) == 1 - inner_integralfunction(f, proto) + if ndims(lims) == 1 + func, ws = inner_integralfunction(f, x0, p) else - len = segs[end] - segs[begin] + integrand, ws, update!, postsolve = outer_integralfunction(f, x0, p) + proto = get_prototype(integrand, x0, p) a, b, = segs x = (a+b)/2 - next = limit_iterate(lims, state, x) + next = (x0[begin:end-1], limit_iterate(lims, state, x)) kws = NamedTuple(kws) + len = segs[end] - segs[begin] kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - integrand = outer_integralfunction(f, proto) subprob = IntegralProblem(integrand, next, p; kwargs...) - CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1], spec[1:ndims(lims)-1]), _update!, _postsolve, proto*x^(ndims(lims)-1), spec[ndims(lims)]) + func = CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1], spec[1:ndims(lims)-1]), update!, postsolve, proto*x^(ndims(lims)-1), spec[ndims(lims)]) end - prob = IntegralProblem(func, segs, (; p, lims_state=(lims, state)); kws...) + prob = IntegralProblem(func, segs, (; p, lims_state=(lims, state), ws); kws...) return init(prob, algs[ndims(lims)]) # the order of updates is somewhat tricky. I think some could be simplified if instead # we use an IntegralProblem modified to contain lims_state, instead of passing the @@ -121,119 +67,22 @@ function do_integral(f, dom, p, alg::NestedQuad, cacheval; kws...) cacheval.kwargs = (; cacheval.kwargs..., kws...) return solve!(cacheval) end -function inner_integralfunction(f::IntegralFunction, proto) - IntegralFunction(proto) do x, (; p, lims_state) +function inner_integralfunction(f::IntegralFunction, x0, p) + proto = get_prototype(f, x0, p) + func = IntegralFunction(proto) do x, (; p, lims_state) f.f(limit_iterate(lims_state..., x), p) end + ws = nothing + return func, ws end -function outer_integralfunction(f::IntegralFunction, proto) - return IntegralFunction(f.f, proto) -end -#= -function init_nest(f::F, fxx, dom, p,lims, state, algs, cacheval; kws_...) where {F} - kws = NamedTuple(kws_) - xx = float(oneunit(eltype(dom))) - FX = typeof(fxx/xx) - TX = typeof(xx) - TP = Tuple{typeof(p),typeof(lims),typeof(state)} - if algs isa Tuple{} # inner integral - if f isa BatchIntegrand - return f - elseif f isa NestedBatchIntegrand - nchunk = length(f.f) - return BatchIntegrand(FunctionWrapper{Nothing,Tuple{typeof(f.y),typeof(f.x),TP}}() do y, x, (p, lims, state) - Threads.@threads for ichunk in 1:min(nchunk, length(x)) - for (i, j) in zip(getchunk(x, ichunk, nchunk, :scatter), getchunk(y, ichunk, nchunk, :scatter)) - xi = x[i] - y[j] = f.f[ichunk](limit_iterate(lims, state, xi), p) - end - end - return nothing - end, f.y, f.x, max_batch=f.max_batch) - elseif f isa InplaceIntegrand - return InplaceIntegrand(FunctionWrapper{Nothing,Tuple{FX,TX,TP}}() do y, x, (p, lims, state) - f.f!(y, limit_iterate(lims, state, x), p) - return nothing - end, f.I) - else - return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) - return f(limit_iterate(lims, state, x), p) - end - end - else - if f isa InplaceIntegrand - return InplaceIntegrand(FunctionWrapper{Nothing,Tuple{FX,TX,TP}}() do y, x, (p, lims, state) - segs, lims_, state_ = limit_iterate(lims, state, x) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - do_solve(InplaceIntegrand(f.f!, y), lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval; kwargs...) - return nothing - end, f.I) - elseif f isa NestedBatchIntegrand - nchunks = length(f.f) - return BatchIntegrand(FunctionWrapper{Nothing,Tuple{typeof(f.y),typeof(f.x),TP}}() do y, x, (p, lims, state) - Threads.@threads for ichunk in 1:min(nchunks, length(x)) - for (i, j) in zip(getchunk(x, ichunk, nchunks, :scatter), getchunk(y, ichunk, nchunks, :scatter)) - xi = x[i] - segs, lims_, state_ = limit_iterate(lims, state, xi) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - y[j] = do_solve(f.f[ichunk], lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval[ichunk]; kwargs...).u - end - end - return nothing - end, f.y, f.x, max_batch=f.max_batch) - else - return FunctionWrapper{FX,Tuple{TX,TP}}() do x, (p, lims, state) - segs, lims_, state_ = limit_iterate(lims, state, x) - len = segs[end] - segs[1] - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - sol = do_solve(f, lims_, NestState(p, segs, state_), NestedQuad(algs), cacheval; kwargs...) - return sol.u - end - end - end -end - -struct NestState{P,G,S} - p::P - segs::G - state::S +function outer_integralfunction(f::IntegralFunction, x0, p) + proto = get_prototype(f, x0, p) + func = IntegralFunction(f.f, proto) + ws = nothing + return func, ws, _update!, _postsolve end +# TODO implement for InplaceIntegralFunction, InplaceBatchIntegralFunction, CommonSolveIntegralFunction -function do_solve(f::F, lims::AbstractIteratedLimits, p_, alg::NestedQuad, cacheval; kws...) where {F} - g, p, segs, state = if p_ isa NestState - gg = if f isa NestedBatchIntegrand - fx = eltype(f.x) === Nothing ? float(eltype(p_.segs))[] : f.x - NestedBatchIntegrand(f.f, f.y, fx, max_batch=f.max_batch) - else - f - end - gg, p_.p, p_.segs, p_.state - else - seg, lim, sta = limit_iterate(lims) - gg = if f isa BatchIntegrand - fx = eltype(f.x) === Nothing ? typeof(interior_point(lims))[] : f.x - BatchIntegrand(f.y, similar(f.x, eltype(eltype(fx))), max_batch=f.max_batch) do y, xs, (p, lims, state) - resize!(fx, length(xs)) - f.f!(y, map!(x -> limit_iterate(lims, state, x), fx, xs), p) - end - elseif f isa NestedBatchIntegrand - # this should be done recursively at the outermost level, but it is lazy. - fx = eltype(f.x) === Nothing ? float(eltype(seg))[] : f.x - NestedBatchIntegrand(f.f, f.y, fx, max_batch=f.max_batch) - else - f - end - gg, p_, seg, sta - end - dom = PuncturedInterval(segs) - dim = ndims(lims) # constant propagation :) - algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(dim)) : alg.algs - nest = init_nest(g, cacheval[3], dom, p, lims, state, algs[1:dim-1], cacheval[1]; kws...) - return do_solve(nest, dom, (p, lims, state), algs[dim], cacheval[2]; kws...) -end -=# """ AbsoluteEstimate(est_alg, abs_alg; kws...) diff --git a/src/fourier.jl b/src/fourier.jl index e695908..9275a78 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -43,6 +43,66 @@ function get_prototype(f::FourierIntegralFunction, x, ws, p) end get_prototype(f::FourierIntegralFunction, x, p) = get_prototype(f, x, f.s, p) +""" + FourierCommonSolveIntegralFunction(prob, alg, update!, postsolve, s, [prototype, specialize]; alias=false, kws...) + +Constructor for an integrand that solves a problem defined with the CommonSolve.jl +interface, `prob`, which is instantiated using `init(prob, alg; kws...)`. Helper functions +include: `update!(cache, x, s(x), p)` is called before +`solve!(cache)`, followed by `postsolve(sol, x, s(x), p)`, which should return the value of the solution. +The `prototype` argument can help control how much to `specialize` on the type of the +problem, which defaults to `FullSpecialize()` so that run times are improved. However +`FunctionWrapperSpecialize()` may help reduce compile times. +""" +struct FourierCommonSolveIntegralFunction{P,A,S,K,U,PS,T,M<:AbstractSpecialization} <: AbstractIntegralFunction + prob::P + alg::A + s::S + kwargs::K + update!::U + postsolve::PS + prototype::T + specialize::M + alias::Bool +end +function FourierCommonSolveIntegralFunction(prob, alg, update!, postsolve, s, prototype=nothing, specialize=FullSpecialize(); alias=false, kws...) + return FourierCommonSolveIntegralFunction(prob, alg, s, NamedTuple(kws), update!, postsolve, prototype, specialize, alias) +end + +function do_solve!(cache, f::FourierCommonSolveIntegralFunction, x, s, p) + f.update!(cache, x, s, p) + sol = solve!(cache) + return f.postsolve(sol, x, s, p) +end +function get_prototype(f::FourierCommonSolveIntegralFunction, x, ws, p) + if isnothing(f.prototype) + cache = init(f.prob, f.alg; f.kwargs...) + do_solve!(cache, f, x, ws(x), p) + else + f.prototype + end +end +get_prototype(f::FourierCommonSolveIntegralFunction, x, p) = get_prototype(f, x, f.s, p) + +function init_specialized_fourierintegrand(cache, f, dom, p; x=get_prototype(dom), s = f.s(x), prototype=f.prototype) + proto = prototype === nothing ? do_solve!(cache, f, x, s, p) : prototype + func = (x, s, p) -> do_solve!(cache, f, x, s, p) + integrand = if f.specialize isa FullSpecialize + func + elseif f.specialize isa FunctionWrapperSpecialize + FunctionWrapper{typeof(prototype), typeof((x, s, p))}(func) + else + throw(ArgumentError("$(f.specialize) is not implemented")) + end + return integrand, proto +end +function _init_fouriercommonsolvefunction(f, dom, p; kws...) + cache = init(f.prob, f.alg; f.kwargs...) + integrand, prototype = init_specialized_fourierintegrand(cache, f, dom, p; kws...) + return cache, integrand, prototype +end + + # similar to workspace_allocate, but more type-stable because of loop unrolling and vector types function workspace_allocate_vec(s::AbstractFourierSeries{N}, x::NTuple{N,Any}, len::NTuple{N,Integer}=ntuple(one,Val(N))) where {N} # Only the top-level workspace has an AbstractFourierSeries in the series field @@ -108,46 +168,28 @@ function autosymptr_integrand(f::FourierIntegralFunction, p, segs, cacheval) x -> x isa FourierValue ? f.f(x.x, x.s, p) : f.f(x, ws(x), p) end -function _fourier_update!(cache, x, (; p, ws, lims_state)) - segs, lims, state = limit_iterate(lims_state..., x) - s = workspace_contract!(ws, x) - len = segs[end] - segs[begin] - kws = cache.kwargs - cache.p = p - cache.cacheval.dom = segs - cache.cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - cache.cacheval.p = (; cache.cacheval.p..., ws=s, lims_state=(lims, state)) +function _fourier_update!(cache, x, p) + _update!(cache, x, p) + s = workspace_contract!(p.ws, x) + cache.cacheval.p = (; cache.cacheval.p..., ws=s) return end -function init_cacheval(f::FourierIntegralFunction, nextdom, p, alg::NestedQuad; kws...) - x0, (segs, lims, state) = if nextdom isa AbstractIteratedLimits - interior_point(nextdom), limit_iterate(nextdom) - else - nothing, nextdom - end - algs = alg.algs isa IntegralAlgorithm ? ntuple(i -> alg.algs, Val(ndims(lims))) : alg.algs - spec = alg.specialize isa AbstractSpecialization ? ntuple(i -> alg.specialize, Val(ndims(lims))) : alg.specialize +function inner_integralfunction(f::FourierIntegralFunction, x0, p) ws = f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) proto = get_prototype(f, x0, ws, p) - func = if ndims(lims) == 1 - IntegralFunction(proto) do x, (; p, ws, lims_state) - f.f(limit_iterate(lims_state..., x), workspace_evaluate!(ws, x), p) - end - else - len = segs[end] - segs[begin] - a, b, = segs - x = (a+b)/2 - next = limit_iterate(lims, state, x) - s = workspace_contract!(ws, x) - kws = NamedTuple(kws) - kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol/len,)) : kws - integrand = FourierIntegralFunction(f.f, s, proto; alias=true) - subprob = IntegralProblem(integrand, next, p; kwargs...) - CommonSolveIntegralFunction(subprob, NestedQuad(algs[1:ndims(lims)-1], spec[1:ndims(lims)-1]), _fourier_update!, _postsolve, proto*x^(ndims(lims)-1), spec[ndims(lims)]) + func = IntegralFunction(proto) do x, (; p, ws, lims_state) + f.f(limit_iterate(lims_state..., x), workspace_evaluate!(ws, x), p) end - prob = IntegralProblem(func, segs, (; p, ws, lims_state=(lims, state)); kws...) - return init(prob, algs[ndims(lims)]) + return func, ws +end +function outer_integralfunction(f::FourierIntegralFunction, x0, p) + ws = f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + proto = get_prototype(f, x0, ws, p) + s = workspace_contract!(ws, x0[end]) + func = FourierIntegralFunction(f.f, s, proto; alias=true) + return func, ws, _fourier_update!, _postsolve end +# TODO implement FourierCommonSolveIntegralFunction # PTR rules diff --git a/src/interfaces.jl b/src/interfaces.jl index 8372b97..56a5478 100644 --- a/src/interfaces.jl +++ b/src/interfaces.jl @@ -89,7 +89,7 @@ function CommonSolveIntegralFunction(prob, alg, update!, postsolve, prototype=no return CommonSolveIntegralFunction(prob, alg, NamedTuple(kws), update!, postsolve, prototype, specialize) end -function do_solve!(cache, f, x, p) +function do_solve!(cache, f::CommonSolveIntegralFunction, x, p) f.update!(cache, x, p) sol = solve!(cache) return f.postsolve(sol, x, p) From 5882b1223eaff292b039d0b3bd2c6202005c959e Mon Sep 17 00:00:00 2001 From: lxvm Date: Thu, 13 Jun 2024 22:56:34 -0400 Subject: [PATCH 05/17] restore autoptr --- aps_example/aps_example.jl | 11 +++++----- src/algorithms_autosymptr.jl | 5 +++-- src/brillouin.jl | 40 +++++++++++++++++++++--------------- src/domains.jl | 2 ++ src/fourier.jl | 7 ++++++- 5 files changed, 39 insertions(+), 26 deletions(-) diff --git a/aps_example/aps_example.jl b/aps_example/aps_example.jl index 43f3b2f..c0c2fcf 100644 --- a/aps_example/aps_example.jl +++ b/aps_example/aps_example.jl @@ -72,8 +72,8 @@ end dos_solver_iai = dos_solver(prob_dos, IAI(QuadGKJL())) @time greens_iai = hchebinterp(dos_solver_iai, ω_min, ω_max; atol=1e-2, order=cheb_order) -# dos_solver_ptr = dos_solver(prob_dos, PTR(; npt=100)) -# @time greens_ptr = hchebinterp(dos_solver_ptr, ω_min, ω_max; atol=1e-2, order=cheb_order) +dos_solver_ptr = dos_solver(prob_dos, PTR(; npt=100)) +@time greens_ptr = hchebinterp(dos_solver_ptr, ω_min, ω_max; atol=1e-2, order=cheb_order) using CairoMakie @@ -84,10 +84,9 @@ ax1 = Axis(fig1[1,1], limits=((10,15), (0,6)), xlabel="ω (eV)", ylabel="SVO DOS p1 = lines!(ax1, 10:η/100:15, ω -> -imag(greens_iai(ω))/pi/det(bz.B); label="IAI, η=$η") axislegend(ax1) save("iai_svo_dos.pdf", fig1) -#= + fig2 = Figure() -ax2 = Axis(fig2[1,1], limits=((10,15), (0,det(bz.B)*6)), xlabel="ω (eV)", ylabel="SVO DOS (eV⁻¹ Å⁻³)") -p2 = lines!(ax2, 10:η/100:15, ω -> -imag(greens_ptr(ω))/pi; label="PTR, η=$η") +ax2 = Axis(fig2[1,1], limits=((10,15), (0,6)), xlabel="ω (eV)", ylabel="SVO DOS (eV⁻¹)") +p2 = lines!(ax2, 10:η/100:15, ω -> -imag(greens_ptr(ω))/pi/det(bz.B); label="PTR, η=$η") axislegend(ax2) save("ptr_svo_dos.pdf", fig2) -=# diff --git a/src/algorithms_autosymptr.jl b/src/algorithms_autosymptr.jl index f9c8347..d19a8b1 100644 --- a/src/algorithms_autosymptr.jl +++ b/src/algorithms_autosymptr.jl @@ -82,11 +82,12 @@ function init_cacheval(f::IntegralFunction, dom, p, alg::AutoSymPTRJL; kws...) return (rule=rule, cache=cache, buffer=buffer) end -function do_solve(f, dom, p, alg::AutoSymPTRJL, cacheval; +function do_integral(f, dom, p, alg::AutoSymPTRJL, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) g = autosymptr_integrand(f, p, dom, cacheval) - value, error = autosymptr(g, dom; syms = alg.syms, rule = cacheval.rule, cache = cacheval.cache, keepmost = alg.keepmost, + bas = get_basis(dom) + value, error = autosymptr(g, bas; syms = alg.syms, rule = cacheval.rule, cache = cacheval.cache, keepmost = alg.keepmost, abstol = abstol, reltol = reltol, maxevals = maxiters, norm=alg.norm, buffer=cacheval.buffer) retcode = error < max(something(abstol, zero(error)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(eltype(a))) : abstol)) ? Success : Failure stats = (; error) diff --git a/src/brillouin.jl b/src/brillouin.jl index e1fcdb2..722d67c 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -47,6 +47,8 @@ nsyms(::FullBZ) = 1 Base.summary(bz::SymmetricBZ) = string(checksquare(bz.A), "-dimensional Brillouin zone with ", bz isa FullBZ ? "trivial" : nsyms(bz), " symmetries") Base.show(io::IO, bz::SymmetricBZ) = print(io, summary(bz)) +Base.ndims(::SymmetricBZ{S,L,d}) where {S,L,d} = d +Base.eltype(::Type{<:SymmetricBZ{S,L,d,TA,TB}}) where {S,L,d,TA,TB} = TB # Define traits for symmetrization based on symmetry representations @@ -391,10 +393,6 @@ function init_cacheval(rep, f, bz::SymmetricBZ, p, bzalg::AutoBZAlgorithm; kws.. return init(prob, alg) end -function do_solve(f, bz::SymmetricBZ, p, bzalg::AutoBZAlgorithm, cacheval; kws...) - do_solve_autobz(bz_to_standard, f, bz, p, bzalg, cacheval; kws...) -end - function do_solve_autobz(rep, f, bz, p, bzalg::AutoBZAlgorithm, cacheval; _kws...) j = abs(det(bz.B)) # rescale tolerance to (I)BZ coordinate and get the right number of digits kws = NamedTuple(_kws) @@ -472,19 +470,28 @@ function AutoPTR(; norm=norm, a=1.0, nmin=50, nmax=1000, n₀=6.0, Δn=log(10), return AutoPTR(norm, a, nmin, nmax, n₀, Δn, keepmost, nthreads) end -function bz_to_standard(f, bz, p, alg::AutoPTR; kws...) - return IntegralProblem(f, canonical_ptr_basis(bz.B), p; kws...), AutoSymPTRJL(norm=alg.norm, a=alg.a, nmin=alg.nmin, nmax=alg.nmax, n₀=alg.n₀, Δn=alg.Δn, keepmost=alg.keepmost, syms=bz.syms, nthreads=alg.nthreads) + +struct RepBZ{R,B} + rep::R + bz::B end +Base.ndims(dom::RepBZ) = ndims(dom.bz) +Base.eltype(::Type{RepBZ{R,B}}) where {R,B} = eltype(B) -function init_cacheval(f, bz::SymmetricBZ, p, bzalg::AutoPTR) - bz_, dom, alg = bz_to_standard(bz, bzalg) - f isa NestedBatchIntegrand && throw(ArgumentError("AutoSymPTRJL doesn't support nested batching")) - rule = SymmetricRuleDef(init_rule(dom, alg), SymRep(f), bz_) - cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) - buffer = init_buffer(f, alg.nthreads) - return (rule=rule, cache=cache, buffer=buffer) +function init_cacheval(rep, f, bz::SymmetricBZ, p, bzalg::AutoPTR; kws...) + prob = IntegralProblem(f, RepBZ(rep, bz), p; kws...) + alg = AutoSymPTRJL(norm=bzalg.norm, a=bzalg.a, nmin=bzalg.nmin, nmax=bzalg.nmax, n₀=bzalg.n₀, Δn=bzalg.Δn, keepmost=bzalg.keepmost, syms=bz.syms, nthreads=bzalg.nthreads) + return init(prob, alg) +end +get_basis(dom::RepBZ) = canonical_ptr_basis(dom.bz.B) +function init_rule(dom::RepBZ, alg::AutoSymPTRJL) + B = get_basis(dom) + rule = init_rule(B, alg) + return SymmetricRuleDef(rule, dom.rep, dom.bz) end -function do_solve_autobz(bz_to_standard, f, bz, p, bzalg::AutoPTR, cacheval; _kws...) +# The spectral convergence of the PTR for integrands with non-trivial symmetry action +# requires symmetrizing inside the quadrature +function do_solve_autobz(rep, f, bz, p, bzalg::AutoPTR, cacheval; _kws...) j = abs(det(bz.B)) # rescale tolerance to (I)BZ coordinate and get the right number of digits kws = NamedTuple(_kws) cacheval.f = f @@ -493,11 +500,11 @@ function do_solve_autobz(bz_to_standard, f, bz, p, bzalg::AutoPTR, cacheval; _kw sol = solve!(cacheval) value = j*sol.value - error = j*sol.stats.error - stats = (; sol.stats..., error) + stats = (; sol.stats..., error=sol.stats.error*j) return IntegralSolution(value, sol.retcode, stats) end + """ TAI(; norm=norm, initdivs=1) @@ -516,7 +523,6 @@ function bz_to_standard(f, bz, p, alg::TAI; kws...) return IntegralProblem(f, HyperCube(bz.lims.a, bz.lims.b), p; kws...), HCubatureJL(norm=alg.norm, initdiv = alg.initdiv) end - #= """ PTR_IAI(; ptr=PTR(), iai=IAI()) diff --git a/src/domains.jl b/src/domains.jl index 190870a..d011be4 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -6,6 +6,8 @@ function get_prototype(dom) return (a+b)/2 end +get_basis(B::Basis) = B + """ PuncturedInterval(s) diff --git a/src/fourier.jl b/src/fourier.jl index 9275a78..0cffd39 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -411,10 +411,15 @@ function init_cacheval(f::FourierIntegralFunction, dom , p, alg::MonkhorstPack; return (; rule, buffer=nothing, ws) end -function init_fourier_rule(w::FourierWorkspace, dom::Basis, alg::AutoSymPTRJL) +function init_fourier_rule(w::FourierWorkspace, dom, alg::AutoSymPTRJL) @assert ndims(w.series) == ndims(dom) return FourierMonkhorstPackRule(w, alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) end +function init_fourier_rule(w::FourierWorkspace, dom::RepBZ, alg::AutoSymPTRJL) + B = get_basis(dom) + rule = init_fourier_rule(w, B, alg) + return SymmetricRuleDef(rule, dom.rep, dom.bz) +end function init_cacheval(f::FourierIntegralFunction, dom, p, alg::AutoSymPTRJL; kws...) ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) rule = init_fourier_rule(ws, dom, alg) From 727b58e477c12206a7b4cd7de45e223a51c1a749 Mon Sep 17 00:00:00 2001 From: lxvm Date: Thu, 13 Jun 2024 23:21:36 -0400 Subject: [PATCH 06/17] update example --- aps_example/Project.toml | 1 + aps_example/README.md | 4 +++- aps_example/aps_example.jl | 8 ++++---- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/aps_example/Project.toml b/aps_example/Project.toml index fd3a1f7..acc2155 100644 --- a/aps_example/Project.toml +++ b/aps_example/Project.toml @@ -1,6 +1,7 @@ [deps] AutoBZCore = "66bd3e16-1600-45cf-8f55-0b550710682b" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" +FourierSeriesEvaluators = "2a892dea-6eef-4bb5-9d1c-de966c9f6db5" HChebInterp = "78faba9b-a54b-441f-8118-62407cbe4e59" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" diff --git a/aps_example/README.md b/aps_example/README.md index cae6563..38d6fbe 100644 --- a/aps_example/README.md +++ b/aps_example/README.md @@ -1,4 +1,6 @@ -[Scripts updated for AutoBZCore v0.3] +[Scripts in development for AutoBZCore v0.4] + +Warning: this example may not work well using Julia 1.10. Hello, to reproduce the code example in these [slides](https://web.mit.edu/lxvm/www/slides/autobz_aps.pdf) follow these steps: diff --git a/aps_example/aps_example.jl b/aps_example/aps_example.jl index c0c2fcf..eb203d2 100644 --- a/aps_example/aps_example.jl +++ b/aps_example/aps_example.jl @@ -1,6 +1,6 @@ -# import Pkg -# Pkg.activate(".") # reproducible environment included -# Pkg.instantiate() # install dependencies +import Pkg +Pkg.activate(".") # reproducible environment included +Pkg.instantiate() # install dependencies using WannierIO @@ -24,7 +24,7 @@ using FourierSeriesEvaluators, LinearAlgebra h = FourierSeries(H_R, period=1.0) -η = 5e-2 # 10 meV (scattering amplitude) +η = 1e-2 # 10 meV (scattering amplitude) ω_min = 10 ω_max = 15 p0 = (; η, ω=(ω_min + ω_max)/2) # initial parameters From 940fee5c8f8f5f961384f51ab58abd31e76d6b05 Mon Sep 17 00:00:00 2001 From: lxvm Date: Wed, 19 Jun 2024 17:42:48 -0400 Subject: [PATCH 07/17] remove hdf5 interface --- Project.toml | 16 ++--- ext/HDF5Ext.jl | 160 ----------------------------------------------- test/hdf5ext.jl | 56 ----------------- test/runtests.jl | 1 - 4 files changed, 5 insertions(+), 228 deletions(-) delete mode 100644 ext/HDF5Ext.jl delete mode 100644 test/hdf5ext.jl diff --git a/Project.toml b/Project.toml index 6615069..71cbe12 100644 --- a/Project.toml +++ b/Project.toml @@ -1,24 +1,22 @@ name = "AutoBZCore" uuid = "66bd3e16-1600-45cf-8f55-0b550710682b" authors = ["lxvm "] -version = "0.3.8" +version = "0.4.0-DEV" [deps] AutoSymPTR = "78a0c066-08f1-49a8-82f0-b29cd485e1d3" ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e" +CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" FourierSeriesEvaluators = "2a892dea-6eef-4bb5-9d1c-de966c9f6db5" FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e" HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49" IteratedIntegration = "3ecdc4d6-ee34-4049-885a-a4e3631db98b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" -Reexport = "189a3867-3050-52da-a836-e630ba90ab69" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [weakdeps] AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a" -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" Polyhedra = "67491407-f73d-577b-9b50-8179a7c68029" SymmetryReduceBZ = "49a35663-c880-4242-bebb-1ec8c0fa8046" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" @@ -26,7 +24,6 @@ WannierIO = "cb1bc77f-5443-4951-af9f-05b616a3e422" [extensions] AtomsBaseExt = "AtomsBase" -HDF5Ext = "HDF5" SymmetryReduceBZExt = ["SymmetryReduceBZ", "Polyhedra"] UnitfulExt = "Unitful" WannierIOExt = "WannierIO" @@ -36,17 +33,15 @@ Aqua = "0.7" AtomsBase = "0.3" AutoSymPTR = "0.4" ChunkSplitters = "2" +CommonSolve = "0.2" Elliptic = "1" FourierSeriesEvaluators = "1" FunctionWrappers = "1" GeneralizedGaussianQuadrature = "0.1" HCubature = "1.4" -HDF5 = "0.16.15" IteratedIntegration = "0.5" LinearAlgebra = "1.9" -Printf = "1.9" QuadGK = "2.6" -Reexport = "1" StaticArrays = "1" SymmetryReduceBZ = "0.2" WannierIO = "0.1,0.2" @@ -56,9 +51,8 @@ julia = "1.9" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a" Elliptic = "b305315f-e792-5b7a-8f41-49f472929428" -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" -LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" GeneralizedGaussianQuadrature = "958e0c08-f14d-42e8-a0ab-84193b3783f2" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" Polyhedra = "67491407-f73d-577b-9b50-8179a7c68029" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -68,4 +62,4 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" WannierIO = "cb1bc77f-5443-4951-af9f-05b616a3e422" [targets] -test = ["Aqua", "Elliptic", "LinearAlgebra", "GeneralizedGaussianQuadrature", "Test", "HDF5", "StaticArrays", "OffsetArrays", "SymmetryReduceBZ", "Polyhedra"] +test = ["Aqua", "Elliptic", "LinearAlgebra", "GeneralizedGaussianQuadrature", "Test", "StaticArrays", "OffsetArrays", "SymmetryReduceBZ", "Polyhedra"] diff --git a/ext/HDF5Ext.jl b/ext/HDF5Ext.jl deleted file mode 100644 index 40feaf8..0000000 --- a/ext/HDF5Ext.jl +++ /dev/null @@ -1,160 +0,0 @@ -module HDF5Ext - -using Printf: @sprintf -using LinearAlgebra: norm - -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 - - -""" - read_h5_to_nt(filename) - -Loads the h5 archive from `filename` and reads its datasets into a `NamedTuple` -and its groups into `NamedTuple`s recursively. -""" -read_h5_to_nt(filename) = h5open(read_h5_to_nt_, filename, "r") -read_h5_to_nt_(h5) = NamedTuple([Pair(Symbol(key), ((val = h5[key]) isa Group) ? read_h5_to_nt_(val) : read_dataset(h5, key)) for key in keys(h5)]) - -""" - write_nt_to_h5(nt::NamedTuple, filename) - -Takes a `NamedTuple` and writes its values, which must be arrays, into an h5 -archive at `filename` with dataset names corresponding to the tuple names. -If a value is a `NamedTuple`, its datasets are written to h5 groups recursively. -""" -write_nt_to_h5(nt::NamedTuple, filename) = h5open(filename, "w") do h5 - write_nt_to_h5_(nt, h5) -end -function write_nt_to_h5_(nt::NamedTuple, h5) - for key in keys(nt) - if (val = nt[key]) isa NamedTuple - write_nt_to_h5_(val, create_group(h5, string(key))) - else - write_dataset(h5, string(key), val) - end - end -end - -# parallelization - -# returns (dset, ax) to allow for array-valued data types -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,A}}, dims) where {T, A} - # 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", A, dims) - return (gval, gaux), (axval, axaux) -end - -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) - create_dataset(g, string(i), S, dims) - end - return g -end -function param_group(parent, ::Type{MixedParameters{T,NamedTuple{K,V}}}, dims) where {T,K,V} - g = create_group(parent, "args") - for (i, S) in enumerate(T.parameters) - create_dataset(g, string(i), S, dims) - end - q = create_group(parent, "kwargs") - for (key, val) in zip(K,V.parameters) - create_dataset(q, string(key), val, dims) - end - return (g,q) -end -function param_record(group, p, i) - set_value(group, i, p) - return nothing -end -function param_record(group, p::Tuple, i) - for (j,e) in enumerate(p) - 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)) - set_value(g[string(j)], i, e) - end - for (k,v) in pairs(getfield(p, :kwargs)) - set_value(q[string(k)], i, v) - end - return nothing -end - -const liblock = ReentrantLock() - -""" - batchsolve(h5::H5DataStore, f::IntegralSolver, ps, [T]; verb=true, nthreads=Threads.nthreads()) - -Solves the integral `f` at all parameters `ps` in a multi-threaded fashion, using `nthreads` -threads to parallelize over `ps`, and writes the results to the `h5` archive. -""" -function batchsolve(h5::H5DataStore, f::IntegralSolver, ps::AbstractArray, T=solver_type(f, ps[begin]); flush=true, 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) - gn = create_dataset(h5, "numevals", Int, dims) - gp = param_group(h5, eltype(ps), dims) - flush && Base.flush(h5) - - function h5callback(_, i, n, p, sol, t) - lock(liblock) - try - 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)) - set_value(gn, i, sol.numevals) - param_record(gp, p, i) - flush && Base.flush(h5) - finally - unlock(liblock) - end - end - - verb && @info "Started parameter sweep" - t = time() - sol = batchsolve(f, ps, T; callback=h5callback, nthreads=nthreads) - t = time()-t - 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 diff --git a/test/hdf5ext.jl b/test/hdf5ext.jl deleted file mode 100644 index 565aae9..0000000 --- a/test/hdf5ext.jl +++ /dev/null @@ -1,56 +0,0 @@ -using Test -using LinearAlgebra -using StaticArrays -using HDF5 -using AutoBZCore - -fn = tempname() -h5open(fn, "w") do io - g1 = create_group(io, "Number") - @testset "Number" begin - prob = IntegralProblem((x, p) -> p, 0, 1) - solver = IntegralSolver(prob, QuadGKJL()) - params = 1:1.0:10 - values = batchsolve(g1, solver, params, verb=false) - @test g1["I"][:] ≈ values - end - g2 = create_group(io, "SArray") - @testset "SArray" begin - f(x,p) = ((s,c) = sincos(p*x) ; SHermitianCompact{2,Float64,3}([s, c, -s])) - prob = IntegralProblem(f, 0.0, 1pi) - solver = IntegralSolver(prob, QuadGKJL()) - params = [0.8, 0.9, 1.0] - values = batchsolve(g2, solver, params, verb=false) - # Arrays of SArray are flattened to multidimensional arrays - @test g2["I"][:,:,:] ≈ reshape(collect(Iterators.flatten(values)), 2, 2, :) - end - g3 = create_group(io, "AuxValue") - @testset "AuxValue" begin - f(x,p) = ((re,im) = reim(inv(complex(cos(x), p))) ; AuxValue(re, im)) - prob = IntegralProblem(f, 0.0, 2pi) - solver = IntegralSolver(prob, QuadGKJL(), abstol=1e-3) - params = [2.0, 1.0, 0.5] - values = batchsolve(g3, solver, params, verb=false) - @test g3["I/val"][:] ≈ getproperty.(values, :val) - @test g3["I/aux"][:] ≈ getproperty.(values, :aux) - end - g4 = create_group(io, "0d") - g5 = create_group(io, "3d") - @testset "parameter dimensions" begin - prob = IntegralProblem((x, p) -> p[1] + p[2] + p[3], 0, 1) - solver = IntegralSolver(prob, QuadGKJL()) - # 0-d parameters - params = paramzip(0, 1, 2) - values = batchsolve(g4, solver, params, verb=false) - @test g4["I"][] ≈ values[] ≈ 3 - @test g4["args/1"][] ≈ 0 - @test g4["args/2"][] ≈ 1 - @test g4["args/3"][] ≈ 2 - # 3-d parameters - params = paramproduct(1:2, 1:2, 1:2) - values = batchsolve(g5, solver, params, verb=false) - @test g5["I"][1,1,1] ≈ values[1] ≈ 3 - @test g5["I"][2,2,2] ≈ values[8] ≈ 6 - end -end -rm(fn) diff --git a/test/runtests.jl b/test/runtests.jl index d4318b4..538b731 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,7 +5,6 @@ include("utils.jl") @testset "interface" include("interface_tests.jl") @testset "brillouin" include("brillouin.jl") @testset "fourier" include("fourier.jl") -@testset "HDF5Ext" include("hdf5ext.jl") @testset "SymmetryReduceBZExt" include("test_ibz.jl") # @testset "AtomsBaseExt" include("atomsbaseext.jl") # @testset "WannierIOExt" include("wannierioext.jl") From 6883b236f4ab9e4d953b999fc7e0e7266a4ebb6f Mon Sep 17 00:00:00 2001 From: lxvm Date: Wed, 19 Jun 2024 17:43:58 -0400 Subject: [PATCH 08/17] additional algorithms with tests --- src/AutoBZCore.jl | 4 +- src/algorithms.jl | 37 ++++---- src/algorithms_autosymptr.jl | 31 ++++--- src/algorithms_iterated.jl | 161 +++++++++++++++++++++++------------ src/algorithms_meta.jl | 1 - src/domains.jl | 5 ++ src/fourier.jl | 139 +++++++++++++++++++++++++----- src/interfaces.jl | 13 ++- test/fourier.jl | 54 ++++++++++-- test/interface_tests.jl | 55 +++++++----- 10 files changed, 360 insertions(+), 140 deletions(-) diff --git a/src/AutoBZCore.jl b/src/AutoBZCore.jl index 03ba301..5c1a29d 100644 --- a/src/AutoBZCore.jl +++ b/src/AutoBZCore.jl @@ -75,7 +75,7 @@ include("interfaces.jl") export QuadGKJL, HCubatureJL, QuadratureFunction include("algorithms.jl") export AuxQuadGKJL, ContQuadGKJL, MeroQuadGKJL -# include("algorithms_iterated.jl") +include("algorithms_iterated.jl") export MonkhorstPack, AutoSymPTRJL include("algorithms_autosymptr.jl") export NestedQuad, AbsoluteEstimate, EvalCounter @@ -88,7 +88,7 @@ export AutoBZProblem export IAI, PTR, AutoPTR, TAI, PTR_IAI, AutoPTR_IAI include("brillouin.jl") -export FourierIntegralFunction +export FourierIntegralFunction, CommonSolveFourierIntegralFunction include("fourier.jl") export DOSProblem diff --git a/src/algorithms.jl b/src/algorithms.jl index 252e5de..9162cc6 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -97,7 +97,7 @@ end function call_quadgk(f::CommonSolveIntegralFunction, p, u, usegs, cacheval; kws...) # cache = cacheval[2] could call do_solve!(cache, f, x, p) to fully specialize integrand = cacheval[3] - quadgk(x -> integrand(x, p), usegs...; kws..., segbuf=cacheval[1]) + quadgk(x -> integrand(u * x, p), usegs...; kws..., segbuf=cacheval[1]) end @@ -175,28 +175,35 @@ struct QuadratureFunction{F} <: IntegralAlgorithm end QuadratureFunction(; fun=trapz, npt=50, nthreads=1) = QuadratureFunction(fun, npt, nthreads) -function init_cacheval(f::IntegralFunction, dom, p, alg::QuadratureFunction; kws...) +function init_rule(dom, alg::QuadratureFunction) x, w = alg.fun(alg.npt) - return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing) + return [(w,x) for (w,x) in zip(w,x)] end -function init_cacheval(f::InplaceIntegralFunction, dom, p, alg::QuadratureFunction; kws...) - x, w = alg.fun(alg.npt) - proto = get_prototype(f, first(x), p) +function init_autosymptr_cache(f::IntegralFunction, dom, p, bufsize; kws...) + return (; buffer=nothing) +end +function init_autosymptr_cache(f::InplaceIntegralFunction, dom, p, bufsize; kws...) + x = get_prototype(dom) + proto = get_prototype(f, x, p) y = similar(proto) ytmp = similar(proto) - I = y * first(w) + I = y * prod(x) Itmp = similar(I) - return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing, I, Itmp, y, ytmp) + return (; buffer=nothing, I, Itmp, y, ytmp) end -function init_cacheval(f::InplaceBatchIntegralFunction, dom, p, alg::QuadratureFunction; kws...) - x, w = alg.fun(alg.npt) - proto=get_prototype(f, first(x), p) - return (rule=[(w,x) for (w,x) in zip(w,x)], buffer=similar(proto, len), y=similar(proto, len), x=Vector{float(eltype(dom))}(undef, len)) +function init_autosymptr_cache(f::InplaceBatchIntegralFunction, dom, p, bufsize; kws...) + x0 = get_prototype(dom) + proto=get_prototype(f, x0, p) + return (; buffer=similar(proto, bufsize), y=similar(proto, bufsize), x=Vector{typeof(x0)}(undef, bufsize)) end -function init_cacheval(f::CommonSolveIntegralFunction, dom, p, alg::QuadratureFunction; kws...) - x, w = alg.fun(alg.npt) +function init_autosymptr_cache(f::CommonSolveIntegralFunction, dom, p, bufsize; kws...) cache, integrand, = _init_commonsolvefunction(f, dom, p) - return (; rule=[(w,x) for (w,x) in zip(w,x)], buffer=nothing, cache, integrand) + return (; buffer=nothing, cache, integrand) +end +function init_cacheval(f, dom, p, alg::QuadratureFunction; kws...) + rule = init_rule(dom, alg) + cache = init_autosymptr_cache(f, dom, p, alg.npt; kws...) + return (; rule, cache...) end function do_integral(f, dom, p, alg::QuadratureFunction, cacheval; diff --git a/src/algorithms_autosymptr.jl b/src/algorithms_autosymptr.jl index d19a8b1..feb7981 100644 --- a/src/algorithms_autosymptr.jl +++ b/src/algorithms_autosymptr.jl @@ -31,17 +31,17 @@ end rule_type(::AutoSymPTR.PTR{N,T}) where {N,T} = SVector{N,T} rule_type(::AutoSymPTR.MonkhorstPack{N,T}) where {N,T} = SVector{N,T} -function init_cacheval(f::IntegralFunction, dom, p, alg::MonkhorstPack; kws...) - rule = init_rule(dom, alg) - buf = init_buffer(f, alg.nthreads) - return (rule=rule, buffer=buf) +function init_cacheval(f, dom, p, alg::MonkhorstPack; kws...) + b = get_basis(dom) + rule = init_rule(b, alg) + cache = init_autosymptr_cache(f, b, p, alg.nthreads; kws...) + return (; rule, cache...) end - function do_integral(f, dom, p, alg::MonkhorstPack, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - g = autosymptr_integrand(f, p, dom, cacheval) - # TODO convert the domain to a Basis - value = cacheval.rule(g, dom, cacheval.buffer) + b = get_basis(dom) + g = autosymptr_integrand(f, p, b, cacheval) + value = cacheval.rule(g, b, cacheval.buffer) retcode = Success stats = (; numevals=length(cacheval.rule)) return IntegralSolution(value, retcode, stats) @@ -75,11 +75,14 @@ end function init_rule(dom, alg::AutoSymPTRJL) return AutoSymPTR.MonkhorstPackRule(alg.syms, alg.a, alg.nmin, alg.nmax, alg.n₀, alg.Δn) end -function init_cacheval(f::IntegralFunction, dom, p, alg::AutoSymPTRJL; kws...) - rule = init_rule(dom, alg) - cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) - buffer = init_buffer(f, alg.nthreads) - return (rule=rule, cache=cache, buffer=buffer) + + +function init_cacheval(f, dom, p, alg::AutoSymPTRJL; kws...) + b = get_basis(dom) + rule = init_rule(b, alg) + rule_cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) + cache = init_autosymptr_cache(f, b, p, alg.nthreads; kws...) + return (; rule, rule_cache, cache...) end function do_integral(f, dom, p, alg::AutoSymPTRJL, cacheval; @@ -87,7 +90,7 @@ function do_integral(f, dom, p, alg::AutoSymPTRJL, cacheval; g = autosymptr_integrand(f, p, dom, cacheval) bas = get_basis(dom) - value, error = autosymptr(g, bas; syms = alg.syms, rule = cacheval.rule, cache = cacheval.cache, keepmost = alg.keepmost, + value, error = autosymptr(g, bas; syms = alg.syms, rule = cacheval.rule, cache = cacheval.rule_cache, keepmost = alg.keepmost, abstol = abstol, reltol = reltol, maxevals = maxiters, norm=alg.norm, buffer=cacheval.buffer) retcode = error < max(something(abstol, zero(error)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(eltype(a))) : abstol)) ? Success : Failure stats = (; error) diff --git a/src/algorithms_iterated.jl b/src/algorithms_iterated.jl index 7ca6d5e..6bd1a69 100644 --- a/src/algorithms_iterated.jl +++ b/src/algorithms_iterated.jl @@ -1,5 +1,3 @@ -# TODO move these to an extension if SymmetricBZ can be simplified - """ AuxQuadGKJL(; order = 7, norm = norm) @@ -15,36 +13,60 @@ function AuxQuadGKJL(; order = 7, norm = norm) return AuxQuadGKJL(order, norm) end -function init_cacheval(f, dom, p, alg::AuxQuadGKJL) - f isa NestedBatchIntegrand && throw(ArgumentError("AuxQuadGKJL doesn't support nested batching")) - return init_segbuf(f, dom, p, alg.norm) +function init_cacheval(f::IntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) + segs = PuncturedInterval(dom) + prototype = get_prototype(f, get_prototype(segs), p) + return init_segbuf(prototype, segs, alg) end - -function do_solve(f, dom, p, alg::AuxQuadGKJL, cacheval; +function init_cacheval(f::InplaceIntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) + segs = PuncturedInterval(dom) + prototype = get_prototype(f, get_prototype(segs), p) + return init_segbuf(prototype, segs, alg), similar(prototype) +end +function init_cacheval(f::InplaceBatchIntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) + segs = PuncturedInterval(dom) + pt = get_prototype(segs) + prototype = get_prototype(f, pt, p) + prototype isa AbstractVector || throw(ArgumentError("QuadGKJL only supports batch integrands with vector outputs")) + pts = zeros(typeof(pt), 2*alg.order+1) + upts = pts / pt + return init_segbuf(first(prototype), segs, alg), similar(prototype), pts, upts +end +function init_cacheval(f::CommonSolveIntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) + segs = PuncturedInterval(dom) + x = get_prototype(segs) + cache, integrand, prototype = _init_commonsolvefunction(f, dom, p; x) + return init_segbuf(prototype, segs, alg), cache, integrand +end +function do_integral(f, dom, p, alg::AuxQuadGKJL, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - - segs = segments(dom) + # we need to strip units from the limits since infinity transformations change the units + # of the limits, which can break the segbuf u = oneunit(eltype(dom)) - usegs = map(x -> x/u, segs) - if f isa InplaceIntegrand - g! = (y, x) -> f.f!(y, u*x, p) - result = f.I / u - val, err = auxquadgk!(g!, result, usegs, maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(f.I .= u .* val, u*err, true, -1) - elseif f isa BatchIntegrand - xx = eltype(f.x) === Nothing ? typeof((segs[1]+segs[end])/2)[] : f.x - g_ = (y, x) -> (resize!(xx, length(x)); f.f!(y, xx .= u .* x, p)) - g = IteratedIntegration.AuxQuadGK.BatchIntegrand(g_, f.y, xx/u, max_batch=f.max_batch) - val, err = auxquadgk(g, usegs, maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(u*val, u*err, true, -1) - else - g = x -> f(u*x, p) - val, err = auxquadgk(g, usegs, maxevals = maxiters, - rtol = reltol, atol = isnothing(abstol) ? abstol : abstol/u, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(u*val, u*err, true, -1) - end + usegs = map(x -> x/u, dom) + atol = isnothing(abstol) ? abstol : abstol/u + val, err = call_auxquadgk(f, p, u, usegs, cacheval; maxevals = maxiters, rtol = reltol, atol, order = alg.order, norm = alg.norm) + value = u*val + retcode = err < max(something(atol, zero(err)), alg.norm(val)*something(reltol, isnothing(atol) ? sqrt(eps(one(eltype(usegs)))) : 0)) ? Success : Failure + stats = (; error=u*err) + return IntegralSolution(value, retcode, stats) +end +function call_auxquadgk(f::IntegralFunction, p, u, usegs, cacheval; kws...) + auxquadgk(x -> f.f(u*x, p), usegs...; kws..., segbuf=cacheval) +end +function call_auxquadgk(f::InplaceIntegralFunction, p, u, usegs, cacheval; kws...) + # TODO allocate everything in the AuxQuadGK.InplaceIntegrand in the cacheval + auxquadgk!((y, x) -> f.f!(y, u*x, p), cacheval[2], usegs...; kws..., segbuf=cacheval[1]) +end +function call_auxquadgk(f::InplaceBatchIntegralFunction, p, u, usegs, cacheval; kws...) + pts = cacheval[3] + g = IteratedIntegration.AuxQuadGK.BatchIntegrand((y, x) -> f.f!(y, resize!(pts, length(x)) .= u .* x, p), cacheval[2], cacheval[4]; max_batch=f.max_batch) + auxquadgk(g, usegs...; kws..., segbuf=cacheval[1]) +end +function call_auxquadgk(f::CommonSolveIntegralFunction, p, u, usegs, cacheval; kws...) + # cache = cacheval[2] could call do_solve!(cache, f, x, p) to fully specialize + integrand = cacheval[3] + auxquadgk(x -> integrand(u*x, p), usegs...; kws..., segbuf=cacheval[1]) end """ @@ -66,34 +88,50 @@ function ContQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIn return ContQuadGKJL(order, norm, rho, rootmeth) end -function init_cacheval(f, dom, p, alg::ContQuadGKJL) - f isa NestedBatchIntegrand && throw(ArgumentError("ContQuadGK doesn't support nested batching")) - f isa BatchIntegrand && throw(ArgumentError("ContQuadGK doesn't support batching")) - f isa InplaceIntegrand && throw(ArgumentError("ContQuadGK doesn't support inplace integrands")) - - a, b = endpoints(dom) +function init_csegbuf(prototype, dom, alg::ContQuadGKJL) + segs = PuncturedInterval(dom) + a, b = endpoints(segs) x, s = (a+b)/2, (b-a)/2 TX = typeof(x) + convert(ComplexF64, prototype) fx_s = one(ComplexF64) * s # currently the integrand is forcibly written to a ComplexF64 buffer TI = typeof(fx_s) TE = typeof(alg.norm(fx_s)) r_segbuf = IteratedIntegration.ContQuadGK.PoleSegment{TX,TI,TE}[] - fc_s = f(complex(x), p) * complex(s) # the regular evalrule is used on complex segments + fc_s = prototype * complex(s) # the regular evalrule is used on complex segments TCX = typeof(complex(x)) TCI = typeof(fc_s) TCE = typeof(alg.norm(fc_s)) c_segbuf = IteratedIntegration.ContQuadGK.Segment{TCX,TCI,TCE}[] return (r=r_segbuf, c=c_segbuf) end +function init_cacheval(f::IntegralFunction, dom, p, alg::ContQuadGKJL; kws...) + segs = PuncturedInterval(dom) + prototype = get_prototype(f, get_prototype(segs), p) + init_csegbuf(prototype, dom, alg) +end +function init_cacheval(f::CommonSolveIntegralFunction, dom, p, alg::ContQuadGKJL; kws...) + segs = PuncturedInterval(dom) + cache, integrand, prototype = _init_commonsolvefunction(f, dom, p; x=get_prototype(segs)) + segbufs = init_csegbuf(prototype, dom, alg) + return (; segbufs..., cache, integrand) +end -function do_solve(f, dom, p, alg::ContQuadGKJL, cacheval; +function do_integral(f, dom, p, alg::ContQuadGKJL, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - segs = segments(dom) - g = x -> f(x, p) - val, err = contquadgk(g, segs, maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, + value, err = call_contquadgk(f, p, dom, cacheval; maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, rtol = reltol, atol = abstol, order = alg.order, norm = alg.norm, r_segbuf=cacheval.r, c_segbuf=cacheval.c) - return IntegralSolution(val, err, true, -1) + retcode = err < max(something(abstol, zero(err)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(one(eltype(dom)))) : 0)) ? Success : Failure + stats = (; error=err) + return IntegralSolution(value, retcode, stats) +end +function call_contquadgk(f::IntegralFunction, p, segs, cacheval; kws...) + contquadgk(x -> f.f(x, p), segs; kws...) +end +function call_contquadgk(f::CommonSolveIntegralFunction, p, segs, cacheval; kws...) + integrand = cacheval.integrand + contquadgk(x -> integrand(x, p), segs...; kws...) end """ @@ -114,23 +152,40 @@ function MeroQuadGKJL(; order = 7, norm = norm, rho = 1.0, rootmeth = IteratedIn return MeroQuadGKJL(order, norm, rho, rootmeth) end -function init_cacheval(f, dom, p, alg::MeroQuadGKJL) - f isa NestedBatchIntegrand && throw(ArgumentError("MeroQuadGK doesn't support nested batching")) - f isa BatchIntegrand && throw(ArgumentError("MeroQuadGK doesn't support batching")) - f isa InplaceIntegrand && throw(ArgumentError("MeroQuadGK doesn't support inplace integrands")) - a, b = endpoints(dom) +function init_msegbuf(prototype, dom, alg::MeroQuadGKJL; kws...) + segs = PuncturedInterval(dom) + a, b = endpoints(segs) x, s = (a + b)/2, (b-a)/2 + convert(ComplexF64, prototype) fx_s = one(ComplexF64) * s # ignore the actual integrand since it is written to CF64 array err = alg.norm(fx_s) return IteratedIntegration.alloc_segbuf(typeof(x), typeof(fx_s), typeof(err)) end +function init_cacheval(f::IntegralFunction, dom, p, alg::MeroQuadGKJL; kws...) + segs = PuncturedInterval(dom) + prototype = get_prototype(f, get_prototype(segs), p) + segbuf = init_msegbuf(prototype, dom, alg) + return (; segbuf) +end +function init_cacheval(f::CommonSolveIntegralFunction, dom, p, alg::MeroQuadGKJL; kws...) + segs = PuncturedInterval(dom) + cache, integrand, prototype = _init_commonsolvefunction(f, dom, p; x=get_prototype(segs)) + segbuf = init_msegbuf(prototype, dom, alg) + return (; segbuf, cache, integrand) +end -function do_solve(f, dom, p, alg::MeroQuadGKJL, cacheval; +function do_integral(f, dom, p, alg::MeroQuadGKJL, cacheval; reltol = nothing, abstol = nothing, maxiters = typemax(Int)) - - segs = segments(dom) - g = x -> f(x, p) - val, err = meroquadgk(g, segs, maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, - rtol = reltol, atol = abstol, order = alg.order, norm = alg.norm, segbuf=cacheval) - return IntegralSolution(val, err, true, -1) + value, err = call_meroquadgk(f, p, dom, cacheval; maxevals = maxiters, rho = alg.rho, rootmeth = alg.rootmeth, + rtol = reltol, atol = abstol, order = alg.order, norm = alg.norm, segbuf=cacheval.segbuf) + retcode = err < max(something(abstol, zero(err)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(one(eltype(dom)))) : 0)) ? Success : Failure + stats = (; error=err) + return IntegralSolution(value, retcode, stats) +end +function call_meroquadgk(f::IntegralFunction, p, segs, cacheval; kws...) + meroquadgk(x -> f.f(x, p), segs; kws...) +end +function call_meroquadgk(f::CommonSolveIntegralFunction, p, segs, cacheval; kws...) + integrand = cacheval.integrand + meroquadgk(x -> integrand(x, p), segs...; kws...) end diff --git a/src/algorithms_meta.jl b/src/algorithms_meta.jl index aacd3ba..e491b22 100644 --- a/src/algorithms_meta.jl +++ b/src/algorithms_meta.jl @@ -81,7 +81,6 @@ function outer_integralfunction(f::IntegralFunction, x0, p) ws = nothing return func, ws, _update!, _postsolve end -# TODO implement for InplaceIntegralFunction, InplaceBatchIntegralFunction, CommonSolveIntegralFunction """ AbsoluteEstimate(est_alg, abs_alg; kws...) diff --git a/src/domains.jl b/src/domains.jl index d011be4..cc725b0 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -6,7 +6,12 @@ function get_prototype(dom) return (a+b)/2 end +function get_prototype(B::Basis) + return B * zero(SVector{ndims(B),float(eltype(B))}) +end + get_basis(B::Basis) = B +get_basis(B::AbstractMatrix) = Basis(B) """ PuncturedInterval(s) diff --git a/src/fourier.jl b/src/fourier.jl index 0cffd39..8a5b441 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -21,6 +21,8 @@ # the nested batched integrand is optional, but when included it allows for thread-safe # parallelization +abstract type AbstractFourierIntegralFunction <: AbstractIntegralFunction end + """ FourierIntegralFunction(f, s, [prototype=nothing]; alias=false) @@ -30,7 +32,7 @@ - `prototype`: - `alias::Bool`: whether to `deepcopy` the series (false) or use the series as-is (true) """ -struct FourierIntegralFunction{F,S,P} <: AbstractIntegralFunction +struct FourierIntegralFunction{F,S,P} <: AbstractFourierIntegralFunction f::F s::S prototype::P @@ -43,8 +45,10 @@ function get_prototype(f::FourierIntegralFunction, x, ws, p) end get_prototype(f::FourierIntegralFunction, x, p) = get_prototype(f, x, f.s, p) +# TODO implement FourierInplaceIntegrand FourierInplaceBatchIntegrand + """ - FourierCommonSolveIntegralFunction(prob, alg, update!, postsolve, s, [prototype, specialize]; alias=false, kws...) + CommonSolveFourierIntegralFunction(prob, alg, update!, postsolve, s, [prototype, specialize]; alias=false, kws...) Constructor for an integrand that solves a problem defined with the CommonSolve.jl interface, `prob`, which is instantiated using `init(prob, alg; kws...)`. Helper functions @@ -54,7 +58,7 @@ The `prototype` argument can help control how much to `specialize` on the type o problem, which defaults to `FullSpecialize()` so that run times are improved. However `FunctionWrapperSpecialize()` may help reduce compile times. """ -struct FourierCommonSolveIntegralFunction{P,A,S,K,U,PS,T,M<:AbstractSpecialization} <: AbstractIntegralFunction +struct CommonSolveFourierIntegralFunction{P,A,S,K,U,PS,T,M<:AbstractSpecialization} <: AbstractFourierIntegralFunction prob::P alg::A s::S @@ -65,16 +69,16 @@ struct FourierCommonSolveIntegralFunction{P,A,S,K,U,PS,T,M<:AbstractSpecializati specialize::M alias::Bool end -function FourierCommonSolveIntegralFunction(prob, alg, update!, postsolve, s, prototype=nothing, specialize=FullSpecialize(); alias=false, kws...) - return FourierCommonSolveIntegralFunction(prob, alg, s, NamedTuple(kws), update!, postsolve, prototype, specialize, alias) +function CommonSolveFourierIntegralFunction(prob, alg, update!, postsolve, s, prototype=nothing, specialize=FullSpecialize(); alias=false, kws...) + return CommonSolveFourierIntegralFunction(prob, alg, s, NamedTuple(kws), update!, postsolve, prototype, specialize, alias) end -function do_solve!(cache, f::FourierCommonSolveIntegralFunction, x, s, p) +function do_solve!(cache, f::CommonSolveFourierIntegralFunction, x, s, p) f.update!(cache, x, s, p) sol = solve!(cache) return f.postsolve(sol, x, s, p) end -function get_prototype(f::FourierCommonSolveIntegralFunction, x, ws, p) +function get_prototype(f::CommonSolveFourierIntegralFunction, x, ws, p) if isnothing(f.prototype) cache = init(f.prob, f.alg; f.kwargs...) do_solve!(cache, f, x, ws(x), p) @@ -82,9 +86,9 @@ function get_prototype(f::FourierCommonSolveIntegralFunction, x, ws, p) f.prototype end end -get_prototype(f::FourierCommonSolveIntegralFunction, x, p) = get_prototype(f, x, f.s, p) +get_prototype(f::CommonSolveFourierIntegralFunction, x, p) = get_prototype(f, x, f.s, p) -function init_specialized_fourierintegrand(cache, f, dom, p; x=get_prototype(dom), s = f.s(x), prototype=f.prototype) +function init_specialized_fourierintegrand(cache, f, dom, p; x=get_prototype(dom), ws=f.s, s = ws(x), prototype=f.prototype) proto = prototype === nothing ? do_solve!(cache, f, x, s, p) : prototype func = (x, s, p) -> do_solve!(cache, f, x, s, p) integrand = if f.specialize isa FullSpecialize @@ -96,12 +100,13 @@ function init_specialized_fourierintegrand(cache, f, dom, p; x=get_prototype(dom end return integrand, proto end -function _init_fouriercommonsolvefunction(f, dom, p; kws...) +function _init_commonsolvefourierfunction(f, dom, p; kws...) cache = init(f.prob, f.alg; f.kwargs...) integrand, prototype = init_specialized_fourierintegrand(cache, f, dom, p; kws...) return cache, integrand, prototype end +# TODO implement CommonSolveFourierInplaceIntegrand CommonSolveFourierInplaceBatchIntegrand # similar to workspace_allocate, but more type-stable because of loop unrolling and vector types function workspace_allocate_vec(s::AbstractFourierSeries{N}, x::NTuple{N,Any}, len::NTuple{N,Integer}=ntuple(one,Val(N))) where {N} @@ -143,9 +148,21 @@ function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadGKJL; kws... prototype = get_prototype(f, get_prototype(segs), ws, p) return init_segbuf(prototype, segs, alg), ws end +function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::QuadGKJL; kws...) + segs = PuncturedInterval(dom) + x = get_prototype(segs) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; x, ws) + return init_segbuf(prototype, segs, alg), ws, cache, integrand +end function call_quadgk(f::FourierIntegralFunction, p, u, usegs, cacheval; kws...) segbuf, ws = cacheval - quadgk(x -> (ux = u*x; f.f(ux, ws(ux), p)), usegs...; kws..., segbuf) + quadgk(x -> (ux = + u*x; f.f(ux, ws(ux), p)), usegs...; kws..., segbuf) +end +function call_quadgk(f::CommonSolveFourierIntegralFunction, p, u, usegs, cacheval; kws...) + segbuf, ws, _, integrand = cacheval + quadgk(x -> (ux = u*x; integrand(ux, ws(ux), p)), usegs...; kws..., segbuf) end function init_cacheval(f::FourierIntegralFunction, dom, p, ::HCubatureJL; kws...) @@ -157,16 +174,97 @@ function hcubature_integrand(f::FourierIntegralFunction, p, a, b, ws) x -> f.f(x, ws(x), p) end -function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadratureFunction; kws...) - x, w = alg.fun(alg.npt) +function init_autosymptr_cache(f::FourierIntegralFunction, dom, p, bufsize; kws...) ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) - # TODO: pre-allocate, later evaluate the series on the - return (rule=[(w, x) for (w,x) in zip(w,x)], buffer=nothing, ws) + return (; buffer=nothing, ws) +end +function init_autosymptr_cache(f::CommonSolveFourierIntegralFunction, dom, p, bufsize; kws...) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + cache, integrand, = _init_commonsolvefourierfunction(f, dom, p; ws) + return (; buffer=nothing, ws, cache, integrand) end function autosymptr_integrand(f::FourierIntegralFunction, p, segs, cacheval) ws = cacheval.ws x -> x isa FourierValue ? f.f(x.x, x.s, p) : f.f(x, ws(x), p) end +function autosymptr_integrand(f::CommonSolveFourierIntegralFunction, p, segs, cacheval) + integrand = cacheval.integrand + ws = cacheval.ws + return x -> x isa FourierValue ? integrand(x.x, x.s, p) : integrand(x, ws(x), p) +end + + +function init_cacheval(f::FourierIntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) + segs = PuncturedInterval(dom) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + prototype = get_prototype(f, get_prototype(segs), ws, p) + return init_segbuf(prototype, segs, alg), ws +end +function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) + segs = PuncturedInterval(dom) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + x = get_prototype(segs) + cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; x, ws) + return init_segbuf(prototype, segs, alg), ws, cache, integrand +end +function call_auxquadgk(f::FourierIntegralFunction, p, u, usegs, cacheval; kws...) + segbuf, ws = cacheval + auxquadgk(x -> (ux=u*x; f.f(ux, ws(ux), p)), usegs...; kws..., segbuf) +end +function call_auxquadgk(f::CommonSolveFourierIntegralFunction, p, u, usegs, cacheval; kws...) + # cache = cacheval[2] could call do_solve!(cache, f, x, p) to fully specialize + segbuf, ws, _, integrand = cacheval + auxquadgk(x -> (ux=u*x; integrand(ux, ws(ux), p)), usegs...; kws..., segbuf) +end + + +function init_cacheval(f::FourierIntegralFunction, dom, p, alg::ContQuadGKJL; kws...) + segs = PuncturedInterval(dom) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + prototype = get_prototype(f, get_prototype(segs), ws, p) + segbufs = init_csegbuf(prototype, dom, alg) + return (; segbufs..., ws) +end +function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::ContQuadGKJL; kws...) + segs = PuncturedInterval(dom) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; ws, x=get_prototype(segs)) + segbufs = init_csegbuf(prototype, dom, alg) + return (; segbufs..., ws, cache, integrand) +end +function call_contquadgk(f::FourierIntegralFunction, p, segs, cacheval; kws...) + ws = cacheval.ws + contquadgk(x -> f.f(x, ws(x), p), segs; kws...) +end +function call_contquadgk(f::CommonSolveFourierIntegralFunction, p, segs, cacheval; kws...) + integrand = cacheval.integrand + ws = cacheval.ws + contquadgk(x -> integrand(x, ws(x), p), segs...; kws...) +end + +function init_cacheval(f::FourierIntegralFunction, dom, p, alg::MeroQuadGKJL; kws...) + segs = PuncturedInterval(dom) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + prototype = get_prototype(f, get_prototype(segs), ws, p) + segbuf = init_msegbuf(prototype, dom, alg) + return (; segbuf, ws) +end +function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::MeroQuadGKJL; kws...) + segs = PuncturedInterval(dom) + ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; ws, x=get_prototype(segs)) + segbuf = init_msegbuf(prototype, dom, alg) + return (; segbuf, ws, cache, integrand) +end +function call_meroquadgk(f::FourierIntegralFunction, p, segs, cacheval; kws...) + ws = cacheval.ws + meroquadgk(x -> f.f(x, ws(x), p), segs; kws...) +end +function call_meroquadgk(f::CommonSolveFourierIntegralFunction, p, segs, cacheval; kws...) + integrand = cacheval.integrand + ws = cacheval.ws + meroquadgk(x -> integrand(x, ws(x), p), segs...; kws...) +end function _fourier_update!(cache, x, p) _update!(cache, x, p) @@ -189,7 +287,7 @@ function outer_integralfunction(f::FourierIntegralFunction, x0, p) func = FourierIntegralFunction(f.f, s, proto; alias=true) return func, ws, _fourier_update!, _postsolve end -# TODO implement FourierCommonSolveIntegralFunction +# TODO implement CommonSolveFourierIntegralFunction # PTR rules @@ -420,9 +518,10 @@ function init_fourier_rule(w::FourierWorkspace, dom::RepBZ, alg::AutoSymPTRJL) rule = init_fourier_rule(w, B, alg) return SymmetricRuleDef(rule, dom.rep, dom.bz) end -function init_cacheval(f::FourierIntegralFunction, dom, p, alg::AutoSymPTRJL; kws...) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) +function init_cacheval(f::AbstractFourierIntegralFunction, dom, p, alg::AutoSymPTRJL; kws...) + cache = init_autosymptr_cache(f, dom, p, alg.nthreads; kws...) + ws = cache.ws rule = init_fourier_rule(ws, dom, alg) - cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) - return (; rule, cache, buffer=nothing, ws) + rule_cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) + return (; rule, rule_cache, cache...) end diff --git a/src/interfaces.jl b/src/interfaces.jl index 56a5478..df8d748 100644 --- a/src/interfaces.jl +++ b/src/interfaces.jl @@ -120,8 +120,7 @@ function _init_commonsolvefunction(f, dom, p; kws...) return cache, integrand, prototype end -# TODO reimplement CommonSolveIntegralFunction as a CommonSolveIntegralProblem OR -# (why? so that FunctionWrappers can be used more systematically) +# TODO add InplaceCommonSolveIntegralFunction and InplaceBatchCommonSolveIntegralFunction # TODO add ThreadedCommonSolveIntegralFunction and DistributedCommonSolveIntegralFunction """ @@ -162,7 +161,7 @@ function IntegralProblem(f, dom, p=NullParameters(); kws...) return IntegralProblem(IntegralFunction(f), dom, p; kws...) end -mutable struct IntegralCache{F,D,P,A,C,K} +mutable struct IntegralSolver{F,D,P,A,C,K} f::F dom::D p::P @@ -179,7 +178,7 @@ function checkkwargs(kwargs) end """ - init(::IntegralProblem, ::IntegralAlgorithm; kws...)::IntegralCache + init(::IntegralProblem, ::IntegralAlgorithm; kws...)::IntegralSolver Construct a cache for an [`IntegralProblem`](@ref), [`IntegralAlgorithm`](@ref), and the keyword arguments to the solver (i.e. `abstol`, `reltol`, or `maxiters`) that can be reused @@ -190,7 +189,7 @@ function init(prob::IntegralProblem, alg::IntegralAlgorithm; kwargs...) kws = (; prob.kwargs..., kwargs...) checkkwargs(kws) cacheval = init_cacheval(f, dom, p, alg; kws...) - return IntegralCache(f, dom, p, alg, cacheval, kws) + return IntegralSolver(f, dom, p, alg, cacheval, kws) end """ @@ -218,11 +217,11 @@ smaller than the absolute tolerance. solve(prob::IntegralProblem, alg::IntegralAlgorithm; kwargs...) """ - solve!(::IntegralCache)::IntegralSolution + solve!(::IntegralSolver)::IntegralSolution Compute the solution to an [`IntegralProblem`](@ref) constructed from [`init`](@ref). """ -function solve!(c::IntegralCache) +function solve!(c::IntegralSolver) return do_integral(c.f, c.dom, c.p, c.alg, c.cacheval; c.kwargs...) end diff --git a/test/fourier.jl b/test/fourier.jl index 9f237a4..1d5abc4 100644 --- a/test/fourier.jl +++ b/test/fourier.jl @@ -17,23 +17,65 @@ using AutoBZCore: PuncturedInterval, HyperCube, segments, endpoints ref = (b-a)*p + t*(b*sin(b/t*2pi) + t*cos(b/t*2pi) - (a*sin(a/t*2pi) + t*cos(a/t*2pi))) abstol = 1e-5 prob = IntegralProblem(FourierIntegralFunction(int, s), (a, b), p; abstol) - for alg in (QuadGKJL(), QuadratureFunction()) + for alg in (QuadGKJL(), QuadratureFunction(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) @test solve(prob, alg).value ≈ ref atol=abstol end end + @testset "commonproblem" begin + a = 0 + b = 1 + p = 0.0 + t = 1.0 + update! = (cache, x, s, p) -> cache.p = (x, s, p) + postsolve = (sol, x, s, p) -> sol.value + s = FourierSeries([1, 0, 1]/2; period=t, offset=-2) + f = (x, (y, s, p)) -> x * s + p + y + subprob = IntegralProblem(f, (a, b), ((a+b)/2, s((a+b)/2), 1.0)) + abstol = 1e-5 + prob = IntegralProblem(CommonSolveFourierIntegralFunction(subprob, QuadGKJL(), update!, postsolve, s), (a, b), p; abstol) + for alg in (QuadGKJL(), QuadratureFunction(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) + cache = init(prob, alg) + for p in [3.0, 4.0] + ref = (b-a)*(t*p + (b-a)^2/2) + (b-a)^2/2*t*(b*sin(b/t*2pi) + t*cos(b/t*2pi) - (a*sin(a/t*2pi) + t*cos(a/t*2pi))) + cache.p = p + sol = solve!(cache) + @test ref ≈ sol.value atol=abstol + end + end + f = (x, (y, s, p)) -> x * s + p + subprob = IntegralProblem(f, (a, b), ([(a+b)/2], s((a+b)/2), 1.0)) + abstol = 1e-5 + prob = IntegralProblem(CommonSolveFourierIntegralFunction(subprob, QuadGKJL(), update!, postsolve, s), AutoBZCore.Basis(t*I(1)), p; abstol) + for alg in (MonkhorstPack(), AutoSymPTRJL(),) + cache = init(prob, alg) + for p in [3.0, 4.0] + ref = (b-a)*(t*p) + (b-a)^2/2*t*(b*sin(b/t*2pi) + t*cos(b/t*2pi) - (a*sin(a/t*2pi) + t*cos(a/t*2pi))) + cache.p = p + sol = solve!(cache) + @test ref ≈ sol.value atol=abstol + end + end + end @testset "cubature" for dim in 2:3 a = zeros(dim) b = ones(dim) p = 0.0 t = 1.0 s = FourierSeries([prod(x) for x in Iterators.product([(0.1, 0.5, 0.3) for i in 1:dim]...)]; period=t, offset=-2) - int(x, s, p) = prod(x) * s + p + f = (x, s, p) -> prod(x) * s + p abstol = 1e-4 - prob = IntegralProblem(FourierIntegralFunction(int, s), (a, b), p; abstol) - refprob = IntegralProblem(IntegralFunction((x, p) -> int(x, s(x), p)), (a, b), p; abstol) + prob = IntegralProblem(FourierIntegralFunction(f, s), (a, b), p; abstol) + refprob = IntegralProblem(IntegralFunction(let f=f; (x, p) -> f(x, s(x), p); end), (a, b), p; abstol) for alg in (HCubatureJL(),) @test solve(prob, alg).value ≈ solve(refprob, alg).value atol=abstol end + p=1.3 + f = (x, s, p) -> inv(s + im*p) + prob = IntegralProblem(FourierIntegralFunction(f, s), AutoBZCore.Basis(t*I(dim)), p; abstol) + refprob = IntegralProblem(IntegralFunction(let f=f; (x, p) -> f(x, s(x), p); end), AutoBZCore.Basis(t*I(dim)), p; abstol) + for alg in (MonkhorstPack(), AutoSymPTRJL(),) + @test solve(prob, alg).value ≈ solve(refprob, alg).value atol=abstol + end end @testset "meta-algorithms" for dim in 2:3 # NestedQuad @@ -46,7 +88,7 @@ using AutoBZCore: PuncturedInterval, HyperCube, segments, endpoints abstol = 1e-4 prob = IntegralProblem(FourierIntegralFunction(int, s), CubicLimits(a, b), p0; abstol) refprob = IntegralProblem(FourierIntegralFunction(int, s), (a, b), p0; abstol) - for alg in (QuadGKJL(),) + for alg in (QuadGKJL(), AuxQuadGKJL()) cache = init(prob, NestedQuad(alg)) refcache = init(refprob, HCubatureJL()) for p in [5.0, 6.0] @@ -55,6 +97,8 @@ using AutoBZCore: PuncturedInterval, HyperCube, segments, endpoints @test solve!(cache).value ≈ solve!(refcache).value atol=abstol end end + # TODO implement CommonSolveFourierInplaceIntegralFunction + # TODO implement CommonSolveFourierInplaceBatchIntegralFunction end end diff --git a/test/interface_tests.jl b/test/interface_tests.jl index 77e2c31..f6f511b 100644 --- a/test/interface_tests.jl +++ b/test/interface_tests.jl @@ -36,7 +36,7 @@ end ((x,p) -> inv(p-cos(x)), (b-a)/sqrt(p^2-1)), ) prob = IntegralProblem(f, (a, b), p; abstol) - for alg in (QuadratureFunction(), QuadGKJL(),)#(QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) + for alg in (QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) sol = solve(prob, alg) @test ref ≈ sol.value atol=abstol end @@ -55,7 +55,7 @@ end subprob = IntegralProblem(f, (a, b), ((a+b)/2, p0); abstol) integrand = CommonSolveIntegralFunction(subprob, QuadGKJL(), update!, postsolve) prob = IntegralProblem(integrand, (a, b), p0; abstol) - for alg in (QuadratureFunction(), QuadGKJL(), HCubatureJL()) #(QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) + for alg in (QuadratureFunction(), QuadGKJL(), HCubatureJL(), AuxQuadGKJL(), ContQuadGKJL(), MeroQuadGKJL()) cache = init(prob, alg) for p in [3.0, 4.0] ref = p*(b-a)*(b^2-a^2) @@ -64,6 +64,19 @@ end @test ref ≈ sol.value atol=abstol end end + f = (x, (y, p)) -> p*(sin(only(y))^2 + x) + subprob = IntegralProblem(f, (a, b), ([b/2], p0); abstol) + integrand = CommonSolveIntegralFunction(subprob, QuadGKJL(), update!, postsolve) + prob = IntegralProblem(integrand, AutoBZCore.Basis(b*I(1)), p0; abstol) + for alg in (MonkhorstPack(), AutoSymPTRJL(),) + cache = init(prob, alg) + for p in [3.0, 4.0] + ref = p*((b-a)+(b^2-a^2))*b/2 + cache.p = p + sol = solve!(cache) + @test ref ≈ sol.value atol=abstol + end + end end @testset "cubature" begin @@ -81,12 +94,10 @@ end for alg in (HCubatureJL(),) @test ref ≈ solve(prob, alg).value atol=abstol end - #= - prob = IntegralProblem(f, Basis(b*I(dim)), p) + prob = IntegralProblem(f, AutoBZCore.Basis(b*I(dim)), p; abstol) for alg in (MonkhorstPack(), AutoSymPTRJL(),) - @test ref ≈ solve(prob, alg, abstol=abstol).value atol=abstol + @test ref ≈ solve(prob, alg).value atol=abstol end - =# end end @@ -103,15 +114,13 @@ end ) integrand = InplaceIntegralFunction(f, [0.0]) inplaceprob = IntegralProblem(integrand, (a, b), p; abstol) - for alg in (QuadGKJL(),)#(QuadratureFunction(), QuadGKJL(), AuxQuadGKJL(), HCubatureJL(),) + for alg in (QuadGKJL(), QuadratureFunction(), QuadGKJL(), AuxQuadGKJL()) @test ref ≈ solve(inplaceprob, alg).value atol=abstol end - #= - inplaceprob = IntegralProblem(integrand, Basis([b;;]), p) + inplaceprob = IntegralProblem(integrand, AutoBZCore.Basis([b;;]), p; abstol) for alg in (MonkhorstPack(), AutoSymPTRJL()) - @test ref ≈ solve(inplaceprob, alg, abstol=abstol).value atol=abstol + @test ref ≈ solve(inplaceprob, alg).value atol=abstol end - =# end end @@ -128,16 +137,14 @@ end ((y,x,p) -> y .= inv.(p .- cos.(only.(x))), (b-a)/sqrt(p^2-1)), ) integrand = InplaceBatchIntegralFunction(f, zeros(1)) - batchprob = IntegralProblem(integrand, (a, b), p) - for alg in (QuadGKJL(),)#(QuadratureFunction(), AuxQuadGKJL()) - @test ref ≈ solve(batchprob, alg, abstol=abstol).value atol=abstol + batchprob = IntegralProblem(integrand, (a, b), p; abstol) + for alg in (QuadGKJL(), QuadratureFunction(), AuxQuadGKJL()) + @test ref ≈ solve(batchprob, alg).value atol=abstol end - #= - batchprob = IntegralProblem(integrand, Basis([b;;]), p) + batchprob = IntegralProblem(integrand, AutoBZCore.Basis([b;;]), p; abstol) for alg in (MonkhorstPack(), AutoSymPTRJL()) - @test ref ≈ solve(batchprob, alg, abstol=abstol).value atol=abstol + @test ref ≈ solve(batchprob, alg).value atol=abstol end - =# end end @testset "multi-algorithms" begin @@ -145,7 +152,7 @@ end f(x, p) = 1.0 + p*sum(abs2 ∘ cos, x) abstol=1e-3 p0 = 0.0 - for dim in 1:3, alg in (QuadratureFunction(), QuadGKJL())# AuxQuadGKJL()) + for dim in 1:3, alg in (QuadratureFunction(), QuadGKJL(), AuxQuadGKJL()) dom = CubicLimits(zeros(dim), 2pi*ones(dim)) prob = IntegralProblem(f, dom, p0; abstol) ndalg = NestedQuad(alg) @@ -154,11 +161,13 @@ end cache.p = p ref = (2pi)^dim + dim*p*pi*(2pi)^(dim-1) @test ref ≈ solve!(cache).value atol=abstol + # TODO implement CommonSolveInplaceIntegralFunction + inplaceprob = IntegralProblem(InplaceIntegralFunction((y,x,p) -> y .= f(x,p), [0.0]), dom, p) + @test_broken [ref] ≈ solve(inplaceprob, ndalg, abstol=abstol).value atol=abstol + # TODO implement CommonSolveInplaceBatchIntegralFunction + batchprob = IntegralProblem(InplaceBatchIntegralFunction((y,x,p) -> y .= f.(x,Ref(p)), zeros(Float64, 1)), dom, p) + @test_broken ref ≈ solve(batchprob, ndalg, abstol=abstol).value atol=abstol end - # inplaceprob = IntegralProblem(InplaceIntegrand((y,x,p) -> y .= f(x,p), [0.0]), dom, p) - # @test [ref] ≈ solve(inplaceprob, ndalg, abstol=abstol).value atol=abstol - # batchprob = IntegralProblem(BatchIntegrand((y,x,p) -> y .= f.(x,Ref(p)), Float64), dom, p) - # @test ref ≈ solve(batchprob, ndalg, abstol=abstol).value atol=abstol end #= # AbsoluteEstimate From 4f487888dbf3395a5b41a547e7dd54cf8ab31976 Mon Sep 17 00:00:00 2001 From: lxvm Date: Wed, 19 Jun 2024 17:47:00 -0400 Subject: [PATCH 09/17] rename cache to solver --- aps_example/aps_example.jl | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/aps_example/aps_example.jl b/aps_example/aps_example.jl index eb203d2..cf65d93 100644 --- a/aps_example/aps_example.jl +++ b/aps_example/aps_example.jl @@ -47,22 +47,22 @@ using HChebInterp cheb_order = 15 function dos_solver(prob, alg) - cache = init(prob, alg) + solver = init(prob, alg) ω -> begin - cache.p = (; cache.p..., ω) - solve!(cache).value + solver.p = (; solver.p..., ω) + solve!(solver).value end end function threaded_dos_solver(prob, alg; nthreads=min(cheb_order, Threads.nthreads())) - caches = [init(prob, alg) for _ in 1:nthreads] + solvers = [init(prob, alg) for _ in 1:nthreads] BatchFunction() do ωs out = Vector{typeof(prototype)}(undef, length(ωs)) Threads.@threads for i in 1:nthreads - cache = caches[i] + solver = solvers[i] for j in i:nthreads:length(ωs) ω = ωs[j] - cache.p = (; cache.p..., ω) - out[j] = solve!(cache).value + solver.p = (; solver.p..., ω) + out[j] = solve!(solver).value end end return out From 21204c1f7528188ea2284b1a75eb5e3e17f3930e Mon Sep 17 00:00:00 2001 From: lxvm Date: Sat, 22 Jun 2024 22:39:58 -0400 Subject: [PATCH 10/17] remove frankentuples --- src/parameters.jl | 120 ---------------------------------------------- 1 file changed, 120 deletions(-) delete mode 100644 src/parameters.jl diff --git a/src/parameters.jl b/src/parameters.jl deleted file mode 100644 index c9d52ce..0000000 --- a/src/parameters.jl +++ /dev/null @@ -1,120 +0,0 @@ -""" - MixedParameters(args::Tuple, kwargs::NamedTuple) - MixedParameters(args...; kwargs...) - -A struct to store the arguments and keyword arguments to a function. Indicies access -`args`, i.e. `MixedParameters(args...; kwargs...)[i] == args[i]` and properties access -`kwargs`, i.e. `MixedParameters(args...; kwargs...).name == kwargs.name`. - -Used internally to store partially complete collections of function arguments or parameters. -""" -struct MixedParameters{A<:Tuple,K<:NamedTuple} - args::A - kwargs::K -end -MixedParameters(args...; kwargs...) = MixedParameters(args, NamedTuple(kwargs)) - -# parameter fusion and iteration utilities - -Base.getindex(p::MixedParameters, i::Int) = getindex(getfield(p, :args), i) -Base.getproperty(p::MixedParameters, name::Symbol) = getproperty(getfield(p, :kwargs), name) - -Base.merge(p::MixedParameters, q) = - MixedParameters((getfield(p, :args)..., q), getfield(p, :kwargs)) -Base.merge(q, p::MixedParameters) = - MixedParameters((q, getfield(p, :args)...), getfield(p, :kwargs)) -Base.merge(p::MixedParameters, q::NamedTuple) = - MixedParameters(getfield(p, :args), merge(getfield(p, :kwargs), q)) -Base.merge(q::NamedTuple, p::MixedParameters) = - MixedParameters(getfield(p, :args), merge(q, getfield(p, :kwargs))) -Base.merge(p::MixedParameters, q::Tuple) = - MixedParameters((getfield(p, :args)..., q...), getfield(p, :kwargs)) -Base.merge(q::Tuple, p::MixedParameters) = - MixedParameters((q..., getfield(p, :args)...), getfield(p, :kwargs)) -Base.merge(p::MixedParameters, q::MixedParameters) = - MixedParameters((getfield(p, :args)..., getfield(q, :args)...), merge(getfield(p, :kwargs), getfield(q, :kwargs))) - -function paramzip_(::Tuple{}, ::NamedTuple{(), Tuple{}}) - MixedParameters{Tuple{},NamedTuple{(), Tuple{}}}[] -end -function paramzip_(args::Tuple, ::NamedTuple{(), Tuple{}}) - [MixedParameters(arg, NamedTuple()) for arg in zip(args...)] -end -function paramzip_(::Tuple{}, kwargs::NamedTuple) - [MixedParameters((), NamedTuple{keys(kwargs)}(val)) for val in zip(values(kwargs)...)] -end -function paramzip_(args::Tuple, kwargs::NamedTuple) - [MixedParameters(arg, NamedTuple{keys(kwargs)}(val)) for (arg, val) in zip(zip(args...), zip(values(kwargs)...))] -end -""" - paramzip(args...; kwargs...) - -Behaves similarly to `zip(zip(args...), zip(kwargs...))` with [`MixedParameters`](@ref) -return values so that `paramzip(args...; kwargs...)[i][j] == args[j][i]` and -`paramzip(args...; kwargs...)[i].name == kwargs.name[i]`. -""" -paramzip(args...; kwargs...) = paramzip_(args, NamedTuple(kwargs)) - -function paramproduct_(args::Tuple, kwargs::NamedTuple) - [MixedParameters(item[1:length(args)], NamedTuple{keys(kwargs)}(item[length(args)+1:end])) for item in Iterators.product(args..., values(kwargs)...)] -end - -""" - paramproduct(args...; kwargs...) - -Behaves similarly to `product(args..., kwargs...)` with [`MixedParameters`](@ref) return -values so that `paramzip(args...; kwargs...)[i1, ...,ij, il, ...in] == -MixedParameters(args[begin][i1], ..., args[end][ij]; kwargs[begin][il], ..., kwargs[end][in])` -""" -paramproduct(args...; kwargs...) = paramproduct_(args, NamedTuple(kwargs)) - -""" - ParameterIntegrand(f, args...; kwargs...) - -Represent an integrand with a partial collection of parameters `p`. When the -`ParameterIntegrand` is invoked with one argument, e.g. `int(x)`, it evaluates `f(x, -p...; kwargs...)`. However when invoked with two arguments, as in an `IntegralProblem`, -e.g. `int(x, p2)`, it evaluates the union of parameters `f(x, p..., p2...; kwargs...)`. -This allows for convenient parametrization of the integrand. -""" -struct ParameterIntegrand{F,P} - f::F - p::P - function ParameterIntegrand{F}(f::F, p::P) where {F,P<:MixedParameters} - return new{F,P}(f, p) - end -end - -function ParameterIntegrand(f, args...; kwargs...) - p = MixedParameters(args...; kwargs...) - return ParameterIntegrand{typeof(f)}(f, p) -end - -# provide Integrals.jl interface -function (f::ParameterIntegrand)(x, q) - p = merge(f.p, q) - return f.f(x, getfield(p, :args)...; getfield(p, :kwargs)...) -end -(f::ParameterIntegrand)(x, ::NullParameters) = f(x, ()) - -# move all parameters from f.p to p for convenience -remake_integrand_cache(args...) = IntegralCache(args...) -function remake_cache(f::ParameterIntegrand, dom, p, alg, cacheval, kwargs) - new = ParameterIntegrand(f.f) - return remake_integrand_cache(new, dom, merge(f.p, p), alg, cacheval, kwargs) -end - -function do_solve(f::ParameterIntegrand, dom, p, alg::EvalCounter, cacheval; kws...) - n::Int = 0 - function g(x, args...; kwargs...) - n += 1 - return f.f(x, args...; kwargs...) - end - sol = do_solve(ParameterIntegrand{typeof(g)}(g, f.p), dom, p, alg.alg, cacheval; kws...) - return IntegralSolution(sol.u, sol.resid, true, n) -end - -# ambiguity -function do_solve(f::ParameterIntegrand, bz::SymmetricBZ, p, alg::EvalCounter{<:AutoBZAlgorithm}, cacheval; kws...) - return do_solve(f, bz, p, AutoBZEvalCounter(bz_to_standard(bz, alg.alg)...), cacheval; kws...) -end From 93396424883c7a8d15e25e0328962f035bf48a07 Mon Sep 17 00:00:00 2001 From: lxvm Date: Tue, 25 Jun 2024 09:28:16 -0400 Subject: [PATCH 11/17] minor fixes --- src/algorithms_autosymptr.jl | 3 --- src/brillouin.jl | 6 ++++-- src/fourier.jl | 10 ++++------ 3 files changed, 8 insertions(+), 11 deletions(-) diff --git a/src/algorithms_autosymptr.jl b/src/algorithms_autosymptr.jl index feb7981..470d0d1 100644 --- a/src/algorithms_autosymptr.jl +++ b/src/algorithms_autosymptr.jl @@ -28,9 +28,6 @@ function init_rule(dom, alg::MonkhorstPack) end end -rule_type(::AutoSymPTR.PTR{N,T}) where {N,T} = SVector{N,T} -rule_type(::AutoSymPTR.MonkhorstPack{N,T}) where {N,T} = SVector{N,T} - function init_cacheval(f, dom, p, alg::MonkhorstPack; kws...) b = get_basis(dom) rule = init_rule(b, alg) diff --git a/src/brillouin.jl b/src/brillouin.jl index 722d67c..527d996 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -49,6 +49,7 @@ Base.summary(bz::SymmetricBZ) = string(checksquare(bz.A), "-dimensional Brilloui Base.show(io::IO, bz::SymmetricBZ) = print(io, summary(bz)) Base.ndims(::SymmetricBZ{S,L,d}) where {S,L,d} = d Base.eltype(::Type{<:SymmetricBZ{S,L,d,TA,TB}}) where {S,L,d,TA,TB} = TB +get_prototype(bz::SymmetricBZ) = interior_point(bz.lims) # Define traits for symmetrization based on symmetry representations @@ -113,7 +114,6 @@ Base.getindex(r::SymmetricRule, i) = getindex(r.rule, i) Base.eltype(::Type{SymmetricRule{R,U,B}}) where {R,U,B} = eltype(R) Base.length(r::SymmetricRule) = length(r.rule) Base.iterate(r::SymmetricRule, args...) = iterate(r.rule, args...) -rule_type(r::SymmetricRule) = rule_type(r.rule) function (r::SymmetricRule)(f::F, args...) where {F} out = r.rule(f, args...) return symmetrize(r.rep, r.bz, out) @@ -340,7 +340,7 @@ implementing an AbstractSymRep for your type. """ function AutoBZProblem(rep::AbstractSymRep, f::AbstractIntegralFunction, bz::SymmetricBZ, p=NullParameters(); kws...) - proto = get_prototype(f, interior_point(bz.lims), p) + proto = get_prototype(f, get_prototype(bz), p) if rep isa UnknownRep && !(bz isa FullBZ) && !(proto isa TrivialRepType) @warn WARN_UNKNOWN_SYMMETRY fbz = SymmetricBZ(bz.A, bz.B, lattice_bz_limits(bz.B), nothing) @@ -477,6 +477,8 @@ struct RepBZ{R,B} end Base.ndims(dom::RepBZ) = ndims(dom.bz) Base.eltype(::Type{RepBZ{R,B}}) where {R,B} = eltype(B) +get_prototype(dom::RepBZ) = get_prototype(dom.bz) + function init_cacheval(rep, f, bz::SymmetricBZ, p, bzalg::AutoPTR; kws...) prob = IntegralProblem(f, RepBZ(rep, bz), p; kws...) diff --git a/src/fourier.jl b/src/fourier.jl index 8a5b441..e1db166 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -452,8 +452,6 @@ Base.eltype(::Type{FourierMonkhorstPack{d,W,T,S}}) where {d,W,T,S} = Tuple{W,Fou Base.length(r::FourierMonkhorstPack) = length(r.wxs) Base.iterate(rule::FourierMonkhorstPack, args...) = iterate(rule.wxs, args...) -rule_type(::FourierMonkhorstPack{d,W,T,S}) where {d,W,T,S} = FourierValue{SVector{d,T},S} - function (rule::FourierMonkhorstPack{d})(f::F, B::Basis, buffer=nothing) where {d,F} arule = AutoSymPTR.AffineQuad(rule, B) return AutoSymPTR.quadsum(arule, f, arule.vol / (rule.npt^d * rule.nsyms), buffer) @@ -494,7 +492,6 @@ end # return f.nest isa NestedBatchIntegrand ? Vector{eltype(f.nest.y)}(undef, len) : nothing # end -rule_type(::FourierPTR{N,T,S}) where {N,T,S} = FourierValue{SVector{N,T},S} function init_fourier_rule(w::FourierWorkspace, dom, alg::MonkhorstPack) @assert ndims(w.series) == ndims(dom) if alg.syms === nothing @@ -503,10 +500,11 @@ function init_fourier_rule(w::FourierWorkspace, dom, alg::MonkhorstPack) return FourierMonkhorstPack(w, eltype(dom), Val(ndims(dom)), alg.npt, alg.syms) end end -function init_cacheval(f::FourierIntegralFunction, dom , p, alg::MonkhorstPack; kws...) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) +function init_cacheval(f::AbstractFourierIntegralFunction, dom , p, alg::MonkhorstPack; kws...) + cache = init_autosymptr_cache(f, dom, p, alg.nthreads; kws...) + ws = cache.ws rule = init_fourier_rule(ws, dom, alg) - return (; rule, buffer=nothing, ws) + return (; rule, buffer=nothing, ws, cache...) end function init_fourier_rule(w::FourierWorkspace, dom, alg::AutoSymPTRJL) From 4174670a187025d319855b3253954a5355cece43 Mon Sep 17 00:00:00 2001 From: lxvm Date: Wed, 26 Jun 2024 20:18:51 -0400 Subject: [PATCH 12/17] IAI for CommonSolveFourierIntegralFunction --- src/fourier.jl | 52 ++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/src/fourier.jl b/src/fourier.jl index e1db166..4bc072c 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -45,6 +45,10 @@ function get_prototype(f::FourierIntegralFunction, x, ws, p) end get_prototype(f::FourierIntegralFunction, x, p) = get_prototype(f, x, f.s, p) +function get_fourierworkspace(f::AbstractFourierIntegralFunction) + f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) +end + # TODO implement FourierInplaceIntegrand FourierInplaceBatchIntegrand """ @@ -144,14 +148,14 @@ end function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) prototype = get_prototype(f, get_prototype(segs), ws, p) return init_segbuf(prototype, segs, alg), ws end function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::QuadGKJL; kws...) segs = PuncturedInterval(dom) x = get_prototype(segs) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; x, ws) return init_segbuf(prototype, segs, alg), ws, cache, integrand end @@ -167,7 +171,7 @@ end function init_cacheval(f::FourierIntegralFunction, dom, p, ::HCubatureJL; kws...) # TODO utilize hcubature_buffer - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) return ws end function hcubature_integrand(f::FourierIntegralFunction, p, a, b, ws) @@ -175,11 +179,11 @@ function hcubature_integrand(f::FourierIntegralFunction, p, a, b, ws) end function init_autosymptr_cache(f::FourierIntegralFunction, dom, p, bufsize; kws...) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) return (; buffer=nothing, ws) end function init_autosymptr_cache(f::CommonSolveFourierIntegralFunction, dom, p, bufsize; kws...) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) cache, integrand, = _init_commonsolvefourierfunction(f, dom, p; ws) return (; buffer=nothing, ws, cache, integrand) end @@ -196,13 +200,13 @@ end function init_cacheval(f::FourierIntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) prototype = get_prototype(f, get_prototype(segs), ws, p) return init_segbuf(prototype, segs, alg), ws end function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::AuxQuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) x = get_prototype(segs) cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; x, ws) return init_segbuf(prototype, segs, alg), ws, cache, integrand @@ -220,14 +224,14 @@ end function init_cacheval(f::FourierIntegralFunction, dom, p, alg::ContQuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) prototype = get_prototype(f, get_prototype(segs), ws, p) segbufs = init_csegbuf(prototype, dom, alg) return (; segbufs..., ws) end function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::ContQuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; ws, x=get_prototype(segs)) segbufs = init_csegbuf(prototype, dom, alg) return (; segbufs..., ws, cache, integrand) @@ -244,14 +248,14 @@ end function init_cacheval(f::FourierIntegralFunction, dom, p, alg::MeroQuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) prototype = get_prototype(f, get_prototype(segs), ws, p) segbuf = init_msegbuf(prototype, dom, alg) return (; segbuf, ws) end function init_cacheval(f::CommonSolveFourierIntegralFunction, dom, p, alg::MeroQuadGKJL; kws...) segs = PuncturedInterval(dom) - ws = FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) cache, integrand, prototype = _init_commonsolvefourierfunction(f, dom, p; ws, x=get_prototype(segs)) segbuf = init_msegbuf(prototype, dom, alg) return (; segbuf, ws, cache, integrand) @@ -273,7 +277,7 @@ function _fourier_update!(cache, x, p) return end function inner_integralfunction(f::FourierIntegralFunction, x0, p) - ws = f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) proto = get_prototype(f, x0, ws, p) func = IntegralFunction(proto) do x, (; p, ws, lims_state) f.f(limit_iterate(lims_state..., x), workspace_evaluate!(ws, x), p) @@ -281,13 +285,33 @@ function inner_integralfunction(f::FourierIntegralFunction, x0, p) return func, ws end function outer_integralfunction(f::FourierIntegralFunction, x0, p) - ws = f.s isa FourierWorkspace ? f.s : FourierSeriesEvaluators.workspace_allocate(f.alias ? f.s : deepcopy(f.s), FourierSeriesEvaluators.period(f.s)) + ws = get_fourierworkspace(f) proto = get_prototype(f, x0, ws, p) s = workspace_contract!(ws, x0[end]) func = FourierIntegralFunction(f.f, s, proto; alias=true) return func, ws, _fourier_update!, _postsolve end -# TODO implement CommonSolveFourierIntegralFunction +# TODO it would be desirable to allow the inner integralfunction to be of the +# same type as f, which requires moving workspace out of the parameters into +# some kind of mutable storage +function inner_integralfunction(f::CommonSolveFourierIntegralFunction, x0, p) + ws = get_fourierworkspace(f) + proto = get_prototype(f, x0, ws, p) + cache = init(f.prob, f.alg; f.kwargs...) + func = IntegralFunction(proto) do x, (; p, ws, lims_state) + y = limit_iterate(lims_state..., x) + s = workspace_evaluate!(ws, x) + do_solve!(cache, f, y, s, p) + end + return func, ws +end +function outer_integralfunction(f::CommonSolveFourierIntegralFunction, x0, p) + ws = get_fourierworkspace(f) + proto = get_prototype(f, x0, ws, p) + s = workspace_contract!(ws, x0[end]) + func = CommonSolveFourierIntegralFunction(f.prob, f.alg, f.update!, f.postsolve, s, proto, f.specialize; alias=true, f.kwargs...) + return func, ws, _fourier_update!, _postsolve +end # PTR rules From 020a1bd6e769b1a32ae5b41c875c9a5caae76d68 Mon Sep 17 00:00:00 2001 From: lxvm Date: Mon, 15 Jul 2024 15:41:53 -0400 Subject: [PATCH 13/17] pass tests --- src/AutoBZCore.jl | 4 +- src/algorithms_autosymptr.jl | 4 +- src/brillouin.jl | 11 +++--- src/fourier.jl | 1 + test/brillouin.jl | 74 ++---------------------------------- test/dos.jl | 1 + test/interface_tests.jl | 1 + test/test_ibz.jl | 1 + 8 files changed, 17 insertions(+), 80 deletions(-) diff --git a/src/AutoBZCore.jl b/src/AutoBZCore.jl index 5c1a29d..098dd94 100644 --- a/src/AutoBZCore.jl +++ b/src/AutoBZCore.jl @@ -50,7 +50,7 @@ module AutoBZCore using LinearAlgebra: I, norm, det, checksquare, isdiag, Diagonal, tr, diag, eigen, Hermitian -using StaticArrays: SVector, SMatrix, pushfirst, sacollect +using StaticArrays: SVector, SMatrix, sacollect using FunctionWrappers: FunctionWrapper using ChunkSplitters: chunks, getchunk using AutoSymPTR @@ -85,7 +85,7 @@ export SymmetricBZ, nsyms export load_bz, FBZ, IBZ, InversionSymIBZ, CubicSymIBZ export AbstractSymRep, UnknownRep, TrivialRep export AutoBZProblem -export IAI, PTR, AutoPTR, TAI, PTR_IAI, AutoPTR_IAI +export IAI, PTR, AutoPTR, TAI#, PTR_IAI, AutoPTR_IAI include("brillouin.jl") export FourierIntegralFunction, CommonSolveFourierIntegralFunction diff --git a/src/algorithms_autosymptr.jl b/src/algorithms_autosymptr.jl index 470d0d1..1be5501 100644 --- a/src/algorithms_autosymptr.jl +++ b/src/algorithms_autosymptr.jl @@ -76,7 +76,7 @@ end function init_cacheval(f, dom, p, alg::AutoSymPTRJL; kws...) b = get_basis(dom) - rule = init_rule(b, alg) + rule = init_rule(dom, alg) rule_cache = AutoSymPTR.alloc_cache(eltype(dom), Val(ndims(dom)), rule) cache = init_autosymptr_cache(f, b, p, alg.nthreads; kws...) return (; rule, rule_cache, cache...) @@ -89,7 +89,7 @@ function do_integral(f, dom, p, alg::AutoSymPTRJL, cacheval; bas = get_basis(dom) value, error = autosymptr(g, bas; syms = alg.syms, rule = cacheval.rule, cache = cacheval.rule_cache, keepmost = alg.keepmost, abstol = abstol, reltol = reltol, maxevals = maxiters, norm=alg.norm, buffer=cacheval.buffer) - retcode = error < max(something(abstol, zero(error)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(eltype(a))) : abstol)) ? Success : Failure + retcode = error < max(something(abstol, zero(error)), alg.norm(value)*something(reltol, isnothing(abstol) ? sqrt(eps(eltype(bas))) : abstol)) ? Success : Failure stats = (; error) return IntegralSolution(value, retcode, stats) end diff --git a/src/brillouin.jl b/src/brillouin.jl index 527d996..51da90b 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -116,7 +116,8 @@ Base.length(r::SymmetricRule) = length(r.rule) Base.iterate(r::SymmetricRule, args...) = iterate(r.rule, args...) function (r::SymmetricRule)(f::F, args...) where {F} out = r.rule(f, args...) - return symmetrize(r.rep, r.bz, out) + val = symmetrize(r.rep, r.bz, out) + return val end struct SymmetricRuleDef{R,U,B} @@ -349,11 +350,11 @@ function AutoBZProblem(rep::AbstractSymRep, f::AbstractIntegralFunction, bz::Sym return AutoBZProblem(rep, f, bz, p, NamedTuple(kws)) end end -function AutoBZProblem(rep::AbstractSymRep, f, bz::SymmetricBZ, p=NullParameters(); kws...) - return AutoBZProblem(IntegralFunction(f), bz, p; kws...) +function AutoBZProblem(f::AbstractIntegralFunction, bz::SymmetricBZ, p=NullParameters(); kws...) + return AutoBZProblem(UnknownRep(), f, bz, p; kws...) end function AutoBZProblem(f, bz::SymmetricBZ, p=NullParameters(); kws...) - return AutoBZProblem(UnknownRep(), f, bz, p; kws...) + return AutoBZProblem(IntegralFunction(f), bz, p; kws...) end mutable struct AutoBZCache{R,F,BZ,P,A,C,K} @@ -401,7 +402,7 @@ function do_solve_autobz(rep, f, bz, p, bzalg::AutoBZAlgorithm, cacheval; _kws.. cacheval.kwargs = haskey(kws, :abstol) ? merge(kws, (abstol=kws.abstol / (j * nsyms(bz)),)) : kws sol = solve!(cacheval) - value = j*symmetrize_(rep, bz, sol.value) + value = j*symmetrize(rep, bz, sol.value) stats = (; sol.stats...) # err = sol.resid === nothing ? nothing : j*symmetrize(f, bz_, sol.resid) return IntegralSolution(value, sol.retcode, stats) diff --git a/src/fourier.jl b/src/fourier.jl index 4bc072c..f5d28da 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -145,6 +145,7 @@ struct FourierValue{X,S} s::S end @inline AutoSymPTR.mymul(w, x::FourierValue) = FourierValue(AutoSymPTR.mymul(w, x.x), x.s) +@inline AutoSymPTR.mymul(::AutoSymPTR.One, x::FourierValue) = x function init_cacheval(f::FourierIntegralFunction, dom, p, alg::QuadGKJL; kws...) segs = PuncturedInterval(dom) diff --git a/test/brillouin.jl b/test/brillouin.jl index 5745092..3cbbd80 100644 --- a/test/brillouin.jl +++ b/test/brillouin.jl @@ -1,7 +1,6 @@ using Test using LinearAlgebra using AutoBZCore -using AutoBZCore: IntegralProblem, solve, MixedParameters using AutoBZCore: PuncturedInterval, HyperCube, segments, endpoints @testset "domains" begin @@ -35,77 +34,10 @@ end A = I(dims) vol = (2π)^dims for bz in (load_bz(FBZ(), A), load_bz(InversionSymIBZ(), A)) - ip = IntegralProblem((x,p) -> 1.0, bz) # unit measure + ip = AutoBZProblem((x,p) -> 1.0, bz) # unit measure for alg in (IAI(), TAI(), PTR(), AutoPTR()) - @test @inferred(solve(ip, alg)).u ≈ vol + solver = init(ip, alg) + @test @inferred(solve!(solver)).value ≈ vol end - @test @inferred(solve(IntegralProblem((x, p) -> exp(-x^2), -Inf, Inf), QuadGKJL())).u ≈ sqrt(pi) - end -end - -@testset "interfaces" begin - @testset "MixedParameters" begin - args = (1, 2) - kwargs = (a = "a", b = "b") - p = MixedParameters(args...) - q = MixedParameters(; kwargs...) - # we only merge non-MixedParameters in the right argument - for pq = (merge(p, q), merge(p, kwargs), merge(q, args)) - @test args[1] == pq[1] - @test args[2] == pq[2] - @test kwargs.a == pq.a - @test kwargs.b == pq.b - end - @test 3 == merge(p, 3)[3] == merge(q, 3)[1] # generic types are appended - @test "c" == merge(p, (a="c",)).a == merge(q, (a="c",)).a # keywords overwritten - end - @testset "IntegralSolver" begin - f = (x,p) -> p - p = 0.81 - # SciML interface: (a, b, alg) - a = 0; b = 1 - prob = IntegralProblem(f, a, b, 33) # ordinary integrands get parameters replaced - solver = IntegralSolver(prob, QuadGKJL()) - @test solver(p) == solve(IntegralProblem(f, a, b, p), QuadGKJL()).u - # AutoBZ interface: (bz, alg) - dims = 3 - A = I(dims) - B = AutoBZCore.canonical_reciprocal_basis(A) - bz = load_bz(FBZ(), A, B) - prob = IntegralProblem(f, bz, p) - solver = IntegralSolver(IntegralProblem(f, bz), IAI()) - @test solver(p) == solve(prob, IAI()).u # use the plain interface - # g = (x,p) -> sum(x)*p[1]+p.a - # solver2 = IntegralSolver(ParameterIntegrand(g), bz, IAI()) # use the MixedParameters interface - # prob2 = IntegralProblem(g, bz, MixedParameters(12.0, a=1.0)) - # @test solver2(12.0, a=1.0) == solve(prob2, IAI()).u - end - @testset "ParameterIntegrand" begin - # AutoBZ interface user function: f(x, args...; kwargs...) where args & kwargs - # stored in MixedParameters - f(x, a; b) = a*x+b - # SciML interface for ParameterIntegrand: f(x, p) (# and parameters can be preloaded and - # p is merged with MixedParameters) - @test f(6.7, 1.3, b=4.2) == ParameterIntegrand(f, 1.3, b=4.2)(6.7, AutoBZCore.NullParameters()) == ParameterIntegrand(f)(6.7, MixedParameters(1.3, b=4.2)) - # A ParameterIntegrand merges its parameters with the problem's - prob = IntegralProblem(ParameterIntegrand(f, 1.3, b=4.2), 0, 1) - u = IntegralSolver(prob, QuadGKJL())() - v = IntegralSolver(ParameterIntegrand(f), 0, 1, QuadGKJL())(1.3, b=4.2) - w = IntegralSolver(ParameterIntegrand(f, b=4.2), 0, 1, QuadGKJL())(1.3) - @test u == v == w - @test solve(prob, EvalCounter(QuadGKJL(order=7))).numevals == 15 - end - @testset "batchsolve" begin - # SciML interface: iterable of parameters - prob = IntegralProblem((x, p) -> p, 0, 1) - solver = IntegralSolver(prob, QuadGKJL()) - params = range(1, 2, length=3) - @test [solver(p) for p in params] == batchsolve(solver, params) - # AutoBZ interface: array of MixedParameters - f(x, a; b) = a*x+b - solver = IntegralSolver(ParameterIntegrand(f), 0, 1, QuadGKJL()) - as = rand(3); bs = rand(3) - @test [solver(a, b=b) for (a,b) in Iterators.zip(as, bs)] == batchsolve(solver, paramzip(as, b=bs)) - @test [solver(a, b=b) for (a,b) in Iterators.product(as, bs)] == batchsolve(solver, paramproduct(as, b=bs)) end end diff --git a/test/dos.jl b/test/dos.jl index 357b48f..7e5708a 100644 --- a/test/dos.jl +++ b/test/dos.jl @@ -1,5 +1,6 @@ using Test, AutoBZCore, LinearAlgebra, StaticArrays, OffsetArrays, Elliptic using GeneralizedGaussianQuadrature: generalizedquadrature +using FourierSeriesEvaluators, QuadGK # test set of known DOS examples diff --git a/test/interface_tests.jl b/test/interface_tests.jl index f6f511b..cc99c8f 100644 --- a/test/interface_tests.jl +++ b/test/interface_tests.jl @@ -41,6 +41,7 @@ end @test ref ≈ sol.value atol=abstol end end + @test @inferred(solve(IntegralProblem((x, p) -> exp(-x^2), (-Inf, Inf)), QuadGKJL())).value ≈ sqrt(pi) end @testset "commonproblem" begin diff --git a/test/test_ibz.jl b/test/test_ibz.jl index e8ce4f1..4661e15 100644 --- a/test/test_ibz.jl +++ b/test/test_ibz.jl @@ -5,6 +5,7 @@ import SymmetryReduceBZ.Symmetry: calc_bz, calc_ibz import SymmetryReduceBZ.Utilities: volume, vertices, get_uniquefacets using Polyhedra: Polyhedron using AutoBZCore +using IteratedIntegration: nested_quad using LinearAlgebra using Test From 83efe958375373763676abe0cc0fbada76581187 Mon Sep 17 00:00:00 2001 From: lxvm Date: Mon, 15 Jul 2024 23:39:15 -0400 Subject: [PATCH 14/17] update docs --- docs/Project.toml | 3 + docs/make.jl | 1 - docs/src/algorithms.md | 35 +++++---- docs/src/dos.md | 63 ---------------- docs/src/examples.md | 141 +++++++++++++++++++++++------------ docs/src/extensions.md | 6 -- docs/src/index.md | 3 - docs/src/integrands.md | 39 ++-------- docs/src/problems.md | 87 +++++++++++---------- docs/src/reference.md | 22 ++---- src/AutoBZCore.jl | 3 +- src/algorithms.jl | 2 +- src/algorithms_autosymptr.jl | 4 +- src/algorithms_meta.jl | 8 +- src/brillouin.jl | 37 +++------ src/dos_ggr.jl | 2 +- src/dos_interfaces.jl | 9 +-- 17 files changed, 205 insertions(+), 260 deletions(-) delete mode 100644 docs/src/dos.md diff --git a/docs/Project.toml b/docs/Project.toml index dfa65cd..9b8f31f 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,2 +1,5 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +FourierSeriesEvaluators = "2a892dea-6eef-4bb5-9d1c-de966c9f6db5" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" diff --git a/docs/make.jl b/docs/make.jl index abee2bd..50d83c4 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -23,7 +23,6 @@ makedocs( "Algorithms" => "algorithms.md", "Reference" => "reference.md", "Extensions" => "extensions.md", - "Density of States" => "dos.md", ], ) diff --git a/docs/src/algorithms.md b/docs/src/algorithms.md index 41f0303..ab620af 100644 --- a/docs/src/algorithms.md +++ b/docs/src/algorithms.md @@ -1,10 +1,12 @@ -# Integral algorithms +# Algorithms + +## `IntegralProblem` algorithms ```@docs AutoBZCore.IntegralAlgorithm ``` -## Quadrature +### Quadrature ```@docs AutoBZCore.QuadratureFunction @@ -14,7 +16,7 @@ AutoBZCore.ContQuadGKJL AutoBZCore.MeroQuadGKJL ``` -## Cubature +### Cubature ```@docs AutoBZCore.HCubatureJL @@ -22,20 +24,17 @@ AutoBZCore.MonkhorstPack AutoBZCore.AutoSymPTRJL ``` -## Meta-algorithms +### Meta-algorithms ```@docs AutoBZCore.NestedQuad -AutoBZCore.EvalCounter -AutoBZCore.AbsoluteEstimate ``` -# BZ-specific integral algorithms +## `AutoBZProblem` algorithms In order to make algorithms domain-agnostic, the BZ loaded from -[`load_bz`](@ref) can be called with the algorithms below, which are wrappers -for algorithms above with the additional capability of mapping integrals over -the IBZ to the FBZ. +[`load_bz`](@ref) can be called with the algorithms below, which are aliases +for algorithms above ```@docs AutoBZCore.AutoBZAlgorithm @@ -43,6 +42,16 @@ AutoBZCore.IAI AutoBZCore.TAI AutoBZCore.PTR AutoBZCore.AutoPTR -AutoBZCore.PTR_IAI -AutoBZCore.AutoPTR_IAI -``` \ No newline at end of file +``` + +## `DOSProblem` algorithms + +Currently the available algorithms are an initial release and we would like to include +the following reference algorithms that are also common in the literature in a future release: +- (Linear) Tetrahedron Method +- Adaptive Gaussian broadening + +```@docs +AutoBZCore.DOSAlgorithm +AutoBZCore.GGR +``` diff --git a/docs/src/dos.md b/docs/src/dos.md deleted file mode 100644 index 9dfd240..0000000 --- a/docs/src/dos.md +++ /dev/null @@ -1,63 +0,0 @@ -# Density of States - -Computing the density of states (DOS) of a self-adjoint, or Hermitian, operator is a -related, but distinct problem to the integrals also presented in this package. -In fact, many DOS algorithms will compute integrals to approximate the DOS of an -operator that depends on a parameter such as crystal momentum. To solve these -types of problems, this package defines the following problem type: - -```@docs -AutoBZCore.DOSProblem -``` - -## Algorithms - -Currently the available algorithms are experimental and we would like to include -the following reference algorithms that are more common in the literature in a future release: -- (Linear) Tetrahedron Method -- Adaptive Gaussian broadening - -```@docs -AutoBZCore.DOSAlgorithm -AutoBZCore.GGR -``` - -## Caching interface - -Using the [`AutoBZCore.init`](@ref) and [`AutoBZCore.solve!`](@ref) functions, it is possible to -construct a cache to solve a [`DOSProblem`](@ref) for several energies or -several Hamiltonians. As an example of solving for several energies, -```julia -using AutoBZCore, StaticArrays -h = FourierSeries(SMatrix{1,1,Float64,1}.([0.5, 0.0, 0.5]), period=1.0, offset=-2) -E = 0.3 -bz = load_bz(FBZ(), [2pi;;]) -prob = DOSProblem(h, E, bz) -alg = GGR(; npt=100) -cache = AutoBZCore.init(prob, alg) -Es = range(-2, 2, length=1000) -data = map(Es) do E - cache.domain = E - AutoBZCore.solve!(cache).u -end -``` - -As an example of interchanging Hamiltonians, which must remain the same type, -```julia -using AutoBZCore, StaticArrays - -h = FourierSeries(SMatrix{1,1,Float64,1}.([0.5, 0.0, 0.5]), period=1.0, offset=-2) -bz = load_bz(FBZ(), [2pi;;]) -prob = DOSProblem(h, 0.0, bz) -alg = GGR() - -cache = AutoBZCore.init(prob, alg) -sol1 = AutoBZCore.solve!(cache) - -h.c .*= 2 -cache.isfresh = true - -sol2 = AutoBZCore.solve!(cache) - -sol1.u*2 ≈ sol2.u -``` \ No newline at end of file diff --git a/docs/src/examples.md b/docs/src/examples.md index a2e5f1d..a64be6d 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -2,9 +2,9 @@ The following are several examples of how to use the algorithms and integrands provided by AutoBZCore.jl. -For background on defining integrals see the [Problem definitions](@ref) page +For background on the essential interface see the [Problem definitions](@ref) page -## Green's function +## Green's function integration A common integral appearing in [Dynamical mean-field theory](https://en.wikipedia.org/wiki/Dynamical_mean-field_theory) is that of @@ -14,11 +14,13 @@ G(\omega) = \int d^d \vec{k}\ \operatorname{Tr} \left[ \left( \omega - H \left( ``` For simplicity, we take ``\Sigma(\omega) = -i\eta``. We can define the integrand -as a function of ``\vec{k}`` and ``H`` and (required) parameters ``\eta, \omega``. -```julia +as a function of ``\vec{k}`` and ``H`` and parameters ``\eta, \omega``. +```@example gloc using LinearAlgebra -gloc_integrand(k, h; η, ω) = inv(complex(ω,η)*I-h(k)) +gloc_integrand(k, (; h, η, ω)) = tr(inv(complex(ω,η)*I-h(k))) ``` +Here we use named tuple destructuring syntax to unpack a named tuple of +parameters in the function definition. Commonly, ``H(\vec{k})`` is evaluated using Wannier interpolation, i.e. as a Fourier series. For a simple tight-binding model, the integer lattice, the @@ -31,33 +33,44 @@ complex Fourier series it becomes easier to use the representation in terms of Fourier coefficients. Using the package [FourierSeriesEvaluators.jl](https://github.com/lxvm/FourierSeriesEvaluators.jl), we can define ``H(k) = \cos(2\pi k)`` by the following: -```julia +```@example gloc using FourierSeriesEvaluators h = FourierSeries([0.5, 0.0, 0.5]; period=1, offset=-2) ``` The coefficient values of ``1/2`` can be determined from Euler's formula, as -used in the expansion of ``cos`` above, and the value of `offset` is chosen to +used in the expansion of ``\cos`` above, and the value of `offset` is chosen to offset the coefficient array indices, `1:3` since Julia has 1-based indexing, to correspond to values of ``n`` in the phase factors ``e^{2\pi i n k}`` used in -the Fourier series above, i.e. `-1:1`. Now we proceed to the define the integral problem -```julia +the Fourier series above, i.e. `-1:1`. Now we proceed to the define the +[`IntegralProblem`](@ref) and solve it with a generic adaptive +integration scheme, [`QuadGKJL`](@ref) +```@example gloc using AutoBZCore -using AutoBZCore: IntegralProblem -integrand = ParameterIntegrand(gloc_integrand, h, η=0.1) -prob = IntegralProblem(integrand, 0, 1) -``` -Here we wrapped our function with two of its arguments, `h, η` as a -[`ParameterIntegrand`](@ref) that allows us to provide partial arguments so that -we can solve the integral as a function of the remaining parameters, in this -case `ω`. We also created an [`AutoBZCore.IntegralProblem`](@ref) to integrate our function -over its period ``[0,1]``. To solve this problem, we pick any of the package's -[Integral algorithms](@ref) and the -tolerance to which we would like the solution. Then we make an -[`IntegralSolver`](@ref) to evaluate ``G(\omega)`` as a function. -```julia +dom = (0, 1) +p = (; h, η=0.1, ω=0.0) +prob = IntegralProblem(gloc_integrand, dom, p) alg = QuadGKJL() -gloc = IntegralSolver(prob, alg; abstol=1e-3) -gloc(ω=0.0) # -2.7755575615628914e-17 - 0.9950375451895513im +solve(prob, alg; abstol=1e-3).value +``` + +## BZ integration + +To perform integration over a Brillouin zone, we can load one using the +[`load_bz`](@ref) function and then construct an [`AutoBZProblem`](@ref) to +solve. Since the Brillouin zone may be reduced using point group symmetries, a +common optimization, it is also required to specify the symmetry representation +of the integrand. Continuing the previous example, the trace of the Green's +function has no band/orbital degrees of freedom and transforms trivially under +the point group, so it is a [`TrivialRep`](@ref). The previous calculation can +be replicated as: + +```@example gloc +using AutoBZCore +bz = load_bz(FBZ(), 2pi*I(1)) +p = (; h, η=0.1, ω=0.0) +prob = AutoBZProblem(TrivialRep(), IntegralFunction(gloc_integrand), bz, p) +alg = TAI() +solve(prob, alg; abstol=1e-3).value ``` Now we proceed to multi-dimensional integrals. In this case, Wannier @@ -76,37 +89,73 @@ which is comparable to the computational complexity of a [multi-dimensional FFT](https://en.wikipedia.org/wiki/Fast_Fourier_transform#Multidimensional_FFTs). Since the constants of a FFT may not be trivial, this scheme is competitive. -This capability is provided by [`FourierIntegrand`](@ref). Let's use this with a Fourier series corresponding to ``H(\vec{k}) = \cos(2\pi k_{x}) + \cos(2\pi k_{y})`` -```julia +and define a new method of `gloc_integrand` that accepts the (efficiently) +evaluated Fourier series in the second argument +```@example gloc h = FourierSeries([0.0; 0.5; 0.0;; 0.5; 0.0; 0.5;; 0.0; 0.5; 0.0]; period=1, offset=-2) -integrand = FourierIntegrand(gloc_integrand, h, η=0.1) -``` -However, since [`FourierIntegrand`](@ref) evaluates ``H(k)`` for us and gives it -as a `FourierValue` together with ``k``, we need to define another method for -our integrand to comply with the interface -```julia -gloc_integrand(h_k::FourierValue; η, ω) = inv(complex(ω,η)*I-h_k.s) +gloc_integrand(k, h_k, (; η, ω)) = tr(inv(complex(ω,η)*I-h_k)) ``` Similar to before, we construct an [`AutoBZCore.IntegralProblem`](@ref) and this time we take the integration domain to correspond to the full Brillouin zone of a square -lattice with lattice vectors `2pi*I(2)`. (See the -[Reference](@ref) for more details on constructing BZs.) -```julia +lattice with lattice vectors `2pi*I(2)`. +```@example gloc +integrand = FourierIntegralFunction(gloc_integrand, h) bz = load_bz(FBZ(2), 2pi*I(2)) -prob = IntegralProblem(integrand, bz) -``` -This package provides several [BZ-specific integral algorithms](@ref) that we -can use to solve the multidimensional integral. -```julia +p = (; η=0.1, ω=0.0) +prob = AutoBZProblem(TrivialRep(), integrand, bz, p) alg = IAI() -gloc = IntegralSolver(prob, alg; abstol=1e-3) -gloc(ω=0.0) # 1.5265566588595902e-16 - 1.3941704019631334im +solve(prob, alg).value ``` - -## Density of states +This package provides several [`AutoBZProblem` algorithms](@ref) that we +can use to solve the multidimensional integral. The [repo's demo](https://github.com/lxvm/AutoBZCore.jl/tree/main/aps_example) on density of states provides a complete example of how to compute and -interpolate an integral as a function of its parameters. +interpolate an integral as a function of its parameters using the [`init` and +`solve!`](@ref) interface + + +## Density of States + +Computing the density of states (DOS) of a self-adjoint, or Hermitian, operator is a +related, but distinct problem to the integrals also presented in this package. +In fact, many DOS algorithms will compute integrals to approximate the DOS of an +operator by introducing an artificial broadening. +To handle the ``T=0^{+}`` limit of the broadening, we implement the well-known +[Gilat-Raubenheimer method](https://arxiv.org/abs/1711.07993) as an algorithm +for the [`AutoBZCore.DOSProblem`](@ref) + +Using the [`AutoBZCore.init`](@ref) and [`AutoBZCore.solve!`](@ref) functions, it is possible to +construct a cache to solve a [`DOSProblem`](@ref) for several energies or +several Hamiltonians. As an example of solving for several energies, +```@example dos +using AutoBZCore, FourierSeriesEvaluators, StaticArrays +h = FourierSeries(SMatrix{1,1,Float64,1}.([0.5, 0.0, 0.5]), period=1.0, offset=-2) +E = 0.3 +bz = load_bz(FBZ(), [2pi;;]) +prob = DOSProblem(h, E, bz) +alg = GGR(; npt=100) +cache = init(prob, alg) +Es = range(-1, 1, length=10) * 1.1 +data = map(Es) do E + cache.domain = E + solve!(cache).value +end +``` + +As an example of interchanging Hamiltonians, which must remain the same type, we +can double the energies, which will halve the DOS +```@example dos +cache.domain = E +sol1 = AutoBZCore.solve!(cache) + +h.c .*= 2 +cache.isfresh = true +cache.domain = 2E + +sol2 = AutoBZCore.solve!(cache) + +sol1.value ≈ 2sol2.value +``` \ No newline at end of file diff --git a/docs/src/extensions.md b/docs/src/extensions.md index 64b5015..956c14d 100644 --- a/docs/src/extensions.md +++ b/docs/src/extensions.md @@ -6,12 +6,6 @@ Loading [SymmetryReduceBZ.jl](https://github.com/jerjorg/SymmetryReduceBZ.jl) provides a specialized method of [`load_bz`](@ref) that when provided atom species and positions can compute the [`IBZ`](@ref). -## HDF5.jl - -Loading [HDF5.jl](https://github.com/JuliaIO/HDF5.jl) provides a specialized -method of [`batchsolve`](@ref) that accepts an H5 archive or group in the first -argument and writes the integral results and timings to a dataset. - ## WannierIO.jl Loading [WannierIO.jl](https://github.com/qiaojunfeng/WannierIO.jl) provides a diff --git a/docs/src/index.md b/docs/src/index.md index 654a97b..1619965 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,8 +1,5 @@ # AutoBZCore.jl -!!! note "Work in progress" - We apologize for the inconvenience as we continue improving the documentation. - ```@docs AutoBZCore ``` \ No newline at end of file diff --git a/docs/src/integrands.md b/docs/src/integrands.md index a57abfd..f99f4e5 100644 --- a/docs/src/integrands.md +++ b/docs/src/integrands.md @@ -3,39 +3,12 @@ The design of AutoBZCore.jl uses multiple dispatch to provide multiple interfaces for user integrands that allow various optimizations to be compatible with a common interface for solvers. -Unfortunately, not all algorithms support all integrands, since the underlying -libraries must support the same interface. - -# Functions - -A user can pass an integrand of the form `f(x,p)` in the same way as in Integrals.jl - -# `ParameterIntegrand` - -```@docs -AutoBZCore.ParameterIntegrand -``` - -# `InplaceIntegrand` - -```@docs -AutoBZCore.InplaceIntegrand -``` - -# `BatchIntegrand` - -```@docs -AutoBZCore.BatchIntegrand -``` - -# `FourierIntegrand` - -```@docs -AutoBZCore.FourierIntegrand -``` - -# `NestedBatchIntegrand` ```@docs -AutoBZCore.NestedBatchIntegrand +AutoBZCore.IntegralFunction +AutoBZCore.InplaceIntegralFunction +AutoBZCore.InplaceBatchIntegralFunction +AutoBZCore.CommonSolveIntegralFunction +AutoBZCore.FourierIntegralFunction +AutoBZCore.CommonSolveFourierIntegralFunction ``` \ No newline at end of file diff --git a/docs/src/problems.md b/docs/src/problems.md index 6397a9e..346d59b 100644 --- a/docs/src/problems.md +++ b/docs/src/problems.md @@ -1,64 +1,71 @@ # Problem definitions -The design of AutoBZCore.jl is heavily influenced by the -[SciML](https://sciml.ai/) package -[Integrals.jl](https://github.com/SciML/Integrals.jl) -and may eventually become implemented in it. +The design of AutoBZCore.jl is heavily influenced by +[SciML](https://sciml.ai/) packages and uses the +[CommonSolve.jl](https://github.com/SciML/CommonSolve.jl) +interface. Eventually, this package may contribute to +[Integrals.jl](https://github.com/SciML/Integrals.jl). -## SciML interface +## Problem interface -AutoBZCore.jl replicates the Integrals.jl interface, but does not export it in -order to avoid name conflicts with other SciML packages. +AutoBZCore.jl replicates the Integrals.jl interface, using an +[`IntegralProblem`](@ref) type to setup an integral from an +integrand, a domain, and parameters. -### Quickstart +```@example prob +using AutoBZCore -```julia -using AutoBZCore: IntegralProblem, init, solve! - -prob = IntegralProblem((x,p) -> sin(p*x), 0, 1, 0.3) -cache = init(prob, QuadGKJL()) -solve!(cache) # 0.14887836958131329 - -# solve again at a new parameter -cache.p = 0.4 -solve!(cache) # 0.1973475149927873 +f = (x,p) -> sin(p*x) +dom = (0, 1) +p = 0.3 +prob = IntegralProblem(f, dom, p) ``` -### Reference - ```@docs AutoBZCore.IntegralProblem -AutoBZCore.solve -AutoBZCore.init -AutoBZCore.solve! AutoBZCore.NullParameters ``` -## Functor interface +## `solve` + +To solve an integral problem, pick an algorithm and call [`solve`](@ref) -As shown in the quickstart of the [`AutoBZCore`](@ref) page, AutoBZCore.jl also defines -a functor interface to solving integrals +```@example prob +alg = QuadGKJL() +solve(prob, alg) +``` ```@docs -AutoBZCore.IntegralSolver +AutoBZCore.solve ``` -The functor interface is also extended by [`ParameterIntegrand`](@ref) and -[`FourierIntegrand`](@ref) to -allow a more flexible interface for passing (partial) positional and keyword -arguments to user-defined integrands. +## `init` and `solve!` -```@docs -AutoBZCore.MixedParameters -AutoBZCore.paramzip -AutoBZCore.paramproduct +To solve many problems with the same integrand but different domains or +parameters, use [`init`](@ref) to allocate a solver and +[`solve!`](@ref) to get the solution + +```@example prob +solver = init(prob, alg) +solve!(solver).value ``` -## Batched evaluation +To solve again, update the parameters of the solver in place and `solve!` again +```@example prob +# solve again at a new parameter +solver.p = 0.4 +solve!(solver).value +``` + + +```@docs +AutoBZCore.init +AutoBZCore.solve! +``` -The routine [`batchsolve`](@ref) allows multi-threaded evaluation of an -[`IntegralSolver`](@ref) at many parameter points. +## Additional problems ```@docs -AutoBZCore.batchsolve -``` \ No newline at end of file +AutoBZCore.AutoBZProblem +AutoBZCore.DOSProblem +``` diff --git a/docs/src/reference.md b/docs/src/reference.md index 4d74657..9b84c2b 100644 --- a/docs/src/reference.md +++ b/docs/src/reference.md @@ -2,18 +2,11 @@ The following symbols are exported by AutoBZCore.jl -## Domains - -```@docs -AutoBZCore.PuncturedInterval -AutoBZCore.HyperCube -AutoBZCore.load_bz -AutoBZCore.SymmetricBZ -``` - ## Brillouin-zone kinds ```@docs +AutoBZCore.load_bz +AutoBZCore.load_bz(::IBZ, A, B, species, positions) AutoBZCore.AbstractBZ AutoBZCore.FBZ AutoBZCore.IBZ @@ -25,21 +18,20 @@ AutoBZCore.CubicSymIBZ ```@docs AutoBZCore.AbstractSymRep -AutoBZCore.SymRep AutoBZCore.TrivialRep AutoBZCore.UnknownRep -AutoBZCore.symmetrize_ +AutoBZCore.symmetrize ``` ## Internal -The following docstrings belong to internal functions that may change between +The following docstrings belong to internal types and functions that may change between versions of AutoBZCore. ```@docs +AutoBZCore.PuncturedInterval +AutoBZCore.HyperCube +AutoBZCore.SymmetricBZ AutoBZCore.trapz AutoBZCore.cube_automorphisms -AutoBZCore.batchparam -AutoBZCore.symmetrize -AutoBZCore.FourierValue ``` \ No newline at end of file diff --git a/src/AutoBZCore.jl b/src/AutoBZCore.jl index 098dd94..8be08b3 100644 --- a/src/AutoBZCore.jl +++ b/src/AutoBZCore.jl @@ -39,7 +39,6 @@ Quadrature algorithms: Meta-Algorithms: - Iterated integration: [`NestedQuad`](@ref) -- Integrand evaluation counter: [`EvalCounter`](@ref) # Extended help @@ -85,7 +84,7 @@ export SymmetricBZ, nsyms export load_bz, FBZ, IBZ, InversionSymIBZ, CubicSymIBZ export AbstractSymRep, UnknownRep, TrivialRep export AutoBZProblem -export IAI, PTR, AutoPTR, TAI#, PTR_IAI, AutoPTR_IAI +export IAI, PTR, AutoPTR, TAI include("brillouin.jl") export FourierIntegralFunction, CommonSolveFourierIntegralFunction diff --git a/src/algorithms.jl b/src/algorithms.jl index 9162cc6..4c70209 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -165,7 +165,7 @@ The default quadrature rule is [`trapz`](@ref), although other packages provide alg = QuadratureFunction(fun=gausslegendre, npt=100) `nthreads` sets the numbers of threads used to parallelize the quadrature only when the -integrand is a [`BatchIntegrand`](@ref), in which case the user must parallelize the +integrand is a , in which case the user must parallelize the integrand evaluations. For no threading set `nthreads=1`. """ struct QuadratureFunction{F} <: IntegralAlgorithm diff --git a/src/algorithms_autosymptr.jl b/src/algorithms_autosymptr.jl index 1be5501..19aa389 100644 --- a/src/algorithms_autosymptr.jl +++ b/src/algorithms_autosymptr.jl @@ -8,7 +8,7 @@ Periodic trapezoidal rule with a fixed number of k-points per dimension, `npt`, using the `PTR` rule from [AutoSymPTR.jl](https://github.com/lxvm/AutoSymPTR.jl). `nthreads` sets the numbers of threads used to parallelize the quadrature only when the -integrand is a [`BatchIntegrand`](@ref), in which case the user must parallelize the +integrand is a , in which case the user must parallelize the integrand evaluations. For no threading set `nthreads=1`. **The caller should check that the integral is converged w.r.t. `npt`**. """ @@ -51,7 +51,7 @@ Periodic trapezoidal rule with automatic convergence to tolerances passed to the solver with respect to `norm` using the routine `autosymptr` from [AutoSymPTR.jl](https://github.com/lxvm/AutoSymPTR.jl). `nthreads` sets the numbers of threads used to parallelize the quadrature only when the -integrand is a [`BatchIntegrand`](@ref), in which case the user must parallelize the +integrand is a in which case the user must parallelize the integrand evaluations. For no threading set `nthreads=1`. **This algorithm is the most efficient for smooth integrands**. """ diff --git a/src/algorithms_meta.jl b/src/algorithms_meta.jl index e491b22..c769c06 100644 --- a/src/algorithms_meta.jl +++ b/src/algorithms_meta.jl @@ -81,7 +81,7 @@ function outer_integralfunction(f::IntegralFunction, x0, p) ws = nothing return func, ws, _update!, _postsolve end - +#= """ AbsoluteEstimate(est_alg, abs_alg; kws...) @@ -107,7 +107,7 @@ function AbsoluteEstimate(est_alg, abs_alg; norm=norm, kwargs...) checkkwargs(kws) return AbsoluteEstimate(est_alg, abs_alg, norm, kws) end -#= + function init_cacheval(f, dom, p, alg::AbsoluteEstimate) return (est=init_cacheval(f, dom, p, alg.est_alg), abs=init_cacheval(f, dom, p, alg.abs_alg)) @@ -122,7 +122,7 @@ function do_solve(f, dom, p, alg::AbsoluteEstimate, cacheval; return do_solve(f, dom, p, alg.abs_alg, cacheval.abs; abstol=atol, reltol=zero(rtol), maxiters=maxiters) end -=# + """ EvalCounter(::IntegralAlgorithm) @@ -133,7 +133,7 @@ The count is stored in the `sol.numevals` field. struct EvalCounter{T<:IntegralAlgorithm} <: IntegralAlgorithm alg::T end -#= + function init_cacheval(f, dom, p, alg::EvalCounter) return init_cacheval(f, dom, p, alg.alg) end diff --git a/src/brillouin.jl b/src/brillouin.jl index 51da90b..ac3a8d3 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -64,6 +64,7 @@ abstract type AbstractSymRep end UnknownRep() Fallback symmetry representation for array types without a user-defined `SymRep`. +Will perform FBZ integration regardless of available BZ symmetries. """ struct UnknownRep <: AbstractSymRep end @@ -74,35 +75,20 @@ Symmetry representation of objects with trivial transformation under the group. """ struct TrivialRep <: AbstractSymRep end -const TrivialRepType = Union{Number,AbstractArray{<:Any,0}} - -""" - symmetrize(f, ::SymmetricBZ, xs...) - symmetrize(f, ::SymmetricBZ, x::Union{Number,AbstractArray{<:Any,0}}) - -Transform `x` by the symmetries of the parametrization used to reduce the -domain, thus mapping the value of `x` on the parametrization to the full domain. """ -symmetrize(f, bz, xs...) = map(x -> symmetrize(f, bz, x), xs) -symmetrize(f, bz, x) = symmetrize_(f isa AbstractSymRep ? f : SymRep(f), bz, x) -symmetrize(f, bz, x::TrivialRepType) = - symmetrize_(TrivialRep(), bz, x) + symmetrize(rep::AbstractSymRep, ::SymmetricBZ, x) +Transform `x` by the representation of the symmetries of the point group used to reduce the +domain, thus mapping the value of `x` on to the full Brillouin zone. """ - symmetrize_(rep::AbstractSymRep, bz::SymmetricBZ, x) - -Transform `x` under representation `rep` using the symmetries in `bz` to obtain -the result of an integral on the FBZ from `x`, which was calculated on the IBZ. -""" -symmetrize_(::TrivialRep, bz::SymmetricBZ, x) = nsyms(bz)*x -symmetrize_(::UnknownRep, ::SymmetricBZ, x) = x -symmetrize_(::UnknownRep, bz::SymmetricBZ, x::TrivialRepType) = symmetrize_(TrivialRep(), bz, x) - +symmetrize(rep, bz::SymmetricBZ, x) = symmetrize_(rep, bz, x) symmetrize(_, ::FullBZ, x) = x -symmetrize(_, ::FullBZ, x::TrivialRepType) = x -symmetrize(f, bz, x::AuxValue) = AuxValue(symmetrize(f, bz, x.val, x.aux)...) -symmetrize(_, ::FullBZ, x::AuxValue) = x +symmetrize_(rep, bz, x) = symmetrize__(rep, bz, x) +symmetrize_(rep, bz, x::AuxValue) = AuxValue(symmetrize__(rep, bz, x.val), symmetrize__(rep, bz, x.aux)) + +symmetrize__(::TrivialRep, bz, x) = nsyms(bz)*x +symmetrize__(::UnknownRep, bz, x) = error("unknown representation cannot be symmetrized") struct SymmetricRule{R,U,B} rule::R @@ -158,7 +144,8 @@ Interface to loading Brillouin zones. - `B::AbstractMatrix`: a ``d \\times d`` matrix whose columns are the reciprocal-space lattice vectors of a ``d``-dimensional Brillouin zone (default: `A' \\ 2πI`) -!!! note "Assumptions" `AutoBZCore` assumes that all calculations occur in the reciprocal +!!! 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. We also assume that the integrands are cheap to evaluate, which is why we provide adaptive methods in the first diff --git a/src/dos_ggr.jl b/src/dos_ggr.jl index d3384ab..2a4f6e6 100644 --- a/src/dos_ggr.jl +++ b/src/dos_ggr.jl @@ -52,7 +52,7 @@ function dos_solve(h, domain, p, alg::GGR, cacheval; A = sum_ggr(ndims(bz.lims), alg.npt, E, cacheval...) - return DOSSolution(A, nothing, true, -1) + return DOSSolution(A, Success, (;)) end function sum_ggr(ndim, npt, E, weights, energies, velocities) diff --git a/src/dos_interfaces.jl b/src/dos_interfaces.jl index 56be1f5..cf47f9c 100644 --- a/src/dos_interfaces.jl +++ b/src/dos_interfaces.jl @@ -40,11 +40,10 @@ function DOSProblem(H, domain, p=NullParameters(); kws...) return DOSProblem(H, domain, p, NamedTuple(kws)) end -struct DOSSolution{U,E} - u::U - err::E - retcode::Bool - numevals::Int +struct DOSSolution{V,S} + value::V + retcode::ReturnCode + stats::S end # store the data in a mutable cache so that the user can update the cache and From 7d19010ca54af358c5b50c2af460a6e09b02b5b7 Mon Sep 17 00:00:00 2001 From: lxvm Date: Tue, 16 Jul 2024 07:46:18 -0400 Subject: [PATCH 15/17] set v0.4 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 71cbe12..88eeacf 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "AutoBZCore" uuid = "66bd3e16-1600-45cf-8f55-0b550710682b" authors = ["lxvm "] -version = "0.4.0-DEV" +version = "0.4.0" [deps] AutoSymPTR = "78a0c066-08f1-49a8-82f0-b29cd485e1d3" From fed82239f0d00e3c31b71e2fcb9c453cdb52415e Mon Sep 17 00:00:00 2001 From: lxvm Date: Tue, 16 Jul 2024 07:49:34 -0400 Subject: [PATCH 16/17] fix undefined export --- src/AutoBZCore.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AutoBZCore.jl b/src/AutoBZCore.jl index 8be08b3..fc5d408 100644 --- a/src/AutoBZCore.jl +++ b/src/AutoBZCore.jl @@ -77,7 +77,7 @@ export AuxQuadGKJL, ContQuadGKJL, MeroQuadGKJL include("algorithms_iterated.jl") export MonkhorstPack, AutoSymPTRJL include("algorithms_autosymptr.jl") -export NestedQuad, AbsoluteEstimate, EvalCounter +export NestedQuad#, AbsoluteEstimate, EvalCounter include("algorithms_meta.jl") export SymmetricBZ, nsyms From a8b7f005532bfe321319e799992af51c9aa7789f Mon Sep 17 00:00:00 2001 From: lxvm Date: Tue, 16 Jul 2024 08:25:23 -0400 Subject: [PATCH 17/17] wrapup --- src/brillouin.jl | 7 +++---- test/dos.jl | 12 +++++++----- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/brillouin.jl b/src/brillouin.jl index ac3a8d3..8472624 100644 --- a/src/brillouin.jl +++ b/src/brillouin.jl @@ -322,14 +322,13 @@ end const WARN_UNKNOWN_SYMMETRY = """ A symmetric BZ was used with an integrand whose symmetry representation is unknown. -For correctness, the calculation will proceed on the full BZ. -However, it is better either to integrate without symmetries or to use symmetries by -implementing an AbstractSymRep for your type. +For correctness, the calculation will proceed on the full BZ, i.e. without symmetry. +To integrate with symmetry, define an AbstractSymRep for your integrand. """ function AutoBZProblem(rep::AbstractSymRep, f::AbstractIntegralFunction, bz::SymmetricBZ, p=NullParameters(); kws...) proto = get_prototype(f, get_prototype(bz), p) - if rep isa UnknownRep && !(bz isa FullBZ) && !(proto isa TrivialRepType) + if rep isa UnknownRep && !(bz isa FullBZ) @warn WARN_UNKNOWN_SYMMETRY fbz = SymmetricBZ(bz.A, bz.B, lattice_bz_limits(bz.B), nothing) return AutoBZProblem(rep, f, fbz, p, NamedTuple(kws)) diff --git a/test/dos.jl b/test/dos.jl index 7e5708a..87cb9bc 100644 --- a/test/dos.jl +++ b/test/dos.jl @@ -106,15 +106,15 @@ for (model, solution, bandwidth, bzkind) in ( cache = AutoBZCore.init(prob, alg) for e in E cache.domain = e - @test AutoBZCore.solve!(cache).u ≈ solution(e) atol=1e-2 + @test AutoBZCore.solve!(cache).value ≈ solution(e) atol=1e-2 end end end # test caching behavior -let h = FourierSeries(SMatrix{1,1,Float64,1}.([0.5, 0.0, 0.5]), period=1.0, offset=-2) +let h = FourierSeries(SMatrix{1,1,Float64,1}.([0.5, 0.0, 0.5]), period=1.0, offset=-2), E = 0.1 bz = load_bz(FBZ(), [2pi;;]) - prob = DOSProblem(h, 0.0, bz) + prob = DOSProblem(h, E, bz) alg = GGR() cache = AutoBZCore.init(prob, alg) @@ -122,12 +122,14 @@ let h = FourierSeries(SMatrix{1,1,Float64,1}.([0.5, 0.0, 0.5]), period=1.0, offs h.c .*= 2 cache.isfresh = true + cache.domain = 2E sol2 = AutoBZCore.solve!(cache) - @test sol1.u*2 ≈ sol2.u + @test sol1.value ≈ 2sol2.value cache.H = FourierSeries(2*h.c, period=h.t, offset=h.o) + cache.domain = 4E sol3 = AutoBZCore.solve!(cache) - @test sol2.u*2 ≈ sol3.u + @test sol2.value ≈ 2sol3.value end