Add internal types for division

Author: asterycs

src/forward.jl

@@ -883,6 +883,10 @@ function evaluate(op::BinaryOperation{Mult})
     evaluate(Mult(), evaluate(op.arg1), evaluate(op.arg2))
 end
 
+function evaluate(op::BinaryOperation{Div})
+    BinaryOperation{Div}(evaluate(op.arg1), evaluate(op.arg2))
+end
+
 function evaluate(op::BinaryOperation{Op}) where {Op<:AdditiveOperation}
     evaluate(Op(), evaluate(op.arg1), evaluate(op.arg2))
 end

src/ir.jl

@@ -55,6 +55,11 @@ struct Product <: IR
     r::IR
 end
 
+struct Quotient <: IR
+    num::IR
+    den::IR
+end
+
 struct HadamardProduct <: IR
     l::IR
     r::IR
@@ -500,6 +505,10 @@ function to_ir(arg::BinaryOperation{Mult})
     throw_not_std(arg)
 end
 
+function to_ir(arg::BinaryOperation{Div})
+    return ir.Quotient(to_ir(arg.arg1), to_ir(arg.arg2))
+end
+
 function to_ir(arg::Power)
     return ir.Power(to_ir(arg.base), arg.exponent)
 end

src/ricci.jl

@@ -109,6 +109,7 @@ abstract type AdditiveOperation end
 struct Add <: AdditiveOperation end
 struct Sub <: AdditiveOperation end
 struct Mult end
+struct Div end
 
 struct Power <: Tensor
     base::Value
@@ -237,6 +238,10 @@ function get_indices(arg::BinaryOperation{Mult})
     return [get_indices(arg.arg1); get_indices(arg.arg2)]
 end
 
+function get_indices(arg::BinaryOperation{Div})
+    return [get_indices(arg.arg1); get_indices(arg.arg2)]
+end
+
 function get_indices(arg::Power)
     return get_indices(arg.base)
 end
@@ -849,6 +854,10 @@ function to_string(arg::BinaryOperation{Mult})
     return parenthesize(arg.arg1) * parenthesize(arg.arg2)
 end
 
+function to_string(arg::BinaryOperation{Div})
+    return parenthesize(arg.arg1) * "/" * parenthesize(arg.arg2)
+end
+
 function to_string(arg::BinaryOperation{Add})
     return to_string(arg.arg1) * " + " * to_string(arg.arg2)
 end

src/simplify.jl

@@ -165,6 +165,12 @@ function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Variable)
     return BinaryOperation{Mult}(arg1, arg2)
 end
 
+function simplify(::Div, arg1::Value, arg2::Value)
+    return evaluate(
+        BinaryOperation{Div}(simplify(evaluate(arg1)), simplify(evaluate(arg2))),
+    )
+end
+
 function simplify(::Mult, arg1::Value, arg2::Value)
     return evaluate(
         BinaryOperation{Mult}(simplify(evaluate(arg1)), simplify(evaluate(arg2))),

src/std.jl

@@ -327,6 +327,10 @@ function to_standard(arg::BinaryOperation{Mult})
     throw_not_std(arg)
 end
 
+function to_standard(arg::BinaryOperation{Div})
+    return BinaryOperation{Div}(to_standard(arg.arg1), to_standard(arg.arg2))
+end
+
 function standardize(arg)
     arg = simplify(arg)
     free_indices = unique(get_free_indices(arg))

src/stdstr.jl

@@ -99,6 +99,10 @@ function to_std_str(arg::ir.Product)
     return parenthesize(to_std_str, arg.l) * parenthesize(to_std_str, arg.r)
 end
 
+function to_std_str(arg::ir.Quotient)
+    return parenthesize(to_std_str, arg.num) * " ⊘ " * parenthesize(to_std_str, arg.den)
+end
+
 function to_std_str(arg::ir.HadamardProduct)
     return to_std_str(arg.l) * " ⊙ " * to_std_str(arg.r)
 end

test/StdTest.jl

@@ -315,6 +315,24 @@ end
     @test to_std(dc.UnaryOperation{dc.Sgn}(mul(A, x))) == "sgn(Ax)"
 end
 
+@testset "to_std output is correct with quotient" begin
+    x = Variable("x", Upper(1))
+    y = Variable("y", Upper(1))
+    z = Variable("z", Upper(2))
+    a = Variable("a")
+    lv = Literal(2, Upper(1))
+
+    function div(l, r)
+        return dc.BinaryOperation{dc.Div}(l, r)
+    end
+
+    @test to_std(div(x, y)) == "x ⊘ y"
+    @test to_std(div(lv, x)) == "vec(2) ⊘ x"
+    @test to_std(div(x, lv)) == "x ⊘ vec(2)"
+    # @test_throws to_std(div(a, x))
+    # @test_throws to_std(div(x, z))
+end
+
 @testset "to_std output is correct with KrD-KrD and one free index" begin
     l = KrD(Upper(1), Lower(2))
     u = KrD(Upper(2), Lower(1))