Add abs and sign

Author: asterycs

README.md

@@ -50,7 +50,7 @@ standard notation.
 
 ### Supported operators
 
-`tr`, `sum`, `sin`, `cos`, `+`, `-`, `'`, `*`, `.*`, `.^`, `^`
+`tr`, `sum`, `sin`, `cos`, `+`, `-`, `'`, `*`, `.*`, `.^`, `^`, `abs`
 
 ### Installation
 

src/forward.jl

@@ -31,6 +31,10 @@ function diff(arg::Real, wrt::Variable)
     return Zero([flip(i) for i ∈ wrt.indices]...)
 end
 
+function diff(arg::UnaryOperation{Abs}, wrt::Variable)
+    return BinaryOperation{Mult}(UnaryOperation{Sgn}(arg.arg), diff(arg.arg, wrt))
+end
+
 function diff(arg::UnaryOperation{Sin}, wrt::Variable)
     return BinaryOperation{Mult}(UnaryOperation{Cos}(arg.arg), diff(arg.arg, wrt))
 end

src/ir.jl

@@ -24,6 +24,14 @@ end
 
 struct Identity <: IR end
 
+struct Abs <: IR
+    arg::IR
+end
+
+struct Sgn <: IR
+    arg::IR
+end
+
 struct Sin <: IR
     arg::IR
 end
@@ -242,6 +250,14 @@ function to_ir(arg::Real)
     return ir.Scal(ir.Const(arg))
 end
 
+function to_ir(arg::UnaryOperation{Abs})
+    return ir.Abs(to_ir(arg.arg))
+end
+
+function to_ir(arg::UnaryOperation{Sgn})
+    return ir.Sgn(to_ir(arg.arg))
+end
+
 function to_ir(arg::UnaryOperation{Sin})
     return ir.Sin(to_ir(arg.arg))
 end

src/ricci.jl

@@ -103,6 +103,8 @@ end
 
 Base.hash(op::UnaryOperation{Op}, h::UInt) where {Op} = hash(op.arg, hash(Op, h))
 
+struct Abs end
+struct Sgn end
 struct Sin end
 struct Cos end
 
@@ -265,6 +267,14 @@ function Base.sum(arg::Tensor)
     return BinaryOperation{Mult}(arg, KrD(first(free_ids), flip(first(free_ids))))
 end
 
+function Base.abs(arg::Tensor)
+    return UnaryOperation{Abs}(arg)
+end
+
+function Base.sign(arg::Tensor)
+    return UnaryOperation{Sgn}(arg)
+end
+
 function Base.broadcasted(::typeof(*), arg1::Tensor, arg2::Tensor)
     arg1_free_indices = get_free_indices(arg1)
     arg2_free_indices = get_free_indices(arg2)
@@ -687,6 +697,14 @@ function to_string(arg::Zero)
     return "0" * join(scripts)
 end
 
+function to_string(arg::UnaryOperation{Abs})
+    return "|$(arg.arg)|"
+end
+
+function to_string(arg::UnaryOperation{Sgn})
+    return "sgn($(arg.arg))"
+end
+
 function to_string(arg::UnaryOperation{Sin})
     return "sin($(arg.arg))"
 end

src/std.jl

@@ -354,6 +354,14 @@ function to_std_str(arg::ir.Identity)
     return "I"
 end
 
+function to_std_str(arg::ir.Abs)
+    return "abs(" * to_std_str(arg.arg) * ")"
+end
+
+function to_std_str(arg::ir.Sgn)
+    return "sgn(" * to_std_str(arg.arg) * ")"
+end
+
 function to_std_str(arg::ir.Sin)
     return "sin(" * to_std_str(arg.arg) * ")"
 end

test/ForwardTest.jl

@@ -626,6 +626,19 @@ end
     @test equivalent(evaluate(D), KrD(Upper(1), Lower(2)))
 end
 
+@testset "diff abs" begin
+    x = Variable("x", Upper(2))
+
+    op = abs(x)
+
+    D = dc.diff(op, Variable("x", Upper(3)))
+
+    @test equivalent(
+        D,
+        dc.BinaryOperation{dc.Mult}(dc.UnaryOperation{dc.Sgn}(x), KrD(Upper(2), Lower(3))),
+    )
+end
+
 @testset "diff sin" begin
     x = Variable("x", Upper(2))
 

test/RicciTest.jl

@@ -70,6 +70,26 @@ end
     @test left != Zero()
 end
 
+@testset "Abs constructor" begin
+    a = KrD(Upper(1), Lower(2))
+    b = Variable("b", Upper(2))
+
+    op = abs(a * b)
+
+    @test typeof(op) == dc.UnaryOperation{dc.Abs}
+    @test typeof(op.arg) == dc.BinaryOperation{dc.Mult}
+end
+
+@testset "Sgn constructor" begin
+    a = KrD(Upper(1), Lower(2))
+    b = Variable("b", Upper(2))
+
+    op = sign(a * b)
+
+    @test typeof(op) == dc.UnaryOperation{dc.Sgn}
+    @test typeof(op.arg) == dc.BinaryOperation{dc.Mult}
+end
+
 @testset "Sin constructor" begin
     a = KrD(Upper(1), Lower(2))
     b = Variable("b", Upper(2))

test/StdStrTest.jl

@@ -12,6 +12,7 @@
     @test to_std(gradient(sin(tr(x * x')), x)) == "cos(xᵀx)2x"
     @test to_std(gradient(cos(tr(x * x')), x)) == "(-1)sin(xᵀx)2x"
     @test to_std(gradient(tr(A), x)) == "vec(0)"
+    @test to_std(gradient(abs(x' * x), x)) == "sgn(xᵀx)2x"
     @test to_std(gradient(x' * B' * A * A * x, x)) == "AᵀAᵀBx + BᵀAAx"
     @test to_std(gradient((A' * B * x)' * A * x, x)) == "AᵀAᵀBx + BᵀAAx"
     @test to_std(gradient(a * sin(y)' * x, x)) == "asin(y)"
@@ -48,6 +49,7 @@ end
 
     @test to_std(jacobian(A * x, x)) == "A"
     @test to_std(jacobian(A' * x, x)) == "Aᵀ"
+    @test to_std(jacobian(abs(x), x)) == "diag(sgn(x))I"
     @test to_std(jacobian(sin(A * x + y), x)) == "diag(cos(Ax + y))A"
     @test to_std(jacobian(((A .* B) * C * x)' * x * x, x)) ==
           "xᵀCᵀ(Aᵀ ⊙ Bᵀ)xI + x(xᵀCᵀ(Aᵀ ⊙ Bᵀ) + xᵀ(A ⊙ B)C)"

test/StdTest.jl

@@ -263,6 +263,32 @@ end
     @test to_std(sum(y)) == "sum(yᵀ)"
 end
 
+@testset "to_std output is correct with abs" begin
+    x = Variable("x", Upper(2))
+    A = Variable("A", Upper(1), Lower(2))
+
+    function mul(l, r)
+        return dc.BinaryOperation{dc.Mult}(l, r)
+    end
+
+    @test to_std(dc.UnaryOperation{dc.Abs}(1)) == "abs(1)"
+    @test to_std(dc.UnaryOperation{dc.Abs}(x)) == "abs(x)"
+    @test to_std(dc.UnaryOperation{dc.Abs}(mul(A, x))) == "abs(Ax)"
+end
+
+@testset "to_std output is correct with sgn" begin
+    x = Variable("x", Upper(2))
+    A = Variable("A", Upper(1), Lower(2))
+
+    function mul(l, r)
+        return dc.BinaryOperation{dc.Mult}(l, r)
+    end
+
+    @test to_std(dc.UnaryOperation{dc.Sgn}(1)) == "sgn(1)"
+    @test to_std(dc.UnaryOperation{dc.Sgn}(x)) == "sgn(x)"
+    @test to_std(dc.UnaryOperation{dc.Sgn}(mul(A, x))) == "sgn(Ax)"
+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))