Merge pull request #48 from MilesCranmer/torch-export

.github/workflows/CI.yml

@@ -61,17 +61,23 @@ jobs:
             python setup.py install
       - name: "Install Coverage tool"
         run: pip install coverage coveralls
+      - name: "Run tests"
+        run: coverage run --source=pysr --omit='*/feynman_problems.py' -m unittest test.test
+        shell: bash
       - name: "Install JAX"
         if: matrix.os != 'windows-latest'
         run: pip install jax jaxlib # (optional import)
         shell: bash
-      - name: "Run tests"
-        run: coverage run --source=pysr --omit='*/feynman_problems.py' -m unittest test.test
-        shell: bash
       - name: "Run JAX tests"
         if: matrix.os != 'windows-latest'
         run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_jax
         shell: bash
+      - name: "Install Torch"
+        run: pip install torch # (optional import)
+        shell: bash
+      - name: "Run Torch tests"
+        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_torch
+        shell: bash
       - name: Coveralls
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

pysr/__init__.py

@@ -1,3 +1,4 @@
 from .sr import pysr, get_hof, best, best_tex, best_callable, best_row
 from .feynman_problems import Problem, FeynmanProblem
-from .export import sympy2jax
+from .export_jax import sympy2jax
+from .export_torch import sympy2torch

pysr/{export.py → export_jax.py}

@@ -3,59 +3,53 @@ import sympy
 import string
 import random
 
-try:
-    import jax
-    from jax import numpy as jnp
-    from jax.scipy import special as jsp
-
 # Special since need to reduce arguments.
-    MUL = 0
-    ADD = 1
-
-    _jnp_func_lookup = {
-        sympy.Mul: MUL,
-        sympy.Add: ADD,
-        sympy.div: "jnp.div",
-        sympy.Abs: "jnp.abs",
-        sympy.sign: "jnp.sign",
-        # Note: May raise error for ints.
-        sympy.ceiling: "jnp.ceil",
-        sympy.floor: "jnp.floor",
-        sympy.log: "jnp.log",
-        sympy.exp: "jnp.exp",
-        sympy.sqrt: "jnp.sqrt",
-        sympy.cos: "jnp.cos",
-        sympy.acos: "jnp.acos",
-        sympy.sin: "jnp.sin",
-        sympy.asin: "jnp.asin",
-        sympy.tan: "jnp.tan",
-        sympy.atan: "jnp.atan",
-        sympy.atan2: "jnp.atan2",
-        # Note: Also may give NaN for complex results.
-        sympy.cosh: "jnp.cosh",
-        sympy.acosh: "jnp.acosh",
-        sympy.sinh: "jnp.sinh",
-        sympy.asinh: "jnp.asinh",
-        sympy.tanh: "jnp.tanh",
-        sympy.atanh: "jnp.atanh",
-        sympy.Pow: "jnp.power",
-        sympy.re: "jnp.real",
-        sympy.im: "jnp.imag",
-        sympy.arg: "jnp.angle",
-        # Note: May raise error for ints and complexes
-        sympy.erf: "jsp.erf",
-        sympy.erfc: "jsp.erfc",
-        sympy.LessThan: "jnp.less",
-        sympy.GreaterThan: "jnp.greater",
-        sympy.And: "jnp.logical_and",
-        sympy.Or: "jnp.logical_or",
-        sympy.Not: "jnp.logical_not",
-        sympy.Max: "jnp.max",
-        sympy.Min: "jnp.min",
-        sympy.Mod: "jnp.mod",
-    }
-except ImportError:
-    ...
+MUL = 0
+ADD = 1
+
+_jnp_func_lookup = {
+    sympy.Mul: MUL,
+    sympy.Add: ADD,
+    sympy.div: "jnp.div",
+    sympy.Abs: "jnp.abs",
+    sympy.sign: "jnp.sign",
+    # Note: May raise error for ints.
+    sympy.ceiling: "jnp.ceil",
+    sympy.floor: "jnp.floor",
+    sympy.log: "jnp.log",
+    sympy.exp: "jnp.exp",
+    sympy.sqrt: "jnp.sqrt",
+    sympy.cos: "jnp.cos",
+    sympy.acos: "jnp.acos",
+    sympy.sin: "jnp.sin",
+    sympy.asin: "jnp.asin",
+    sympy.tan: "jnp.tan",
+    sympy.atan: "jnp.atan",
+    sympy.atan2: "jnp.atan2",
+    # Note: Also may give NaN for complex results.
+    sympy.cosh: "jnp.cosh",
+    sympy.acosh: "jnp.acosh",
+    sympy.sinh: "jnp.sinh",
+    sympy.asinh: "jnp.asinh",
+    sympy.tanh: "jnp.tanh",
+    sympy.atanh: "jnp.atanh",
+    sympy.Pow: "jnp.power",
+    sympy.re: "jnp.real",
+    sympy.im: "jnp.imag",
+    sympy.arg: "jnp.angle",
+    # Note: May raise error for ints and complexes
+    sympy.erf: "jsp.erf",
+    sympy.erfc: "jsp.erfc",
+    sympy.LessThan: "jnp.less",
+    sympy.GreaterThan: "jnp.greater",
+    sympy.And: "jnp.logical_and",
+    sympy.Or: "jnp.logical_or",
+    sympy.Not: "jnp.logical_not",
+    sympy.Max: "jnp.max",
+    sympy.Min: "jnp.min",
+    sympy.Mod: "jnp.mod",
+}
+
 
 def sympy2jaxtext(expr, parameters, symbols_in):
     if issubclass(expr.func, sympy.Float):
@@ -75,7 +69,28 @@ def sympy2jaxtext(expr, parameters, symbols_in):
         else:
             return f'{_func}({", ".join(args)})'
 
-def sympy2jax(equation, symbols_in):
+
+jax_initialized = False
+jax = None
+jnp = None
+jsp = None
+
+def _initialize_jax():
+    global jax_initialized
+    global jax
+    global jnp
+    global jsp
+
+    if not jax_initialized:
+        import jax as _jax
+        from jax import numpy as _jnp
+        from jax.scipy import special as _jsp
+        jax = _jax
+        jnp = _jnp
+        jsp = _jsp
+
+
+def sympy2jax(expression, symbols_in):
     """Returns a function f and its parameters;
     the function takes an input matrix, and a list of arguments:
             f(X, parameters)
@@ -146,9 +161,15 @@ def sympy2jax(equation, symbols_in):
         #                3.5427954 , -2.7479894 ], dtype=float32)
         ```
     """
+    _initialize_jax()
+    global jax_initialized
+    global jax
+    global jnp
+    global jsp
+
     parameters = []
-    functional_form_text = sympy2jaxtext(equation, parameters, symbols_in)
-    hash_string = 'A_' + str(abs(hash(str(equation) + str(symbols_in))))
+    functional_form_text = sympy2jaxtext(expression, parameters, symbols_in)
+    hash_string = 'A_' + str(abs(hash(str(expression) + str(symbols_in))))
     text = f"def {hash_string}(X, parameters):\n"
     text += "    return "
     text += functional_form_text

pysr/export_torch.py

@@ -0,0 +1,171 @@
+#####
+# From https://github.com/patrick-kidger/sympytorch
+# Copied here to allow PySR-specific tweaks
+#####
+
+import collections as co
+import functools as ft
+import sympy
+
+def _reduce(fn):
+    def fn_(*args):
+        return ft.reduce(fn, args)
+    return fn_
+
+torch_initialized = False
+torch = None
+_global_func_lookup = None
+_Node = None
+SingleSymPyModule = None
+
+def _initialize_torch():
+    global torch_initialized
+    global torch
+    global _global_func_lookup
+    global _Node
+    global SingleSymPyModule
+
+    # Way to lazy load torch, only if this is called,
+    # but still allow this module to be loaded in __init__
+    if not torch_initialized:
+        import torch as _torch
+        torch = _torch
+
+        _global_func_lookup = {
+            sympy.Mul: _reduce(torch.mul),
+            sympy.Add: _reduce(torch.add),
+            sympy.div: torch.div,
+            sympy.Abs: torch.abs,
+            sympy.sign: torch.sign,
+            # Note: May raise error for ints.
+            sympy.ceiling: torch.ceil,
+            sympy.floor: torch.floor,
+            sympy.log: torch.log,
+            sympy.exp: torch.exp,
+            sympy.sqrt: torch.sqrt,
+            sympy.cos: torch.cos,
+            sympy.acos: torch.acos,
+            sympy.sin: torch.sin,
+            sympy.asin: torch.asin,
+            sympy.tan: torch.tan,
+            sympy.atan: torch.atan,
+            sympy.atan2: torch.atan2,
+            # Note: May give NaN for complex results.
+            sympy.cosh: torch.cosh,
+            sympy.acosh: torch.acosh,
+            sympy.sinh: torch.sinh,
+            sympy.asinh: torch.asinh,
+            sympy.tanh: torch.tanh,
+            sympy.atanh: torch.atanh,
+            sympy.Pow: torch.pow,
+            sympy.re: torch.real,
+            sympy.im: torch.imag,
+            sympy.arg: torch.angle,
+            # Note: May raise error for ints and complexes
+            sympy.erf: torch.erf,
+            sympy.loggamma: torch.lgamma,
+            sympy.Eq: torch.eq,
+            sympy.Ne: torch.ne,
+            sympy.StrictGreaterThan: torch.gt,
+            sympy.StrictLessThan: torch.lt,
+            sympy.LessThan: torch.le,
+            sympy.GreaterThan: torch.ge,
+            sympy.And: torch.logical_and,
+            sympy.Or: torch.logical_or,
+            sympy.Not: torch.logical_not,
+            sympy.Max: torch.max,
+            sympy.Min: torch.min,
+            # Matrices
+            sympy.MatAdd: torch.add,
+            sympy.HadamardProduct: torch.mul,
+            sympy.Trace: torch.trace,
+            # Note: May raise error for integer matrices.
+            sympy.Determinant: torch.det,
+        }
+
+        class _Node(torch.nn.Module):
+            """SympyTorch code from https://github.com/patrick-kidger/sympytorch"""
+            def __init__(self, *, expr, _memodict, _func_lookup, **kwargs):
+                super().__init__(**kwargs)
+
+                self._sympy_func = expr.func
+
+                if issubclass(expr.func, sympy.Float):
+                    self._value = torch.nn.Parameter(torch.tensor(float(expr)))
+                    self._torch_func = lambda: self._value
+                    self._args = ()
+                elif issubclass(expr.func, sympy.UnevaluatedExpr):
+                    if len(expr.args) != 1 or not issubclass(expr.args[0].func, sympy.Float):
+                        raise ValueError("UnevaluatedExpr should only be used to wrap floats.")
+                    self.register_buffer('_value', torch.tensor(float(expr.args[0])))
+                    self._torch_func = lambda: self._value
+                    self._args = ()
+                elif issubclass(expr.func, sympy.Integer):
+                    # Can get here if expr is one of the Integer special cases,
+                    # e.g. NegativeOne
+                    self._value = int(expr)
+                    self._torch_func = lambda: self._value
+                    self._args = ()
+                elif issubclass(expr.func, sympy.Symbol):
+                    self._name = expr.name
+                    self._torch_func = lambda value: value
+                    self._args = ((lambda memodict: memodict[expr.name]),)
+                else:
+                    self._torch_func = _func_lookup[expr.func]
+                    args = []
+                    for arg in expr.args:
+                        try:
+                            arg_ = _memodict[arg]
+                        except KeyError:
+                            arg_ = type(self)(expr=arg, _memodict=_memodict, _func_lookup=_func_lookup, **kwargs)
+                            _memodict[arg] = arg_
+                        args.append(arg_)
+                    self._args = torch.nn.ModuleList(args)
+
+            def forward(self, memodict):
+                args = []
+                for arg in self._args:
+                    try:
+                        arg_ = memodict[arg]
+                    except KeyError:
+                        arg_ = arg(memodict)
+                        memodict[arg] = arg_
+                    args.append(arg_)
+                return self._torch_func(*args)
+
+
+        class SingleSymPyModule(torch.nn.Module):
+            """SympyTorch code from https://github.com/patrick-kidger/sympytorch"""
+            def __init__(self, expression, symbols_in,
+                    extra_funcs=None, **kwargs):
+                super().__init__(**kwargs)
+                
+                if extra_funcs is None:
+                    extra_funcs = {}
+                _func_lookup = co.ChainMap(_global_func_lookup, extra_funcs)
+
+                _memodict = {}
+                self._node = _Node(expr=expression, _memodict=_memodict, _func_lookup=_func_lookup)
+                self._expression_string = str(expression)
+                self.symbols_in = [str(symbol) for symbol in symbols_in]
+
+            def __repr__(self):
+                return f"{type(self).__name__}(expression={self._expression_string})"
+
+            def forward(self, X):
+                symbols = {symbol: X[:, i]
+                           for i, symbol in enumerate(self.symbols_in)}
+                return self._node(symbols)
+
+
+def sympy2torch(expression, symbols_in, extra_torch_mappings=None):
+    """Returns a module for a given sympy expression with trainable parameters;
+
+    This function will assume the input to the module is a matrix X, where
+        each column corresponds to each symbol you pass in `symbols_in`.
+    """
+    global SingleSymPyModule
+
+    _initialize_torch()
+
+    return SingleSymPyModule(expression, symbols_in, extra_funcs=extra_torch_mappings)

pysr/sr.py

@@ -13,7 +13,6 @@ import shutil
 from pathlib import Path
 from datetime import datetime
 import warnings
-from .export import sympy2jax
 
 global_equation_file = 'hall_of_fame.csv'
 global_n_features = None
@@ -103,6 +102,8 @@ def pysr(X, y, weights=None,
          perturbationFactor=1.0,
          timeout=None,
          extra_sympy_mappings=None,
+         extra_torch_mappings=None,
+         extra_jax_mappings=None,
          equation_file=None,
          verbosity=1e9,
          progress=True,
@@ -124,11 +125,12 @@ def pysr(X, y, weights=None,
          update=True,
          temp_equation_file=False,
          output_jax_format=False,
+         output_torch_format=False,
          optimizer_algorithm="BFGS",
          optimizer_nrestarts=3,
          optimize_probability=1.0,
-         optimizer_iterations=10,
-        ):
+         optimizer_iterations=10
+         ):
     """Run symbolic regression to fit f(X[i, :]) ~ y[i] for all i.
     Note: most default parameters have been tuned over several example
     equations, but you should adjust `niterations`,
@@ -241,6 +243,8 @@ def pysr(X, y, weights=None,
         delete_tempfiles argument.
     :param output_jax_format: Whether to create a 'jax_format' column in the output,
         containing jax-callable functions and the default parameters in a jax array.
+    :param output_torch_format: Whether to create a 'torch_format' column in the output,
+        containing a torch module with trainable parameters.
     :returns: pd.DataFrame or list, Results dataframe,
         giving complexity, MSE, and equations (as strings), as well as functional
         forms. If list, each element corresponds to a dataframe of equations
@@ -334,8 +338,11 @@ def pysr(X, y, weights=None,
                  weightSimplify=weightSimplify,
                  constraints=constraints,
                  extra_sympy_mappings=extra_sympy_mappings,
+                 extra_jax_mappings=extra_jax_mappings,
+                 extra_torch_mappings=extra_torch_mappings,
                  julia_project=julia_project, loss=loss,
                  output_jax_format=output_jax_format,
+                 output_torch_format=output_torch_format,
                  multioutput=multioutput, nout=nout)
 
     kwargs = {**_set_paths(tempdir), **kwargs}
@@ -715,10 +722,10 @@ def run_feature_selection(X, y, select_k_features):
         the k most important features in X, returning indices for those
         features as output."""
 
-    from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
+    from sklearn.ensemble import RandomForestRegressor
     from sklearn.feature_selection import SelectFromModel, SelectKBest
 
-    clf = GradientBoostingRegressor(n_estimators=100, learning_rate=0.1, max_depth=1, random_state=0, loss='ls') #RandomForestRegressor()
+    clf = RandomForestRegressor(n_estimators=100, max_depth=3, random_state=0)
     clf.fit(X, y)
     selector = SelectFromModel(clf, threshold=-np.inf,
             max_features=select_k_features, prefit=True)
@@ -726,6 +733,8 @@ def run_feature_selection(X, y, select_k_features):
 
 def get_hof(equation_file=None, n_features=None, variable_names=None,
             extra_sympy_mappings=None, output_jax_format=False,
+            output_torch_format=False,
+            extra_jax_mappings=None, extra_torch_mappings=None,
             multioutput=None, nout=None, **kwargs):
     """Get the equations from a hall of fame file. If no arguments
     entered, the ones used previously from a call to PySR will be used."""
@@ -770,6 +779,8 @@ def get_hof(equation_file=None, n_features=None, variable_names=None,
         lambda_format = []
         if output_jax_format:
             jax_format = []
+        if output_torch_format:
+            torch_format = []
         use_custom_variable_names = (len(variable_names) != 0)
         local_sympy_mappings = {
                 **extra_sympy_mappings,
@@ -784,11 +795,22 @@ def get_hof(equation_file=None, n_features=None, variable_names=None,
         for i in range(len(output)):
             eqn = sympify(output.loc[i, 'Equation'], locals=local_sympy_mappings)
             sympy_format.append(eqn)
+
+            # Numpy:
+            lambda_format.append(CallableEquation(sympy_symbols, eqn))
+
+            # JAX:
             if output_jax_format:
+                from .export_jax import sympy2jax
                 func, params = sympy2jax(eqn, sympy_symbols)
                 jax_format.append({'callable': func, 'parameters': params})
 
-            lambda_format.append(CallableEquation(sympy_symbols, eqn))
+            # Torch:
+            if output_torch_format:
+                from .export_torch import sympy2torch
+                module = sympy2torch(eqn, sympy_symbols)
+                torch_format.append(module)
+
             curMSE = output.loc[i, 'MSE']
             curComplexity = output.loc[i, 'Complexity']
 
@@ -808,6 +830,9 @@ def get_hof(equation_file=None, n_features=None, variable_names=None,
         if output_jax_format:
             output_cols += ['jax_format']
             output['jax_format'] = jax_format
+        if output_torch_format:
+            output_cols += ['torch_format']
+            output['torch_format'] = torch_format
 
         ret_outputs.append(output[output_cols])
 

test/test_jax.py

@@ -1,6 +1,7 @@
 import unittest
 import numpy as np
-from pysr import pysr, sympy2jax
+from pysr import sympy2jax, get_hof
+import pandas as pd
 from jax import numpy as jnp
 from jax import random
 from jax import grad
@@ -15,3 +16,24 @@ class TestJAX(unittest.TestCase):
         true = 1.0 * jnp.cos(X[:, 0]) + X[:, 1]
         f, params = sympy2jax(cosx, [x, y, z])
         self.assertTrue(jnp.all(jnp.isclose(f(X, params), true)).item())
+    def test_pipeline(self):
+        X = np.random.randn(100, 2)
+        equations = pd.DataFrame({
+            'Equation': ['1.0', 'cos(x0)', 'square(cos(x0))'],
+            'MSE': [1.0, 0.1, 1e-5],
+            'Complexity': [1, 2, 3]
+            })
+
+        equations['Complexity MSE Equation'.split(' ')].to_csv(
+                'equation_file.csv.bkup', sep='|')
+
+        equations = get_hof(
+                'equation_file.csv', n_features=2, variables_names='x0 x1'.split(' '),
+                extra_sympy_mappings={}, output_jax_format=True,
+                multioutput=False, nout=1)
+
+        jformat = equations.iloc[-1].jax_format
+        np.testing.assert_almost_equal(
+                np.array(jformat['callable'](jnp.array(X), jformat['parameters'])),
+                np.square(np.cos(X[:, 0]))
+        )

test/test_torch.py

@@ -0,0 +1,38 @@
+import unittest
+import numpy as np
+import pandas as pd
+from pysr import sympy2torch, get_hof
+import torch
+import sympy
+
+class TestTorch(unittest.TestCase):
+    def test_sympy2torch(self):
+        x, y, z = sympy.symbols('x y z')
+        cosx = 1.0 * sympy.cos(x) + y
+        X = torch.randn((1000, 3))
+        true = 1.0 * torch.cos(X[:, 0]) + X[:, 1]
+        torch_module = sympy2torch(cosx, [x, y, z])
+        self.assertTrue(
+                np.all(np.isclose(torch_module(X).detach().numpy(), true.detach().numpy()))
+        )
+    def test_pipeline(self):
+        X = np.random.randn(100, 2)
+        equations = pd.DataFrame({
+            'Equation': ['1.0', 'cos(x0)', 'square(cos(x0))'],
+            'MSE': [1.0, 0.1, 1e-5],
+            'Complexity': [1, 2, 3]
+            })
+
+        equations['Complexity MSE Equation'.split(' ')].to_csv(
+                'equation_file.csv.bkup', sep='|')
+
+        equations = get_hof(
+                'equation_file.csv', n_features=2, variables_names='x0 x1'.split(' '),
+                extra_sympy_mappings={}, output_torch_format=True,
+                multioutput=False, nout=1)
+
+        tformat = equations.iloc[-1].torch_format
+        np.testing.assert_almost_equal(
+                tformat(torch.tensor(X)).detach().numpy(),
+                np.square(np.cos(X[:, 0]))
+        )