From 5b5a5badeb529b3144f2f4336435e43f5f20e87f Mon Sep 17 00:00:00 2001
From: Luke Morris <lukelukemorrismorris@gmail.com>
Date: Mon, 24 Jun 2024 14:00:02 -0400
Subject: [PATCH] Do not always compose intermediates by default

---
 src/composition.jl  | 28 +++++++++++++++++++---------
 test/composition.jl | 29 +++++++++++++++++++++++++++--
 2 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/src/composition.jl b/src/composition.jl
index 2745af0..587c1ad 100644
--- a/src/composition.jl
+++ b/src/composition.jl
@@ -220,31 +220,41 @@ oapply(r::RelationDiagram, pode::OpenSummationDecapode) = oapply(r, [pode])
 # Default composition
 # -------------------
 
+# TODO: Upstream this to Catlab?
+function construct_relation_diagram(boxes::Vector{Symbol}, junctions::Vector{Vector{Symbol}})
+  tables = map(boxes, junctions) do b, j
+    Expr(:call, b, j...)
+  end
+  quote @relation () begin $(tables...) end end |> eval
+end
+
 # TODO: Add a macro which provides names for boxes via the Symbol of the Decapode.
 """    function default_composition_diagram(podes::Vector{D}, names::Vector{Symbol}) where {D<:SummationDecapode}
 
 Given a list of Decapodes and their names, return a composition diagram which assumes that variables sharing the same name ought to be composed.
 
-Only variables which are found in multiple Decapodes are exposed. No Literals are exposed.
+No Literals are exposed. Use [`unique_lits!`](@ref) after composing.
 
 Throw an error if any individual Decapode already contains a repeated name (except for Literals).
 
+If `only_states_terminals` is `true`, only expose state and terminal variables. Defaults to `false`.
+
 Note that composing immediately with [`oapply`](@ref) will fail if types do not match (e.g. (:infer, :Form0) or (:Form0, :Form1)). So, use the function [`infer_types_from_diagram!`](@ref) after this if needed.
 """
-function default_composition_diagram(podes::Vector{D}, names::Vector{Symbol}) where {D<:SummationDecapode}
+function default_composition_diagram(podes::Vector{D}, names::Vector{Symbol}, only_states_terminals=false) where {D<:SummationDecapode}
   non_lit_names = map(podes) do pode
     pode[findall(!=(:Literal), pode[:type]), :name]
   end
   for (nln, name) in zip(non_lit_names, names)
-    allunique(nln) || error("Decapode $name contains a repeated variable name.")
+    allunique(nln) || error("Decapode $name contains repeated variable names: $nln.")
   end
-  all_names = union(non_lit_names...)
-  tables = map(names, non_lit_names) do name, nln
-    Expr(:call, name, intersect(all_names, nln)...)
+  if only_states_terminals
+    foreach(non_lit_names, podes) do nln, pode
+      outers = infer_state_names(pode) ∪ pode[infer_terminals(pode), :name]
+      filter!(x -> x ∈ outers, nln)
+    end
   end
-  # XXX: The only means of creating a RelationDiagram is via the DSL or
-  # tracking your own indices via @acset/ imperatively.
-  quote @relation () begin $(tables...) end end |> eval
+  construct_relation_diagram(names, non_lit_names)
 end
 
 """    function infer_types_from_diagram!(r::RelationDiagram, podes::Vector{D}) where {D<:SummationDecapode}
diff --git a/test/composition.jl b/test/composition.jl
index 0f42425..290c72b 100644
--- a/test/composition.jl
+++ b/test/composition.jl
@@ -222,8 +222,33 @@ end
       Advection(C,V,ϕ₂)
       Superposition(C,ϕ₁,ϕ₂,Ċ,ϕ)
     end
-  @test default_composition_diagram(
-    [Diffusion, Advection, Superposition], [:Diffusion, :Advection, :Superposition]) == expected
+  @test is_isomorphic(expected, default_composition_diagram(
+          [Diffusion, Advection, Superposition],
+          [:Diffusion, :Advection, :Superposition]))
+
+  # Compose Halfar's equation with Glen's law.
+  GlensLaw = @decapode begin
+    Γ::Form1
+    (A,ρ,g,n)::Constant
+
+    Γ == (2/(n+2))*A*(ρ*g)^n
+  end
+  HalfarsEquation = @decapode begin
+    h::Form0
+    Γ::Form1
+    n::Constant
+
+    ∂ₜ(h) == ∘(⋆, d, ⋆)(Γ  * d(h) ∧ (mag(♯(d(h)))^(n-1)) ∧ (h^(n+2)))
+  end
+  expected =
+    @relation () begin
+      GlensLaw(Γ,A,ρ,g,n)
+      HalfarsEquation(h,Γ,n,ḣ)
+    end
+  @test is_isomorphic(expected, default_composition_diagram(
+          [GlensLaw, HalfarsEquation],
+          [:GlensLaw, :HalfarsEquation],
+          true))
 end
 
 # end