Clean up more unused code

pysr/export_torch.py

@@ -17,16 +17,12 @@ def _reduce(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,
@@ -148,7 +144,7 @@ def _initialize_torch():
                     args.append(arg_)
                 return self._torch_func(*args)
 
-        class SingleSymPyModule(torch.nn.Module):
+        class _SingleSymPyModule(torch.nn.Module):
             """SympyTorch code from https://github.com/patrick-kidger/sympytorch"""
 
             def __init__(
@@ -177,6 +173,8 @@ def _initialize_torch():
                 symbols = {symbol: X[:, i] for i, symbol in enumerate(self.symbols_in)}
                 return self._node(symbols)
 
+        SingleSymPyModule = _SingleSymPyModule
+
 
 def sympy2torch(expression, symbols_in, selection=None, extra_torch_mappings=None):
     """Returns a module for a given sympy expression with trainable parameters;

pysr/feynman_problems.py

@@ -92,7 +92,7 @@ class FeynmanProblem(Problem):
                     ret.append(p)
                 except Exception as e:
                     traceback.print_exc()
-                    print(f"FAILED ON ROW {i}")
+                    print(f"FAILED ON ROW {i} with {e}")
                 ind += 1
         return ret
 
@@ -168,7 +168,6 @@ def do_feynman_experiments(
     problems = FeynmanProblem.mk_problems(
         first=first, gen=True, dp=dp, data_dir=data_dir
     )
-    indx = range(len(problems))
     ids = []
     predictions = []
     true_equations = []

pysr/sr.py

@@ -400,7 +400,6 @@ def pysr(
     kwargs = {**dict(equation_file=equation_file), **kwargs}
 
     pkg_directory = kwargs["pkg_directory"]
-    manifest_file = None
     if kwargs["julia_project"] is not None:
         manifest_filepath = Path(kwargs["julia_project"]) / "Manifest.toml"
     else:

test/test.py

@@ -48,10 +48,10 @@ class TestPipeline(unittest.TestCase):
         w[w >= 0.5] = 1.0
 
         # Double equation when weights are 0:
-        y += (1 - w) * y
+        y = (2 - w) * y
         # Thus, pysr needs to use the weights to find the right equation!
 
-        equations = pysr(
+        pysr(
             self.X,
             y,
             weights=w,
@@ -140,7 +140,7 @@ class TestFeatureSelection(unittest.TestCase):
             X,
             select_k_features=2,
             use_custom_variable_names=True,
-            variable_names=[f"x{i}" for i in range(5)],
+            variable_names=var_names,
             y=y,
         )
         self.assertTrue((2 in selection) and (3 in selection))