Add torch format output; dont import jax/torch by default

pysr/__init__.py

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

pysr/export_jax.py

@@ -2,60 +2,56 @@ import functools as ft
 import sympy
 import string
 import random
-
-try:
-    import jax
-    from jax import numpy as jnp
-    from jax.scipy import special as jsp
+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
+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:
-    ...
+_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):

pysr/export_torch.py

@@ -8,7 +8,6 @@ import functools as ft
 import sympy
 import torch
 
-
 def _reduce(fn):
     def fn_(*args):
         return ft.reduce(fn, args)
@@ -67,7 +66,6 @@ _global_func_lookup = {
     sympy.Determinant: torch.det,
 }
 
-
 class _Node(torch.nn.Module):
     def __init__(self, *, expr, _memodict, _func_lookup, **kwargs):
         super().__init__(**kwargs)

pysr/sr.py

@@ -13,8 +13,6 @@ import shutil
 from pathlib import Path
 from datetime import datetime
 import warnings
-from .export_jax import sympy2jax
-from .export_torch import sympy2torch
 
 global_equation_file = 'hall_of_fame.csv'
 global_n_features = None
@@ -125,11 +123,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`,
@@ -242,6 +241,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
@@ -337,6 +338,7 @@ def pysr(X, y, weights=None,
                  extra_sympy_mappings=extra_sympy_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}
@@ -727,6 +729,7 @@ 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,
             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."""
@@ -771,6 +774,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,
@@ -786,10 +791,19 @@ def get_hof(equation_file=None, n_features=None, variable_names=None,
             eqn = sympify(output.loc[i, 'Equation'], locals=local_sympy_mappings)
             sympy_format.append(eqn)
             if output_jax_format:
+                from .export_jax import sympy2jax
                 func, params = sympy2jax(eqn, sympy_symbols)
                 jax_format.append({'callable': func, 'parameters': params})
+<<<<<<< HEAD
 
             lambda_format.append(CallableEquation(sympy_symbols, eqn))
+=======
+            if output_torch_format:
+                from .export_torch import sympy2torch
+                func, params = sympy2torch(eqn, sympy_symbols)
+                torch_format.append({'callable': func, 'parameters': params})
+            lambda_format.append(lambdify(sympy_symbols, eqn))
+>>>>>>> 6ba697f (Add torch format output; dont import jax/torch by default)
             curMSE = output.loc[i, 'MSE']
             curComplexity = output.loc[i, 'Complexity']
 
@@ -809,6 +823,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])