fix(denoise): learn the signal amplitude, normalize targets, and match Xresampled by name

pysr/denoising.py

@@ -20,7 +20,10 @@ def denoise(
 
     gp_kernel = RBF(np.ones(X.shape[1])) + WhiteKernel(1e-1) + ConstantKernel()
     gpr = GaussianProcessRegressor(
-        kernel=gp_kernel, n_restarts_optimizer=50, random_state=random_state
+        kernel=gp_kernel,
+        normalize_y=True,
+        n_restarts_optimizer=50,
+        random_state=random_state,
     )
     gpr.fit(X, y)
 

pysr/sr.py

@@ -2102,6 +2102,20 @@ def _validate_and_set_fit_params(
                 y_units,
             )
 
+        if Xresampled is not None and self.denoise and weights is not None:
+            raise ValueError(
+                "`weights` cannot be combined with `Xresampled`, as the denoised "
+                "targets are evaluated at the rows of `Xresampled` instead of `X`."
+            )
+        if isinstance(X, pd.DataFrame) and isinstance(Xresampled, pd.DataFrame):
+            missing_columns = X.columns.difference(Xresampled.columns)
+            if len(missing_columns) > 0:
+                raise ValueError(
+                    "`Xresampled` is missing columns of `X`: "
+                    f"{list(missing_columns)}."
+                )
+            Xresampled = Xresampled[X.columns]
+
         if isinstance(X, pd.DataFrame):
             if variable_names:
                 variable_names = None
@@ -2825,12 +2839,14 @@ def fit(
             Resampled training data, of shape (n_resampled, n_features),
             to generate a denoised data on. This
             will be used as the training data, rather than `X`.
+            If both are DataFrames, columns are matched to `X` by name.
         weights : ndarray | pandas.DataFrame
             Weight array of the same shape as `y`.
             Each element is how to weight the mean-square-error loss
             for that particular element of `y`. Alternatively,
             if a custom `loss` was set, it will can be used
-            in arbitrary ways.
+            in arbitrary ways. Cannot be combined with `Xresampled`
+            when `denoise=True`.
         variable_names : list[str]
             A list of names for the variables, rather than "x0", "x1", etc.
             If `X` is a pandas dataframe, the column names will be used

pysr/test/test_main.py

@@ -37,6 +37,7 @@
     jl,
     load_all_packages,
 )
+from pysr.denoising import denoise
 from pysr.export_latex import sympy2latex
 from pysr.export_sympy import pysr2sympy
 from pysr.feature_selection import _handle_feature_selection, run_feature_selection
@@ -1607,6 +1608,44 @@ def test_all_selection_strategies(self):
         self.assertEqual(idx_score, 1)
 
 
+class TestDenoising(unittest.TestCase):
+    def test_denoise_learns_signal_amplitude(self):
+        rstate = np.random.RandomState(0)
+        X = rstate.uniform(-3, 3, size=(200, 1))
+        Xresampled = np.linspace(-2.5, 2.5, 50)[:, None]
+        for amplitude in [1.0, 100.0]:
+            with self.subTest(amplitude=amplitude):
+                noise = 0.1 * amplitude
+                y = amplitude * np.sin(2 * X[:, 0]) + noise * rstate.randn(len(X))
+                _, y_out = denoise(
+                    X, y, Xresampled=Xresampled, random_state=np.random.RandomState(0)
+                )
+                rmse = np.sqrt(
+                    np.mean((y_out - amplitude * np.sin(2 * Xresampled[:, 0])) ** 2)
+                )
+                self.assertLess(rmse, 0.5 * noise)
+
+    def test_xresampled_columns_are_matched_by_name(self):
+        rstate = np.random.RandomState(0)
+        X = pd.DataFrame({"T": rstate.uniform(300, 400, 20), "x": rstate.rand(20)})
+        y = X["T"] * X["x"]
+        Xresampled = pd.DataFrame(
+            {"x": rstate.rand(5), "T": rstate.uniform(300, 400, 5)}
+        )
+        model = PySRRegressor(denoise=True)
+        Xresampled_validated = model._validate_and_set_fit_params(
+            X, y, Xresampled, None, None, None, None, None
+        )[2]
+        np.testing.assert_array_equal(
+            Xresampled_validated, Xresampled[["T", "x"]].to_numpy()
+        )
+
+        with self.assertRaisesRegex(ValueError, "missing columns"):
+            model.fit(X, y, Xresampled=Xresampled[["x"]])
+        with self.assertRaisesRegex(ValueError, "`weights` cannot be combined"):
+            model.fit(X, y, Xresampled=Xresampled, weights=np.ones(len(y)))
+
+
 class TestFeatureSelection(unittest.TestCase):
     def setUp(self):
         self.rstate = np.random.RandomState(0)
@@ -3045,6 +3084,7 @@ def runtests(just_tests=False):
     test_cases = [
         TestPipeline,
         TestBest,
+        TestDenoising,
         TestFeatureSelection,
         TestMiscellaneous,
         TestHelpMessages,