From baaadd09ced08bbbbdf150e6a67cbcbea68205c5 Mon Sep 17 00:00:00 2001 From: Sam Cohen Date: Tue, 9 Jul 2024 12:25:00 -0400 Subject: [PATCH] feat: optimizers with vector functions --- src/FinSetAlgebras.jl | 2 +- src/Objectives.jl | 153 ++------------------------------------ src/Optimizers.jl | 169 +++++++++--------------------------------- test/Objectives.jl | 13 +--- test/Optimizers.jl | 86 +++------------------ 5 files changed, 55 insertions(+), 368 deletions(-) diff --git a/src/FinSetAlgebras.jl b/src/FinSetAlgebras.jl index a173d93..135c791 100644 --- a/src/FinSetAlgebras.jl +++ b/src/FinSetAlgebras.jl @@ -84,7 +84,7 @@ portmap(obj::Open{T}) where T = obj.m # Helper function for when m is identity. function Open{T}(o::T) where T - Open{T}(domain(o), o, id(domain(o))) + Open{T}(dom(o), o, id(dom(o))) end dom(obj::Open{T}) where T = dom(obj.m) diff --git a/src/Objectives.jl b/src/Objectives.jl index 194b0a2..96a6cf4 100644 --- a/src/Objectives.jl +++ b/src/Objectives.jl @@ -1,7 +1,7 @@ module Objectives export PrimalObjective, MinObj, gradient_flow, - SaddleObjective, DualComp, primal_solution, dual_objective, primal_objective, pullback_function + SaddleObjective, DualComp, primal_solution, dual_objective, primal_objective using ..FinSetAlgebras import ..FinSetAlgebras: hom_map, laxator @@ -11,6 +11,8 @@ import Catlab: oapply, dom using ForwardDiff using Optim + + # Primal Minimization Problems and Gradient Descent ################################################### @@ -38,7 +40,7 @@ struct MinObj <: FinSetAlgebra{PrimalObjective} end The morphism map is defined by ϕ ↦ (f ↦ f∘ϕ^*). """ hom_map(::MinObj, ϕ::FinFunction, p::PrimalObjective) = - PrimalObjective(codom(ϕ), x->p(test_pullback_function(ϕ, x))) + PrimalObjective(codom(ϕ), x->p(pullback_function(ϕ, x))) """ laxator(::MinObj, Xs::Vector{PrimalObjective}) @@ -46,7 +48,7 @@ Takes the "disjoint union" of a collection of primal objectives. """ function laxator(::MinObj, Xs::Vector{PrimalObjective}) c = coproduct([dom(X) for X in Xs]) - subproblems = [x -> X(test_pullback_function(l, x)) for (X,l) in zip(Xs, legs(c))] + subproblems = [x -> X(pullback_function(l, x)) for (X,l) in zip(Xs, legs(c))] objective(x) = sum([sp(x) for sp in subproblems]) return PrimalObjective(apex(c), objective) end @@ -96,14 +98,14 @@ struct DualComp <: FinSetAlgebra{SaddleObjective} end # Only "glue" along dual variables hom_map(::DualComp, ϕ::FinFunction, p::SaddleObjective) = SaddleObjective(p.primal_space, codom(ϕ), - (x,λ) -> p(x, test_pullback_function(ϕ, λ))) + (x,λ) -> p(x, pullback_function(ϕ, λ))) # Laxate along both primal and dual variables function laxator(::DualComp, Xs::Vector{SaddleObjective}) c1 = coproduct([X.primal_space for X in Xs]) c2 = coproduct([X.dual_space for X in Xs]) subproblems = [(x,λ) -> - X(test_pullback_function(l1, x), test_pullback_function(l2, λ)) for (X,l1,l2) in zip(Xs, legs(c1), legs(c2))] + X(pullback_function(l1, x), pullback_function(l2, λ)) for (X,l1,l2) in zip(Xs, legs(c1), legs(c2))] objective(x,λ) = sum([sp(x,λ) for sp in subproblems]) return SaddleObjective(apex(c1), apex(c2), objective) end @@ -123,145 +125,4 @@ function gradient_flow(of::Open{SaddleObjective}) λ -> ForwardDiff.gradient(dual_objective(f, x(λ)), λ), of.m) end - - - - - -# New stuff - -struct UnivTypedPrimalObjective - decision_space::FinSet - objective::Function # R^ds -> R NOTE: should be autodifferentiable - type::FinDomFunction # R^ds -> Z+, for our use case. One function to rule them all--this will stay the same across our finsets. FinDomFunction -end - -(p::UnivTypedPrimalObjective)(x::Vector) = p.objective(x) -dom(p::UnivTypedPrimalObjective) = p.decision_space - -""" MinObj - -# Finset-algebra implementing composition of minimization problems by variable sharing. -# """ -struct MinObj <: FinSetAlgebra{PrimalObjective} end - -""" hom_map(::MinObj, ϕ::FinFunction, p::PrimalObjective) - -The morphism map is defined by ϕ ↦ (f ↦ f∘ϕ^*). -""" -hom_map(::MinObj, ϕ::FinFunction, p::UnivTypedPrimalObjective) = # Another version, which wouldn't require a universal type function, would have you pass in a custom type function for your set M. This would require more work in the laxator to take the disjoint union of type functions. - all(p.type(x) == p.type(ϕ(x)) for x in dom(p)) ? - UnivTypedPrimalObjective(codom(ϕ), x -> p(test_pullback_function(ϕ, x)), p.type) : - error("The ϕ provided is not type-preserving.") # throw an error - - - -""" laxator(::MinObj, Xs::Vector{PrimalObjective}) - -Takes the "disjoint union" of a collection of primal objectives. -""" -function laxator(::MinObj, Xs::Vector{UnivTypedPrimalObjective}) - c = coproduct([dom(X) for X in Xs]) - subproblems = [x -> X(test_pullback_function(l, x)) for (X, l) in zip(Xs, legs(c))] - objective(x) = sum([sp(x) for sp in subproblems]) - return UnivTypedPrimalObjective(apex(c), objective, Xs[1].type) # Assuming all have the same type function -end - - - - - - -struct TypedPrimalObjective # Should we put restrictions on the constructor, i.e. check that dom(objective) = dom(type) = decision_space? - decision_space::FinSet - objective::Function # R^ds -> R NOTE: should be autodifferentiable. Make it a FinDomFunction? - type::FinDomFunction -end - -(p::TypedPrimalObjective)(x::Vector) = p.objective(x) -dom(p::TypedPrimalObjective) = p.decision_space - - -hom_map(::MinObj, ϕ::FinFunction, σ::FinDomFunction, τ::FinDomFunction, p::TypedPrimalObjective) = # τ seems completely unnecessary to me - all(p.type(x) == σ(ϕ(x)) && p.type(x) == τ(x) for x in dom(p)) ? - UnivTypedPrimalObjective(codom(ϕ), x -> p(test_pullback_function(ϕ, x)), σ) : # Note: we didn't check but σ must be applicable across all of codom(ϕ) - nothing # throw an error - - - -function laxator(::MinObj, Xs::Vector{TypedPrimalObjective}) - combinedType = copair([X.type for X in Xs]) - c = dom(combinedType) # c is the coproduct of the decision spaces - objective(x) = sum(X(test_pullback_function(l, x)) for (X, l) in zip(Xs, legs(c))) # Calculated the same as before (but simplified onto one line) - return UnivTypedPrimalObjective(apex(c), objective, combinedType) -end - - - - -# Pullback function for a given ϕ and vector v -# function pullback(ϕ::FinFunction, v::Vector) -# output = Vector{eltype(v)}(undef, length(dom(ϕ))) -# for i in 1:length(dom(ϕ)) -# output[i] = v[ϕ(i)] -# end -# return output -# end - -# Optional easier version: return Vector{eltype(v)}(v[ϕ(i)] for i in 1:length(dom(ϕ))) - - - -# Curried version of the pullback function -function curried_pullback(ϕ::FinFunction) - return function (v::Vector) - output = Vector{eltype(v)}(undef, length(dom(ϕ))) - for i in 1:length(dom(ϕ)) - output[i] = v[ϕ(i)] - end - return output - end -end - - - -function pullback_function(ϕ::FinFunction, v::Vector) - return v[ϕ.(1:length(dom(ϕ)))] # Broadcasting with vector of indices - end - - - - -# Not the active one -function typed_pullback_matrix(f::FinFunction) # Modify code - - # Track codomain indices - prefixes = Dict() - lastPrefix = 0 - - for v in codom(f) - prefixes[v] = lastPrefix - lastPrefix += length(v) - end - # lastPrefix now holds the sum of the sizes across all the output vectors - - domLength = 0 - result = [] - for v in dom(f) - domLength += length(v) - for i in 1:length(v) - push!(result, prefixes[f(v)] + i) - end - end - - sparse(1:domLength, result, ones(Int, domLength), domLength, lastPrefix) -end - - - - - - - - end # module \ No newline at end of file diff --git a/src/Optimizers.jl b/src/Optimizers.jl index 5850849..d238bfc 100644 --- a/src/Optimizers.jl +++ b/src/Optimizers.jl @@ -1,30 +1,14 @@ # Implement the cospan-algebra of dynamical systems. module Optimizers -export pullback_matrix, pushforward_matrix, Optimizer, OpenContinuousOpt, OpenDiscreteOpt, Euler, - simulate, typed_pullback_matrix, test_pullback_function, test_pushforward_function +export Optimizer, OpenContinuousOpt, OpenDiscreteOpt, Euler, + simulate, pullback_function, pushforward_function using ..FinSetAlgebras import ..FinSetAlgebras: hom_map, laxator using Catlab import Catlab: oapply, dom -using SparseArrays -""" pullback_matrix(f::FinFunction) - -The pullback of f : n → m is the linear map f^* : Rᵐ → Rⁿ defined by -f^*(y)[i] = y[f(i)]. -""" -function pullback_matrix(f::FinFunction) - n = length(dom(f)) - sparse(1:n, f.(dom(f)), ones(Int, n), dom(f).n, codom(f).n) -end - -""" pushforward_matrix(f::FinFunction) - -The pushforward is the dual of the pullback. -""" -pushforward_matrix(f::FinFunction) = pullback_matrix(f)' """ Optimizer @@ -47,7 +31,7 @@ struct DiscreteOpt <: FinSetAlgebra{Optimizer} end The hom map is defined as ϕ ↦ (s ↦ ϕ_*∘s∘ϕ^*). """ hom_map(::ContinuousOpt, ϕ::FinFunction, s::Optimizer) = - Optimizer(codom(ϕ), x -> test_pushforward_function(ϕ, s(test_pullback_function(ϕ, x)))) + Optimizer(codom(ϕ), x -> pushforward_function(ϕ, s(pullback_function(ϕ, x)))) """ hom_map(::DiscreteOpt, ϕ::FinFunction, s::Optimizer) @@ -55,8 +39,8 @@ The hom map is defined as ϕ ↦ (s ↦ id + ϕ_*∘(s - id)∘ϕ^*). """ hom_map(::DiscreteOpt, ϕ::FinFunction, s::Optimizer) = Optimizer(codom(ϕ), x -> begin - y = test_pullback_function(ϕ, x) - return x + test_pushforward_function(ϕ, (s(y) - y)) + y = pullback_function(ϕ, x) + return x + pushforward_function(ϕ, (s(y) - y)) end) """ laxator(::ContinuousOpt, Xs::Vector{Optimizer}) @@ -65,7 +49,7 @@ Takes the "disjoint union" of a collection of optimizers. """ function laxator(::ContinuousOpt, Xs::Vector{Optimizer}) c = coproduct([dom(X) for X in Xs]) - subsystems = [x -> X(test_pullback_function(l, x)) for (X, l) in zip(Xs, legs(c))] + subsystems = [x -> X(pullback_function(l, x)) for (X, l) in zip(Xs, legs(c))] function parallel_dynamics(x) res = Vector{Vector}(undef, length(Xs)) # Initialize storage for results for i = 1:length(Xs) #=Threads.@threads=# @@ -78,7 +62,14 @@ end # Same as continuous opt laxator(::DiscreteOpt, Xs::Vector{Optimizer}) = laxator(ContinuousOpt(), Xs) + Open{Optimizer}(S::FinSet, v::Function, m::FinFunction) = Open{Optimizer}(S, Optimizer(S, v), m) +Open{Optimizer}(s::Int, v::Function, m::FinFunction) = Open{Optimizer}(FinSet(s), v, m) + +# Special cases: m is an identity +Open{Optimizer}(S::FinSet, v::Function) = Open{Optimizer}(S, Optimizer(S, v), id(S)) +Open{Optimizer}(s::Int, v::Function) = Open{Optimizer}(FinSet(s), v) + # Turn into cospan-algebras. struct OpenContinuousOpt <: CospanAlgebra{Open{Optimizer}} end @@ -95,22 +86,11 @@ end # Euler's method is a natural transformation from continous optimizers to discrete optimizers. function Euler(f::Open{Optimizer}, γ::Float64) return Open{Optimizer}(f.S, - Optimizer(f.S, x -> begin println(x); println(f.o(x)); println("hey there buddy"); println(f.o.dynamics); - x .+ γ .* f.o(x) - end), f.m) + Optimizer(f.S, x -> x .+ γ .* f.o(x)), f.m) end # Run a discrete optimizer the designated number of time-steps. -function simulate(f::Open{Optimizer}, x0::Vector{Float64}, tsteps::Int) - res = x0 - for i in 1:tsteps - res = f.o(res) - end - return res -end - - -function simulate(f::Open{Optimizer}, x0::Vector{Vector{Float64}}, tsteps::Int) +function simulate(f::Open{Optimizer}, x0::Vector, tsteps::Int) res = x0 for i in 1:tsteps res = f.o(res) @@ -119,113 +99,43 @@ function simulate(f::Open{Optimizer}, x0::Vector{Vector{Float64}}, tsteps::Int) end +""" pullback_function(f::FinFunction, v::Vector) - - - -function typed_pullback_matrix(f, domType, codomType) # Modify code - # Track codomain indices - prefixes = Dict() - lastPrefix = 0 - - for v in codom(f) # Assumes codom(f) is a set and has distinct elements - prefixes[v] = lastPrefix - lastPrefix += codomType(v) - end - # lastPrefix now holds the sum of the sizes across all the output vectors - - domLength = 0 - result = [] - for v in dom(f) - domLength += domType(v) - for i in 1:domType(v) - push!(result, prefixes[f(v)] + i) - end - end - - sparse(1:domLength, result, ones(Int, domLength), domLength, lastPrefix) -end - - - -function typed_pullback_matrix(f) # No types provided, so assume everything uses scalars - - # println("inside typed") - - domType = FinFunction(ones(Int, length(dom(f)))) - codomType = FinFunction(ones(Int, length(codom(f)))) - - # Track codomain indices - prefixes = Dict() - lastPrefix = 0 - - for v in codom(f) # Assumes codom(f) is a set and has distinct elements - prefixes[v] = lastPrefix - lastPrefix += codomType(v) - end - # lastPrefix now holds the sum of the sizes across all the output vectors - - domLength = 0 - result = [] - for v in dom(f) - domLength += domType(v) - for i in 1:domType(v) - push!(result, prefixes[f(v)] + i) - end - end - - sparse(1:domLength, result, ones(Int, domLength), domLength, lastPrefix) -end - - -typed_pushforward_matrix(f::FinFunction) = typed_pullback_matrix(f)' -typed_pushforward_matrix(f::FinFunction, domType, codomType) = typed_pullback_matrix(f, domType, codomType)' - - - -function test_pullback_function(f::FinFunction, v::Vector)::Vector +The pullback of f : n → m is the linear map f^* : Rᵐ → Rⁿ defined by +f^*(y)[i] = y[f(i)]. +""" +function pullback_function(f::FinFunction, v::Vector)::Vector return [v[f(i)] for i in 1:length(dom(f))] end -function test_pushforward_function(f::FinFunction, v)::Vector - # output = zeros(length(codom(f))) - output = [[] for _ in 1:length(codom(f))] - - # output = Vector{Vector}(nothing, length(codom(f))) - - # println(output) +""" pushforward_matrix(f::FinFunction, v::Vector{Vector{Float64}}) +The pushforward of f : n → m is the linear map f_* : Rⁿ → Rᵐ defined by +f_*(y)[j] = ∑ y[i] for i ∈ f⁻¹(j). +""" +function pushforward_function(f::FinFunction, v::Vector{Vector{Float64}})::Vector + output = [[] for _ in 1:length(codom(f))] for i in 1:length(dom(f)) - # println("i = ", i) - # println(output) - # println(output[f(i)]) - # println() if isempty(output[f(i)]) output[f(i)] = v[i] else output[f(i)] += v[i] end end - return output end -function test_pushforward_function(f::FinFunction, v::Vector{Float64})::Vector - # output = zeros(length(codom(f))) - output = [0.0 for _ in 1:length(codom(f))] - - - # output = Vector{Vector}(nothing, length(codom(f))) +""" pushforward_matrix(f::FinFunction, v) - # println(output) +The pushforward of f : n → m is the linear map f_* : Rⁿ → Rᵐ defined by +f_*(y)[j] = ∑ y[i] for i ∈ f⁻¹(j). +""" +function pushforward_function(f::FinFunction, v::Vector{Float64})::Vector + output = [0.0 for _ in 1:length(codom(f))] for i in 1:length(dom(f)) - # println("i = ", i) - # println(output) - # println(output[f(i)]) - # println() output[f(i)] += v[i] end @@ -233,17 +143,4 @@ function test_pushforward_function(f::FinFunction, v::Vector{Float64})::Vector end - - - - - - - - - -end # module - - - - +end # module \ No newline at end of file diff --git a/test/Objectives.jl b/test/Objectives.jl index cdd72b1..f7ba52b 100644 --- a/test/Objectives.jl +++ b/test/Objectives.jl @@ -9,6 +9,7 @@ d = @relation (x,y,z) begin h(u,w,z) end + P = rand(-1:0.01:1,5,5) P = P'*P Q = rand(-1:0.01:1,3,3) @@ -42,14 +43,4 @@ tsteps = 1000 r1 = simulate(dc1, x0, tsteps) r2 = simulate(dc2, x0, tsteps) -#println(r1) -#println(r2) - -@test r1 ≈ r2 - - - - -f = FinFunction([1, 2, 2, 3]) -v = [44, 55, 66] -pullback(f, v) \ No newline at end of file +@test r1 ≈ r2 \ No newline at end of file diff --git a/test/Optimizers.jl b/test/Optimizers.jl index 9c83a9e..7198161 100644 --- a/test/Optimizers.jl +++ b/test/Optimizers.jl @@ -3,8 +3,8 @@ using AlgebraicOptimization using Catlab using Test -# Test naturality of Euler -d = @relation (x,y,z,u,w) begin # Do I need to type these variables? Also, why does this low key take forever? +# Test naturality of Euler with scalar functions +d = @relation (x,y,z,u,w) begin f(w,x) g(u,w,y) h(u,w,z) @@ -13,8 +13,6 @@ end A = rand(-1:0.01:1,5,5) B = rand(-1:0.01:1,3,3) C = rand(-1:0.01:1,4,4) -fA = v -> [v[1] * v[1], v[2] + v[3], 4, 4, v[5]/v[4]] - o1 = Open{Optimizer}(FinSet(5), x->A*x, FinFunction([2,4], 5) ) o2 = Open{Optimizer}(FinSet(3), x->B*x, id(FinSet(3))) @@ -31,48 +29,31 @@ do3 = Euler(o3, 0.01) composite_of_discretizations = oapply(OpenDiscreteOpt(), d, [do1,do2,do3]) x0 = repeat([1.0], length(composite_opt.S)) -x2 = repeat([5.0], length(composite_opt.S)) tsteps = 100 r1 = simulate(discretization_of_composites, x0, tsteps) -r3 = simulate(discretization_of_composites, x2, tsteps) - - r2 = simulate(composite_of_discretizations, x0, tsteps) -println(r3 ./ r1) - -#println(r1) -#println(r2) - @test r1 ≈ r2 +# Test naturality of Euler with vector functions +d = @relation (x,y,z,u,w,) begin + f(w,x) + g(u,w,y) + h(u,w,z) +end -# Vector functions testing -mA = v -> v -mB = v -> [ [v[1][1] * v[3][1], 10, v[1][3] / v[2][1]], [v[2][1]], [4, 5] ] +mA = v -> [[v[1][1]], [v[2][1], v[1][1]], [v[3][2], v[3][1], v[4][4]], [v[4][4], v[3][1], v[2][1], v[1][1]], [5, 5, 5, 5, 10]] +mB = v -> [[6, 6, 6, 6, 6, 6], [2, 2], [7, 7, 7, 7, 7, 7, 7]] mC = v -> v -# Why did we lose dimensions? -# Component arrays: typed by symbols (go back and forth between a vector and a collection of vector) - - o1 = Open{Optimizer}(FinSet(5), mA, FinFunction([2,4], 5)) o2 = Open{Optimizer}(FinSet(3), mB, id(FinSet(3))) o3 = Open{Optimizer}(FinSet(4), mC, FinFunction([1,3,4])) - -# Wrong domain sizes! Domain size of first function is # vars in f, codom size is # vars in the finset -# o1 = Open{Optimizer}(FinSet(5), mA, id(FinSet(5))) -# o2 = Open{Optimizer}(FinSet(3), mB, id(FinSet(3))) -# o3 = Open{Optimizer}(FinSet(4), mC, id(FinSet(4))) - - - - composite_opt = oapply(OpenContinuousOpt(), d, [o1,o2,o3]) discretization_of_composites = Euler(composite_opt, 0.01) @@ -81,56 +62,13 @@ do1 = Euler(o1, 0.01) do2 = Euler(o2, 0.01) do3 = Euler(o3, 0.01) + composite_of_discretizations = oapply(OpenDiscreteOpt(), d, [do1,do2,do3]) # x0 = repeat([1.0], length(composite_opt.S)) -x1::Vector{Vector{Float64}} = [[1], [1], [1], [1], [1], [1, 1, 1], [1], [1, 1], [1], [1], [1], [1]] # Checker would be helpful +x1::Vector{Vector{Float64}} = [[1], [1, 2], [1, 2, 3], [1, 2, 3, 4], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8, 9]] # Checker would be helpful tsteps = 100 r1 = simulate(discretization_of_composites, x1, tsteps) r2 = simulate(composite_of_discretizations, x1, tsteps) -r2 - r1 - -#println(r1) -#println(r2) @test r1 ≈ r2 - - -# Various unit tests - - -f = FinFunction([1, 2, 2, 3]) -domType = FinFunction([3, 2, 2, 1]) -codomType = FinFunction([3, 2, 1]) - -domOnes = FinFunction([1, 1, 1, 1]) -codomOnes = FinFunction([1, 1, 1]) - -typed_pullback_matrix(f, domType, codomType) - -typed_pullback_matrix(f) - -typed_pullback_matrix(f, domOnes, codomOnes) - - -f = FinFunction([1, 2, 2, 3]) -v = [[33, 34, 35], [44, 45], [55, 56], [100]] -test_pullback_function(f, v) - -test_pushforward_function(f, v) - - -wTest::Vector{Float64} = [20, 22, 33, 11] -test_pushforward_function(f, wTest) - - -parsed = @relation (x,y,z) where (x::1, y::2, z::3, w::4) begin # Why is it still using the old modules?? - R(x,w) - S(y,w) - T(z,w) -end - -draw(parsed) -draw_types(parsed) - -draw(d) \ No newline at end of file