Merge pull request #545 from MilesCranmer/fix-units

Fix y_units bug

Author: MilesCranmer

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pysr"
-version = "0.17.0"
+version = "0.17.1"
 authors = [
     {name = "Miles Cranmer", email = "miles.cranmer@gmail.com"},
 ]

pysr/sr.py

@@ -1736,7 +1736,9 @@ def _run(self, X, y, mutated_params, weights, seed):
             ),
             y_variable_names=jl_y_variable_names,
             X_units=jl_array(self.X_units_),
-            y_units=jl_array(self.y_units_),
+            y_units=jl_array(self.y_units_)
+            if isinstance(self.y_units_, list)
+            else self.y_units_,
             options=options,
             numprocs=cprocs,
             parallelism=parallelism,

pysr/test/test.py

@@ -1038,7 +1038,7 @@ def test_unit_checks(self):
         valid_units = [
             (np.ones((10, 2)), np.ones(10), ["m/s", "s"], "m"),
             (np.ones((10, 1)), np.ones(10), ["m/s"], None),
-            (np.ones((10, 1)), np.ones(10), None, "m/s"),
+            (np.ones((10, 1)), np.ones(10), None, "km/s"),
             (np.ones((10, 1)), np.ones(10), None, ["m/s"]),
             (np.ones((10, 1)), np.ones((10, 1)), None, ["m/s"]),
             (np.ones((10, 1)), np.ones((10, 2)), None, ["m/s", ""]),
@@ -1053,7 +1053,7 @@ def test_unit_checks(self):
             )
         invalid_units = [
             (np.ones((10, 2)), np.ones(10), ["m/s", "s", "s^2"], None),
-            (np.ones((10, 2)), np.ones(10), ["m/s", "s", "s^2"], "m"),
+            (np.ones((10, 2)), np.ones(10), ["m/s", "s", "s^2"], "km"),
             (np.ones((10, 2)), np.ones((10, 2)), ["m/s", "s"], ["m"]),
             (np.ones((10, 1)), np.ones((10, 1)), "m/s", ["m"]),
         ]