[WIP] experimental way to access expression methods

pysr/__init__.py

@@ -4,6 +4,7 @@
 from .export_torch import sympy2torch
 from .feynman_problems import FeynmanProblem, Problem
 from .julia_helpers import install
+from .pysr_expression import Expression
 from .sr import PySRRegressor
 from .version import __version__
 
@@ -15,6 +16,7 @@
     "Problem",
     "install",
     "PySRRegressor",
+    "Expression",
     "best",
     "best_callable",
     "best_row",

pysr/pysr_expression.py

@@ -0,0 +1,75 @@
+from numbers import Number
+from typing import List, Optional
+
+from .sr import PySRRegressor
+
+
+class Expression:
+    """A wrapper around `SymbolicRegression.Node`"""
+
+    def __init__(
+        self,
+        equation,
+        *,
+        model: PySRRegressor = None,
+        options=None,
+        variable_names: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        # exactly one of model and options is None:
+        assert (model is None) != (
+            options is None
+        ), "Pass exactly one of model and options"
+
+        self.equation = equation
+        self.options = model.sr_options_ if options is None else options
+        self.variable_names = (
+            variable_names
+            if variable_names is not None
+            else (model.feature_names_in_ if model is not None else None)
+        )
+
+        from julia import Main, SymbolicRegression
+
+        self.julia_ = Main
+        self.backend_ = SymbolicRegression
+
+    @classmethod
+    def from_string(
+        cls,
+        s: str,
+        *,
+        model: PySRRegressor = None,
+        options=None,
+        variable_names: Optional[List[str]] = None,
+    ):
+        self = cls(None, model=model, options=options, variable_names=variable_names)
+
+        for i, variable in enumerate(self.variable_names):
+            self.julia_.eval(f"{variable} = Node(feature={i + 1})")
+
+        self.julia_.last_options = self.options
+        self.julia_.eval("SymbolicRegression.@extend_operators last_options")
+
+        equation = self.julia_.eval(s)
+
+        if isinstance(equation, Number):
+            equation = self.julia_.eval(f"Node(val={equation})")
+
+        self.equation = equation
+
+        return self
+
+    def __repr__(self):
+        variable_names = (
+            list(self.variable_names) if self.variable_names is not None else None
+        )
+        return self.backend_.string_tree(
+            self.equation, self.options, variable_names=variable_names
+        )
+
+    def __call__(self, X):
+        return self.equation(X.T, self.options)
+
+    def compute_complexity(self) -> int:
+        return int(self.backend_.compute_complexity(self.equation, self.options))