[Feature] Support external julia libraries for custom operators

[evilmav]

Is your feature request related to a problem? Please describe. I'd like to use a bessel function as a custom binary operator with PySR. It seems that I need to use an external library like SpecialFunctions.jl. However, I can not seem to figure out an easy way to use a library function with binary_operators argument. Everything I've tried so far ended up in JULIA: AssertionError: Your configuration is invalid - one of your operators (besselj) is not well-defined over the real line., although I assume this means a compile rather then mathematical issue underneath.

Describe the solution you'd like A custom parameter with library names to additionally load passed to PySRRegressor could probably be doable?

Additional context I'm not a Julia user, so I might be missing something trivial.

[MilesCranmer]

Great question. Two notes:

• For loading external Julia libraries, you need to manually call a few Julia functions. See example below.
• besselj actually assumes non-negative inputs. PySR assumes all operators are valid over the entire real line, so we need to create a modified operator that works for negative inputs.

Here's a working example:

import numpy as np

# Import bessel function:
import sympy
import scipy.special as sp
import pysr

# Import PyJulia, and import other library:
jl = pysr.julia_helpers.init_julia()

# If necessary:
jl.eval("using Pkg")
jl.eval('Pkg.add("SpecialFunctions")')

# Import the library
jl.eval("using SpecialFunctions")

bessel_function = "besselj_abs(x, y) = besselj(abs(x), abs(y))"

model = pysr.PySRRegressor(
    binary_operators=["+", "*", bessel_function],
    extra_sympy_mappings={
        "besselj_abs": lambda x, y: sympy.besselj(sympy.Abs(x), sympy.Abs(y))
    },
    # No bessel inside bessel:
    nested_constraints={"besselj_abs": {"besselj_abs": 0}},
)

X = np.random.rand(500, 2) * 10

# Actual function we wish to find:
y = sp.jv(0, X[:, 0]) * sp.jv(1, X[:, 1])

model.fit(X, y)

It seems to correctly get the Bessel function for me!

>>> model.sympy()
besselj(1.2737758e-9, Abs(x0))*besselj(1.0000007, Abs(x1))

[MilesCranmer]

besselj seems to be really slow for large inputs (I guess that's a property of bessel functions?). Therefore, you could instead clip the inputs to the range [0, 5] so it is even faster:

bessel_function = """besselj_clip(x::T, y::T) where {T} = begin
    # Same type as x and y:
    zero = T(0)
    five = T(5)

    # Clip x, y to [0, 5]
    x = max(x, zero)
    x = min(x, five)
    y = max(y, zero)
    y = min(y, five)

    # Call function:
    return besselj(x, y)
end
"""

The ::T and where {T} is just templating the function to input type T.

[MilesCranmer]

As a side note, I think I should just make all the SpecialFunctions.jl functions available by default in PySR, since it has a lot of useful functions. What do you think?

[evilmav]

Wow, thank you for the fast response and all the work! This should do the trick =)

As a side note, I think I should just make all the SpecialFunctions.jl functions available by default in PySR, since it has a lot of useful functions. What do you think?

I agree =)

PS I'd be happy to buy you a beer for this amazing project - feel free to drop me a link if you wish.

[MilesCranmer]

Awesome; will add it.

Happy to hear you are enjoying the package! No need for the tip, but thanks for the offer – be sure to spread the word about PySR though :) (Although if we meet in person I will gladly take you up on a beer!)

[MilesCranmer]

This feature should be fully supported now, see the Primes.jl example on https://astroautomata.com/PySR/examples/.

Cheers, Miles