diff --git a/.gitignore b/.gitignore index 29126e47..d4dad176 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,7 @@ docs/site/ # committed for packages, but should be committed for applications that require a static # environment. Manifest.toml + +# development related +.vscode +dev \ No newline at end of file diff --git a/Project.toml b/Project.toml index dc58d615..1e9cbb14 100644 --- a/Project.toml +++ b/Project.toml @@ -12,13 +12,20 @@ Meshes = "eacbb407-ea5a-433e-ab97-5258b1ca43fa" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" +[weakdeps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" + +[extensions] +MeshIntegralsEnzymeExt = "Enzyme" + [compat] CliffordNumbers = "0.1.9" -CoordRefSystems = "0.12, 0.13, 0.14, 0.15, 0.16" +CoordRefSystems = "0.15, 0.16" +Enzyme = "0.13.19" FastGaussQuadrature = "1" HCubature = "1.5" LinearAlgebra = "1" -Meshes = "0.50, 0.51, 0.52" +Meshes = "0.51.20, 0.52" QuadGK = "2.1.1" Unitful = "1.19" julia = "1.9" diff --git a/benchmark/Project.toml b/benchmark/Project.toml index 96f687ac..6eec58b7 100644 --- a/benchmark/Project.toml +++ b/benchmark/Project.toml @@ -1,12 +1,14 @@ [deps] BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Meshes = "eacbb407-ea5a-433e-ab97-5258b1ca43fa" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [compat] BenchmarkTools = "1.5" +Enzyme = "0.13.19" LinearAlgebra = "1" -Meshes = "0.50, 0.51, 0.52" +Meshes = "0.51.20, 0.52" Unitful = "1.19" julia = "1.9" diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 10ec4757..14896a03 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -3,6 +3,7 @@ using LinearAlgebra using Meshes using MeshIntegrals using Unitful +import Enzyme const SUITE = BenchmarkGroup() diff --git a/ext/MeshIntegralsEnzymeExt.jl b/ext/MeshIntegralsEnzymeExt.jl new file mode 100644 index 00000000..339b8b64 --- /dev/null +++ b/ext/MeshIntegralsEnzymeExt.jl @@ -0,0 +1,19 @@ +module MeshIntegralsEnzymeExt + +using MeshIntegrals: MeshIntegrals, AutoEnzyme +using Meshes: Meshes +using Enzyme: Enzyme + +function MeshIntegrals.jacobian( + geometry::Meshes.Geometry, + ts::Union{AbstractVector{T}, Tuple{T, Vararg{T}}}, + ::AutoEnzyme +) where {T <: AbstractFloat} + Dim = Meshes.paramdim(geometry) + if Dim != length(ts) + throw(ArgumentError("ts must have same number of dimensions as geometry.")) + end + return Meshes.to.(Enzyme.jacobian(Enzyme.Forward, geometry, ts...)) +end + +end diff --git a/src/MeshIntegrals.jl b/src/MeshIntegrals.jl index cbc7260f..82da7fc4 100644 --- a/src/MeshIntegrals.jl +++ b/src/MeshIntegrals.jl @@ -10,7 +10,7 @@ import QuadGK import Unitful include("differentiation.jl") -export DifferentiationMethod, FiniteDifference, jacobian +export DifferentiationMethod, FiniteDifference, AutoEnzyme, jacobian include("utils.jl") diff --git a/src/differentiation.jl b/src/differentiation.jl index 4433fd49..58425ad8 100644 --- a/src/differentiation.jl +++ b/src/differentiation.jl @@ -9,7 +9,7 @@ A category of types used to specify the desired method for calculating derivativ Derivatives are used to form Jacobian matrices when calculating the differential element size throughout the integration region. -See also [`FiniteDifference`](@ref). +See also [`FiniteDifference`](@ref), [`AutoEnzyme`](@ref). """ abstract type DifferentiationMethod end @@ -27,8 +27,14 @@ end FiniteDifference{T}() where {T <: AbstractFloat} = FiniteDifference{T}(T(1e-6)) FiniteDifference() = FiniteDifference{Float64}() +""" + AutoEnzyme() + +Use to specify use of the Enzyme.jl for calculating derivatives. +""" +struct AutoEnzyme <: DifferentiationMethod end + # Future Support: -# struct AutoEnzyme <: DifferentiationMethod end # struct AutoZygote <: DifferentiationMethod end ################################################################################ @@ -52,7 +58,7 @@ function jacobian( geometry::G, ts::Union{AbstractVector{T}, Tuple{T, Vararg{T}}} ) where {G <: Geometry, T <: AbstractFloat} - return jacobian(geometry, ts, _default_diff_method(G)) + return jacobian(geometry, ts, _default_diff_method(G, T)) end function jacobian( @@ -68,7 +74,7 @@ function jacobian( # Get the partial derivative along the n'th axis via finite difference # approximation, where ts is the current parametric position function ∂ₙr(ts, n, ε) - # Build left/right parametric coordinates with non-allocating iterators + # Build left/right parametric coordinates with non-allocating iterators left = Iterators.map(((i, t),) -> i == n ? t - ε : t, enumerate(ts)) right = Iterators.map(((i, t),) -> i == n ? t + ε : t, enumerate(ts)) # Select orientation of finite-diff @@ -107,7 +113,7 @@ possible and finite difference approximations otherwise. function differential( geometry::G, ts::Union{AbstractVector{T}, Tuple{T, Vararg{T}}}, - diff_method::DifferentiationMethod = _default_diff_method(G) + diff_method::DifferentiationMethod = _default_diff_method(G, T) ) where {G <: Geometry, T <: AbstractFloat} J = Iterators.map(_KVector, jacobian(geometry, ts, diff_method)) return LinearAlgebra.norm(foldl(∧, J)) diff --git a/src/integral.jl b/src/integral.jl index 34d9695e..1365aa93 100644 --- a/src/integral.jl +++ b/src/integral.jl @@ -3,7 +3,7 @@ ################################################################################ """ - integral(f, geometry[, rule]; diff_method=_default_method(geometry), FP=Float64) + integral(f, geometry[, rule]; diff_method=_default_diff_method(geometry, FP), FP=Float64) Numerically integrate a given function `f(::Point)` over the domain defined by a `geometry` using a particular numerical integration `rule` with floating point @@ -16,7 +16,7 @@ precision of type `FP`. `GaussKronrod()` in 1D and `HAdaptiveCubature()` else) # Keyword Arguments -- `diff_method::DifferentiationMethod = _default_method(geometry)`: the method to +- `diff_method::DifferentiationMethod = _default_diff_method(geometry, FP)`: the method to use for calculating Jacobians that are used to calculate differential elements - `FP = Float64`: the floating point precision desired. """ @@ -42,8 +42,10 @@ function _integral( geometry, rule::GaussKronrod; FP::Type{T} = Float64, - diff_method::DM = _default_diff_method(geometry) + diff_method::DM = _default_diff_method(geometry, FP) ) where {DM <: DifferentiationMethod, T <: AbstractFloat} + _check_diff_method_support(geometry, diff_method) + # Implementation depends on number of parametric dimensions over which to integrate N = Meshes.paramdim(geometry) if N == 1 @@ -70,8 +72,10 @@ function _integral( geometry, rule::GaussLegendre; FP::Type{T} = Float64, - diff_method::DM = _default_diff_method(geometry) + diff_method::DM = _default_diff_method(geometry, FP) ) where {DM <: DifferentiationMethod, T <: AbstractFloat} + _check_diff_method_support(geometry, diff_method) + N = Meshes.paramdim(geometry) # Get Gauss-Legendre nodes and weights of type FP for a region [-1,1]ᴺ @@ -99,8 +103,10 @@ function _integral( geometry, rule::HAdaptiveCubature; FP::Type{T} = Float64, - diff_method::DM = _default_diff_method(geometry) + diff_method::DM = _default_diff_method(geometry, FP) ) where {DM <: DifferentiationMethod, T <: AbstractFloat} + _check_diff_method_support(geometry, diff_method) + N = Meshes.paramdim(geometry) integrand(ts) = f(geometry(ts...)) * differential(geometry, ts, diff_method) diff --git a/src/specializations/BezierCurve.jl b/src/specializations/BezierCurve.jl index cdae2ec6..5a8778d5 100644 --- a/src/specializations/BezierCurve.jl +++ b/src/specializations/BezierCurve.jl @@ -36,13 +36,17 @@ function integral( curve::Meshes.BezierCurve, rule::IntegrationRule; alg::Meshes.BezierEvalMethod = Meshes.Horner(), + FP::Type{T} = Float64, + diff_method::DM = _default_diff_method(curve, FP), kwargs... -) +) where {DM <: DifferentiationMethod, T <: AbstractFloat} + _check_diff_method_support(curve, diff_method) + # Generate a _ParametricGeometry whose parametric function auto-applies the alg kwarg param_curve = _ParametricGeometry(_parametric(curve, alg), Meshes.paramdim(curve)) # Integrate the _ParametricGeometry using the standard methods - return _integral(f, param_curve, rule; kwargs...) + return _integral(f, param_curve, rule; diff_method = diff_method, FP = FP, kwargs...) end ################################################################################ diff --git a/src/specializations/CylinderSurface.jl b/src/specializations/CylinderSurface.jl index 3e79d0be..c359e3af 100644 --- a/src/specializations/CylinderSurface.jl +++ b/src/specializations/CylinderSurface.jl @@ -12,18 +12,22 @@ function integral( f, cyl::Meshes.CylinderSurface, rule::I; + FP::Type{T} = Float64, + diff_method::DM = _default_diff_method(cyl, FP), kwargs... -) where {I <: IntegrationRule} +) where {I <: IntegrationRule, DM <: DifferentiationMethod, T <: AbstractFloat} + _check_diff_method_support(cyl, diff_method) + # The generic method only parametrizes the sides - sides = _integral(f, cyl, rule; kwargs...) + sides = _integral(f, cyl, rule; diff_method = diff_method, FP = FP, kwargs...) # Integrate the Disk at the top disk_top = Meshes.Disk(cyl.top, cyl.radius) - top = _integral(f, disk_top, rule; kwargs...) + top = _integral(f, disk_top, rule; diff_method = diff_method, FP = FP, kwargs...) # Integrate the Disk at the bottom disk_bottom = Meshes.Disk(cyl.bot, cyl.radius) - bottom = _integral(f, disk_bottom, rule; kwargs...) + bottom = _integral(f, disk_bottom, rule; diff_method = diff_method, FP = FP, kwargs...) return sides + top + bottom end diff --git a/src/utils.jl b/src/utils.jl index 6cf5ee14..84ebfdc7 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -12,15 +12,50 @@ end # DifferentiationMethod ################################################################################ -# Return the default DifferentiationMethod instance for a particular geometry type +""" + supports_autoenzyme(geometry) + +Return whether a geometry (or geometry type) has a parametric function that can be +differentiated with Enzyme. See GitHub Issue #154 for more information. +""" +supports_autoenzyme(::Type{<:Meshes.Geometry}) = true +supports_autoenzyme(::Type{<:Meshes.BezierCurve}) = false +supports_autoenzyme(::Type{<:Meshes.CylinderSurface}) = false +supports_autoenzyme(::Type{<:Meshes.Cylinder}) = false +supports_autoenzyme(::Type{<:Meshes.ParametrizedCurve}) = false +supports_autoenzyme(::G) where {G <: Geometry} = supports_autoenzyme(G) + +""" + _check_diff_method_support(::Geometry, ::DifferentiationMethod) -> nothing + +Throw an error if incompatible geometry-diff_method combination detected. +""" +_check_diff_method_support(::Geometry, ::DifferentiationMethod) = nothing +function _check_diff_method_support(geometry::Geometry, ::AutoEnzyme) + if !supports_autoenzyme(geometry) + throw(ArgumentError("AutoEnzyme not supported for this geometry.")) + end +end + +""" + _default_diff_method(geometry, FP) + +Return an instance of the default DifferentiationMethod for a particular geometry +(or geometry type) and floating point type. +""" function _default_diff_method( - g::Type{G} -) where {G <: Geometry} - return FiniteDifference() + g::Type{G}, FP::Type{T} +) where {G <: Geometry, T <: AbstractFloat} + if supports_autoenzyme(g) && FP <: Union{Float32, Float64} + AutoEnzyme() + else + FiniteDifference() + end end -# Return the default DifferentiationMethod instance for a particular geometry instance -_default_diff_method(g::G) where {G <: Geometry} = _default_diff_method(G) +function _default_diff_method(::G, ::Type{T}) where {G <: Geometry, T <: AbstractFloat} + _default_diff_method(G, T) +end ################################################################################ # Numerical Tools diff --git a/test/Project.toml b/test/Project.toml index e38225db..0861ae09 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" CoordRefSystems = "b46f11dc-f210-4604-bfba-323c1ec968cb" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Meshes = "eacbb407-ea5a-433e-ab97-5258b1ca43fa" @@ -12,10 +13,11 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [compat] Aqua = "0.7, 0.8" -CoordRefSystems = "0.12, 0.13, 0.14, 0.15, 0.16" +CoordRefSystems = "0.15, 0.16" +Enzyme = "0.13.19" ExplicitImports = "1.6.0" LinearAlgebra = "1" -Meshes = "0.50, 0.51, 0.52" +Meshes = "0.51.20, 0.52" SpecialFunctions = "2" TestItemRunner = "1" TestItems = "1" diff --git a/test/combinations.jl b/test/combinations.jl index b53420bd..d7db64fb 100644 --- a/test/combinations.jl +++ b/test/combinations.jl @@ -18,6 +18,7 @@ This file includes tests for: using Meshes using MeshIntegrals using Unitful + import Enzyme # Used for testing callable objects as integrand functions struct Callable{F <: Function} @@ -43,27 +44,32 @@ This file includes tests for: gausslegendre::Bool hadaptivecubature::Bool # DifferentiationMethods - # autoenzyme::Bool + autoenzyme::Bool end # Shortcut constructor for geometries with typical support structure - function SupportStatus(geometry::Geometry) - if paramdim(geometry) == 1 + function SupportStatus(g::Geometry, autoenzyme = MeshIntegrals.supports_autoenzyme(g)) + N = Meshes.paramdim(g) + if N == 1 + # line/curve aliases = Bool.((1, 0, 0)) rules = Bool.((1, 1, 1)) - return SupportStatus(aliases..., rules...) - elseif paramdim(geometry) == 2 + return SupportStatus(aliases..., rules..., autoenzyme) + elseif N == 2 + # surface aliases = Bool.((0, 1, 0)) rules = Bool.((1, 1, 1)) - return SupportStatus(aliases..., rules...) - elseif paramdim(geometry) == 3 + return SupportStatus(aliases..., rules..., autoenzyme) + elseif N == 3 + # volume aliases = Bool.((0, 0, 1)) rules = Bool.((0, 1, 1)) - return SupportStatus(aliases..., rules...) + return SupportStatus(aliases..., rules..., autoenzyme) else + # ≥4D aliases = Bool.((0, 0, 0)) rules = Bool.((0, 1, 1)) - return SupportStatus(aliases..., rules...) + return SupportStatus(aliases..., rules..., autoenzyme) end end @@ -110,15 +116,21 @@ This file includes tests for: end end # for - #= iter_diff_methods = ( (supports.autoenzyme, AutoEnzyme()), ) for (supported, diff_method) in iter_diff_methods - @test integral(testable.integrand, testable.geometry; diff_method=diff_method)≈sol rtol=rtol - end - =# + if supported + @test integral( + testable.integrand, testable.geometry; diff_method = diff_method)≈testable.solution rtol=rtol + @test MeshIntegrals.supports_autoenzyme(testable.geometry) == true + else + @test_throws "not supported" integral( + testable.integrand, testable.geometry; diff_method = diff_method) + @test MeshIntegrals.supports_autoenzyme(testable.geometry) == false + end + end # for end # function end #testsnippet @@ -454,32 +466,28 @@ end end @testitem "ParametrizedCurve" setup=[Combinations] begin - # ParametrizedCurve has been added in Meshes v0.51.20 - # If the version is specified as minimal compat bound in the Project.toml, the downgrade test fails - if pkgversion(Meshes) >= v"0.51.20" - using CoordRefSystems: Polar - - # Geometries - # Parametrize a circle centered on origin with specified radius - radius = 4.4 - curve_cart = ParametrizedCurve( - t -> Point(radius * cos(t), radius * sin(t)), (0.0, 2π)) - curve_polar = ParametrizedCurve(t -> Point(Polar(radius, t)), (0.0, 2π)) - - # Integrand & Solution - function integrand(p::Meshes.Point) - ur = norm(to(p)) - r = ustrip(u"m", ur) - exp(-r^2) * u"A" - end - solution = 2π * radius * exp(-radius^2) * u"A*m" + using CoordRefSystems: Polar + + # Geometries + # Parametrize a circle centered on origin with specified radius + radius = 4.4 + curve_cart = ParametrizedCurve( + t -> Point(radius * cos(t), radius * sin(t)), (0.0, 2π)) + curve_polar = ParametrizedCurve(t -> Point(Polar(radius, t)), (0.0, 2π)) - # Package and run tests - testable_cart = TestableGeometry(integrand, curve_cart, solution) - runtests(testable_cart) - testable_polar = TestableGeometry(integrand, curve_polar, solution) - runtests(testable_polar) + # Integrand & Solution + function integrand(p::Meshes.Point) + ur = norm(to(p)) + r = ustrip(u"m", ur) + exp(-r^2) * u"A" end + solution = 2π * radius * exp(-radius^2) * u"A*m" + + # Package and run tests + testable_cart = TestableGeometry(integrand, curve_cart, solution) + runtests(testable_cart) + testable_polar = TestableGeometry(integrand, curve_polar, solution) + runtests(testable_polar) end @testitem "Meshes.Plane" setup=[Combinations] begin diff --git a/test/floating_point_types.jl b/test/floating_point_types.jl index 7af87bd5..d1f6dbe5 100644 --- a/test/floating_point_types.jl +++ b/test/floating_point_types.jl @@ -7,6 +7,7 @@ using LinearAlgebra: norm using Meshes using Unitful + using Enzyme baseatol = Dict( Float32 => 0.01f0, diff --git a/test/utils.jl b/test/utils.jl index 947688a7..5db3d779 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -1,8 +1,10 @@ @testsnippet Utils begin using LinearAlgebra: norm using Meshes + using MeshIntegrals using MeshIntegrals: _default_diff_method, _parametric, _units, _zeros, _ones using Unitful + import Enzyme end @testitem "Utilities" setup=[Utils] begin @@ -26,15 +28,21 @@ end @testitem "Differentiation" setup=[Utils] begin # _default_diff_method sphere = Sphere(Point(0, 0, 0), 1.0) - @test _default_diff_method(Meshes.Sphere) isa FiniteDifference - @test _default_diff_method(sphere) isa FiniteDifference + @test _default_diff_method(Meshes.Sphere, Float64) isa AutoEnzyme + @test _default_diff_method(sphere, Float64) isa AutoEnzyme + @test _default_diff_method(sphere, BigFloat) isa FiniteDifference # FiniteDifference @test FiniteDifference().ε ≈ 1e-6 + # Two-argument jacobian + segment = Segment(Point(0), Point(1)) + @test MeshIntegrals.jacobian(segment, (0.5,)) == (Vec(1),) + # Test jacobian with wrong number of parametric coordinates box = Box(Point(0, 0), Point(1, 1)) - @test_throws ArgumentError jacobian(box, zeros(3)) + @test_throws ArgumentError jacobian(box, zeros(3), FiniteDifference()) + @test_throws ArgumentError jacobian(box, zeros(3), AutoEnzyme()) end @testitem "_ParametricGeometry" setup=[Utils] begin