From d4553b8068c7b0ba0b69fe2aaf8e8a6f60de38ab Mon Sep 17 00:00:00 2001
From: Matt Cuffaro <cuffaro.matt@gmail.com>
Date: Mon, 28 Oct 2024 17:09:36 -0400
Subject: [PATCH] Added steadystates dictionary keyword to `collate` (#83)

Co-authored-by: Luke Morris <70283489+lukem12345@users.noreply.github.com>
Co-authored-by: Luke Morris <lukelukemorrismorris@gmail.com>
---
 src/collages.jl  | 24 +++++++++-------
 test/collages.jl | 74 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 10 deletions(-)

diff --git a/src/collages.jl b/src/collages.jl
index 2761849..42d8148 100644
--- a/src/collages.jl
+++ b/src/collages.jl
@@ -8,12 +8,15 @@ end
 collate(c::Collage) = collate(c.src, c.tgt, c.uwd, c.symbols)
 
 # TODO: This is assuming only "restriction"-type morphisms.
-"""    function collate(equations, boundaries, uwd, symbols)
+"""    function collate(equations, boundaries, uwd, symbols; restrictions::Vector{Symbol})
 
 Create a collage of two Decapodes that simulates with boundary conditions.
-```
+
+Passing a list of symbols into `restrictions` specifies a list of boundary variables that will be restricted.
+
 """
-function collate(equations, boundaries, uwd, symbols)
+function collate(equations, boundaries, uwd, symbols; 
+        restrictions::Dict{Symbol,Symbol}=Dict{Symbol,Symbol}())
   f = SummationDecapode{Any, Any, Symbol}()
   copy_parts!(f, equations, (:Var, :TVar, :Op1, :Op2, :Σ, :Summand))
   for b in boxes(uwd)
@@ -28,9 +31,6 @@ function collate(equations, boundaries, uwd, symbols)
     var = only(incident(f, en, :name))
     en_type = equations[only(incident(equations, en, :name)), :type]
     bn_type = boundaries[only(incident(boundaries, bn, :name)), :type]
-    if en_type != bn_type
-      error("Cannot use $(string(bn)) of type $(string(bn_type)) to bound $(string(en)) of type $(string(en_type)).")
-    end
     # Add a new variable and transfer the children of the original variable to it.
     b_var = add_part!(f, :Var, type=f[var, :type], name=Symbol("r$(b)_" * string(f[var, :name])))
     transfer_children!(f, var, b_var)
@@ -47,14 +47,18 @@ function collate(equations, boundaries, uwd, symbols)
     # Insert the "masking" operation.
     gettype(xs, n) = xs[only(incident(xs, n, :name)), :type]
     newtype = @match (gettype(equations, en), gettype(boundaries, bn)) begin
-        (:infer, _) => :infer
-        (_, :infer) => :infer
-        (:Form0, :Constant) => :Form0
+        (:infer, x) || (x, :infer) => x
+        (x, :Constant) || (:Constant, x) => x
+        (x, :Parameter) || (:Parameter, x) => x
         (x, y) && if x == y end => x
-        (x, y) => error("Type mismatch between $x and $y")
+        _ => error("Cannot use $(string(bn)) of type $(string(bn_type)) to bound $(string(en)) of type $(string(en_type)).")
     end
     s_var = add_part!(f, :Var, type=newtype, name=bn)
     add_part!(f, :Op2, proj1=var, proj2=s_var, res=b_var, op2=uwd[b, :name])
+
+    if bn ∈ keys(restrictions)
+        add_part!(f, :Op1, src=only(incident(equations, restrictions[bn], :name)), tgt=s_var, op1=Symbol("restrict_" * string(bn)))
+    end
   end
 
   f
diff --git a/test/collages.jl b/test/collages.jl
index f01a817..4c9cde9 100644
--- a/test/collages.jl
+++ b/test/collages.jl
@@ -141,3 +141,77 @@ TwiceCollage = DiagrammaticEquations.collate(
   name = [:K, :J, :K̇, :r1_J, :Jb1, :r2_J, :Jb2]
 end
 
+# Test that `restrictions` adds the correct restriction edges
+DiffusionCollage = DiagrammaticEquations.collate(
+  DiffusionDynamics,
+  DiffusionBoundaries,
+  DiffusionMorphism,
+  DiffusionSymbols;
+  restrictions=Dict(:Kb1 => :K))
+
+@test DiffusionCollage == @acset SummationDecapode{Any, Any, Symbol} begin
+  Var = 8
+  TVar = 1
+  Op1 = 3
+  Op2 = 3
+  src = [1, 3, 5]
+  tgt = [7, 2, 4]
+  proj1 = [5, 1, 2]
+  proj2 = [4, 6, 8]
+  res = [3, 5, 7]
+  incl = [2]
+  op1 = Any[:∂ₜ, [:d, :⋆, :d, :⋆], :restrict_Kb1]
+  op2 = [:rb1_leftwall, :rb2_rightwall, :rb3]
+  type = [:Form0, :Form0, :Form0, :Form0, :Form0, :Form0, :Form0, :Form0]
+  name = [:K, :K̇, :r1_K, :Kb1, :r2_K, :Kb2, :r3_K̇, :Null]
+end
+
+# Parameters
+# Test simple boundary masks.
+ParamDiffusionDynamics = infer_types!(@decapode begin
+  K::Form0
+  A::Parameter
+  ∂ₜ(K) == A*(∘(d,⋆,d,⋆)(K))
+end)
+ParamDiffusionBoundaries = @decapode begin
+  (Kb1, Kb2, Null)::Form0
+  Ab::Parameter
+end
+ParamDiffusionMorphism = @relation () begin
+  rb1_leftwall(C, Cb1)
+  rb2_rightwall(C, Cb2)
+  rb3(Ċ, Zero)
+  r0(Param,BoundaryParam)
+end
+ParamDiffusionSymbols = Dict(
+  :C => :K,
+  :Ċ => :K̇,
+  :Param => :A,
+  :BoundaryParam => :Ab,
+  :Cb1 => :Kb1,
+  :Cb2 => :Kb2,
+  :Zero => :Null)
+ParamDiffusionCollage = DiagrammaticEquations.collate(
+  ParamDiffusionDynamics,
+  ParamDiffusionBoundaries,
+  ParamDiffusionMorphism,
+  ParamDiffusionSymbols)
+infer_types!(ParamDiffusionCollage)
+
+@test ParamDiffusionCollage == @acset SummationDecapode{Any, Any, Symbol} begin
+  Var = 12
+  TVar = 1
+  Op1 = 2
+  Op2 = 5
+  src = [1, 5]
+  tgt = [9, 4]
+  proj1 = [11, 7, 1, 3, 2]
+  proj2 = [4, 6, 8, 10, 12]
+  res = [3, 5, 7, 9, 11]
+  incl = [3]
+  op1 = Any[:∂ₜ, [:d, :⋆, :d, :⋆]]
+  op2 = [:*, :rb1_leftwall, :rb2_rightwall, :rb3, :r0]
+  type = [:Form0, :Parameter, :Form0, :infer, :Form0, :Form0, :Form0, :Form0, :Form0, :Form0, :Parameter, :Parameter]
+  name = [:K, :A, :K̇, Symbol("•2"), :r1_K, :Kb1, :r2_K, :Kb2, :r3_K̇, :Null, :r4_A, :Ab]
+end
+