Add Julia code output

Author: asterycs

src/DiffMatic.jl

@@ -10,5 +10,6 @@ include("ir.jl")
 include("simplify.jl")
 include("std.jl")
 include("forward.jl")
+include("julia.jl")
 
 end

src/julia.jl

@@ -0,0 +1,80 @@
+
+function to_julia(arg::ir.Var)
+    return Symbol(arg.id)
+end
+
+function to_julia(arg::ir.Const)
+    return :($(arg.value))
+end
+
+function to_julia(arg::ir.Mat)
+    if arg.id isa ir.Var
+        return to_julia(arg.id)
+    end
+
+    throw(RuntimeError("Unable to generate julia code for constant matrix of unknown size"))
+end
+
+function to_julia(arg::ir.Vec)
+    if arg.id isa ir.Var
+        return to_julia(arg.id)
+    end
+
+    throw(RuntimeError("Unable to generate julia code for constant vector of unknown size"))
+end
+
+function to_julia(arg::ir.Scal)
+    return to_julia(arg.id)
+end
+
+function to_julia(arg::ir.Real)
+    return arg
+end
+
+function to_julia(arg::ir.Identity)
+    return Symbol("I")
+end
+
+function to_julia(arg::ir.Sin)
+    return :(sin.($(to_julia(arg.arg))))
+end
+
+function to_julia(arg::ir.Cos)
+    return :(cos.($(to_julia(arg.arg))))
+end
+
+function to_julia(arg::ir.Add)
+    return :($(to_julia(arg.l)) + $(to_julia(arg.r)))
+end
+
+function to_julia(arg::ir.Sub)
+    return :($(to_julia(arg.l)) - $(to_julia(arg.r)))
+end
+
+function to_julia(arg::ir.Product)
+    return :($(to_julia(arg.l)) * $(to_julia(arg.r)))
+end
+
+function to_julia(arg::ir.HadamardProduct)
+    return :($(to_julia(arg.l)) .* $(to_julia(arg.r)))
+end
+
+function to_julia(arg::ir.Power)
+    return :($(to_julia(arg.base)) .^ $(to_julia(arg.exponent)))
+end
+
+function to_julia(arg::ir.Trace)
+    return :(tr($(to_julia(arg.arg))))
+end
+
+function to_julia(arg::ir.Diag)
+    return :(diagm($(to_julia(arg.arg))))
+end
+
+function to_julia(arg::ir.Transpose)
+    return :(transpose($(to_julia(arg.arg))))
+end
+
+function to_julia(arg::ir.Sum)
+    return :(sum($(to_julia(arg.arg))))
+end

src/std.jl

@@ -306,9 +306,6 @@ function to_standard(arg::BinaryOperation{Mult})
     throw_not_std(arg)
 end
 
-struct Ir end
-struct StdStr end
-
 function to_std_str(arg::ir.Mat)
     if arg.id isa String
         return arg.id
@@ -453,6 +450,10 @@ function standardize(arg)
     return standardized
 end
 
+struct Ir end
+struct StdStr end
+struct Julia end
+
 """
     to_std(expr)
 
@@ -472,14 +473,29 @@ function to_std(arg; format = StdStr())
     return _to_std(format, arg)
 end
 
+function _to_std(format::Ir, arg)
+    standardized = standardize(arg)
+
+    return to_ir(standardized)
+end
+
 function _to_std(format::StdStr, arg)
     standardized = standardize(arg)
 
     return to_std_str(to_ir(standardized))
 end
 
-function _to_std(format::Ir, arg)
+function _to_std(format::Julia, arg)
     standardized = standardize(arg)
 
-    return to_ir(standardized)
+    ir = to_ir(standardized)
+    op = to_julia(ir)
+
+    variables = DiffMatic.ir.get_variables(ir)
+
+    return quote
+        function derivative($(Symbol.(variables)...))
+            return $(op)
+        end
+    end
 end