Feature/componentarray input output

Project.toml

@@ -6,13 +6,15 @@ version = "0.0.1"
 
 [deps]
 Catlab = "134e5e36-593f-5add-ad60-77f754baafbe"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
@@ -21,7 +23,7 @@ Catlab = "0.16.10"
 ForwardDiff = "0.10.36"
 NLsolve = "4.5.1"
 Optim = "1.9.4"
+Reexport = "1.2.2"
 SparseArrays = "1.10.0"
 StatsBase = "0.34.3"
 julia = "1.10"
-Reexport = "1.2.2"

docs/literate/literate_example.jl

@@ -10,4 +10,4 @@ using AlgebraicOptimization
 #
 # We provide the `hello(string)` method which prints "Hello, `string`!"
 
-#hello("World")
+# hello("World")

src/FinSetAlgebras.jl

@@ -2,7 +2,7 @@
 # TODO: upstream into Catlab.jl
 module FinSetAlgebras
 
-export FinSetAlgebra, CospanAlgebra, Open, hom_map, laxator, data, portmap
+export FinSetAlgebra, CospanAlgebra, Open, hom_map, laxator, data, portmap, draw, draw_types
 
 using LinearAlgebra, SparseArrays
 using Catlab
@@ -82,9 +82,13 @@ end
 data(obj::Open{T}) where T = obj.o
 portmap(obj::Open{T}) where T = obj.m
 
-# Helper function for when m is identity.
+# Helper functions 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
+
+function Open{T}(S::FinSet, o::T) where T
+    Open{T}(S, o, id(dom(o)))
 end
 
 function Open{T}(o::T, m::FinFunction) where T
@@ -142,4 +146,14 @@ function oapply(CA::CospanAlgebra{Open{T}}, FA::FinSetAlgebra{T}, d::AbstractUWD
     return oapply(CA, FA, uwd_to_cospan(d), Xs)
 end
 
-end
+
+function draw(uwd)
+    to_graphviz(uwd, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len => ".75"))
+end
+
+function draw_types(uwd)   # Add better typing and error catching for if uwd is untyped
+    to_graphviz(uwd, box_labels=:name, junction_labels=:junction_type, edge_attrs=Dict(:len => ".75"))
+end
+
+
+end  # module

src/Objectives.jl

@@ -10,6 +10,8 @@ using Catlab
 import Catlab: oapply, dom
 using ForwardDiff
 using Optim
+using ComponentArrays
+
 
 # Primal Minimization Problems and Gradient Descent
 ###################################################
@@ -24,7 +26,7 @@ struct PrimalObjective
     decision_space::FinSet
     objective::Function # R^ds -> R NOTE: should be autodifferentiable
 end
-(p::PrimalObjective)(x::Vector) = p.objective(x)
+(p::PrimalObjective)(x) = p.objective(x)    # Removed x::Vector hard typing
 dom(p::PrimalObjective) = p.decision_space
 
 """     MinObj
@@ -38,15 +40,15 @@ struct MinObj <: FinSetAlgebra{PrimalObjective} end
 The morphism map is defined by ϕ ↦ (f ↦ f∘ϕ^*).
 """
 hom_map(::MinObj, ϕ::FinFunction, p::PrimalObjective) = 
-    PrimalObjective(codom(ϕ), x->p(pullback_matrix(ϕ)*x))
+    PrimalObjective(codom(ϕ), x->p(pullback_function(ϕ, x)))
 
 """     laxator(::MinObj, Xs::Vector{PrimalObjective})
 
 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(pullback_matrix(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
@@ -65,7 +67,20 @@ end
 Returns the gradient flow optimizer of a given primal objective.
 """
 function gradient_flow(f::Open{PrimalObjective})
-    return Open{Optimizer}(f.S, x -> -ForwardDiff.gradient(f.o, x), f.m)
+    function f_wrapper(ca::ComponentArray)
+        inputs = [ca[key] for key in keys(ca)]
+        f.o(inputs)    # To spread or not to spread?|
+    end
+
+    function gradient_descent(x)
+        init_conds = ComponentVector(;zip([Symbol(i) for i in eachindex(x)], x)...)
+        grad = -ForwardDiff.gradient(f_wrapper, init_conds)
+        [grad[key] for key in keys(grad)]
+    end
+
+    return Open{Optimizer}(f.S, x -> gradient_descent(x), f.m)
+
+    # return Open{Optimizer}(f.S, x -> -ForwardDiff.gradient(f.o, x), f.m)   # Scalar version
 end
 
 function solve(f::Open{PrimalObjective}, x0::Vector{Float64}, ss::Float64, n_steps::Int)
@@ -82,6 +97,15 @@ struct SaddleObjective
     objective::Function # x × λ → R
 end
 
+# struct SaddleObjective
+#     decision_space::FinSet
+#     type::decision_space -> Bool
+#     objective::Function # x × λ → R
+# end
+
+
+
+
 (p::SaddleObjective)(x,λ) = p.objective(x,λ)
 
 n_primal_vars(p::SaddleObjective) = length(p.primal_space)
@@ -101,14 +125,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, pullback_matrix(ϕ)*λ))
+        (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(pullback_matrix(l1)*x, pullback_matrix(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
@@ -128,4 +152,4 @@ function gradient_flow(of::Open{SaddleObjective})
         λ -> ForwardDiff.gradient(dual_objective(f, x(λ)), λ), of.m)
 end
 
-end
+end  # module

src/OpenFlowGraphs.jl

@@ -49,7 +49,7 @@ struct FG <: FinSetAlgebra{FlowGraph} end
 hom_map(::FG, ϕ::FinFunction, g::FlowGraph) =
     FlowGraph(codom(ϕ), g.edges, 
         g.src⋅ϕ, g.tgt⋅ϕ, 
-        g.edge_costs, pushforward_matrix(ϕ)*g.flows)
+        g.edge_costs, pushforward_function(ϕ)(g.flows))
 
 function laxator(::FG, gs::Vector{FlowGraph})
     laxed_src = reduce(⊕, [g.src for g in gs])