"Custom function my_pow throws MethodError despite NaN handling"

[Hosseinmghasemi]

Description

I'm trying to define a custom function my_pow in PYSR that should handle negative inputs by returning NaN. However, despite various attempts to handle this in the Julia code, PYSR still throws a MethodError when the function is evaluated with negative inputs.
I'd like to use PYSR to test the function such that it tries the power of features to the powers of n, which ranges between 3 to 9.

Code

import numpy as np
import sympy
from pysr import PySRRegressor, 
default_pysr_params = dict(
    populations=80,
    model_selection="best",
)


# Generate some data
np.random.seed(0)
X = np.random.randn(100, 5)
y = np.sum(X**2, axis=1) + 0.1 * np.random.randn(100)


# Define a custom power function class
class my_pow(sympy.Function):
    @classmethod
    def eval(cls, x, n):
        if x.is_negative:
            return sympy.nan
        return x ** n


# Learn equations
model_p = PySRRegressor(
    niterations=60,
    binary_operators=["+", "*","/", "-"],
    # unary_operators=["my_pow(x::T, n::T) where T = (x >= zero(T)) ? x^clamp(n, 3, 9) : T(NaN)"],
    # unary_operators=[ "my_pow(x::T, n::T) where T = abs(x)^clamp(n, 3, 9)"],
    unary_operators=[
        "my_pow(x, n) = x >= 0 ? x^(3 + 6 * (n - 0.5)^2) : convert(typeof(x), NaN)"
    ],    
    extra_sympy_mappings={"my_pow": lambda x, n: np.where(x >= 0, x ** (3 + 6 * (n - 0.5)**2), np.nan)},
    parsimony = 6.4e-8,
    **default_pysr_params,
)

model_p.fit(X2, y)

explanation :
I use Julia's ternary operator ?: to handle invalid inputs.
x >= 0 ? checks if x is non-negative.
If x is non-negative, we compute x^(3 + 6 * (n - 0.5)^2).

This is a more complex expression that ensures n is always between 3 and 9:

(n - 0.5)^2 maps any n to a value between 0 and 0.25.
6 * (n - 0.5)^2 maps this to a value between 0 and 6.
Adding 3 gives us a range of 3 to 9.
If x is negative, we return convert(typeof(x), NaN) to ensure type stability.

But I faced the following issue:
```julia
---------------------------------------------------------------------------
JuliaError                                Traceback (most recent call last)
[<ipython-input-26-d46dfbb57121>](https://localhost:8080/#) in <cell line: 24>()
     22 )
     23 
---> 24 model_p.fit(X2, y)

2 frames
[/usr/local/lib/python3.10/dist-packages/juliacall/__init__.py](https://localhost:8080/#) in __call__(self, *args, **kwargs)
    221             return ValueBase.__dir__(self) + self._jl_callmethod($(pyjl_methodnum(pyjlany_dir)))
    222         def __call__(self, *args, **kwargs):
--> 223             return self._jl_callmethod($(pyjl_methodnum(pyjlany_call)), args, kwargs)
    224         def __bool__(self):
    225             return True

JuliaError: The operator `my_pow` is not well-defined over the real line, as it threw the error `MethodError` when evaluating the input -100.0. You can work around this by returning NaN for invalid inputs. For example, `safe_log(x::T) where {T} = x > 0 ? log(x) : T(NaN)`.

I would appreciate it if any idea how to fix this issue or other alternatives.

[MilesCranmer]

I think it's because you put my_pow in unary operators, but it is actually a binary operator.

(I guess the error message for such an error should be better)

---

Also your extra_sympy_mappings can't use numpy functions – only sympy ones.

[Hosseinmghasemi]

Thank you so much for your reply.

I changed the code to the following :

# Define a custom power function class
class my_pow(sympy.Function):
    @classmethod
    def eval(cls, x, n):
        if x.is_negative:
            return sympy.nan
        return x ** n

def my_pow_mapping(x, n):
    return sympy.Piecewise(
        (sympy.Abs(x) ** (3 + 6 * (n - 0.5)**2), x >= 0),
        (sympy.nan, True)
    )

# Learn equations
model_p = PySRRegressor(
    niterations=60,
    binary_operators=["+", "*","/", "-","my_pow"],
    unary_operators=["exp"],  
    extra_sympy_mappings={"my_pow": my_pow_mapping},
    parsimony = 6.4e-8,
    **default_pysr_params,
)

model_p.fit(X2, y)

but the issue is :
JuliaError: The operator my_pow returned an output of type Float64, when it was given inputs 0.0 and -100.0 of type Float32. Please ensure that your operators return the same type as their inputs.

I tried many solutions based on the documentation, but I could not resolve it.
I appreciate it if you have any idea

[Hosseinmghasemi]

I found the solution and the code works ,
we need to explicitly type x and n as Float64. This is crucial because it tells Julia exactly what type to expect and return.

# Learn equations
model_p = PySRRegressor(
    niterations=60,
    binary_operators=["+", "*","/", "-","pow_n(x::Float64, n::Float64) = x > 0 ? x^(3 + 6*(n-0.5)^2) / (x^2) : NaN"],
    unary_operators=["my_pow(x) = x^2"],
    extra_sympy_mappings={"my_pow": lambda x: x**2,
                          "pow_n": lambda x, n: sympy.Piecewise(
                          ((sympy.Abs(x) ** (3 + 6 * (n - 0.5)**2))/ x**2 , x >= 0),
                          (sympy.nan, True))
},
    parsimony = 6.4e-8,
    **default_pysr_params,
)

model_p.fit(X2, y)

[MilesCranmer]

For this you should set precision=64 as the default is Float32.

Alternatively you can do in the

"pow_n(x::T, n::T) where {T} = x > 0 ? x^(3 + 6*(n-convert(T, 0.5))^2) / (x^2) : convert(T, NaN)"

and it will automatically convert stuff to either Float32 OR Float64, depending on the user input setting.

[Hosseinmghasemi]

Awesome !
Thank you so much for your great help.