From 5af63544c76d9d5c7e112d60cce4dc5f0f36f3f9 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Tue, 8 Jun 2021 16:36:55 -0400
Subject: [PATCH] Clean up more unused code

---
 pysr/export_torch.py     | 8 +++-----
 pysr/feynman_problems.py | 3 +--
 pysr/sr.py               | 1 -
 test/test.py             | 6 +++---
 4 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/pysr/export_torch.py b/pysr/export_torch.py
index bc15651da..107205e54 100644
--- a/pysr/export_torch.py
+++ b/pysr/export_torch.py
@@ -17,16 +17,12 @@ def fn_(*args):
 
 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 forward(self, memodict):
                     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 forward(self, X):
                 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;
diff --git a/pysr/feynman_problems.py b/pysr/feynman_problems.py
index ef6f261bf..f30b02438 100644
--- a/pysr/feynman_problems.py
+++ b/pysr/feynman_problems.py
@@ -92,7 +92,7 @@ def mk_problems(first=100, gen=False, dp=500, data_dir=FEYNMAN_DATASET):
                     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 = []
diff --git a/pysr/sr.py b/pysr/sr.py
index c8bd33294..2feef1113 100644
--- a/pysr/sr.py
+++ b/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:
diff --git a/test/test.py b/test/test.py
index e32e4d48d..87d955149 100644
--- a/test/test.py
+++ b/test/test.py
@@ -48,10 +48,10 @@ def test_multioutput_weighted_with_callable(self):
         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 @@ def test_feature_selection_handler(self):
             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))