Potential bug in parsing LLM suggestions

[jyothis-fp]

I am trying LaSR on a dataset from the LLM-SRBench repository. The ground truth equation is $v = -c \cdot \sqrt{1 - \frac{m_0^2}{m^2}}$. My llm_options is:

  active: True
  weights: {'llm_mutate': 0.0, 'llm_crossover': 0.0, 'llm_gen_random': 0.01}
  num_pareto_context: 3
  api_key: [HIDDEN]
  model: gpt-4.1
  prompt_evol: True
  prompt_concepts: True
  llm_recorder_dir: 'log/'
  prompts_dir: 'prompts/'
  llm_context: 'Find an equation that describes the object's velocity (v), given physical input variables: the relativistic mass of an object (m), the object's rest mass (m_0), the speed of light (c).'
  var_order: {'v': 'v', 'm': 'm', 'm_0': 'm_0', 'c': 'c'}
  idea_threshold: 30
  api_kwargs: {'max_tokens': 4096, 'url': 'https://api.openai.com/v1/'}
  http_kwargs: {'retries': 5, 'readtimeout': 360}

There seems to be a mismatch between LLM output and the equation that LaSR generates as a result. The LLM suggests v = c / sqrt((m / m_0)^2 - 1), but LaSR seems to take it as cos(c / sqrt(((m / m_0) ^ 2.0) * 1.0)). I am attaching the llm_calls log here. The log contains the usual llm input and output, as well as candidate equations (raw expressions from the llm output) and accepted expressions (expressions after parsing). I have added thread id to each item for disambiguation.

Caveat: I am not sure if my var_order is causing the issue.

llm_calls.txt

@MilesCranmer

[MilesCranmer]

@atharvas any idea?

[atharvas]

Hi! Sorry for the delay and thanks for the bug report. We weren't able to find a pre-existing way to parse string equations into julia expressions so had to ship our own. This is implemented in src/Parse.jl. It seems like we weren't supporting y = ... expressions but it is concerning that this wasn't being straight up rejected (and instead an incorrect expression was returned). Adding equation support in a new PR and I'll also add some harder test cases drawn from the llm_calls.txt file. Let me know about any other corner cases and I'll be happy to add it to the list of special edge cases we want to support.

[MilesCranmer]

@atharvas could we use DynamicExpressions.parse_expression? It's used for the new guesses API

[atharvas]

Hi!

#64 should fix this. @MilesCranmer I tried using the DynamicExpressions.parse_expression but I ran into a couple of issues with my roundtrip tests. I can open an issue in DynamicExpressions if this seems like a bug (rather than a LaSR specific issue).

Issue 1: Stack overflow on certain edge cases

========
Parse.jl
========
function parse_expr(
    ::Type{T}, expr_str::String, options::AbstractOptions
) where {T<:DATA_TYPE}
    try
        expr_str = replace(expr_str, r"\*\*" => "^")
        return parse_expression(
            Meta.parse(expr_str),
            operators=options.operators,
            node_type=options.node_type,
            expression_type=options.expression_type,
            variable_names=get_variable_names(options.variable_names)
        )
    catch e
        @info "Failed to parse expression: $expr_str"
        @info "Error: $e"
        @info "Returning a constant node with value 1."
        return _make_constant_node(1, T)
    end
end
====
REPL
====
julia> tree = random_trees[1]
x3 + (((x1 + 0.8001067) + exp(0.14465569)) - -0.36010084)

julia> str_tree = string_tree(tree, options)
"x3 + (((x1 + 0.8001067) + exp(0.14465569)) - -0.36010084)"

julia> output = tree(data, options.operators)

1-element Vector{Float32}:
 2.6800885

julia> str_tree == String(strip(str_tree, [' ', '\n', '"', ',', '.', '[', ']']))
true

julia> expr_tree = parse_expr(Float32, str_tree, options)
x3 + (((x1 + 0.8001067) + exp(0.14465569)) - -0.36010084)

julia> expr_output = expr_tree(data, options.operators)
ERROR: StackOverflowError:
Stacktrace:
     [1] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{…}, cX::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}; kws::@Kwargs{})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:240
     [2] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{…}, cX::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:240
     [3] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{…}, cX::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}; kws::@Kwargs{})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:249
--- the above 2 lines are repeated 20928 more times ---
 [41860] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{…}, cX::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:240
 [41861] (::DynamicExpressions.NodeModule.Node{…})(X::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}; kws::@Kwargs{})
       @ DynamicExpressions.EvaluationHelpersModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/EvaluationHelpers.jl:37
 [41862] (::DynamicExpressions.NodeModule.Node{…})(X::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…})
       @ DynamicExpressions.EvaluationHelpersModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/EvaluationHelpers.jl:36
 [41863] (::DynamicExpressions.ExpressionModule.Expression{…})(X::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}; kws::@Kwargs{})
       @ DynamicExpressions.ExpressionModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Expression.jl:513
 [41864] (::DynamicExpressions.ExpressionModule.Expression{…})(X::Matrix{…}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…})
       @ DynamicExpressions.ExpressionModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Expression.jl:509
Some type information was truncated. Use `show(err)` to see complete types.

julia> show(err)
1-element ExceptionStack:
StackOverflowError:
Stacktrace:
     [1] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{Any}, cX::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}}; kws::@Kwargs{})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:240
     [2] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{Any}, cX::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:240
     [3] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{Any}, cX::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}}; kws::@Kwargs{})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:249
--- the above 2 lines are repeated 20928 more times ---
 [41860] eval_tree_array(tree::DynamicExpressions.NodeModule.Node{String}, cX::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}})
       @ DynamicExpressions.EvaluateModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Evaluate.jl:240
 [41861] (::DynamicExpressions.NodeModule.Node{String})(X::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}}; kws::@Kwargs{})
       @ DynamicExpressions.EvaluationHelpersModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/EvaluationHelpers.jl:37
 [41862] (::DynamicExpressions.NodeModule.Node{String})(X::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}})
       @ DynamicExpressions.EvaluationHelpersModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/EvaluationHelpers.jl:36
 [41863] (::DynamicExpressions.ExpressionModule.Expression{String, DynamicExpressions.NodeModule.Node{String}, @NamedTuple{operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{}, Tuple{}}, variable_names::Nothing}})(X::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}}; kws::@Kwargs{})
       @ DynamicExpressions.ExpressionModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Expression.jl:513
 [41864] (::DynamicExpressions.ExpressionModule.Expression{String, DynamicExpressions.NodeModule.Node{String}, @NamedTuple{operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{}, Tuple{}}, variable_names::Nothing}})(X::Matrix{Float32}, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{Tuple{typeof(+), typeof(*), typeof(safe_pow), typeof(-)}, Tuple{typeof(sin), typeof(cos), typeof(exp)}})
       @ DynamicExpressions.ExpressionModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Expression.jl:509

Issue 2: Inline functions.

(imports)
include("test/test_params.jl")
const T = Float32

@inline sinf(x)    = sin(T(x))::T
@inline cosf(x)    = cos(T(x))::T
@inline expf(x)    = exp(T(clamp(x, -40.0, 40.0)))::T
@inline safelog(x) = log(abs(T(x)) + eps(T))::T
@inline sqr(x)  = (t = T(x); (t*t)::T)
@inline cube(x) = (t = T(x); (t*t*t)::T)
@inline addf(a,b) = (T(a) + T(b))::T
@inline mulf(a,b) = (T(a) * T(b))::T
@inline divf(a,b) = (ta=T(a); tb=T(b); (ta / (abs(tb) + eps(T)))::T)

options = LaSROptions(;
    default_params...,
    binary_operators = [addf, mulf, divf],
    unary_operators  = [sinf, cosf, expf, safelog, sqr, cube],
    nested_constraints = [
        divf    => [divf => 1, expf => 0],
        safelog => [safelog => 0, sinf => 0, cosf => 0, divf => 0, expf => 0],
        sinf    => [sinf => 0, safelog => 0, expf => 0],
        cosf    => [cosf => 0, safelog => 0, expf => 0],
        (expf)    => [(expf) => 0, (sqr) => 0, (cube) => 0, (divf) => 0],
        sqr     => [expf => 0],
        cube    => [expf => 0],
    ],
    constraints = [
        safelog => 5,
        sinf    => 7,
        cosf    => 7,
        expf    => 1,
        sqr     => 3,
        cube    => 2,
    ],
)
(...)

julia> for (i, tree) in enumerate(random_trees)
           output = tree(data, options.operators)
           println("[$i] Testing tree: $tree")
           if any(isnan.(output))
               println("Skipping tree with NaN output: $tree")
               continue
           end
           str_tree = string_tree(tree, options)
           t1 = parse_expression(
               Meta.parse(str_tree),
               operators=options.operators,
               node_type=options.node_type,
               expression_type=options.expression_type,
               variable_names=get_variable_names(options.variable_names)
           )
           # t2 = parse_expression(str_tree)
           # The string might not always be perfectly formatted. Introducing noise.
           str_tree = fuzz_string(str_tree)
           expr_tree = parse_expr(T, str_tree, options)
           expr_output = expr_tree(data, options.operators)
           # @test isapprox(expr_output, output)
           @assert isapprox(expr_output, output) "[$i] Failed for tree: $str_tree"
       end

[1] Testing tree: sinf(divf(sinf(mulf(1.1563058, -0.18177138)), mulf(0.6890417, 0.8001067)))
ERROR: ArgumentError: Tried to interpolate function `sinf` but failed.
Stacktrace:
 [1] _parse_expression(ex::Expr, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}, variable_names::Vector{…}, ::Type{…}, ::Type{…}, evaluate_on::Nothing; kws::@Kwargs{})
   @ DynamicExpressions.ParseModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Parse.jl:267
 [2] _parse_expression
   @ /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Parse.jl:248 [inlined]
 [3] parse_expression(ex::Expr; operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}, binary_operators::Nothing, unary_operators::Nothing, variable_names::Vector{…}, expression_type::Type{…}, node_type::Type{…}, evaluate_on::Nothing, kws::@Kwargs{})
   @ DynamicExpressions.ParseModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Parse.jl:240
 [4] top-level scope
   @ ./REPL[40]:9

caused by: UndefVarError: `sinf` not defined in `DynamicExpressions.ParseModule.EmptyModule`
Suggestion: check for spelling errors or missing imports.
Suggestion: check for spelling errors or missing imports.
Stacktrace:
 [1] top-level scope
   @ :0
 [2] eval
   @ ./boot.jl:430 [inlined]
 [3] _parse_expression(ex::Expr, operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}, variable_names::Vector{…}, ::Type{…}, ::Type{…}, evaluate_on::Nothing; kws::@Kwargs{})
   @ DynamicExpressions.ParseModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Parse.jl:265
 [4] _parse_expression
   @ /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Parse.jl:248 [inlined]
 [5] parse_expression(ex::Expr; operators::DynamicExpressions.OperatorEnumModule.OperatorEnum{…}, binary_operators::Nothing, unary_operators::Nothing, variable_names::Vector{…}, expression_type::Type{…}, node_type::Type{…}, evaluate_on::Nothing, kws::@Kwargs{})
   @ DynamicExpressions.ParseModule /var/local/atharvas/.julia/packages/DynamicExpressions/cYbpm/src/Parse.jl:240
 [6] top-level scope
   @ ./REPL[40]:9
Some type information was truncated. Use `show(err)` to see complete types.

[MilesCranmer]

Thanks, I will look into 1. I think 2. is fixed in the most recent DynamicExpressions.jl, what version are you on?

julia> using DynamicQuantities

julia> const T = Float32
Float32

julia> @inline sinf(x)    = sin(T(x))::T;

julia> @inline cosf(x)    = cos(T(x))::T;

julia> @inline addf(a,b) = (T(a) + T(b))::T;

julia> @inline mulf(a,b) = (T(a) * T(b))::T;

julia> @inline divf(a,b) = (ta=T(a); tb=T(b); (ta / (abs(tb) + eps(T)))::T);

julia> operators = OperatorEnum(2 => [addf, mulf, divf], 1 => [sinf, cosf])
OperatorEnum{Tuple{Tuple{typeof(sinf), typeof(cosf)}, Tuple{typeof(addf), typeof(mulf), typeof(divf)}}}(((sinf, cosf), (addf, mulf, divf)))

julia> parse_expression(Meta.parse("sinf(addf(x1, -0.5))"); operators, variable_names=["x1"])
sinf(addf(x1, -0.5))

In general @inline shouldn't affect execution. It's just a hint to the compiler that it's worth inlining the code. But the compiler won't do so if it changes the result.

[MilesCranmer]

I can't reproduce 1. either. Maybe it's also fixed on the most recent DynamicExpressions.jl?

julia> operators = OperatorEnum(1 => [exp], 2 => [+, -, *])
OperatorEnum{Tuple{Tuple{typeof(exp)}, Tuple{typeof(+), typeof(-), typeof(*)}}}(((exp,), (+, -, *)))

julia> s = "x3 + (((x1 + 0.8001067) + exp(0.14465569)) - -0.36010084)";

julia> expr = parse_expression(Meta.parse(s); variable_names=["x1", "x2", "x3"], operators, node_type=Node{Float32})
x3 + (((x1 + 0.8001067) + exp(0.14465569)) - -0.36010084)

julia> expr(randn(Float32, 3, 10))
10-element Vector{Float32}:
# ...

[atharvas]

Ah. I was on SR.jl 1.12

(LibraryAugmentedSymbolicReg...) pkg> status DynamicExpressions
Project LibraryAugmentedSymbolicRegression v0.3.0
Status `/var/local/atharvas/lasr-experiments/LibraryAugmentedSymbolicRegression.jl/Project.toml`
⌅ [a40a106e] DynamicExpressions v1.10.3
Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated`

Manually checked out 2.0.0 alpha6. I need to resolve some dependency issues but this looks promising! Also, I can't integrate this into LaSR right now because Julia hub only has the versions up to 1.12 here.

[MilesCranmer]

You should be able to do it locally with

[sources]
SymbolicRegression = { url = "https://github.com/MilesCranmer/SymbolicRegression.jl", rev="v2.0.0-alpha.6" }

in your Project.toml.

But yeah this format can't be registered. It's just for local testing

[MilesCranmer]

@atharvas quick ping on this! Let me know if I can help

[atharvas]

Sorry for the delay @MilesCranmer @jyothis-fp . PR #67 integrates v.20 SR.jl and simplifies the Parse.jl functionality. This still isn't a deep integration with the guesses API (which is really cool -- I played around with it). I'm going to be leaving that to v0.4.0 which I'll be working on late September. By the way, let me know if I'm blocking any of y'all on this and happy to move faster.

[MilesCranmer]

@atharvas would you have a chance to work on this relatively soon? @jyothis-fp has shared he is developing a symbolic regression benchmark - I am a bit worried that LaSR won't be on the leaderboard if he can't get it working in time :/

[jyothis-fp]

Hey @atharvas @MilesCranmer, thanks for working on the fix. I was wondering: is it possible to access the latest code from the Python wrapper or is it Julia-only for now?

[jyothis-fp]

When trying the latest code, I encounter a couple of issues:

1. Example: with variable names: Dict("theta" => "theta", "p_d" => "p_d", "T" => "T", "kb" => "kb", "Ef" => "Ef", "n" => "n") and the expression n * cos(theta) * exp((p_d * Ef) / (kb * T)), I get a "failed to parse expression" warning. The expression then gets normalized into n * cos(theta) * exp((p_d * ef) / (kb * t)), which in turn throws an error because of the lower-case ef.
2. For operators like sqrt, I get an error because it's expecting safe_sqrt. If I am not mistaken, this has been fixed by @MilesCranmer in the latest version of DynamicExpressions.jl.

[atharvas]

Hey guys! Sorry for the delay! Super busy with ICLR till 24th of September.

latest code from python wrapper

No this is a julia side fix for now. Sorry I don't have the bandwidth this week to work on this but happy to take a stab at integrating the python interface afterwards!

normalized into ef

Thanks for the analysis. Seems like this line is to blame. This is a quick fix. Give me a couple of minutes.