Add inversion using MadNLP

Author: enigne

data/PIG_testopt_small.mat

@@ -2,30 +2,52 @@ using Enzyme
 using ManualNLPModels
 using MadNLP
 
-function Control_Core(md::model, femmodel::FemModel)
+function Control_Core(md::model, femmodel::FemModel, solutionstring::Symbol)#{{{
 	#independent variable
 	α = md.inversion.independent
 	#initialize derivative as 0
-	∂J_∂α = zero(α)
+	∂J_∂α = make_zero(α)
 	if md.inversion.onlygrad
 		# only compute the gradient
-		ComputeGradient(∂J_∂α, α, femmodel)
+		ComputeGradient!(∂J_∂α, α, femmodel)
 		#Put gradient in results
 		InputUpdateFromVectorx(femmodel, ∂J_∂α, GradientEnum, VertexSIdEnum)
 		RequestedOutputsx(femmodel, [GradientEnum])
 	else
 		# optimization
-		# use user defined grad, errors!
-		#		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())
+		# define cost function and gradient
+		# need to build connection between md and x
+		f(x) = begin
+			fem=DJUICE.ModelProcessor(md, solutionstring)
+			DJUICE.CostFunction(x, fem)
+		end
+
+		g!(gx, x) = begin
+			fem=DJUICE.ModelProcessor(md, solutionstring)
+			DJUICE.ComputeGradient!(gx, x, fem)
+		end
+		nlp = NLPModel(
+							α,
+							f,
+							grad = g!,
+							lvar = md.inversion.min_parameters,
+							uvar = md.inversion.max_parameters,
+							)
+
+		results_qn = madnlp(
+								  nlp;
+								  linear_solver=LapackCPUSolver,
+								  hessian_approximation=MadNLP.CompactLBFGS,
+								  tol=md.inversion.tol,
+								  max_iter=md.inversion.maxiter,
+								  )
+
 		independent_enum = StringToEnum(md.inversion.independent_string)
-		InputUpdateFromVectorx(femmodel, sol.u, independent_enum, VertexSIdEnum)
+		InputUpdateFromVectorx(femmodel, results_qn.solution, independent_enum, VertexSIdEnum)
 		RequestedOutputsx(femmodel, [independent_enum])
 	end
 end#}}}
-
-function ComputeGradient(∂J_∂α::Vector{Float64}, α::Vector{Float64}, femmodel::FemModel) #{{{
+function ComputeGradient!(∂J_∂α::Vector{Float64}, α::Vector{Float64}, femmodel::FemModel) #{{{
 	# zero ALL depth of the model, make sure we get correct gradient
 	dfemmodel = make_zero(Base.Core.Typeof(femmodel), IdDict(), femmodel)
 	# zero the gradient
@@ -54,6 +76,7 @@ function CostFunction(α::Vector{Float64}, femmodel::FemModel) #{{{
 	# compute cost function
 	# TODO: loop through all controls with respect to all the components in the cost function
 	solutionstring = FindParam(Symbol, femmodel.parameters, SolutionTypeEnum)
+	# return J
 	CostFunctionx(femmodel, α, controlvar_enum, VertexSIdEnum, cost_enum_list, Val(solutionstring))
 end#}}}
 function CostFunctionx(femmodel::FemModel, α::Vector{Float64}, controlvar_enum::IssmEnum, SId_enum::IssmEnum, cost_enum_list::Vector{IssmEnum}, ::Val{solutionstring}) where solutionstring #{{{

src/core/control.jl

@@ -44,9 +44,9 @@ function solve(md::model, solution::Symbol) #{{{
 	#Construct FemModel from md
 	femmodel=ModelProcessor(md, solutionkey)
 
-	#Solve (FIXME: to be improved later...)
+	#Solve 
 	if (md.inversion.iscontrol) # solve inverse problem
-		Control_Core(md, femmodel)
+		Control_Core(md, femmodel, solutionkey)
 	else # otherwise forward problem
 		if(solutionkey===:StressbalanceSolution)
 			analysis = StressbalanceAnalysis()

src/core/solve.jl

@@ -272,11 +272,12 @@ mutable struct Inversion
 	max_parameters::Vector{Float64}
 	independent::Vector{Float64}
 	maxiter::Int64
+	tol::Float64
 	independent_string::String
 	dependent_string::Vector{String}
 end
 function Inversion() #{{{
-	return Inversion( false, true, Vector{Float64}(undef,0), Vector{Float64}(undef,0), Vector{Float64}(undef,0), Vector{Float64}(undef,0), Vector{Float64}(undef,0), 0, "Friction", Vector{String}(undef,0))
+	return Inversion( false, true, Vector{Float64}(undef,0), Vector{Float64}(undef,0), Vector{Float64}(undef,0), Vector{Float64}(undef,0), Vector{Float64}(undef,0), 0, 0., "Friction", Vector{String}(undef,0))
 end# }}}
 function Base.show(io::IO, this::Inversion)# {{{
 	IssmStructDisp(io, this)

src/usr/classes.jl

@@ -26,12 +26,13 @@ md.inversion.min_parameters = ones(md.mesh.numberofvertices)*(1.0)
 md.inversion.max_parameters = ones(md.mesh.numberofvertices)*(200)
 md.inversion.independent_string = "FrictionCoefficient"
 md.inversion.dependent_string = ["SurfaceAbsVelMisfit"]
+md.inversion.maxiter = 40
+md.inversion.tol = 1e-5
 
 md.verbose.convergence = false
 
-if onlygrad == 1
-	md = solve(md, :sb)
-end
+#md = solve(md, :sb);
+
 α = md.inversion.independent
 ∂J_∂α = zero(α)
 
@@ -44,20 +45,11 @@ end
 # Enzyme gradient g
 g!(gx, x) = begin
 	femmodel=DJUICE.ModelProcessor(md, :StressbalanceSolution)
-	dfemmodel = Enzyme.Compiler.make_zero(Base.Core.Typeof(femmodel), IdDict(), femmodel)
-	dx = zero(x)
-	Enzyme.autodiff(set_runtime_activity(Enzyme.Reverse), DJUICE.CostFunction, Active, Duplicated(x, gx), Duplicated(femmodel,dfemmodel))
-	@assert typeof(dx) == typeof(gx)
-	gx .= dx
+	DJUICE.ComputeGradient!(gx, x, femmodel)
 end
 
-#g!(gx, x) = begin
-#	femmodel=DJUICE.ModelProcessor(md, :StressbalanceSolution)
-#	DJUICE.ComputeGradient(gx, x, femmodel)
-#end
-
 nlp = NLPModel(
-					α,
+					md.inversion.independent,
 					f,
 					grad = g!,
 					lvar = md.inversion.min_parameters,
@@ -68,7 +60,8 @@ results_qn = madnlp(
 						  nlp;
 						  linear_solver=LapackCPUSolver,
 						  hessian_approximation=MadNLP.CompactLBFGS,
-						  tol=1e-5,
-						  max_iter=40,
+						  tol=md.inversion.tol,
+						  max_iter=md.inversion.maxiter,
 						  )
 
+