Merge pull request #928 from MilesCranmer:negative-losses

Enable negative losses

Author: MilesCranmer

pysr/juliapkg.json

@@ -3,7 +3,7 @@
     "packages": {
         "SymbolicRegression": {
             "uuid": "8254be44-1295-4e6a-a16d-46603ac705cb",
-            "version": "~1.10.0"
+            "version": "~1.11.0"
         },
         "Serialization": {
             "uuid": "9e88b42a-f829-5b0c-bbe9-9e923198166b",

pysr/param_groupings.yml

@@ -14,6 +14,7 @@
     - elementwise_loss
     - loss_function
     - loss_function_expression
+    - loss_scale
     - model_selection
     - dimensional_constraint_penalty
     - dimensionless_constants_only

pysr/sr.py

@@ -389,6 +389,16 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         the innermost `AbstractExpressionNode`. This is useful
         for specifying custom loss functions on `TemplateExpressionSpec`.
         Default is `None`.
+    loss_scale : Literal["log", "linear"]
+        Determines how loss values are scaled when computing scores.
+        "log" (default) uses logarithmic scaling of loss ratios; this mode
+        requires non-negative loss values and is ideal for traditional
+        loss functions that are always non-negative.
+        "linear" uses direct
+        differences between losses; this mode handles any loss values
+        (including negative) and is useful for custom loss functions,
+        especially those based on likelihoods.
+        Default is "log".
     complexity_of_operators : dict[str, int | float]
         If you would like to use a complexity other than 1 for an
         operator, specify the complexity here. For example,
@@ -817,6 +827,7 @@ def __init__(
         elementwise_loss: str | None = None,
         loss_function: str | None = None,
         loss_function_expression: str | None = None,
+        loss_scale: Literal["log", "linear"] = "log",
         complexity_of_operators: dict[str, int | float] | None = None,
         complexity_of_constants: int | float | None = None,
         complexity_of_variables: int | float | list[int | float] | None = None,
@@ -924,6 +935,7 @@ def __init__(
         self.elementwise_loss = elementwise_loss
         self.loss_function = loss_function
         self.loss_function_expression = loss_function_expression
+        self.loss_scale = loss_scale
         self.complexity_of_operators = complexity_of_operators
         self.complexity_of_constants = complexity_of_constants
         self.complexity_of_variables = complexity_of_variables
@@ -1203,7 +1215,11 @@ def __repr__(self) -> str:
             repr_equations = pd.DataFrame(
                 dict(
                     pick=selected,
-                    score=equations["score"],
+                    **(
+                        {"score": equations["score"]}
+                        if "score" in equations.columns
+                        else {}
+                    ),
                     equation=equations["equation"],
                     loss=equations["loss"],
                     complexity=equations["complexity"],
@@ -1993,6 +2009,7 @@ def _run(
             elementwise_loss=custom_loss,
             loss_function=custom_full_objective,
             loss_function_expression=custom_loss_expression,
+            loss_scale=jl.Symbol(self.loss_scale),
             maxsize=int(self.maxsize),
             output_directory=_escape_filename(self.output_directory_),
             npopulations=int(self.populations),
@@ -2644,7 +2661,7 @@ def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]:
             pd.concat(
                 [
                     output,
-                    calculate_scores(output),
+                    *([calculate_scores(output)] if self.loss_scale == "log" else []),
                     self.expression_spec_.create_exports(
                         self, output, search_output, i if self.nout_ > 1 else None
                     ),
@@ -2720,6 +2737,10 @@ def latex_table(
 
 def idx_model_selection(equations: pd.DataFrame, model_selection: str):
     """Select an expression and return its index."""
+
+    # We must default to "accuracy" if no score column is present (like in the case of linear loss_scale)
+    model_selection = model_selection if "score" in equations.columns else "accuracy"
+
     if model_selection == "accuracy":
         chosen_idx = equations["loss"].idxmin()
     elif model_selection == "best":

pysr/test/test_main.py

@@ -765,6 +765,27 @@ def test_tensorboard_logger(self):
                 # Verify model still works as expected
                 self.assertLessEqual(model.get_best()["loss"], 1e-4)
 
+    def test_negative_losses(self):
+        X = self.rstate.rand(100, 3) * 20.0
+        eps = self.rstate.randn(100)
+        y = np.cos(X[:, 0] * 2.1 - 0.5) + X[:, 1] ** 2 + 0.1 * eps
+        spec = TemplateExpressionSpec(
+            expressions=["f_mu", "f_logvar"],
+            variable_names=["x1", "x2", "x3", "y"],
+            combine="mu = f_mu(x1, x2, x3); logvar = f_logvar(x1, x2, x3); 0.5f0 * (logvar + (mu - y)^2 / exp(logvar))",
+        )
+        model = PySRRegressor(
+            **self.default_test_kwargs,
+            expression_spec=spec,
+            binary_operators=["+", "*", "-"],
+            unary_operators=["cos", "log", "exp"],
+            elementwise_loss="(pred, targ) -> pred",
+            loss_scale="linear",
+            early_stop_condition="stop_if_under_n1(loss, complexity) = loss < -1.0",
+        )
+        model.fit(np.column_stack([X, y]), 0 * y)
+        self.assertLessEqual(model.get_best()["loss"], -1.0)
+
     def test_comparison_operator(self):
         X = self.rstate.randn(100, 2)
         y = ((X[:, 0] + X[:, 1]) < (X[:, 0] * X[:, 1])).astype(float)