modified allocations

ranocha · Feb 21, 2023 · b359a9d · b359a9d
1 parent 2233c47
commit b359a9d
Showing 1 changed file with 13 additions and 37 deletions.
diff --git a/src/halley.jl b/src/halley.jl
@@ -42,28 +42,8 @@ function SciMLBase.__solve(prob::NonlinearProblem,
                            maxiters = 1000, kwargs...)
     f = Base.Fix2(prob.f, prob.p)
     x = float(prob.u0)
-    # Defining all derivative expressions in one place before the iterations
     if isa(x, AbstractArray)
-        if alg_autodiff(alg)
-            n = length(x)
-            a_dfdx(x) = ForwardDiff.jacobian(f, x)
-            a_d2fdx(x) = ForwardDiff.jacobian(a_dfdx, x)
-            A = Array{Union{Nothing, Number}}(nothing, n, n)
-            #fx = f(x)
-        else
-            n = length(x)
-            f_dfdx(x) = FiniteDiff.finite_difference_jacobian(f, x, diff_type(alg), eltype(x))
-            f_d2fdx(x) = FiniteDiff.finite_difference_jacobian(f_dfdx, x, diff_type(alg), eltype(x))
-            A = Array{Union{Nothing, Number}}(nothing, n, n)
-        end
-    elseif isa(x, Number)
-        if alg_autodiff(alg)
-            sa_dfdx(x) = ForwardDiff.derivative(f, x)
-            sa_d2fdx(x) = ForwardDiff.derivative(sa_dfdx, x)
-        else
-            sf_dfdx(x) = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg), eltype(x))
-            sf_d2fdx(x) = FiniteDiff.finite_difference_derivative(sf_dfdx, x, diff_type(alg), eltype(x))
-        end
+        n = length(x)
     end
     T = typeof(x)
 
@@ -85,34 +65,30 @@ function SciMLBase.__solve(prob::NonlinearProblem,
         if alg_autodiff(alg)
             if isa(x, Number)
                 fx = f(x)
-                dfx = sa_dfdx(x)
-                d2fx = sa_d2fdx(x)
+                dfx = ForwardDiff.derivative(f, x)
+                d2fx = ForwardDiff.derivative(x -> ForwardDiff.derivative(f, x), x)
             else
                 fx = f(x)
-                dfx = a_dfdx(x)
-                d2fx = reshape(a_d2fdx(x), (n,n,n)) # A 3-dim Hessian Tensor
+                dfx = ForwardDiff.jacobian(f, x)
+                d2fx = ForwardDiff.jacobian(x -> ForwardDiff.jacobian(f, x), x) #  n^2 by n matrix
                 ai = -(dfx \ fx)
-                for j in 1:n
-                    tmp = transpose(d2fx[:, :, j] * ai)
-                    A[j, :] = tmp
-                end
+                A = reshape(d2fx * ai, (n, n))
                 bi = (dfx) \ (A * ai)
                 ci = (ai .* ai) ./ (ai .+ (0.5 .* bi))
             end
         else
             if isa(x, Number)
                 fx = f(x)
-                dfx = sf_dfdx(x)
-                d2fx = sf_d2fdx(x)
+                dfx = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg), eltype(x))
+                d2fx = FiniteDiff.finite_difference_derivative(x -> FiniteDiff.finite_difference_derivative(f, x), x,
+                                                                                            diff_type(alg), eltype(x))
             else
                 fx = f(x)
-                dfx = f_dfdx(x)
-                d2fx = reshape(f_d2fdx(x), (n,n,n)) # A 3-dim Hessian Tensor
+                dfx = FiniteDiff.finite_difference_jacobian(f, x, diff_type(alg), eltype(x))
+                d2fx = FiniteDiff.finite_difference_jacobian(x -> FiniteDiff.finite_difference_jacobian(f, x), x, 
+                                                                        diff_type(alg), eltype(x))
                 ai = -(dfx \ fx)
-                for j in 1:n
-                    tmp = transpose(d2fx[:, :, j] * ai)
-                    A[j, :] = tmp
-                end
+                A = reshape(d2fx * ai, (n, n))
                 bi = (dfx) \ (A * ai)
                 ci = (ai .* ai) ./ (ai .+ (0.5 .* bi))
             end