Expression level rewriting

Project.toml

@@ -2,7 +2,7 @@ name = "DiagrammaticEquations"
 uuid = "6f00c28b-6bed-4403-80fa-30e0dc12f317"
 license = "MIT"
 authors = ["James Fairbanks", "Andrew Baas", "Evan Patterson", "Luke Morris", "George Rauta"]
-version = "0.1.6"
+version = "0.1.7"
 
 [deps]
 ACSets = "227ef7b5-1206-438b-ac65-934d6da304b8"

src/DiagrammaticEquations.jl

@@ -2,6 +2,8 @@
 """
 module DiagrammaticEquations
 
+using Catlab
+
 export
 DerivOp, append_dot, normalize_unicode, infer_states, infer_types!,
 # Deca
@@ -12,6 +14,7 @@ recursive_delete_parents, spacename, varname, unicode!, vec_to_dec!,
 Collage, collate,
 ## composition
 oapply, unique_by, unique_by!, OpenSummationDecapodeOb, OpenSummationDecapode, Open, default_composition_diagram,
+apex, @relation, # Re-exported from Catlab
 ## acset
 SchDecapode, SchNamedDecapode, AbstractDecapode, AbstractNamedDecapode, NamedDecapode, SummationDecapode,
 contract_operators!, contract_operators, add_constant!, add_parameter, fill_names!, dot_rename!, is_expanded, expand_operators, infer_state_names, infer_terminal_names, recognize_types,
@@ -25,12 +28,12 @@ unique_lits!,
 Plus, AppCirc1, Var, Tan, App1, App2,
 ## visualization
 to_graphviz_property_graph, typename, draw_composition,
+to_graphviz, # Re-exported from Catlab
 ## rewrite
 average_rewrite,
 ## openoperators
 transfer_parents!, transfer_children!, replace_op1!, replace_op2!, replace_all_op1s!, replace_all_op2s!
 
-using Catlab
 using Catlab.Theories
 import Catlab.Theories: otimes, oplus, compose, ⊗, ⊕, ⋅, associate, associate_unit, Ob, Hom, dom, codom
 using Catlab.Programs
@@ -62,11 +65,14 @@ include("pretty.jl")
 include("colanguage.jl")
 include("openoperators.jl")
 include("symbolictheoryutils.jl")
+include("graph_traversal.jl")
 include("deca/Deca.jl")
 include("learn/Learn.jl")
 include("SymbolicUtilsInterop.jl")
 
 @reexport using .Deca
 @reexport using .SymbolicUtilsInterop
 
+include("acset2symbolic.jl")
+
 end

src/SymbolicUtilsInterop.jl

@@ -1,6 +1,8 @@
 module SymbolicUtilsInterop
 
-using ..DiagrammaticEquations: AbstractDecapode, Quantity
+using ACSets
+using ..DiagrammaticEquations: AbstractDecapode, Quantity, DerivOp
+using ..DiagrammaticEquations: recognize_types, fill_names!, make_sum_mult_unique!
 import ..DiagrammaticEquations: eval_eq!, SummationDecapode
 using ..decapodes
 using ..Deca
@@ -14,6 +16,7 @@ struct SymbolicEquation{E}
     lhs::E
     rhs::E
 end
+export SymbolicEquation
 
 Base.show(io::IO, e::SymbolicEquation) = begin
     print(io, e.lhs); print(io, " == "); print(io, e.rhs)
@@ -48,7 +51,7 @@ function decapodes.Term(t::SymbolicUtils.BasicSymbolic)
             decapodes.Plus(termargs)
         elseif op == *
             decapodes.Mult(termargs)
-        elseif op == ∂ₜ
+        elseif op ∈ [DerivOp, ∂ₜ]
             decapodes.Tan(only(termargs))
         elseif length(args) == 1
             decapodes.App1(nameof(op, symtype.(args)...), termargs...)
@@ -82,9 +85,9 @@ Example:
   SymbolicUtils.BasicSymbolic(context, Term(a))
 ```
 """
-function SymbolicUtils.BasicSymbolic(context::Dict{Symbol,DataType}, t::decapodes.Term, __module__=@__MODULE__)
+function SymbolicUtils.BasicSymbolic(context::Dict{Symbol,DataType}, t::decapodes.Term)
     # user must import symbols into scope
-    ! = (f -> getfield(__module__, f))
+    ! = (f -> getfield(@__MODULE__, f))
     @match t begin
         Var(name) => SymbolicUtils.Sym{context[name]}(name)
         Lit(v) => Meta.parse(string(v))
@@ -95,17 +98,17 @@ function SymbolicUtils.BasicSymbolic(context::Dict{Symbol,DataType}, t::decapode
             # see test/language.jl
             (f, x) -> (!(f))(x),
             fs;
-            init=BasicSymbolic(context, arg, __module__)
+            init=BasicSymbolic(context, arg)
         )
-        App1(f, x) => (!(f))(BasicSymbolic(context, x, __module__))
-        App2(f, x, y) => (!(f))(BasicSymbolic(context, x, __module__), BasicSymbolic(context, y, __module__))
-        Plus(xs) => +(BasicSymbolic.(Ref(context), xs, Ref(__module__))...)
-        Mult(xs) => *(BasicSymbolic.(Ref(context), xs, Ref(__module__))...)
-        Tan(x) => ∂ₜ(BasicSymbolic(context, x, __module__))
+        App1(f, x) => (!(f))(BasicSymbolic(context, x))
+        App2(f, x, y) => (!(f))(BasicSymbolic(context, x), BasicSymbolic(context, y))
+        Plus(xs) => +(BasicSymbolic.(Ref(context), xs)...)
+        Mult(xs) => *(BasicSymbolic.(Ref(context), xs)...)
+        Tan(x) => (!(DerivOp))(BasicSymbolic(context, x))
     end
 end
 
-function SymbolicContext(d::decapodes.DecaExpr, __module__=@__MODULE__)
+function SymbolicContext(d::decapodes.DecaExpr)
     # associates each var to its sort...
     context = map(d.context) do j
         j.var => symtype(Deca.DECQuantity, j.dim, j.space)
@@ -116,13 +119,13 @@ function SymbolicContext(d::decapodes.DecaExpr, __module__=@__MODULE__)
     end
     context = Dict{Symbol,DataType}(context)
     eqs = map(d.equations) do eq
-        SymbolicEquation{Symbolic}(BasicSymbolic.(Ref(context), [eq.lhs, eq.rhs], Ref(__module__))...)
+        SymbolicEquation{Symbolic}(BasicSymbolic.(Ref(context), [eq.lhs, eq.rhs])...)
     end
     SymbolicContext(vars, eqs)
 end
 
 function eval_eq!(eq::SymbolicEquation, d::AbstractDecapode, syms::Dict{Symbol, Int}, deletions::Vector{Int})
-    eval_eq!(Equation(Term(eq.lhs), Term(eq.rhs)), d, syms, deletions)
+    eval_eq!(Eq(Term(eq.lhs), Term(eq.rhs)), d, syms, deletions)
 end
 
 """    function SummationDecapode(e::SymbolicContext) """
@@ -132,7 +135,7 @@ function SummationDecapode(e::SymbolicContext)
 
     foreach(e.vars) do var
         # convert Sort(var)::PrimalForm0 --> :Form0
-        var_id = add_part!(d, :Var, name=var.name, type=nameof(Sort(var)))
+        var_id = add_part!(d, :Var, name=var.name, type=nameof(symtype(var)))
         symbol_table[var.name] = var_id
     end
 

src/acset.jl

@@ -158,13 +158,16 @@ end
 # A collection of DecaType getters
 # TODO: This should be replaced by using a type hierarchy
 const ALL_TYPES = [:Form0, :Form1, :Form2, :DualForm0, :DualForm1, :DualForm2,
-  :Literal, :Parameter, :Constant, :infer]
+                   :PVF, :DVF,
+                   :Literal, :Parameter, :Constant, :infer]
 
 const FORM_TYPES = [:Form0, :Form1, :Form2, :DualForm0, :DualForm1, :DualForm2]
 const PRIMALFORM_TYPES = [:Form0, :Form1, :Form2]
 const DUALFORM_TYPES = [:DualForm0, :DualForm1, :DualForm2]
 
-const NONFORM_TYPES = [:Constant, :Parameter, :Literal, :infer]
+const VECTORFIELD_TYPES = [:PVF, :DVF]
+
+const NON_EC_TYPES = [:Constant, :Parameter, :Literal, :infer]
 const USER_TYPES = [:Constant, :Parameter]
 const NUMBER_TYPES = [:Literal]
 const INFER_TYPES = [:infer]
@@ -184,6 +187,7 @@ function recognize_types(d::AbstractNamedDecapode)
   isempty(unrecognized_types) ||
   error("Types $unrecognized_types are not recognized. CHECK: $types")
 end
+export recognize_types
 
 """    is_expanded(d::AbstractNamedDecapode)
 
@@ -427,12 +431,12 @@ function safe_modifytype!(d::SummationDecapode, var_idx::Int, new_type::Symbol)
 end
 
 """
-    filterfor_forms(types::AbstractVector{Symbol})
+    filterfor_ec_types(types::AbstractVector{Symbol})
 
-Return any form type symbols.
+Return any form or vector-field type symbols.
 """
-function filterfor_forms(types::AbstractVector{Symbol})
-  conditions = x -> !(x in NONFORM_TYPES)
+function filterfor_ec_types(types::AbstractVector{Symbol})
+  conditions = x -> !(x in NON_EC_TYPES)
   filter(conditions, types)
 end
 
@@ -447,29 +451,26 @@ function infer_sum_types!(d::SummationDecapode, Σ_idx::Int)
   types = d[idxs, :type]
   all(t != :infer for t in types) && return applied # We need not infer
 
-  forms = unique(filterfor_forms(types))
+  ec_types = unique(filterfor_ec_types(types))
 
-  form = @match length(forms) begin
+  ec_type = @match length(ec_types) begin
     0 => return applied # We can not infer
-    1 => only(forms)
-    _ => error("Type mismatch in summation $Σ_idx, all the following forms appear: $forms")
+    1 => only(ec_types)
+    _ => error("Type mismatch in summation $Σ_idx, all the following forms appear: $ec_types")
   end
 
   for idx in idxs
-    applied |= safe_modifytype!(d, idx, form)
+    applied |= safe_modifytype!(d, idx, ec_type)
   end
 
   return applied
 end
 
 function apply_inference_rule_op1!(d::SummationDecapode, op1_id, rule)
-  type_src = d[d[op1_id, :src], :type]
-  type_tgt = d[d[op1_id, :tgt], :type]
+  score_src = (rule.src_type == d[d[op1_id, :src], :type])
+  score_tgt = (rule.tgt_type == d[d[op1_id, :tgt], :type])
 
-  score_src = (rule.src_type == type_src)
-  score_tgt = (rule.tgt_type == type_tgt)
   check_op = (d[op1_id, :op1] in rule.op_names)
-
   if(check_op && (score_src + score_tgt == 1))
     mod_src = safe_modifytype!(d, d[op1_id, :src], rule.src_type)
     mod_tgt = safe_modifytype!(d, d[op1_id, :tgt], rule.tgt_type)
@@ -480,19 +481,15 @@ function apply_inference_rule_op1!(d::SummationDecapode, op1_id, rule)
 end
 
 function apply_inference_rule_op2!(d::SummationDecapode, op2_id, rule)
-  type_proj1 = d[d[op2_id, :proj1], :type]
-  type_proj2 = d[d[op2_id, :proj2], :type]
-  type_res = d[d[op2_id, :res], :type]
+  score_proj1 = (rule.proj1_type == d[d[op2_id, :proj1], :type])
+  score_proj2 = (rule.proj2_type == d[d[op2_id, :proj2], :type])
+  score_res = (rule.res_type == d[d[op2_id, :res], :type])
 
-  score_proj1 = (rule.proj1_type == type_proj1)
-  score_proj2 = (rule.proj2_type == type_proj2)
-  score_res = (rule.res_type == type_res)
   check_op = (d[op2_id, :op2] in rule.op_names)
-
-  if(check_op && (score_proj1 + score_proj2 + score_res == 2))
+  if check_op && (score_proj1 + score_proj2 + score_res == 2)
     mod_proj1 = safe_modifytype!(d, d[op2_id, :proj1], rule.proj1_type)
     mod_proj2 = safe_modifytype!(d, d[op2_id, :proj2], rule.proj2_type)
-    mod_res =   safe_modifytype!(d, d[op2_id, :res], rule.res_type)
+    mod_res   = safe_modifytype!(d, d[op2_id, :res], rule.res_type)
     return mod_proj1 || mod_proj2 || mod_res
   end
 

src/acset2symbolic.jl

@@ -0,0 +1,97 @@
+using DiagrammaticEquations
+using ACSets
+using SymbolicUtils
+using SymbolicUtils: BasicSymbolic, Symbolic
+
+# TODO: Expose only the symbolic_rewriting function
+export symbolics_lookup, extract_symexprs, apply_rewrites, merge_equations, to_acset, symbolic_rewriting
+
+const DECA_EQUALITY_SYMBOL = (==)
+
+function symbolics_lookup(d::SummationDecapode)
+  Dict{Symbol, BasicSymbolic}(map(parts(d, :Var)) do i
+    (d[i, :name], decavar_to_symbolics(d, i))
+  end)
+end
+
+function decavar_to_symbolics(d::SummationDecapode, idx::Int; space = :I)
+  new_type = SymbolicUtils.symtype(Deca.DECQuantity, d[idx, :type], space)
+  SymbolicUtils.Sym{new_type}(d[idx, :name])
+end
+
+function to_symbolics(d::SummationDecapode, symvar_lookup::Dict{Symbol, BasicSymbolic}, op_idx::Int, op_type::Symbol)
+  input_syms = getindex.(Ref(symvar_lookup), d[edge_inputs(d,op_idx,Val(op_type)), :name])
+  output_sym = getindex.(Ref(symvar_lookup), d[edge_output(d,op_idx,Val(op_type)), :name])
+  op_sym = getfield(@__MODULE__, edge_function(d,op_idx,Val(op_type)))
+
+  S = promote_symtype(op_sym, input_syms...)
+  SymbolicEquation{Symbolic}(output_sym, SymbolicUtils.Term{S}(op_sym, input_syms))
+end
+
+function symbolic_rewriting(old_d::SummationDecapode, rewriter = nothing)
+  d = infer_types!(deepcopy(old_d))
+  eqns = merge_equations(d)
+  to_acset(d, apply_rewrites(eqns, rewriter))
+end
+
+apply_rewrites(eqns, rewriter) = isnothing(rewriter) ? eqns : map(rewriter, eqns)
+
+function extract_symexprs(d::SummationDecapode, symvar_lookup::Dict{Symbol, BasicSymbolic})
+  map(topological_sort_edges(d)) do node
+    to_symbolics(d, symvar_lookup, node.index, node.name)
+  end
+end
+
+# XXX SymbolicUtils.substitute swaps the order of multiplication.
+# e.g. @decapode begin
+#   ∂ₜ(G) == κ*u
+# end
+# will have the κ*u term rewritten to u*κ
+function merge_equations(d::SummationDecapode)
+  symvar_lookup = symbolics_lookup(d)
+  symexpr_list = extract_symexprs(d, symvar_lookup)
+
+  eqn_lookup = Dict()
+
+  terminal_vars = infer_terminal_names(d)
+  terminal_eqns = SymbolicEquation{Symbolic}[]
+
+  foreach(symexpr_list) do x
+    push!(eqn_lookup, (x.lhs => SymbolicUtils.substitute(x.rhs, eqn_lookup)))
+    if x.lhs.name in terminal_vars
+      push!(terminal_eqns, SymbolicEquation{Symbolic}(x.lhs, eqn_lookup[x.lhs]))
+    end
+  end
+
+  formed_deca_eqn.(terminal_eqns)
+end
+
+formed_deca_eqn(symeqn::SymbolicEquation{Symbolic}) = SymbolicUtils.Term{Number}(DECA_EQUALITY_SYMBOL, [symeqn.lhs, symeqn.rhs])
+
+function to_acset(d::SummationDecapode, sym_exprs)
+  outer_types = map([infer_states(d)..., infer_terminals(d)...]) do i
+    :($(d[i, :name])::$(d[i, :type]))
+  end
+
+  tangents = map(incident(d, DerivOp, :op1)) do op1
+    :($(d[d[op1, :tgt], :name]) == $DerivOp($(d[d[op1, :src], :name])))
+  end
+
+  #TODO: This step is breaking up summations
+  final_exprs = SymbolicUtils.Code.toexpr.(sym_exprs)
+  reify!(exprs) = foreach(exprs) do x
+    if typeof(x)==Expr && x.head == :call
+      x.args[1] = nameof(x.args[1])
+      reify!(x.args[2:end])
+    end
+  end
+  reify!(final_exprs)
+
+  deca_block = quote
+    $(outer_types...) 
+    $(final_exprs...)
+  end
+
+  ∘(infer_types!, SummationDecapode, parse_decapode)(deca_block)
+end
+