Run in `Main`

pysr/julia_helpers.py

@@ -11,11 +11,8 @@ if os.environ.get("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes") != "yes":
 
 os.environ["PYTHON_JULIACALL_HANDLE_SIGNALS"] = "yes"
 
-import juliacall
 import juliapkg
-
-jl = juliacall.newmodule("PySR")
-
+from juliacall import Main as jl
 from juliacall import convert as jl_convert
 
 jl.seval("using PythonCall: PythonCall")
@@ -53,3 +50,9 @@ def _backend_version_assertion():
 def _load_cluster_manager(cluster_manager):
     jl.seval(f"using ClusterManagers: addprocs_{cluster_manager}")
     return jl.seval(f"addprocs_{cluster_manager}")
+
+
+def jl_array(x):
+    if x is None:
+        return None
+    return jl_convert(jl.Array, x)

pysr/sr.py

@@ -37,6 +37,7 @@ from .julia_helpers import (
     _escape_filename,
     _load_cluster_manager,
     jl,
+    jl_array,
     jl_convert,
 )
 from .utils import (
@@ -1618,8 +1619,8 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         options = SymbolicRegression.Options(
             binary_operators=jl.seval(str(binary_operators).replace("'", "")),
             unary_operators=jl.seval(str(unary_operators).replace("'", "")),
-            bin_constraints=jl_convert(jl.Vector, bin_constraints),
-            una_constraints=jl_convert(jl.Vector, una_constraints),
+            bin_constraints=jl_array(bin_constraints),
+            una_constraints=jl_array(una_constraints),
             complexity_of_operators=complexity_of_operators,
             complexity_of_constants=self.complexity_of_constants,
             complexity_of_variables=self.complexity_of_variables,
@@ -1684,16 +1685,16 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             np_dtype = {32: np.complex64, 64: np.complex128}[self.precision]
 
         # This converts the data into a Julia array:
-        jl_X = jl_convert(jl.Array, np.array(X, dtype=np_dtype).T)
+        jl_X = jl_array(np.array(X, dtype=np_dtype).T)
         if len(y.shape) == 1:
-            jl_y = jl_convert(jl.Vector, np.array(y, dtype=np_dtype))
+            jl_y = jl_array(np.array(y, dtype=np_dtype))
         else:
-            jl_y = jl_convert(jl.Array, np.array(y, dtype=np_dtype).T)
+            jl_y = jl_array(np.array(y, dtype=np_dtype).T)
         if weights is not None:
             if len(weights.shape) == 1:
-                jl_weights = jl_convert(jl.Vector, np.array(weights, dtype=np_dtype))
+                jl_weights = jl_array(np.array(weights, dtype=np_dtype))
             else:
-                jl_weights = jl_convert(jl.Array, np.array(weights, dtype=np_dtype).T)
+                jl_weights = jl_array(np.array(weights, dtype=np_dtype).T)
         else:
             jl_weights = None
 
@@ -1711,17 +1712,12 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         if len(y.shape) > 1:
             # We set these manually so that they respect Python's 0 indexing
             # (by default Julia will use y1, y2...)
-            jl_y_variable_names = jl_convert(
-                jl.Vector, [f"y{_subscriptify(i)}" for i in range(y.shape[1])]
+            jl_y_variable_names = jl_array(
+                [f"y{_subscriptify(i)}" for i in range(y.shape[1])]
             )
         else:
             jl_y_variable_names = None
 
-        jl_feature_names = jl_convert(jl.Vector, self.feature_names_in_.tolist())
-        jl_display_feature_names = jl_convert(
-            jl.Vector, self.display_feature_names_in_.tolist()
-        )
-
         PythonCall.GC.disable()
         # Call to Julia backend.
         # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl
@@ -1730,11 +1726,11 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             jl_y,
             weights=jl_weights,
             niterations=int(self.niterations),
-            variable_names=jl_feature_names,
-            display_variable_names=jl_display_feature_names,
+            variable_names=jl_array(self.feature_names_in_.tolist()),
+            display_variable_names=jl_array(self.display_feature_names_in_.tolist()),
             y_variable_names=jl_y_variable_names,
-            X_units=self.X_units_,
-            y_units=self.y_units_,
+            X_units=jl_array(self.X_units_),
+            y_units=jl_array(self.y_units_),
             options=options,
             numprocs=cprocs,
             parallelism=parallelism,

pysr/test/test.py

@@ -106,11 +106,11 @@ class TestPipeline(unittest.TestCase):
             warm_start=True,
         )
         model.fit(self.X, y)
-        from pysr.sr import Main
+        from pysr.sr import jl
 
         # We should have that the model state is now a Float64 hof:
-        Main.test_state = model.raw_julia_state_
-        self.assertTrue(Main.eval("typeof(test_state[2]).parameters[1] == Float64"))
+        jl.test_state = model.raw_julia_state_
+        self.assertTrue(jl.seval("typeof(test_state[2]).parameters[1] == Float64"))
 
     def test_multioutput_custom_operator_quiet_custom_complexity(self):
         y = self.X[:, [0, 1]] ** 2
@@ -229,11 +229,11 @@ class TestPipeline(unittest.TestCase):
             early_stop_condition=None,
         )
         # Check that the the julia state is saved:
-        from pysr.sr import Main
+        from pysr.sr import jl
 
         # We should have that the model state is now a Float32 hof:
-        Main.test_state = regressor.raw_julia_state_
-        self.assertTrue(Main.eval("typeof(test_state[2]).parameters[1] == Float32"))
+        jl.test_state = regressor.raw_julia_state_
+        self.assertTrue(jl.eval("typeof(test_state[2]).parameters[1] == Float32"))
         # This should exit almost immediately, and use the old equations
         regressor.fit(X, y)