Test optimization

src/core/control.jl

@@ -1,36 +1,43 @@
 using Enzyme
+Enzyme.API.looseTypeAnalysis!(false)
+Enzyme.API.strictAliasing!(false)
+Enzyme.API.typeWarning!(false)
+
 using Optimization, OptimizationOptimJL
 
-#Enzyme.API.looseTypeAnalysis!(false)
-#Enzyme.API.strictAliasing!(false)
-#Enzyme.API.typeWarning!(false)
 Enzyme.Compiler.RunAttributor[] = false
 
 function Control_Core(md::model, femmodel::FemModel) #{{{
 	# solve for optimization
+	# TODO: just a first try, need to add all the features
 	α = md.inversion.independent
+	∂J_∂α = zero(α)
 	n = length(α)
+	# use user defined grad, errors!
+	#optprob = OptimizationFunction(costfunction, Optimization.AutoEnzyme(), grad=computeGradient(∂J_∂α, α, femmodel))
 	optprob = OptimizationFunction(costfunction, Optimization.AutoEnzyme())
 	prob = Optimization.OptimizationProblem(optprob, α, femmodel, lb=md.inversion.min_parameters, ub=md.inversion.max_parameters)
 	sol = Optimization.solve(prob, Optim.LBFGS())
 
-	#TODO: Put the solution back in results according to its Enum
-	#InputUpdateFromVectorx(femmodel, sol.u, GradientEnum, VertexSIdEnum)
-	#RequestedOutputsx(femmodel, [GradientEnum])
+	independent_enum = StringToEnum(md.inversion.independent_string)
+	InputUpdateFromVectorx(femmodel, sol.u, independent_enum, VertexSIdEnum)
+	RequestedOutputsx(femmodel, [independent_enum])
 end#}}}
 function computeGradient(md::model, femmodel::FemModel) #{{{
 	#independent variable
 	α = md.inversion.independent
 	#initialize derivative as 0
 	∂J_∂α = zero(α)
-
-	# zero ALL depth of the model, make sure we get correct gradient
-	dfemmodel = Enzyme.Compiler.make_zero(Base.Core.Typeof(femmodel), IdDict(), femmodel)
-	# compute the gradient
-	#println("CALLING AUTODIFF, prepare to die...")
-	@time autodiff(Enzyme.Reverse, costfunction, Duplicated(femmodel, dfemmodel), Duplicated(α, ∂J_∂α))
+	# Compute Gradient
+	computeGradient(∂J_∂α, α, femmodel)
 
 	#Put gradient in results
 	InputUpdateFromVectorx(femmodel, ∂J_∂α, GradientEnum, VertexSIdEnum)
 	RequestedOutputsx(femmodel, [GradientEnum])
 end#}}}
+function computeGradient(∂J_∂α::Vector{Float64}, α::Vector{Float64}, femmodel::FemModel) #{{{
+	# zero ALL depth of the model, make sure we get correct gradient
+	dfemmodel = Enzyme.Compiler.make_zero(Base.Core.Typeof(femmodel), IdDict(), femmodel)
+	# compute the gradient
+	@time autodiff(Enzyme.Reverse, costfunction, Duplicated(α, ∂J_∂α), Duplicated(femmodel,dfemmodel))
+end#}}}

src/core/costfunctions.jl

@@ -1,5 +1,5 @@
 
-function costfunction(femmodel::FemModel, α::Vector{Float64}) #{{{
+function costfunction(α::Vector{Float64}, femmodel::FemModel) #{{{
 	# get the md.inversion.control_string
 	control_string = FindParam(String, femmodel.parameters, InversionControlParametersEnum)
 	# get the Enum
@@ -9,9 +9,9 @@ function costfunction(femmodel::FemModel, α::Vector{Float64}) #{{{
 	end
 	# compute cost function
 	# TODO: loop through all controls with respect to all the components in the cost function
-	costfunction(femmodel, α,  controlvar_enum, VertexSIdEnum)
+	costfunction(α, femmodel, controlvar_enum, VertexSIdEnum)
 end#}}}
-function costfunction(femmodel::FemModel, α::Vector{Float64}, controlvar_enum::IssmEnum, SId_enum::IssmEnum) #{{{
+function costfunction(α::Vector{Float64}, femmodel::FemModel, controlvar_enum::IssmEnum, SId_enum::IssmEnum) #{{{
 	#Update FemModel accordingly
 	InputUpdateFromVectorx(femmodel, α, controlvar_enum, SId_enum)
 

test/testad.jl

@@ -5,6 +5,7 @@ using MAT
 using Test
 using Enzyme
 
+Enzyme.API.typeWarning!(false)
 Enzyme.Compiler.RunAttributor[] = false
 
 #Load model from MATLAB file
@@ -36,20 +37,20 @@ addJ = md.results["StressbalanceSolution"]["Gradient"]
 
    α = md.inversion.independent
    femmodel=dJUICE.ModelProcessor(md, :StressbalanceSolution)
-   J1 = dJUICE.costfunction(femmodel, α)
+   J1 = dJUICE.costfunction(α, femmodel)
    @test ~isnothing(J1)
 end
 
 @testset "AD gradient calculation for FrictionC" begin
 	α = md.inversion.independent
 	delta = 1e-7
 	femmodel=dJUICE.ModelProcessor(md, :StressbalanceSolution)
-	J1 = dJUICE.costfunction(femmodel, α)
+	J1 = dJUICE.costfunction(α, femmodel)
 	for i in 1:md.mesh.numberofvertices
 		dα = zero(md.friction.coefficient)
 		dα[i] = delta
 		femmodel=dJUICE.ModelProcessor(md, :StressbalanceSolution)
-		J2 = dJUICE.costfunction(femmodel, α+dα)
+		J2 = dJUICE.costfunction(α+dα, femmodel)
 		dJ = (J2-J1)/delta
 
 		@test abs(dJ - addJ[i])< 1e-5

test/testad2.jl

@@ -5,6 +5,7 @@ using MAT
 using Test
 using Enzyme
 
+Enzyme.API.typeWarning!(false)
 Enzyme.Compiler.RunAttributor[] = false
 
 #Load model from MATLAB file
@@ -35,20 +36,20 @@ addJ = md.results["StressbalanceSolution"]["Gradient"]
 
    α = md.inversion.independent
    femmodel=dJUICE.ModelProcessor(md, :StressbalanceSolution)
-   J1 = dJUICE.costfunction(femmodel, α)
+   J1 = dJUICE.costfunction(α, femmodel)
    @test ~isnothing(J1)
 end
 
 @testset "AD results RheologyB" begin
 	α = md.inversion.independent
 	delta = 1e-8
 	femmodel=dJUICE.ModelProcessor(md, :StressbalanceSolution)
-	J1 = dJUICE.costfunction(femmodel, α)
+	J1 = dJUICE.costfunction(α, femmodel)
 	for i in 1:md.mesh.numberofvertices
 		dα = zero(md.friction.coefficient)
 		dα[i] = delta
 		femmodel=dJUICE.ModelProcessor(md, :StressbalanceSolution)
-		J2 = dJUICE.costfunction(femmodel, α+dα)
+		J2 = dJUICE.costfunction(α+dα, femmodel)
 		dJ = (J2-J1)/delta
 
 		@test abs(dJ - addJ[i])< 1e-6

test/testoptimization.jl

@@ -0,0 +1,46 @@
+module enzymeDiff
+
+using dJUICE
+using MAT
+using Test
+using Enzyme
+
+Enzyme.Compiler.RunAttributor[] = false
+using Optimization, OptimizationOptimJL
+
+#Load model from MATLAB file
+#file = matopen(joinpath(@__DIR__, "..", "data","temp12k.mat")) #BIG model
+file = matopen(joinpath(@__DIR__, "..", "data","temp.mat")) #SMALL model (35 elements)
+mat  = read(file, "md")
+close(file)
+md = model(mat)
+
+#make model run faster 
+md.stressbalance.maxiter = 20
+
+#Now call AD!
+md.inversion.iscontrol = 1
+md.inversion.independent = md.friction.coefficient
+md.inversion.min_parameters = ones(md.mesh.numberofvertices)*(0.0)
+md.inversion.max_parameters = ones(md.mesh.numberofvertices)*(1.0e3)
+md.inversion.independent_string = "FrictionCoefficient"
+
+α = md.inversion.independent
+∂J_∂α = zero(α)
+
+femmodel=dJUICE.ModelProcessor(md, :StressbalanceSolution)
+n = length(α)
+# use user defined grad, errors!
+optprob = OptimizationFunction(dJUICE.costfunction, Optimization.AutoEnzyme(), grad=dJUICE.computeGradient)
+#optprob = OptimizationFunction(dJUICE.costfunction, Optimization.AutoEnzyme())
+#prob = Optimization.OptimizationProblem(optprob, α, femmodel, lb=md.inversion.min_parameters, ub=md.inversion.max_parameters)
+prob = Optimization.OptimizationProblem(optprob, α, femmodel)
+sol = Optimization.solve(prob, Optim.LBFGS())
+
+
+#md = solve(md, :sb)
+
+# compute gradient by finite differences at each node
+#addJ = md.results["StressbalanceSolution"]["Gradient"]
+
+end