Fix latexify

Author: YingboMa

src/direct.jl

@@ -231,18 +231,36 @@ function sparsehessian(O, vars::AbstractVector; simplify=false)
     return H
 end
 
-"""
-    canonicalexpr(O) -> (canonical::Bool, expr)
+# `_toexpr` is only used for latexify
+function _toexpr(O; canonicalize=true)
+    if canonicalize
+        canonical, O = canonicalexpr(O)
+        canonical && return O
+    else
+        !istree(O) && return O
+    end
+
+    op = operation(O)
+    args = arguments(O)
+    if op isa Differential
+        ex = _toexpr(args[1]; canonicalize=canonicalize)
+        wrt = _toexpr(op.x; canonicalize=canonicalize)
+        return :(_derivative($ex, $wrt))
+    elseif op isa Sym
+        isempty(args) && return nameof(op)
+        return Expr(:call, _toexpr(op; canonicalize=canonicalize), _toexpr(args; canonicalize=canonicalize)...)
+    end
+    return Expr(:call, op, _toexpr(args; canonicalize=canonicalize)...)
+end
+_toexpr(s::Sym; kw...) = nameof(s)
 
-Canonicalize `O`. Return `canonical` if `expr` is valid code to generate.
-"""
 function canonicalexpr(O)
     !istree(O) && return true, O
     op = operation(O)
     args = arguments(O)
     if op === (^)
         if length(args) == 2 && args[2] isa Number && args[2] < 0
-            ex = toexpr(args[1])
+            ex = _toexpr(args[1])
             if args[2] == -1
                 expr = Expr(:call, inv, ex)
             else
@@ -254,11 +272,15 @@ function canonicalexpr(O)
     return false, O
 end
 
-function toexpr(eq::Equation; kw...)
-    Expr(:(=), toexpr(eq.lhs; kw...), toexpr(eq.rhs; kw...))
-end
+for fun in [:toexpr, :_toexpr]
+    @eval begin
+        function $fun(eq::Equation; kw...)
+            Expr(:(=), $fun(eq.lhs; kw...), $fun(eq.rhs; kw...))
+        end
 
-toexpr(eqs::AbstractArray; kw...) = map(eq->toexpr(eq; kw...), eqs)
-toexpr(x::Integer; kw...) = x
-toexpr(x::AbstractFloat; kw...) = x
-toexpr(x::Num; kw...) = toexpr(value(x); kw...)
+        $fun(eqs::AbstractArray; kw...) = map(eq->$fun(eq; kw...), eqs)
+        $fun(x::Integer; kw...) = x
+        $fun(x::AbstractFloat; kw...) = x
+        $fun(x::Num; kw...) = $fun(value(x); kw...)
+    end
+end

src/latexify_recipes.jl

@@ -11,13 +11,13 @@ prettify_expr(expr::Expr) = Expr(expr.head, prettify_expr.(expr.args)...)
     # that latexify can deal with
 
     rhs = getfield.(eqs, :rhs)
-    rhs = prettify_expr.(toexpr(rhs; canonicalize=false))
+    rhs = prettify_expr.(_toexpr(rhs; canonicalize=false))
     rhs = [postwalk(x -> x isa Expr && length(x.args) == 1 ? x.args[1] : x, eq) for eq in rhs]
     rhs = [postwalk(x -> x isa Expr && x.args[1] == :_derivative && length(x.args[2].args) == 2 ? :($(Symbol(:d, x.args[2]))/($(Symbol(:d, x.args[2].args[2])))) : x, eq) for eq in rhs]
     rhs = [postwalk(x -> x isa Expr && x.args[1] == :_derivative ? "\\frac{d\\left($(Latexify.latexraw(x.args[2]))\\right)}{d$(Latexify.latexraw(x.args[3]))}" : x, eq) for eq in rhs]
 
     lhs = getfield.(eqs, :lhs)
-    lhs = prettify_expr.(toexpr(lhs; canonicalize=false))
+    lhs = prettify_expr.(_toexpr(lhs; canonicalize=false))
     lhs = [postwalk(x -> x isa Expr && length(x.args) == 1 ? x.args[1] : x, eq) for eq in lhs]
     lhs = [postwalk(x -> x isa Expr && x.args[1] == :_derivative && length(x.args[2].args) == 2 ? :($(Symbol(:d, x.args[2]))/($(Symbol(:d, x.args[2].args[2])))) : x, eq) for eq in lhs]
     lhs = [postwalk(x -> x isa Expr && x.args[1] == :_derivative ? "\\frac{d\\left($(Latexify.latexraw(x.args[2]))\\right)}{d$(Latexify.latexraw(x.args[3]))}" : x, eq) for eq in lhs]