Add and fix multimaterial example (#18)

* First sketch

* ignore derivatives

* solved assertion issue

* zygote "foreign call" error

* fix TopOpt example and use Zygote-over-FDM

Co-authored-by: LucasMSpereira <lucas.lmspereira@gmail.com>

Author: mohdibntarek

Project.toml

@@ -7,13 +7,18 @@ version = "0.1.0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+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"
+NonconvexTOBS = "6c0b5230-d4c9-466e-bfd4-b31e6272ab65"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+TopOpt = "53a1e1a5-51bb-58a9-8a02-02056cc81109"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 

src/ReliabilityOptimization.jl

@@ -2,6 +2,7 @@ module ReliabilityOptimization
 
 using ImplicitDifferentiation, Zygote, LinearAlgebra, ChainRulesCore, SparseArrays
 using UnPack, NonconvexIpopt, Statistics, Distributions, Reexport, DistributionsAD
+using FiniteDifferences, StaticArraysCore
 @reexport using LinearAlgebra
 @reexport using Statistics
 export RandomFunction, FORM, RIA, MvNormal
@@ -15,9 +16,14 @@ end
 struct FORM{M}
     method::M
 end
-struct RIA end
+struct RIA{A, O}
+    optim_alg::A
+    optim_options::O
+end
+RIA() = RIA(IpoptAlg(), IpoptOptions(print_level = 0, max_wall_time = 1.0))
 
-function get_forward(f, p, ::FORM{<:RIA})
+function get_forward(f, p, method::FORM{<:RIA})
+    alg, options = method.method.optim_alg, method.method.optim_options
     function forward(x)
         # gets an objective function of p
         obj = pc -> begin
@@ -39,9 +45,9 @@ function get_forward(f, p, ::FORM{<:RIA})
         add_eq_constraint!(innerOptModel, constr)
         result = optimize(
             innerOptModel,
-            IpoptAlg(),
+            alg,
             [mean(p); 0.0],
-            options = IpoptOptions(print_level = 0),
+            options = options,
         )
         return vcat(result.minimizer, result.problem.mult_g[1])
     end
@@ -54,7 +60,7 @@ function get_conditions(f, ::FORM{<:RIA})
         c = pcmult[end-1]
         mult = pcmult[end]
         return vcat(
-            2 * p + Zygote.pullback(p -> f(x, p), p)[2](mult)[1],
+            2 * p + vec(_jacobian(f, x, p)) * mult,
             2c - mult,
             f(x, p) .- c,
         )
@@ -79,23 +85,103 @@ 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')
+function get_identity_vecs(M)
+    return [Vector(sparsevec([i], [1.0], M)) for i in 1:M]
+end
+function ChainRulesCore.rrule(::typeof(get_identity_vecs), M::Int)
+    get_identity_vecs(M), _ -> (NoTangent(), NoTangent())
+end
+reduce_hcat(vs) = reduce(hcat, vs)
+# function ChainRulesCore.rrule(::typeof(reduce_hcat), vs::Vector{<:Vector})
+#     return reduce_hcat(vs), Δ -> begin
+#         return NoTangent(), [Δ[:, i] for i in 1:size(Δ, 2)]
+#     end
+# end
+
+const fdm = FiniteDifferences.central_fdm(5, 1)
+
+function _jacobian(f, x1, x2)
+    # val, pb = Zygote.pullback(f, x1, x2)
+    # if val isa Vector
+    #     M = length(val)
+    #     vecs = get_identity_vecs(M)
+    #     cotangents = map(pb, vecs)
+    #     Jt = reduce_hcat(map(last, cotangents))
+    #     return copy(Jt')
+    # elseif val isa Real
+    #     Jt = last(pb(1.0))
+    #     return copy(Jt')
+    # else
+    #     throw(ArgumentError("Output type not supported."))
+    # end
+    ẏs = map(eachindex(x2)) do n
+        return fdm(zero(eltype(x2))) do ε
+            xn = x2[n]
+            xcopy = vcat(x2[1:n-1], xn + ε, x2[n+1:end])
+            ret = copy(f(x1, xcopy))  # copy required incase `f(x)` returns something that aliases `x`
+            return ret
+        end
+    end
+    return reduce(hcat, ẏs)
 end
+# function ChainRulesCore.rrule(::typeof(_jacobian), f, x1, x2)
+#     (val, pb), _pb_pb = Zygote.pullback(Zygote.pullback, f, x1, x2)
+#     M = length(val)
+#     if val isa Vector
+#         vecs = get_identity_vecs(M)
+#         _pb = (pb, v) -> last(pb(v))
+#         co1, pb_pb = Zygote.pullback(_pb, pb, first(vecs))
+#         cotangents = vcat([co1], last.(map(pb, @view(vecs[2:end]))))
+#         Jt, hcat_pb = Zygote.pullback(reduce_hcat, cotangents)
+#         return copy(Jt'), Δ -> begin
+#             temp = hcat_pb(Δ')[1]
+#             co_pb = map(temp) do t
+#                 first(pb_pb(t))
+#             end
+#             co_f_x = _pb_pb.(tuple.(Ref(nothing), co_pb))
+#             co_f = sum(getindex.(co_f_x, 1))
+#             co_x1 = sum(getindex.(co_f_x, 2))
+#             co_x2 = sum(getindex.(co_f_x, 3))
+#             return NoTangent(), co_f, co_x1, co_x2
+#         end
+#     elseif val isa Real
+#         println(1)
+#         _pb = (pb, v) -> pb(v)[end]
+#         println(2)
+#         @show _pb(pb, 1.0)
+#         Jt, pb_pb = Zygote.pullback(_pb, pb, 1.0)
+#         println(3)
+#         return copy(Jt'), Δ -> begin
+#             println(4)
+#             @show vec(Δ)
+#             @show Δ
+#             @show pb_pb(Δ')
+#             co_pb = first(pb_pb(vec(Δ)))
+#             co_f_x = _pb_pb((nothing, co_pb))
+#             co_f = co_f_x[1]
+#             co_x1 = co_f_x[2]
+#             co_x2 = co_f_x[3]
+#             return NoTangent(), co_f, co_x1, co_x2
+#         end
+#     else
+#         throw(ArgumentError("Output type not supported."))
+#     end
+# end
 
 function (f::RandomFunction)(x)
     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)
+    dfdp0 = _jacobian(f.f, x, p0)
     fp0 = f.f(x, p0)
-    muf = _vec(fp0) + dfdp0 * (mup - p0)
+    muf = _vec(fp0) .+ dfdp0 * (mup - p0)
     covf = dfdp0 * covp * dfdp0'
     return MvNormal(muf, covf)
 end
 
+# necessary type piracy FiniteDifferences._estimate_magnitudes uses this constructor which Zygote struggles to differentiate on its own
+function ChainRulesCore.rrule(::typeof(StaticArraysCore.SVector{3}), x1::T, x2::T, x3::T) where {T}
+    StaticArraysCore.SVector{3}(x1, x2, x3), Δ -> (NoTangent(), Δ[1], Δ[2], Δ[3])
 end
+
+end

test/example.jl

@@ -0,0 +1,54 @@
+using ReliabilityOptimization, Test, NonconvexTOBS, ChainRulesCore, TopOpt, Zygote, FiniteDifferences
+
+const densities = [0.0, 0.5, 1.0] # for mass calculation
+const nmats = 3 # number of materials
+const v = 0.3 # Poisson’s ratio
+const f = 1.0 # downward force
+const problemSize = (4, 4) # size of rectangular mesh
+const elSize = (1.0, 1.0) # size of QUAD4 elements
+# Point load cantilever problem to be solved
+problem = PointLoadCantilever(Val{:Linear}, problemSize, elSize, 1.0, v, f)
+ncells = TopOpt.getncells(problem) # Number of elements in mesh
+solver = FEASolver(Direct, problem; xmin = 0.0)
+filter = DensityFilter(solver; rmin = 3.0) # filter to avoid checkerboarding
+M = 1 / nmats / 2 # mass fraction
+comp = Compliance(solver) # function that returns compliance
+penalty = TopOpt.PowerPenalty(3.0) # SIMP penalty
+# Young’s modulii of air (void) + 2 materials
+const avgEs = [1e-6, 0.5, 2.0]
+# since first E is close to zero,
+# can the deviation make the final value negative?
+logEs = MvNormal(log.(avgEs), Matrix(Diagonal(0.1 .* abs.(log.(avgEs)))))
+# 'Original' function. At least one input is random.
+# In this example, Es is the random input.
+function uncertainComp(x, logEs)
+  Es = exp.(logEs)
+  # interpolation of properties between materials
+  interp = MaterialInterpolation(Es, penalty)
+  MultiMaterialVariables(x, nmats) |> interp |> filter |> comp
+  # return sum(x) + sum(Es)
+end
+# wrap original function in RandomFunction struct
+rf = RandomFunction(uncertainComp, logEs, FORM(RIA()))
+# initial homogeneous distribution of pseudo-densities
+x0 = fill(M, ncells * (length(logEs) - 1))
+# call wrapper with example input
+# (Returns propability distribution of the objective for current point)
+d = rf(x0)
+# mass constraint
+constr = x -> begin
+    ρs = PseudoDensities(MultiMaterialVariables(x, nmats))
+    return sum(element_densities(ρs, densities)) / ncells - 0.3 # unit element volume
+end
+function obj(x) # objective for TO problem
+  dist = rf(x)
+  mean(dist)[1] + 2 * sqrt(cov(dist)[1, 1])
+end
+obj(x0)
+Zygote.gradient(obj, x0)
+FiniteDifferences.grad(central_fdm(5, 1), obj, x0)[1]
+
+m = Model(obj) # create optimization model
+addvar!(m, zeros(length(x0)), ones(length(x0))) # setup optimization variables
+Nonconvex.add_ineq_constraint!(m, constr) # setup volume inequality constraint
+@time r = Nonconvex.optimize(m, TOBSAlg(), x0; options = TOBSOptions())