Add sorting option to get_variables function

src/utils.jl

@@ -20,31 +20,40 @@ function build_expr(head::Symbol, args)
 end
 
 """
-    get_variables(O) -> Vector{BasicSymbolic}
+    get_variables(e, varlist = nothing; sort::Bool = false)
 
-Returns the variables in the expression. Note that the returned variables are
-not wrapped in the `Num` type.
+Return a vector of variables appearing in `e`, optionally restricting to variables in `varlist`.
+
+Note that the returned variables are not wrapped in the `Num` type.
 
 # Examples
 ```julia
-julia> @variables t x y z(t)
-4-element Vector{Num}:
-    t
-    x
-    y
- z(t)
+julia> @variables t x y z(t);
 
-julia> ex = x + y + sin(z)
-(x + y) + sin(z(t))
-
-julia> Symbolics.get_variables(ex)
-3-element Vector{Any}:
+julia> Symbolics.get_variables(x + y + sin(z); sort = true)
+3-element Vector{SymbolicUtils.BasicSymbolic}:
  x
  y
  z(t)
+
+julia> Symbolics.get_variables(x - y; sort = true)
+2-element Vector{SymbolicUtils.BasicSymbolic}:
+ x
+ y
 ```
 """
-get_variables(e::Num, varlist=nothing) = get_variables(value(e), varlist)
+function get_variables(e::Num, varlist = nothing; sort::Bool = false)
+    get_variables(value(e), varlist; sort)
+end
+function get_variables(e, varlist = nothing; sort::Bool = false)
+    vars = Vector{BasicSymbolic}()
+    get_variables!(vars, e, varlist)
+    if sort
+        sort!(vars; by = SymbolicUtils.get_degrees)
+    end
+    vars
+end
+
 get_variables!(vars, e::Num, varlist=nothing) = get_variables!(vars, value(e), varlist)
 get_variables!(vars, e, varlist=nothing) = vars
 
@@ -76,8 +85,6 @@ function get_variables!(vars, e::Equation, varlist=nothing)
   get_variables!(vars, e.rhs, varlist)
 end
 
-get_variables(e, varlist=nothing) = get_variables!([], e, varlist)
-
 # Sym / Term --> Symbol
 Base.Symbol(x::Union{Num,Symbolic}) = tosymbol(x)
 tosymbol(t::Num; kwargs...) = tosymbol(value(t); kwargs...)

test/runtests.jl

@@ -53,6 +53,7 @@ if GROUP == "All" || GROUP == "Core"
         @safetestset "LuxCore extensions Test" begin include("extensions/lux.jl") end
         @safetestset "Registration without using Test" begin include("registration_without_using.jl") end
         @safetestset "Show Test" begin include("show.jl") end
+        @safetestset "Utility Function Test" begin include("utils.jl") end
     end
 end
 

test/utils.jl

@@ -0,0 +1,21 @@
+using Symbolics
+
+@testset "get_variables" begin
+    @variables t x y z(t)
+
+    ex1 = x + y + sin(z)
+    vars1 = Symbolics.get_variables(ex1)
+    @test length(vars1) == 3
+    @test allunique(vars1)
+
+    sorted_vars1 = Symbolics.get_variables(ex1; sort = true)
+    @test isequal(sorted_vars1, [x, y, z])
+
+    ex2 = x - y
+    vars2 = Symbolics.get_variables(ex2)
+    @test length(vars2) == 2
+    @test allunique(vars2)
+
+    sorted_vars2 = Symbolics.get_variables(ex2; sort = true)
+    @test isequal(sorted_vars2, [x, y])
+end