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)
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)
Sym
s and callable Sym
s
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 Term
x(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