Fix implementation and tests

Project.toml

@@ -6,23 +6,30 @@ version = "0.1.0"
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ImplicitDifferentiation = "57b37032-215b-411a-8a7c-41a003a55207"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Nonconvex = "01bcebdf-4d21-426d-b5c4-6132c1619978"
 NonconvexIpopt = "bf347577-a06d-49ad-a669-8c0e005493b8"
-Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-TopOpt = "53a1e1a5-51bb-58a9-8a02-02056cc81109"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+ChainRulesCore = "1"
+Distributions = "0.25"
+DistributionsAD = "0.6"
+ImplicitDifferentiation = "0.2"
+NonconvexIpopt = "0.4"
+Reexport = "1"
+UnPack = "1"
+Zygote = "0.6"
 julia = "1"
 
 [extras]
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["FiniteDifferences", "Test"]

src/Uncertainty.jl

@@ -1,61 +1,95 @@
 module Uncertainty
 
-using ImplicitDifferentiation, Zygote, TopOpt, ChainRulesCore, LinearAlgebra
-using UnPack, Nonconvex, Statistics, GLMakie, Distributions, Reexport
-Nonconvex.@load Ipopt
+using ImplicitDifferentiation, Zygote, LinearAlgebra, ChainRulesCore, ForwardDiff
+using UnPack, NonconvexIpopt, Statistics, Distributions, Reexport, DistributionsAD
 @reexport using LinearAlgebra
+@reexport using Statistics
+export RandomFunction, FORM, RIA, MvNormal
 
-struct RandomFunction{F,P}
+struct RandomFunction{F,P,M}
   f::F
   p::P
-	method
+  method::M
 end
 
-function getp0(randFunc::RandomFunction, x)
-  @unpack f, p, method = randFunc
-	lag_multipliers = 0.0
+struct FORM{M}
+	method::M
+end
+struct RIA end
+
+function get_forward(f, p, ::FORM{<:RIA})
 	function forward(x)
-		obj = p -> p'*p # gets an objective function of p
-		constr = p -> f(x, p) # gets the constraint on p
-		# solve the RIA problem to find p0
-		innerOptModel = Nonconvex.Model(obj)
-		addvar!(innerOptModel, zeros(size(p)), fill(10.0, size(p)))
-		add_eq_constraint!(innerOptModel, constr)
-		result = optimize(innerOptModel, IpoptAlg(), mean(p), options = IpoptOptions())
-		lag_multipliers = result # get the Lagrangian multipliers
-		return result.minimizer
-	end
-	function kkt_conditions(x, p)
-		constr = p -> f(x, p) # gets the constraint on p
-		return 2 * p + Zygote.gradient(constr, p)[1] * lag_multipliers
+		# gets an objective function of p
+		obj = pc -> begin
+		   _p = pc[1:end-1]
+		   c = pc[end]
+		   return _p'*_p + c^2
+	   end
+	   # gets the constraint on p
+	   constr = pc -> begin
+		   _p = pc[1:end-1]
+		   c = pc[end]
+		   f(x, _p) .- c
+	   end
+	   # solve the RIA problem to find p0
+	   # should use and reuse the VecModel
+	   innerOptModel = Model(obj)
+	   n = size(p)[1]
+	   addvar!(innerOptModel, fill(-Inf, n + 1), fill(Inf, n + 1))
+	   add_eq_constraint!(innerOptModel, constr)
+	   result = optimize(innerOptModel, IpoptAlg(), [mean(p); 0.0], options = IpoptOptions(print_level = 0))
+	   return vcat(result.minimizer, result.problem.mult_g[1])
+   end
+   return forward
+end
+
+# this is inefficient for multiple reasons
+# first is the use of ForwardDiff because Zygote over Zygote fails
+# second is that we are taking the jacobian wrt to both x and p because
+# of a limitation in Zygote over ForwardDiff
+function jac2(f, x, p)
+	ForwardDiff.jacobian(
+		xp -> f(xp[1:length(x)], xp[length(x)+1:end]),
+		vcat(x, p),
+	)[:,length(x)+1:end]
+end
+
+function get_conditions(f, ::FORM{<:RIA})
+	function kkt_conditions(x, pcmult)
+		p = pcmult[1:end-2]
+		c = pcmult[end-1]
+		mult = pcmult[end]
+		return vcat(2 * p + jac2(f, x, p)' * mult, 2c - mult, f(x, p) .- c)
 	end
-	implicit_f = ImplicitFunction(forward, kkt_conditions)
-	return implicit_f(x)
 end
 
-struct FORM{M}
-	method::M
+function get_implicit(f, p, method)
+	forward = get_forward(f, p, method)
+	kkt_conditions = get_conditions(f, method)
+	return ImplicitFunction(forward, kkt_conditions)
 end
-struct RIA end
 
-function RandomFunction(f, p)
-	return RandomFunction(f, p, FORM(RIA()))
+function getp0(f, x, p, method::FORM{<:RIA})
+	implicit_f = get_implicit(f, p, method)
+	return implicit_f(x)[1:size(p)[1]]
 end
-RandomFunction(f, p; method) = RandomFunction(f, p, method)
+
+function RandomFunction(f, p; method = FORM(RIA()))
+	return RandomFunction(f, p, method)
+end
+
+_vec(x::Real) = [x]
+_vec(x) = x
 
 function (f::RandomFunction)(x)
 	mup = mean(f.p)
 	covp = cov(f.p)
-	p0 = getp0(f, x) # should use ImplicitFunctions to be differentiable and efficient
-	dfdp0 = Zygote.jacobian(p -> f.f(x, p), p0)[1]
+	p0 = getp0(f.f, x, f.p, f.method)
+	dfdp0 = jac2(f.f, x, p0)
 	fp0 = f.f(x, p0)
-	muf = fp0 + dfdp0 * (mup - p0)
+	muf = _vec(fp0) + dfdp0 * (mup - p0)
 	covf = dfdp0 * covp * dfdp0'
 	return MvNormal(muf, covf)
 end
 
-export RandomFunction
-export FORM
-export MvNormal
-
-end
+end

test/runtests.jl

@@ -1,21 +1,21 @@
-using Uncertainty, Test
+using Uncertainty, Test, FiniteDifferences, Zygote
 
-# @testset "Uncertainty.jl" begin
-    # function with at least one random input
-    pol(x, y) = norm(x)^2 + norm(y)^2
-    # define deterministic inputs, in case there are any
-    const y = [4; 5; 8]
+@testset "Uncertainty.jl" begin
+    # test function - y is random
+    pol(x, y) = [norm(x + y)^2]
     # input value to be used in example
-    _x = [2; 3; 6]
+    x = [2.0, 3.0, 6.0]
     # wrap original function in RandomFunction struct
-    rf = RandomFunction(
-        pol, MvNormal(_x, Diagonal(ones(3))), method = FORM(:RIA)
-    )
+    y = MvNormal(zeros(3), Diagonal(ones(3)))
+    rf = RandomFunction(pol, y, FORM(RIA()))
     # call wrapper with example input
-    d = rf(_x)
+    d = rf(x)
     function obj(x)
         dist = rf(x)
-        mean(dist) + 2 * std(dist)
+        mean(dist)[1] + 2 * sqrt(cov(dist)[1,1])
     end
-
-# end
+    obj(x)
+    g1 = FiniteDifferences.grad(central_fdm(5, 1), obj, x)[1]
+    g2 = Zygote.gradient(obj, x)[1]
+    @test norm(g1 - g2) < 1e-7
+end