[BUG]: Cannot load model that has been fitted with TemplateExpressionSpec

[ibengtsson]

What happened?

When trying to load a model (from_file(...)) that has been fitted with a template function, I get a pickle error.

Version

1.5.8

Operating System

macOS

Package Manager

pip

Interface

Script (i.e., python my_script.py)

Relevant log output

Traceback (most recent call last):
  File "/Users/isakbe/Dev/modelling/il-sr/il_sr/scripts/minimal_example.py", line 32, in <module>
    main()
  File "/Users/isakbe/Dev/modelling/il-sr/il_sr/scripts/minimal_example.py", line 28, in main
    loaded_model = PySRRegressor().from_file(run_directory="./outputs/test")
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/isakbe/Library/Caches/pypoetry/virtualenvs/il-sr-9TFUWRsR-py3.11/lib/python3.11/site-packages/pysr/sr.py", line 1140, in from_file
    model = cast("PySRRegressor", pkl.load(f))
                                  ^^^^^^^^^^^
Exception: error deserializing this value

Extra Info

The model used to produce the error was defined as:

    x = np.random.uniform(low=[-1, 3, 10], high=[6, 10, 100], size=(3000, 3))
    y = np.sin(x[:, 0] - x[:, 2] ** 3) * np.exp(x[:, 1] / 5 + x[:, 0]) 
    

    template = TemplateExpressionSpec(
        expressions=["f", "g"],
        variable_names=["x1", "x2", "x3"],
        combine="sin(f(x1, x3)) * exp(g(x1, x2))",
    )

    model = PySRRegressor(
        binary_operators=["+", "-", "*", "/", "^"],
        expression_spec=template,
        constraints={"^": (1, 9)},
        run_id="test",
    )

[MilesCranmer]

Thanks for posting this!

So, unfortunately, I am not sure this will be possible at the moment. The tricky part about template expressions compared to regular expressions is that they contain live Julia objects. And this doesn't play nice with Python pickling.

Regular expressions are fine because they convert to SymPy and are stored in SymPy format. But this is hard to do in a general way for template expressions, because there are so many ways you could write one.

I'm not sure. Maybe it is possible. Because I guess the template expression is defined by the string... Just need a way to serialize the expression objects safely. I guess we could do the normal sympy conversion for both f and g? But we would need a way to go back to Julia after unpickling... Which is going to be super hard. So I'm not sure. Maybe we need a better integration with Julia's Serialization.jl.

[ibengtsson]

Ah, then I understand the root cause... Yes, going back to the template definition in Julia based on a pickled Python object seems hard! I'll work with the .csv-file instead, that'll probably work fine in my case! Appreciate the explanation for why it is hard!

[MilesCranmer]

Just messing around with this a bit more in Julia. Seems like it fails to serialize a closure when you have enclosed variables:

julia> using Serialization

julia> c = 1
1

julia> f = (x, y, z) -> c + x
#15 (generic function with 1 method)

julia> open("tmp.jls", "w") do io
           serialize(io, f)
       end

and in a new session:

julia> using Serialization

julia> f = open("tmp.jls", "r") do io
           deserialize(io)
       end
#15 (generic function with 1 method)

julia> f(1, 2, 3)
ERROR: UndefVarError: `c` not defined in `Main`
Suggestion: check for spelling errors or missing imports.
Stacktrace:
 [1] (::Serialization.__deserialized_types__.var"#15#16")(x::Int64, y::Int64, z::Int64)
   @ Main ./REPL[4]:1
 [2] top-level scope
   @ REPL[3]:1

So I'm not sure what to do for this kind of situation. Maybe there's some other Julia package that can correctly trace the variables...

[GoldenGoldy]

I have exactly the same problem. However, this used to work fine before. I recently upgraded to the latest version of PySR (1.5.8), and now I have the problem described in this issue. And that is when attempting to load certain models that used to load fine before, when using an older version of PySR. I just tested this by downgrading to PySR version 1.0.0, and surely after this I can load models with template expressions again using "from_file". After this, I re-upgraded to 1.5.8 and now exactly the same model doesn't load with "from_file", resulting in the same error as reported in this issue.
Unfortunately, I don't know what is the latest version of PySR with which it does load correctly, I went back to 1.0.0 just as a test to go back as far as possible to see if it would work again.

It seems like something has changed between the mentioned versions that results in the issue? Between these two PySR versions, there are a lot of changes also in the backend/Julia it seems, so changes on that side might result in the issue?
Would this information possibly help to fix it also for the latest version of PySR?

UPDATE:
It actually looks like, at least in my case, the error occurs when trying to load older models (from before updating PySR) with the new version of PySR, or vice versa, when trying to load newer models (generated with the new version of PySR) with the older version of PySR. Is this perhaps related to the changed format/structure for using template expressions?

[MilesCranmer]

Ah, I see. Yes you cannot load checkpoints across versions of PySR. A checkpoint should only be loaded within the same version.

It's because checkpoints may store different information for different versions. Also, if one of the internal types changes, trying to load the binary representation for the old type could result in correctness issues or worse.

[GoldenGoldy]

Right, I wasn't aware of that, but that makes sense, thanks!

[MilesCranmer]

I think we should store the version in the checkpoint so this can be warned about automatically. Would be useful!

[GoldenGoldy]

Sounds like a useful change indeed! And in general it can be useful to be able to see with what version a model was created.

[GongJr0]

Hi @MilesCranmer, @GoldenGoldy

I've been trying to get a model using TemplateExpressionSpec pickled in multiple approaches. I chose to focus on somehow using the PySRRegressor.get_best's output as the pickled object. If we would be able to detach the selected expression from the model, substituting the pysr.expression_specs.CallableJuliaExpression with a python native equivalent (if possible) would likely resolve the issue. I noticed the sub-expressions receive positional identifiers at fit-time in julia. The output is something like:

 4.172e+03  5.126e-03  ╭ f = (#1 - ((-6.3987 - tan(#1)) / 0.82737)) - (tan(sqrt(#1)) ...
                                  / 0.50056)
                                  ╰ w = [1.2479, -0.077092, 0.034817, ..., 0.06366]

where #1 marks the position to place the expression with positional-id = 1. Could trying to evaluate the expression as a string by substituting the variables and arguments in a try/except block be a crude fix for this? In output I've shown for example, I'm confident the expression is lambda compatible since it is a simple AR(n) [i.e. $w_1X_1+...+w_nX_n$ ] process being passed to the regressor.

This is the definition of the template I use:

    def get_expr(self) -> TemplateExpressionSpec:
        var_set = [col for col in self.X.columns if col.startswith(self.target)]
        params = {"w": len(var_set)}

        sub_expr = " + ".join([f"w[{n+1}] * {col}" for n, col in enumerate(var_set)])  # n+1 to match julia's 1-based indexing
        return TemplateExpressionSpec(
            expressions=["f"],
            variable_names=var_set,
            parameters=params,
            combine=f"f({sub_expr})"

having a check for lambdaify compatibility in string form may resolve the pickling issue if we substitute the callble julia object with the with the evaluated version inside PySRRegressor.equations_. I know this may effect the prediction process since (as far as I know) the predictions are being made by selecting the equation from the DataFrame instead of having the best eq stored separately. If so, a flag to swap the julia callable with the lambda format can be used specifically when the user calls get_best and the callable can remain in the equations_ frame for prediction functionality. Would something like this be in scope for PySR?