Add diff rules for log

Author: asterycs

README.md

@@ -51,9 +51,9 @@ standard notation.
 
 ### Supported functions and operators
 
-`+`, `-`, `'`, `*`, `^`, `abs`, `sin`, `cos`
+`+`, `-`, `'`, `*`, `^`, `abs`, `sin`, `cos`, `log`
 
-Element-wise operations `sin.`, `cos.`, `abs.`, `.*` and `.^` are supported.  
+Element-wise operations `sin.`, `cos.`, `abs.`, `.*`, `.^` and `log.` are supported.  
 Vector 1-norm and 2-norm can be computed with `LinearAlgebra.norm(..., 1)` and `LinearAlgebra.norm(..., 2)`.  
 Sums of vectors can be computed with `sum`.  
 Matrix traces can be computed with `LinearAlgebra.tr`.

docs/src/examples.md

@@ -118,6 +118,21 @@ expr = sum(A * x)
 
 A¹₄x⁴1₁
 ```
+#### Log and element-wise log
+```jldoctest usage
+expr = log(x' * y)
+
+# output
+
+log(x₃y³)
+```
+```jldoctest usage
+expr = log.(x)' * y
+
+# output
+
+log(x₃)y³
+```
 #### Vector Norms
 ```jldoctest usage
 expr = norm(A * x, 2)

docs/src/usage.md

@@ -66,8 +66,8 @@ c * A
 cA¹₂
 ```
 Supported operators and functions when creating expressions:
-  - Basic operators `+`, `-`, `'`, `*`, `^`, `abs`, `sin` and `cos`
-  - Element-wise operators `sin.`, `cos.`, `abs.`, `.*` and `.^`
+  - Basic operators `+`, `-`, `'`, `*`, `^`, `abs`, `sin`, `cos` and `log`
+  - Element-wise operators `sin.`, `cos.`, `abs.`, `.*`, `.^` and `log.`
   - Vector 1-norm and 2-norm can be computed with `norm1` and `norm2`
   - Sums of vectors can be computed with `sum`.
   - Matrix traces can be computed with `tr`.

src/forward.jl

@@ -49,6 +49,13 @@ function diff(arg::UnaryOperation{Cos}, wrt::Variable)
     return BinaryOperation{Mult}(-UnaryOperation{Sin}(arg.arg), diff(arg.arg, wrt))
 end
 
+function diff(arg::Log, wrt::Variable)
+    return BinaryOperation{Mult}(
+        BinaryOperation{Div}(Literal(1, get_free_indices(arg)...), arg.arg),
+        diff(arg.arg, wrt),
+    )
+end
+
 function diff(arg::Power, wrt::Variable)
     outer = replace_bound_letters(arg.base, wrt)
 
@@ -287,6 +294,20 @@ function evaluate(::Mult, arg1::KrD, arg2::BinaryOperation{Mult})
     return evaluate(Mult(), arg2, arg1)
 end
 
+function evaluate(::Mult, arg1::BinaryOperation{Div}, arg2::KrD)
+    attempt = BinaryOperation{Div}(
+        evaluate(Mult(), evaluate(arg1.arg1), evaluate(arg2)),
+        evaluate(Mult(), evaluate(arg1.arg2), evaluate(arg2)),
+    )
+
+    # This ensures that arg1.base and arg2 can contract and that the contraction is simple
+    if length(get_free_indices(attempt)) == length(get_free_indices(arg1))
+        return attempt
+    end
+
+    return BinaryOperation{Mult}(arg1, arg2)
+end
+
 function evaluate(::Mult, arg1::BinaryOperation{Mult}, arg2::KrD)
     ci = indices_in_common(arg1.arg1, arg1.arg2)
 
@@ -377,6 +398,17 @@ function evaluate(::Mult, arg1::KrD, arg2::UnaryOp) where {UnaryOp<:UnaryOperati
     return BinaryOperation{Mult}(evaluate(arg2), evaluate(arg1))
 end
 
+function evaluate(::Mult, arg1::Log, arg2::KrD)
+    attempt = Log(evaluate(Mult(), evaluate(arg1.arg), evaluate(arg2)))
+
+    # This ensures that arg1.base and arg2 can contract and that the contraction is simple
+    if length(get_free_indices(attempt)) == length(get_free_indices(arg1))
+        return attempt
+    end
+
+    return BinaryOperation{Mult}(evaluate(arg1), evaluate(arg2))
+end
+
 function evaluate(::Mult, arg1::Power, arg2::KrD)
     attempt = Power(evaluate(Mult(), evaluate(arg1.base), evaluate(arg2)), arg1.exponent)
 
@@ -874,6 +906,10 @@ function _sub_from_product(arg1::BinaryOperation{Mult}, arg2::Value)
     return BinaryOperation{Sub}(evaluate(arg1), evaluate(arg2))
 end
 
+function evaluate(op::Log)
+    return Log(evaluate(op.arg))
+end
+
 function evaluate(op::Power)
     if op.exponent == 1
         return evaluate(op.base)

src/ir.jl

@@ -65,6 +65,10 @@ struct HadamardProduct <: IR
     r::IR
 end
 
+struct Log <: IR
+    arg::IR
+end
+
 struct Power <: IR
     base::IR
     exponent::Union{Int,Rational{Int}}
@@ -509,6 +513,10 @@ function to_ir(arg::BinaryOperation{Div})
     return ir.Quotient(to_ir(arg.arg1), to_ir(arg.arg2))
 end
 
+function to_ir(arg::Log)
+    return ir.Log(to_ir(arg.arg))
+end
+
 function to_ir(arg::Power)
     return ir.Power(to_ir(arg.base), arg.exponent)
 end

src/ricci.jl

@@ -118,6 +118,12 @@ end
 
 Base.hash(op::Power, h::UInt) = hash(op.exponent, hash(op.base, hash(Power, h)))
 
+struct Log <: Tensor
+    arg::Value
+end
+
+Base.hash(op::Log, h::UInt) = hash(op.arg, hash(Log, h))
+
 struct UnaryOperation{Op} <: Tensor where {Op}
     arg::Value
 end
@@ -246,6 +252,10 @@ function get_indices(arg::Power)
     return get_indices(arg.base)
 end
 
+function get_indices(arg::Log)
+    return get_indices(arg.arg)
+end
+
 function get_indices(arg::BinaryOperation{Op}) where {Op<:AdditiveOperation}
     arg1_free_ids, arg2_free_ids = get_free_indices.((arg.arg1, arg.arg2))
 
@@ -432,13 +442,29 @@ function Base.literal_pow(f::typeof(^), base::Tensor, exponent::Val{E}) where {E
     return Power(base, E)
 end
 
+function Base.log(arg::Tensor)
+    if !isempty(get_free_indices(arg))
+        throw(DomainError(arg, "Argument is not a scalar, use log. for element-wise log."))
+    end
+
+    return Log(arg)
+end
+
+function Base.broadcasted(::typeof(log), arg::Tensor)
+    return Log(arg)
+end
+
 function replace_letters(arg::BinaryOperation{Mult}, letter_map::Dict)
     return BinaryOperation{Mult}(
         replace_letters(arg.arg1, letter_map),
         replace_letters(arg.arg2, letter_map),
     )
 end
 
+function replace_letters(arg::Log, letter_map::Dict)
+    return Log(replace_letters(arg.arg, letter_map))
+end
+
 function replace_letters(arg::Power, letter_map::Dict)
     return Power(replace_letters(arg.base, letter_map), arg.exponent)
 end
@@ -697,6 +723,10 @@ function Base.adjoint(arg::T) where {T<:UnaryOperation}
     return T(arg.arg')
 end
 
+function Base.adjoint(arg::Log)
+    return Log(adjoint(arg.arg))
+end
+
 function Base.adjoint(arg::Power)
     return Power(adjoint(arg.base), arg.exponent)
 end
@@ -846,6 +876,10 @@ function to_string(arg::Power)
     return b * ".^" * parenthesize(arg.exponent)
 end
 
+function to_string(arg::Log)
+    return "log(" * to_string(arg.arg) * ")"
+end
+
 function to_string(arg::BinaryOperation{Mult})
     if arg.arg1 == -1
         return "-" * parenthesize(arg.arg2)

src/simplify.jl

@@ -111,7 +111,7 @@ function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Literal)
     return BinaryOperation{Mult}(arg1, arg2)
 end
 
-function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Variable)
+function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
     if is_diag(arg1) && !is_elementwise_multiplication(arg1, arg2)
         d = get_diag_delta(arg1)
 

src/std.jl

@@ -217,6 +217,10 @@ function to_standard(term::UnaryOperation{Op}) where {Op}
     return UnaryOperation{Op}(to_standard(term.arg))
 end
 
+function to_standard(arg::Log)
+    return Log(to_standard(arg.arg))
+end
+
 function to_standard(term::Power)
     return Power(to_standard(term.base), term.exponent)
 end

src/stdstr.jl

@@ -107,6 +107,12 @@ function to_std_str(arg::ir.HadamardProduct)
     return to_std_str(arg.l) * " ⊙ " * to_std_str(arg.r)
 end
 
+function to_std_str(arg::ir.Log)
+    out = to_std_str(arg.arg)
+
+    return "log(" * out * ")"
+end
+
 function to_std_str(arg::ir.Power)
     out = to_std_str(arg.base)
 

test/StdStrTest.jl

@@ -36,6 +36,9 @@
     @test to_std(gradient(sum((x + y) .^ 2), x)) == "2(x + y)"
     @test to_std(gradient(sum((x .* y) .^ 2), x)) == "2(x ⊙ y ⊙ y)"
     @test to_std(gradient(sum((A * x - y) .^ 2), x)) == "2Aᵀ(Ax - y)"
+    @test to_std(gradient(log.(x)'*x, x)) == "(vec(1) ⊘ x ⊙ x) + log(x)" # TODO: Add simplification rule for quotients
+    @test to_std(gradient(log.(x)'*log.(x), x)) ==
+          "diag(vec(1)ᵀ ⊘ xᵀ)Iᵀlog(x) + (vec(1) ⊘ x ⊙ log(x))"
     @test to_std(gradient((x' * A * x) ^ (-2), x)) == "(-2)(xᵀAᵀx)^(-3)(Aᵀx + Ax)"
     @test to_std(gradient((x' * A * x) ^ 2, x)) == "2xᵀAᵀx(Aᵀx + Ax)"
     @test to_std(gradient(((A .* B) * C * x)' * x, x)) == "(A ⊙ B)Cx + Cᵀ(Aᵀ ⊙ Bᵀ)x"
@@ -52,6 +55,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(log.(x), x)) == "diag(vec(1) ⊘ 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)"