Switch to using sympytorch

pysr/export_torch.py

@@ -14,151 +14,46 @@ def _reduce(fn):
 
 torch_initialized = False
 torch = None
-_global_func_lookup = None
-_Node = None
-SingleSymPyModule = None
+sympytorch = None
+PySRTorchModule = None
 
 def _initialize_torch():
     global torch_initialized
     global torch
-    global _global_func_lookup
-    global _Node
-    global SingleSymPyModule
+    global sympytorch
+    global PySRTorchModule
 
     # 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
+        try:
+            import torch
+            import sympytorch
+        except ImportError:
+            raise ImportError("You need to pip install `torch` and `sympytorch` before exporting to pytorch.")
+        torch_initialized = True
 
-        _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):
+        class PySRTorchModule(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,
+            def __init__(self, *, expression, symbols_in,
                          selection=None, 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._module = sympytorch.SymPyModule(
+                        expressions=[expression],
+                        extra_funcs=extra_funcs)
                 self._selection = selection
-                self.symbols_in = [str(symbol) for symbol in symbols_in]
-
+                self._symbols = symbols_in
+                
             def __repr__(self):
                 return f"{type(self).__name__}(expression={self._expression_string})"
 
             def forward(self, X):
                 if self._selection is not None:
                     X = X[:, self._selection]
-                symbols = {symbol: X[:, i]
-                           for i, symbol in enumerate(self.symbols_in)}
-                return self._node(symbols)
+                symbols = {str(symbol): X[:, i]
+                           for i, symbol in enumerate(self._symbols)}
+                return self._module(**symbols)[..., 0]
 
 
 def sympy2torch(expression, symbols_in,
@@ -168,10 +63,11 @@ def sympy2torch(expression, symbols_in,
     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
+    global PySRTorchModule
 
     _initialize_torch()
 
-    return SingleSymPyModule(expression, symbols_in,
-                             selection=selection,
-                             extra_funcs=extra_torch_mappings)
+    return PySRTorchModule(expression=expression,
+                           symbols_in=symbols_in,
+                           selection=selection,
+                           extra_funcs=extra_torch_mappings)

setup.py

@@ -1,8 +1,19 @@
+import importlib.util
 import setuptools
 
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
+extra_installs = []
+
+torch_installed = (importlib.util.find_spec('torch') is not None)
+install_sympytorch = torch_installed
+
+if install_sympytorch:
+    extra_installs.append('sympytorch')
+
+print(extra_installs)
+
 setuptools.setup(
     name="pysr", # Replace with your own username
     version="0.6.0rc1",
@@ -12,11 +23,11 @@ setuptools.setup(
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/MilesCranmer/pysr",
-    install_requires=[
+    install_requires=([
         "numpy",
         "pandas",
         "sympy"
-        ],
+        ] + extra_installs),
     packages=setuptools.find_packages(),
     package_data={
         'pysr': ['../Project.toml', '../datasets/*']
@@ -26,5 +37,5 @@ setuptools.setup(
         "Programming Language :: Python :: 3",
         "Operating System :: OS Independent",
     ],
-    python_requires='>=3.3',
+    python_requires='>=3.7',
 )