Symbolic Calculations and Building Fast Parallel Functions

Symbolics.jl is a symbolic modeling language. The way to define symbolic variables is via the @variables macro:

@variables x y

After defining variables as symbolic, symbolic expressions, which we call a istree object, can be generated by utilizing Julia expressions. For example:

z = x^2 + y

Here, z is an expression tree for "square x and add y". To make an array of symbolic expressions, simply make an array of symbolic expressions:

A = [x^2 + y 0 2x
     0       0 2y
     y^2 + x 0 0]

3×3 Matrix{Num}:
 y + x^2  0  2x
       0  0  2y
 x + y^2  0   0

Note that by default, @variables returns Sym or istree objects wrapped in Num in order to make them behave like subtypes of Real. Any operation on these Num objects will return a new Num object, wrapping the result of computing symbolically on the underlying values.

To better view the results, we can use Latexify.jl. Symbolics.jl comes with Latexify recipes so it works automatically:

using Latexify
latexify(A)
\[\begin{array}{ccc} (x ^ 2) + y & 0 & 2 * x \\ 0 & 0 & 2 * y \\ (y ^ 2) + x & 0 & 0 \\ \end{array}\]

Normal Julia functions work on Symbolics expressions, so if we want to create the sparse version of A we would just call sparse:

using SparseArrays
spA = sparse(A)

3×3 SparseMatrixCSC{Num, Int64} with 4 stored entries:
 y + x^2  ⋅  2x
       ⋅  ⋅  2y
 x + y^2  ⋅   ⋅

We can thus use normal Julia functions as generators for sparse expressions. For example, here we will define

function f(u)
  [u[1] - u[3], u[1]^2 - u[2], u[3] + u[2]]
end
f([x, y, z]) # Recall that z = x^2 + y

3-element Vector{Num}:
 x - y - (x^2)
       x^2 - y
      x^2 + 2y

Or we can build array variables and use these to trace:

@variables u[1:3]
f(u)

3-element Array{Num,1}:
   u₁ - u₃
 u₁^2 - u₂
   u₂ + u₃

Building Functions

The function for building functions from symbolic expressions is the aptly-named build_function. The first argument is the symbolic expression or the array of symbolic expressions to compile, and the trailing arguments are the arguments for the function. For example:

to_compute = [x^2 + y, y^2 + x]
f_expr = build_function(to_compute, [x, y])

gives back two codes. The first is a function f([x, y]) that computes and builds an output vector [x^2 + y, y^2 + x]. Because this tool was made to be used by all the cool kids writing fast Julia codes, it is specialized to Julia and supports features like StaticArrays. For example:

using StaticArrays
myf = eval(f_expr[1])
myf(SA[2.0, 3.0])

2-element SArray{Tuple{2},Float64,1,2} with indices SOneTo(2):
  7.0
 11.0

The second function is an in-place non-allocating mutating function which mutates its first value:

Base.remove_linenums!(f_expr[2])

:(function (var"##out#259", var"##arg#258")
      let x = var"##arg#258"[1], y = var"##arg#258"[2]
          @inbounds begin
                  var"##out#259"[1] = (+)(y, (^)(x, 2))
                  var"##out#259"[2] = (+)(x, (^)(y, 2))
                  nothing
          end
      end
  end)

Thus we'd use it like the following:

myf! = eval(f_expr[2])
out = zeros(2)
myf!(out, [2.0, 3.0])
out

2-element Array{Float64,1}:
  7.0
 11.0

To save the symbolic calculations for later, we can take this expression and save it out to a file:

write("function.jl", string(f_expr[2]))

Note that if we need to avoid eval, for example to avoid world-age issues, one could do expression = Val{false}:

build_function(to_compute, [x, y], expression=Val{false})

which will use RuntimeGeneratedFunctions.jl to build Julia functions which avoid world-age issues.

Building Non-Allocating Parallel Functions for Sparse Matrices

Now let's show off a little bit. build_function is kind of like if lambdify ate its spinach. To show this, let's build a non-allocating function that fills sparse matrices in a multithreaded manner. To do this, we just have to represent the operation that we're turning into a function via a sparse matrix. For example:

using LinearAlgebra
N = 8
A = sparse(Tridiagonal([x^i for i in 1:N-1], [x^i * y^(8-i) for i in 1:N], [y^i for i in 1:N-1]))

8×8 SparseMatrixCSC{Num, Int64} with 22 stored entries:
 x*(y^7)            y            ⋅  …            ⋅            ⋅        ⋅    ⋅
       x  (x^2)*(y^6)          y^2               ⋅            ⋅        ⋅    ⋅
       ⋅          x^2  (x^3)*(y^5)               ⋅            ⋅        ⋅    ⋅
       ⋅            ⋅          x^3             y^4            ⋅        ⋅    ⋅
       ⋅            ⋅            ⋅     (x^5)*(y^3)          y^5        ⋅    ⋅
       ⋅            ⋅            ⋅  …          x^5  (x^6)*(y^2)      y^6    ⋅
       ⋅            ⋅            ⋅               ⋅          x^6  y*(x^7)  y^7
       ⋅            ⋅            ⋅               ⋅            ⋅      x^7  x^8

Now we call build_function:

build_function(A,[x,y],parallel=Symbolics.MultithreadedForm())[2]

which generates the code:

(var"##MTIIPVar#317", var"##MTKArg#315")->begin
        @inbounds begin
                @sync begin
                        let (x, y) = (var"##MTKArg#315"[1], var"##MTKArg#315"[2])
                            begin
                                Threads.@spawn begin
                                        (var"##MTIIPVar#317").nzval[1] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 1), (getproperty(Base, :^))(y, 7))
                                        (var"##MTIIPVar#317").nzval[2] = (getproperty(Base, :^))(x, 1)
                                        (var"##MTIIPVar#317").nzval[3] = (getproperty(Base, :^))(y, 1)
                                        (var"##MTIIPVar#317").nzval[4] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 2), (getproperty(Base, :^))(y, 6))
                                        (var"##MTIIPVar#317").nzval[5] = (getproperty(Base, :^))(x, 2)
                                        (var"##MTIIPVar#317").nzval[6] = (getproperty(Base, :^))(y, 2)
                                    end
                            end
                            begin
                                Threads.@spawn begin
                                        (var"##MTIIPVar#317").nzval[7] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 3), (getproperty(Base, :^))(y, 5))
                                        (var"##MTIIPVar#317").nzval[8] = (getproperty(Base, :^))(x, 3)
                                        (var"##MTIIPVar#317").nzval[9] = (getproperty(Base, :^))(y, 3)
                                        (var"##MTIIPVar#317").nzval[10] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 4), (getproperty(Base, :^))(y, 4))
                                        (var"##MTIIPVar#317").nzval[11] = (getproperty(Base, :^))(x, 4)
                                        (var"##MTIIPVar#317").nzval[12] = (getproperty(Base, :^))(y, 4)
                                    end
                            end
                            begin
                                Threads.@spawn begin
                                        (var"##MTIIPVar#317").nzval[13] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 5), (getproperty(Base, :^))(y, 3))
                                        (var"##MTIIPVar#317").nzval[14] = (getproperty(Base, :^))(x, 5)
                                        (var"##MTIIPVar#317").nzval[15] = (getproperty(Base, :^))(y, 5)
                                        (var"##MTIIPVar#317").nzval[16] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 6), (getproperty(Base, :^))(y, 2))
                                        (var"##MTIIPVar#317").nzval[17] = (getproperty(Base, :^))(x, 6)
                                        (var"##MTIIPVar#317").nzval[18] = (getproperty(Base, :^))(y, 6)
                                    end
                            end
                            begin
                                Threads.@spawn begin
                                        (var"##MTIIPVar#317").nzval[19] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 7), (getproperty(Base, :^))(y, 1))
                                        (var"##MTIIPVar#317").nzval[20] = (getproperty(Base, :^))(x, 7)
                                        (var"##MTIIPVar#317").nzval[21] = (getproperty(Base, :^))(y, 7)
                                        (var"##MTIIPVar#317").nzval[22] = (getproperty(Base, :*))((getproperty(Base, :^))(x, 8), (getproperty(Base, :^))(y, 0))
                                    end
                            end
                        end
                    end
            end
        nothing
    end

Let's unpack what that's doing. It's using A.nzval in order to linearly index through the sparse matrix, avoiding the A[i,j] form because that is a more expensive way to index a sparse matrix if you know exactly the order that the data is stored. Then, it's splitting up the calculation into Julia threads so they can be operated on in parallel. It synchronizes after spawning all of the tasks so the computation is ensured to be complete before moving on. And it does this with all in-place operations to the original matrix instead of generating arrays. This is the fanciest way you could fill a sparse matrix, and Symbolics makes this dead simple.

(Note: this example may be slower with multithreading due to the thread spawning overhead, but the full version was not included in the documentation for brevity. It will be the faster version if N is sufficiently large!)

Derivatives

One common thing to compute in a symbolic system is derivatives. In Symbolics.jl, derivatives are represented lazily via operations, just like any other function. To build a differential operator, use Differential like:

@variables t
D = Differential(t)

This is the differential operator $D = \frac{\partial}{\partial t}$. We can compose the differential operator by *, e.g. Differential(t) * Differential(x) or Differential(t)^2. Now let's write down the derivative of some expression:

z = t + t^2
D(z) # Symbolics.derivative(t + t^2, t)

Notice that this hasn't computed anything yet: D is a lazy operator because it lets us symbolically represent "The derivative of $z$ with respect to $t$", which is useful for example when representing our favorite thing in the world, differential equations. However, if we want to expand the derivative operators, we'd use expand_derivatives:

expand_derivatives(D(z)) # 1 + 2t

To get the variable that you are taking the derivative with respect to is accessed with:

D.x # t

We can also have simplified functions for multivariable calculus. For example, we can compute the Jacobian of an array of expressions like:

Symbolics.jacobian([x + x*y, x^2 + y], [x, y])

2×2 Matrix{Num}:
 1 + y  x
    2x  1

and similarly we can do Hessians, gradients, and define whatever other derivatives you want.

Simplification and Substitution

Symbolics interfaces with SymbolicUtils.jl to allow for simplifying symbolic expressions. This is done simply through the simplify command:

simplify(2x + 2y) # 2(x + y)

This can be applied to arrays by using Julia's broadcast mechanism:

B = simplify.([t + t^2 + t + t^2  2t + 4t
               x + y + y + 2t     x^2 - x^2 + y^2])

2×2 Matrix{Num}:
   2(t + t^2)   6t
 x + 2(t + y)  y^2

We can then use substitute to change values of an expression around:

simplify.(substitute.(B, (Dict(x => y^2),)))

2×2 Matrix{Num}:
     2(t + t^2)   6t
 y^2 + 2(t + y)  y^2

and we can use this to interactively evaluate expressions without generating and compiling Julia functions:

V = substitute.(B, (Dict(x => 2.0, y => 3.0, t => 4.0),))

2×2 Matrix{Num}:
 40.0  24.0
 16.0   9.0

Where we can reference the values via:

Symbolics.value.(V)

2×2 Matrix{Float64}:
 40.0  24.0
 16.0   9.0

Non-Interactive Development (No Macro Version)

Note that the macros are for the high-level case where you're doing symbolic computation on your own code. If you want to do symbolic computation on someone else's code, like in a macro, you may not want to do @variables x because you might want the name "x" to come from the user's code. For these cases, Symbolics.jl allows for fully macro-free usage. For example:

using Symbolics: Sym

x = Num(Sym{Float64}(:x))
y = Num(Sym{Float64}(:y))
x + y^2.0 # isa Num

α = Num(Variable(:α))
σ = Num(Variable{Symbolics.FnType{Tuple{Any}, Real}}(:σ)) # left uncalled, since it is used as a function
w = Num(Variable{Symbolics.FnType{Tuple{Any}, Real}}(:w)) # unknown, left uncalled
x = Num(Variable{Symbolics.FnType{Tuple{Any}, Real}}(:x))(t)  # unknown, depends on `t`
y = Num(Variable(:y))   # unknown, no dependents
# Line below throw an error since \alpha is not defined
z = Num(Variable{Symbolics.FnType{NTuple{3, Any}, Real}}(:z))(t, α, x)  # unknown, multiple arguments
β₁ = Num(Variable(:β, 1)) # with index 1
β₂ = Num(Variable(:β, 2)) # with index 2

expr = β₁ * x + y^α + σ(3) * (z - t) - β₂ * w(t - 1)

Does what you'd expect. Note that Variable is simply a convenient function for making variables with indices that always returns Sym. The reference documentation shows how to define any of the quantities in such a way that the names can come from runtime values.

If we need to use this to generate new Julia code, we can simply convert the output to an Expr:

Symbolics.toexpr(x + y^2)

Syms and callable Syms

In the definition

@variables t x(t) y(t)

t is of type Sym{Real} but the name x refers to an object that represents the Termx(t). The operation of this expression is itself the object Sym{FnType{Tuple{Real}, Real}}(:x). The type Sym{FnType{...}} represents a callable object. In this case specifically it's a function that takes 1 Real argument (noted by Tuple{Real}) and returns a Real result. You can call such a callable Sym with either a number or a symbolic expression of a permissible type.

This expression also defines t as a independent variable while x(t) and y(t) are dependent variables. This is accounted for in differentiation:

z = x + y*t
expand_derivatives(D(z)) # derivative(x(t), t) + y(t) + derivative(y(t), t) * t

Since x and y are time-dependent, they are not automatically eliminated from the expression and thus the D(x) and D(y) operations still exist in the expanded derivatives for correctness.

We can also define unrestricted functions:

@variables g(..)

Here g is a variable that is a function of other variables. Any time that we reference g we have to utilize it as a function:

z = g(x) + g(y)

Registering Functions

One of the benefits of a one-language Julia symbolic stack is that the primitives are all written in Julia, and therefore it's trivially extendible from Julia itself. By default, new functions are traced to the primitives and the symbolic expressions are written on the primitives. However, we can expand the allowed primitives by registering new functions. For example, let's register a new function h:

h(x, y) = x^2 + y
@register h(x, y)

Now when we use h(x, y), it is a symbolic expression and doesn't expand:

julia> h(x, y) + y^2
h(x(t), y(t)) + (y(t))^2

In order to use it with the differentiation system, we need to register its derivatives. We would do it like this:

# Derivative w.r.t. the first argument
Symbolics.derivative(::typeof(h), args::NTuple{2,Any}, ::Val{1}) = 2args[1]
# Derivative w.r.t. the second argument
Symbolics.derivative(::typeof(h), args::NTuple{2,Any}, ::Val{2}) = 1

and now it works with the rest of the system:

Symbolics.derivative(h(x, y) + y^2, x) # 2x
Symbolics.derivative(h(x, y) + y^2, y) # 1 + 2y