minor code refactoring for Gradient

bcube-project · Dec 11, 2023 · fed0a38 · fed0a38
1 parent 94267de
commit fed0a38
Showing 1 changed file with 5 additions and 10 deletions.
diff --git a/src/algebra/gradient.jl b/src/algebra/gradient.jl
@@ -103,24 +103,19 @@ function Gradient(
     MapOver(grad_shape_functionsNA(fs, n, ctype, cnodes, ξ))
 end
 
-function Gradient(
-    cellFunction::AbstractCellFunction{<:PhysicalDomain, 1},
-    cPoint::CellPoint,
-)
-    f = get_function(cellFunction)
-    x = change_domain(cPoint, PhysicalDomain()) # transparent if already in PhysicalDomain
-    return ForwardDiff.gradient(f, get_coord(x))
-end
-
 function Gradient(
     cellFunction::AbstractCellFunction{<:PhysicalDomain, S},
     cPoint::CellPoint,
 ) where {S}
     f = get_function(cellFunction)
     x = change_domain(cPoint, PhysicalDomain()) # transparent if already in PhysicalDomain
-    return ForwardDiff.jacobian(f, get_coord(x))
+    return _gradient_or_jacobian(Val(S), f, get_coord(x))
 end
 
+# dispatch on codomain size :
+_gradient_or_jacobian(::Val{1}, f, x) = ForwardDiff.gradient(f, x)
+_gradient_or_jacobian(::Val{S}, f, x) where {S} = ForwardDiff.jacobian(f, x)
+
 function grad_shape_functionsNA(fs::AbstractFunctionSpace, n::Val{1}, ctype, cnodes, ξ)
     grad = grad_shape_functions(fs, n, ctype, cnodes, ξ)
     _reshape_gradient_shape_function_impl(grad, n)