Merge pull request #13 from LucasMSpereira/mt/fix_zygote

Use Zygote-over-Zygote and format

Author: LucasMSpereira

.JuliaFormatter.toml

@@ -0,0 +1,3 @@
+remove_extra_newlines = true
+always_for_in = true
+conditional_to_if = true

Project.toml

@@ -7,11 +7,12 @@ version = "0.1.0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ImplicitDifferentiation = "57b37032-215b-411a-8a7c-41a003a55207"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NonconvexIpopt = "bf347577-a06d-49ad-a669-8c0e005493b8"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
@@ -20,8 +21,8 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ChainRulesCore = "1"
 Distributions = "0.25"
 DistributionsAD = "0.6"
-ForwardDiff = "0.10"
 ImplicitDifferentiation = "0.2"
+JuliaFormatter = "1"
 NonconvexIpopt = "0.4"
 Reexport = "1"
 UnPack = "1"

src/Uncertainty.jl

@@ -1,95 +1,101 @@
 module Uncertainty
 
-using ImplicitDifferentiation, Zygote, LinearAlgebra, ChainRulesCore, ForwardDiff
+using ImplicitDifferentiation, Zygote, LinearAlgebra, ChainRulesCore, SparseArrays
 using UnPack, NonconvexIpopt, Statistics, Distributions, Reexport, DistributionsAD
 @reexport using LinearAlgebra
 @reexport using Statistics
 export RandomFunction, FORM, RIA, MvNormal
 
 struct RandomFunction{F,P,M}
-  f::F
-  p::P
-  method::M
+    f::F
+    p::P
+    method::M
 end
 
 struct FORM{M}
-	method::M
+    method::M
 end
 struct RIA end
 
 function get_forward(f, p, ::FORM{<:RIA})
-	function forward(x)
-		# 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]
+    function forward(x)
+        # 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
 
 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
+    function kkt_conditions(x, pcmult)
+        p = pcmult[1:end-2]
+        c = pcmult[end-1]
+        mult = pcmult[end]
+        return vcat(
+            2 * p + Zygote.pullback(p -> f(x, p), p)[2](mult)[1],
+            2c - mult,
+            f(x, p) .- c,
+        )
+    end
 end
 
 function get_implicit(f, p, method)
-	forward = get_forward(f, p, method)
-	kkt_conditions = get_conditions(f, method)
-	return ImplicitFunction(forward, kkt_conditions)
+    forward = get_forward(f, p, method)
+    kkt_conditions = get_conditions(f, method)
+    return ImplicitFunction(forward, kkt_conditions)
 end
 
 function getp0(f, x, p, method::FORM{<:RIA})
-	implicit_f = get_implicit(f, p, method)
-	return implicit_f(x)[1:size(p)[1]]
+    implicit_f = get_implicit(f, p, method)
+    return implicit_f(x)[1:size(p)[1]]
 end
 
 function RandomFunction(f, p; method = FORM(RIA()))
-	return RandomFunction(f, p, method)
+    return RandomFunction(f, p, method)
 end
 
 _vec(x::Real) = [x]
 _vec(x) = x
 
+function _jacobian(f, x)
+    val, pb = Zygote.pullback(f, x)
+    M = length(val)
+    vecs = [Vector(sparsevec([i], [true], M)) for i in 1:M]
+    Jt = reduce(hcat, first.(pb.(vecs)))
+    return copy(Jt')
+end
+
 function (f::RandomFunction)(x)
-	mup = mean(f.p)
-	covp = cov(f.p)
-	p0 = getp0(f.f, x, f.p, f.method)
-	dfdp0 = jac2(f.f, x, p0)
-	fp0 = f.f(x, p0)
-	muf = _vec(fp0) + dfdp0 * (mup - p0)
-	covf = dfdp0 * covp * dfdp0'
-	return MvNormal(muf, covf)
+    mup = mean(f.p)
+    covp = cov(f.p)
+    p0 = getp0(f.f, x, f.p, f.method)
+    dfdp0 = _jacobian(p -> f.f(x, p), p0)
+    fp0 = f.f(x, p0)
+    muf = _vec(fp0) + dfdp0 * (mup - p0)
+    covf = dfdp0 * covp * dfdp0'
+    return MvNormal(muf, covf)
 end
 
 end

src/results.jl

@@ -2,7 +2,7 @@ using Nonconvex
 Nonconvex.@load Ipopt
 
 f(x) = sqrt(x[2])
-g(x, a, b) = (a*x[1] + b)^3 - x[2]
+g(x, a, b) = (a * x[1] + b)^3 - x[2]
 
 model = Model(f)
 addvar!(model, [0.0, 0.0], [10.0, 10.0])
@@ -16,4 +16,4 @@ propertynames(r)
 println(propertynames(r.problem))
 typeof(r.problem.intermediate)
 propertynames(r.problem.intermediate)
-r.problem
+r.problem

test/runtests.jl

@@ -12,7 +12,7 @@ using Uncertainty, Test, FiniteDifferences, Zygote
     d = rf(x)
     function obj(x)
         dist = rf(x)
-        mean(dist)[1] + 2 * sqrt(cov(dist)[1,1])
+        mean(dist)[1] + 2 * sqrt(cov(dist)[1, 1])
     end
     obj(x)
     g1 = FiniteDifferences.grad(central_fdm(5, 1), obj, x)[1]