Fix bug where integral fails when Enzyme ext not loaded

CHANGELOG.md

@@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Implemented a more efficient internal parametric transformation for `Meshes.Tetrahedron`, resulting in about an 80% integral performance improvement.
 
+### Fixed
+
+- Fixed a bug where `integral` would default to `diff_method=AutoEnzyme()` even when the Enzyme extension isn't loaded.
+
 
 ## [0.16.0] - 2024-12-14
 

ext/MeshIntegralsEnzymeExt.jl

@@ -16,4 +16,12 @@ function MeshIntegrals.jacobian(
     return Meshes.to.(Enzyme.jacobian(Enzyme.Forward, geometry, ts...))
 end
 
+# Supports all geometries except for those that throw errors
+# See GitHub Issue #154 for more information
+MeshIntegrals.supports_autoenzyme(::Type{<:Meshes.Geometry}) = true
+MeshIntegrals.supports_autoenzyme(::Type{<:Meshes.BezierCurve}) = false
+MeshIntegrals.supports_autoenzyme(::Type{<:Meshes.CylinderSurface}) = false
+MeshIntegrals.supports_autoenzyme(::Type{<:Meshes.Cylinder}) = false
+MeshIntegrals.supports_autoenzyme(::Type{<:Meshes.ParametrizedCurve}) = false
+
 end

src/utils.jl

@@ -13,48 +13,63 @@ end
 ################################################################################
 
 """
-    supports_autoenzyme(geometry)
+    supports_autoenzyme(geometry::Geometry)
+    supports_autoenzyme(type::Type{<:Geometry})
 
-Return whether a geometry (or geometry type) has a parametric function that can be
+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
+function supports_autoenzyme end
+
+# Returns false for all geometries when Enzyme extension is not loaded
+supports_autoenzyme(::Type{<:Any}) = false
+
+# If provided a geometry instance, re-run with the type as argument
 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.
+Throw an error if incompatible combination {geometry, diff_method} detected.
 """
-_check_diff_method_support(::Geometry, ::DifferentiationMethod) = nothing
+function _check_diff_method_support end
+
+# If diff_method == Enzyme, then perform check
 function _check_diff_method_support(geometry::Geometry, ::AutoEnzyme)
     if !supports_autoenzyme(geometry)
         throw(ArgumentError("AutoEnzyme not supported for this geometry."))
     end
 end
 
+# If diff_method != AutoEnzyme, then do nothing
+_check_diff_method_support(::Geometry, ::DifferentiationMethod) = nothing
+
 """
     _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}, FP::Type{T}
-) where {G <: Geometry, T <: AbstractFloat}
-    if supports_autoenzyme(g) && FP <: Union{Float32, Float64}
-        AutoEnzyme()
+        ::Type{G},
+        ::Type{FP}
+) where {G <: Geometry, FP <: AbstractFloat}
+    # Enzyme only works with these FP types
+    uses_Enzyme_supported_FP_type = (FP <: Union{Float32, Float64})
+
+    if supports_autoenzyme(G) && uses_Enzyme_supported_FP_type
+        return AutoEnzyme()
     else
-        FiniteDifference()
+        return FiniteDifference()
     end
 end
 
-function _default_diff_method(::G, ::Type{T}) where {G <: Geometry, T <: AbstractFloat}
-    _default_diff_method(G, T)
+# If provided a geometry instance, re-run with the type as argument
+function _default_diff_method(
+        ::G,
+        ::Type{FP}
+) where {G <: Geometry, FP <: AbstractFloat}
+    return _default_diff_method(G, FP)
 end
 
 ################################################################################

test/combinations.jl

@@ -48,8 +48,12 @@ This file includes tests for:
     end
 
     # Shortcut constructor for geometries with typical support structure
-    function SupportStatus(g::Geometry, autoenzyme = MeshIntegrals.supports_autoenzyme(g))
-        N = Meshes.paramdim(g)
+    function SupportStatus(geometry::G;) where {G <: Geometry}
+        # Check whether AutoEnzyme should be supported, i.e. not on blacklist
+        unsupported_Gs = Union{BezierCurve, Cylinder, CylinderSurface, ParametrizedCurve}
+        autoenzyme = !(G <: unsupported_Gs)
+
+        N = Meshes.paramdim(geometry)
         if N == 1
             # line/curve
             aliases = Bool.((1, 0, 0))
@@ -70,14 +74,17 @@ This file includes tests for:
             aliases = Bool.((0, 0, 0))
             rules = Bool.((0, 1, 1))
             return SupportStatus(aliases..., rules..., autoenzyme)
-        end
-    end
+        end #if
+    end # function
+
+    # Generate applicable tests for this geometry
+    function runtests(testable::TestableGeometry; rtol = sqrt(eps()))
+        # Determine support matrix for this geometry
+        supports = SupportStatus(testable.geometry)
+
+        # Ensure consistency of SupportStatus with supports_autoenzyme
+        @test MeshIntegrals.supports_autoenzyme(testable.geometry) == supports.autoenzyme
 
-    function runtests(
-            testable::TestableGeometry,
-            supports::SupportStatus = SupportStatus(testable.geometry);
-            rtol = sqrt(eps())
-    )
         # Test alias functions
         for alias in (lineintegral, surfaceintegral, volumeintegral)
             # if supports.alias
@@ -86,15 +93,14 @@ This file includes tests for:
             else
                 @test_throws "not supported" alias(testable.integrand, testable.geometry)
             end
-        end
+        end # for
 
+        # Iteratively test all IntegrationRules
         iter_rules = (
             (supports.gausskronrod, GaussKronrod()),
             (supports.gausslegendre, GaussLegendre(100)),
             (supports.hadaptivecubature, HAdaptiveCubature())
         )
-
-        # Test rules
         for (supported, rule) in iter_rules
             if supported
                 # Scalar integrand
@@ -116,19 +122,21 @@ This file includes tests for:
             end
         end # for
 
+        # Iteratively test all DifferentiationMethods
         iter_diff_methods = (
-            (supports.autoenzyme, AutoEnzyme()),
+            (true, FiniteDifference()),
+            (supports.autoenzyme, AutoEnzyme())
         )
+        for (supported, method) in iter_diff_methods
+            # Aliases for improved code readability
+            f = testable.integrand
+            geometry = testable.geometry
+            sol = testable.solution
 
-        for (supported, diff_method) in iter_diff_methods
             if supported
-                @test integral(
-                    testable.integrand, testable.geometry; diff_method = diff_method)≈testable.solution rtol=rtol
-                @test MeshIntegrals.supports_autoenzyme(testable.geometry) == true
+                @test integral(f, geometry; diff_method = method)≈sol rtol=rtol
             else
-                @test_throws "not supported" integral(
-                    testable.integrand, testable.geometry; diff_method = diff_method)
-                @test MeshIntegrals.supports_autoenzyme(testable.geometry) == false
+                @test_throws "not supported" integral(f, geometry; diff_method = method)
             end
         end # for
     end # function
@@ -465,7 +473,7 @@ end
     runtests(testable)
 end
 
-@testitem "ParametrizedCurve" setup=[Combinations] begin
+@testitem "Meshes.ParametrizedCurve" setup=[Combinations] begin
     using CoordRefSystems: Polar
 
     # Geometries

test/utils.jl

@@ -25,12 +25,25 @@ end
     @test _ones(Float32, 2) == (1.0f0, 1.0f0)
 end
 
-@testitem "Differentiation" setup=[Utils] begin
-    # _default_diff_method
-    sphere = Sphere(Point(0, 0, 0), 1.0)
-    @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
+@testitem "Differentiation (EnzymeExt loaded)" setup=[Utils] begin
+    # supports_autoenzyme(::Type{<:Any})
+    @test MeshIntegrals.supports_autoenzyme(Nothing) == false
+
+    # _default_diff_method -- using type or instance, Enzyme-supported combination
+    let sphere = Sphere(Point(0, 0, 0), 1.0)
+        @test _default_diff_method(Meshes.Sphere, Float64) isa AutoEnzyme
+        @test _default_diff_method(sphere, Float64) isa AutoEnzyme
+    end
+
+    # _default_diff_method -- Enzyme-unsupported FP types
+    @test _default_diff_method(Meshes.Sphere, Float16) isa FiniteDifference
+    @test _default_diff_method(Meshes.Sphere, BigFloat) isa FiniteDifference
+
+    # _default_diff_method -- geometries that currently error with AutoEnzyme
+    @test _default_diff_method(Meshes.BezierCurve, Float64) isa FiniteDifference
+    @test _default_diff_method(Meshes.CylinderSurface, Float64) isa FiniteDifference
+    @test _default_diff_method(Meshes.Cylinder, Float64) isa FiniteDifference
+    @test _default_diff_method(Meshes.ParametrizedCurve, Float64) isa FiniteDifference
 
     # FiniteDifference
     @test FiniteDifference().ε ≈ 1e-6