Add missing PyTorch/JAX export for logical_or, logical_and, and relu

pysr/export_jax.py

@@ -69,7 +69,7 @@ def sympy2jaxtext(expr, parameters, symbols_in, extra_jax_mappings=None):
         _func = {**_jnp_func_lookup, **extra_jax_mappings}[expr.func]
     except KeyError:
         raise KeyError(
-            f"Function {expr.func} was not found in JAX function mappings."
+            f"Function {expr.func} was not found in JAX function mappings. "
             "Please add it to extra_jax_mappings in the format, e.g., "
             "{sympy.sqrt: 'jnp.sqrt'}."
         )

pysr/export_torch.py

@@ -33,6 +33,68 @@ def _initialize_torch():
 
         torch = _torch
 
+        # Allows PyTorch to map Piecewise functions:
+        def expr_cond_pair(expr, cond):
+            if isinstance(cond, torch.Tensor) and not isinstance(expr, torch.Tensor):
+                expr = torch.tensor(expr, dtype=cond.dtype, device=cond.device)
+            elif isinstance(expr, torch.Tensor) and not isinstance(cond, torch.Tensor):
+                cond = torch.tensor(cond, dtype=expr.dtype, device=expr.device)
+            else:
+                return expr, cond
+
+            # First, make sure expr and cond are same size:
+            if expr.shape != cond.shape:
+                if len(expr.shape) == 0:
+                    expr = expr.expand(cond.shape)
+                elif len(cond.shape) == 0:
+                    cond = cond.expand(expr.shape)
+                else:
+                    raise ValueError(
+                        "expr and cond must have same shape, or one must be a scalar."
+                    )
+            return expr, cond
+
+def if_then_else(*conds):
+    a, b, c = conds
+    return torch.where(a, torch.where(b, True, False), torch.where(c, True, False))
+
+        def piecewise(*expr_conds):
+            output = None
+            already_used = None
+            for expr, cond in expr_conds:
+                if not isinstance(cond, torch.Tensor) and not isinstance(
+                    expr, torch.Tensor
+                ):
+                    # When we just have scalars, have to do this a bit more complicated
+                    # due to the fact that we need to evaluate on the correct device.
+                    if output is None:
+                        already_used = cond
+                        output = expr if cond else 0.0
+                    else:
+                        if not isinstance(output, torch.Tensor):
+                            output += expr if cond and not already_used else 0.0
+                            already_used = already_used or cond
+                        else:
+                            expr = torch.tensor(
+                                expr, dtype=output.dtype, device=output.device
+                            ).expand(output.shape)
+                            output += torch.where(
+                                cond & ~already_used, expr, torch.zeros_like(expr)
+                            )
+                            already_used = already_used | cond
+                else:
+                    if output is None:
+                        already_used = cond
+                        output = torch.where(cond, expr, torch.zeros_like(expr))
+                    else:
+                        output += torch.where(
+                            cond.bool() & ~already_used, expr, torch.zeros_like(expr)
+                        )
+                        already_used = already_used | cond.bool()
+            return output
+
+        # TODO: Add test that makes sure tensors are on the same device
+
         _global_func_lookup = {
             sympy.Mul: _reduce(torch.mul),
             sympy.Add: _reduce(torch.add),
@@ -81,6 +143,12 @@ def _initialize_torch():
             sympy.Heaviside: torch.heaviside,
             sympy.core.numbers.Half: (lambda: 0.5),
             sympy.core.numbers.One: (lambda: 1.0),
+            sympy.logic.boolalg.Boolean: lambda x: x,
+            sympy.logic.boolalg.BooleanTrue: (lambda: True),
+            sympy.logic.boolalg.BooleanFalse: (lambda: False),
+            sympy.functions.elementary.piecewise.ExprCondPair: expr_cond_pair,
+            sympy.Piecewise: piecewise,
+            sympy.logic.boolalg.ITE: if_then_else,
         }
 
         class _Node(torch.nn.Module):
@@ -125,7 +193,7 @@ def __init__(self, *, expr, _memodict, _func_lookup, **kwargs):
                         self._torch_func = _func_lookup[expr.func]
                     except KeyError:
                         raise KeyError(
-                            f"Function {expr.func} was not found in Torch function mappings."
+                            f"Function {expr.func} was not found in Torch function mappings. "
                             "Please add it to extra_torch_mappings in the format, e.g., "
                             "{sympy.sqrt: torch.sqrt}."
                         )
@@ -153,7 +221,13 @@ def forward(self, memodict):
                         arg_ = arg(memodict)
                         memodict[arg] = arg_
                     args.append(arg_)
-                return self._torch_func(*args)
+                try:
+                    return self._torch_func(*args)
+                except Exception as err:
+                    # Add information about the current node to the error:
+                    raise type(err)(
+                        f"Error occurred in node {self._sympy_func} with args {args}"
+                    )
 
         class _SingleSymPyModule(torch.nn.Module):
             """SympyTorch code from https://github.com/patrick-kidger/sympytorch"""

pysr/sr.py

@@ -603,8 +603,12 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         Path to the temporary equations directory.
     equation_file_ : str
         Output equation file name produced by the julia backend.
-    raw_julia_state_ : tuple[list[PyCall.jlwrap], PyCall.jlwrap]
+    sr_state_ : tuple[list[PyCall.jlwrap], PyCall.jlwrap]
         The state for the julia SymbolicRegression.jl backend post fitting.
+    sr_options_ : PyCall.jlwrap
+        The options used by `SymbolicRegression.jl`, created during
+        a call to `.fit`. You may use this to manually call functions
+        in `SymbolicRegression` which take an `::Options` argument.
     equation_file_contents_ : list[pandas.DataFrame]
         Contents of the equation file output by the Julia backend.
     show_pickle_warnings_ : bool
@@ -1031,7 +1035,7 @@ def __getstate__(self):
         serialization.
 
         Thus, for `PySRRegressor` to support pickle serialization, the
-        `raw_julia_state_` attribute must be hidden from pickle. This will
+        `sr_state_` attribute must be hidden from pickle. This will
         prevent the `warm_start` of any model that is loaded via `pickle.loads()`,
         but does allow all other attributes of a fitted `PySRRegressor` estimator
         to be serialized. Note: Jax and Torch format equations are also removed
@@ -1041,9 +1045,9 @@ def __getstate__(self):
         show_pickle_warning = not (
             "show_pickle_warnings_" in state and not state["show_pickle_warnings_"]
         )
-        if "raw_julia_state_" in state and show_pickle_warning:
+        if ("sr_state_" in state or "sr_options_" in state) and show_pickle_warning:
             warnings.warn(
-                "raw_julia_state_ cannot be pickled and will be removed from the "
+                "sr_state_ and sr_options_ cannot be pickled and will be removed from the "
                 "serialized instance. This will prevent a `warm_start` fit of any "
                 "model that is deserialized via `pickle.load()`."
             )
@@ -1055,7 +1059,10 @@ def __getstate__(self):
                     "serialized instance. When loading the model, please redefine "
                     f"`{state_key}` at runtime."
                 )
-        state_keys_to_clear = ["raw_julia_state_"] + state_keys_containing_lambdas
+        state_keys_to_clear = [
+            "sr_state_",
+            "sr_options_",
+        ] + state_keys_containing_lambdas
         pickled_state = {
             key: (None if key in state_keys_to_clear else value)
             for key, value in state.items()
@@ -1105,6 +1112,14 @@ def equations(self):  # pragma: no cover
         )
         return self.equations_
 
+    @property
+    def raw_julia_state_(self):  # pragma: no cover
+        warnings.warn(
+            "PySRRegressor.raw_julia_state_ is now deprecated. "
+            "Please use PySRRegressor.sr_state_ instead.",
+        )
+        return self.sr_state_
+
     def get_best(self, index=None):
         """
         Get best equation using `model_selection`.
@@ -1605,7 +1620,7 @@ def _run(self, X, y, mutated_params, weights, seed):
 
         # Call to Julia backend.
         # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/OptionsStruct.jl
-        options = SymbolicRegression.Options(
+        self.sr_options_ = SymbolicRegression.Options(
             binary_operators=Main.eval(str(binary_operators).replace("'", "")),
             unary_operators=Main.eval(str(unary_operators).replace("'", "")),
             bin_constraints=bin_constraints,
@@ -1704,7 +1719,7 @@ def _run(self, X, y, mutated_params, weights, seed):
 
         # Call to Julia backend.
         # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl
-        self.raw_julia_state_ = SymbolicRegression.equation_search(
+        self.sr_state_ = SymbolicRegression.equation_search(
             Main.X,
             Main.y,
             weights=Main.weights,
@@ -1714,10 +1729,10 @@ def _run(self, X, y, mutated_params, weights, seed):
             y_variable_names=y_variable_names,
             X_units=self.X_units_,
             y_units=self.y_units_,
-            options=options,
+            options=self.sr_options_,
             numprocs=cprocs,
             parallelism=parallelism,
-            saved_state=self.raw_julia_state_,
+            saved_state=self.sr_state_,
             return_state=True,
             addprocs_function=cluster_manager,
             progress=progress and self.verbosity > 0 and len(y.shape) == 1,
@@ -1786,10 +1801,10 @@ def fit(
             Fitted estimator.
         """
         # Init attributes that are not specified in BaseEstimator
-        if self.warm_start and hasattr(self, "raw_julia_state_"):
+        if self.warm_start and hasattr(self, "sr_state_"):
             pass
         else:
-            if hasattr(self, "raw_julia_state_"):
+            if hasattr(self, "sr_state_"):
                 warnings.warn(
                     "The discovered expressions are being reset. "
                     "Please set `warm_start=True` if you wish to continue "
@@ -1799,7 +1814,8 @@ def fit(
             self.equations_ = None
             self.nout_ = 1
             self.selection_mask_ = None
-            self.raw_julia_state_ = None
+            self.sr_state_ = None
+            self.sr_options_ = None
             self.X_units_ = None
             self.y_units_ = None
 

pysr/test/test.py

@@ -109,7 +109,7 @@ def test_high_precision_search_custom_loss(self):
         from pysr.sr import Main
 
         # We should have that the model state is now a Float64 hof:
-        Main.test_state = model.raw_julia_state_
+        Main.test_state = model.sr_state_
         self.assertTrue(Main.eval("typeof(test_state[2]).parameters[1] == Float64"))
 
     def test_multioutput_custom_operator_quiet_custom_complexity(self):
@@ -232,7 +232,7 @@ def test_empty_operators_single_input_warm_start(self):
         from pysr.sr import Main
 
         # We should have that the model state is now a Float32 hof:
-        Main.test_state = regressor.raw_julia_state_
+        Main.test_state = regressor.sr_state_
         self.assertTrue(Main.eval("typeof(test_state[2]).parameters[1] == Float32"))
         # This should exit almost immediately, and use the old equations
         regressor.fit(X, y)