feat: raise friendly error when loss functions have bad signatures (#1138)

* fix: validate loss_function arity and raise friendly error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: cover loss_function arity validation branches

* test: add docstrings for #982 validation tests

* refactor: extract loss_function validation helper

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor: move loss_function validator near class

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor: place loss_function validator near other checks

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor: avoid custom closures

* fix: improve loss function validation and errors

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: make loss validators beartype-safe

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: DRY loss validators and revert docstring edits

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: remove blanket exception in jl.isnothing helper

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* fix: validate elementwise_loss arity based on weights

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* test: cover loss validation branches

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* fix: apply review suggestions for loss validation

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* fix: remove _jl_is_nothing helper

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* fix: clarify loss validator arity + error text

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* test: cover signature error branch in loss validators

* test: avoid julia function name collisions in loss validation

* refactor: hardcode elementwise_loss suggestion

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: 3 people

pysr/sr.py

@@ -235,6 +235,84 @@ def _check_assertions(
             )
 
 
+def _validate_elementwise_loss(custom_loss, *, has_weights: bool) -> None:
+    """Validate that a Julia `elementwise_loss` is callable.
+
+    We require exactly 2 args unless the user passed `weights=` to fit,
+    in which case we require 3 args.
+    """
+
+    # This can be either a LossFunctions.jl object (e.g. `L2DistLoss()`) or a Julia function.
+    # Only validate arity when the evaluated object is actually a function.
+    if not jl_is_function(custom_loss):
+        return
+
+    if has_weights:
+        ok = bool(jl.applicable(custom_loss, 1.0, 1.0, 1.0))
+        if not ok:
+            raise ValueError(
+                "`elementwise_loss` must accept (prediction, target, weight) when `weights` is passed to `fit`."
+            )
+    else:
+        ok = bool(jl.applicable(custom_loss, 1.0, 1.0))
+        if not ok:
+            raise ValueError(
+                "`elementwise_loss` must accept (prediction, target). If you intended a full objective, use "
+                "`loss_function` or `loss_function_expression`."
+            )
+
+
+def _validate_custom_objective(
+    custom_objective,
+    *,
+    knob,
+    signature,
+    other_alternative=None,
+) -> None:
+    if not jl_is_function(custom_objective):
+        raise ValueError(f"`{knob}` must evaluate to a callable Julia function.")
+
+    methods = jl.collect(jl.methods(custom_objective))
+
+    def _accepts_npos(m, npos: int) -> bool:
+        required_npos = int(m.nargs) - 1
+        if bool(m.isva):
+            return required_npos <= npos
+        return required_npos == npos
+
+    accepts_three_args = any(_accepts_npos(m, 3) for m in methods)
+    accepts_two_args = any(_accepts_npos(m, 2) for m in methods)
+
+    if not accepts_three_args and accepts_two_args:
+        msg = (
+            f"`{knob}` must have signature like {signature}. "
+            "If you intended an elementwise loss, use `elementwise_loss`."
+        )
+        if other_alternative is not None:
+            msg += f" If you intended the other full-objective mode, use `{other_alternative}`."
+        raise ValueError(msg)
+
+    if not accepts_three_args:
+        raise ValueError(f"`{knob}` must have signature {signature}.")
+
+
+def _validate_custom_full_objective(custom_full_objective) -> None:
+    _validate_custom_objective(
+        custom_full_objective,
+        knob="loss_function",
+        signature="(tree, dataset, options)",
+        other_alternative="loss_function_expression",
+    )
+
+
+def _validate_custom_expression_objective(custom_loss_expression) -> None:
+    _validate_custom_objective(
+        custom_loss_expression,
+        knob="loss_function_expression",
+        signature="(expression, dataset, options)",
+    )
+
+
 def _validate_export_mappings(extra_jax_mappings, extra_torch_mappings):
     # It is expected extra_jax/torch_mappings will be updated after fit.
     # Thus, validation is performed here instead of in _validate_init_params
@@ -2036,14 +2114,22 @@ def _run(
             if self.elementwise_loss is not None
             else "nothing"
         )
+        if self.elementwise_loss is not None:
+            _validate_elementwise_loss(custom_loss, has_weights=weights is not None)
+
         custom_full_objective = jl.seval(
             str(self.loss_function) if self.loss_function is not None else "nothing"
         )
+        if self.loss_function is not None:
+            _validate_custom_full_objective(custom_full_objective)
+
         custom_loss_expression = jl.seval(
             str(self.loss_function_expression)
             if self.loss_function_expression is not None
             else "nothing"
         )
+        if self.loss_function_expression is not None:
+            _validate_custom_expression_objective(custom_loss_expression)
 
         early_stop_condition = jl.seval(
             str(self.early_stop_condition)

pysr/test/test_main.py

@@ -41,6 +41,8 @@
     _check_assertions,
     _process_constraints,
     _suggest_keywords,
+    _validate_elementwise_loss,
+    _validate_export_mappings,
     idx_model_selection,
 )
 
@@ -183,6 +185,200 @@ def test_high_precision_search_custom_loss(self):
             jl.seval("((::Val{x}) where x) -> x")(model.julia_options_.turbo), False
         )
 
+    def test_loss_function_with_elementwise_signature_errors_early(self):
+        """Issue #982: elementwise (prediction, target) loss passed via loss_function errors."""
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function="bad_full_objective(prediction, target) = (prediction - target)^2",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y)
+        self.assertIn("elementwise_loss", str(cm.exception))
+
+    def test_loss_function_noncallable_errors_early(self):
+        """Issue #982: non-callable loss_function (e.g. '1.0') errors early."""
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function="1.0",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y)
+        self.assertIn("callable", str(cm.exception))
+
+    def test_loss_function_valid_full_objective_runs(self):
+        """Issue #982: a valid (tree, dataset, options) objective is accepted."""
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function="""
+            begin
+                goodloss(tree, dataset, options) = zero(eltype(dataset.y))
+                goodloss
+            end
+            """,
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        model.fit(X, y)
+
+    def test_loss_function_varargs_objective_runs(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function="""
+            begin
+                goodvarloss(tree, dataset, options...) = zero(eltype(dataset.y))
+                goodvarloss
+            end
+            """,
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        model.fit(X, y)
+
+    def test_elementwise_loss_wrong_signature_errors_early(self):
+        """Validate `elementwise_loss` signature (prediction, target[, weights])."""
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            elementwise_loss="myloss_bad_arity(a) = a",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y)
+        self.assertIn("elementwise_loss", str(cm.exception))
+
+    def test_elementwise_loss_with_weights_requires_three_args(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            elementwise_loss="myloss2(prediction, target) = (prediction - target)^2",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        weights = np.array([1.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y, weights=weights)
+        self.assertIn("elementwise_loss", str(cm.exception))
+        self.assertIn("weights", str(cm.exception))
+
+    def test_elementwise_loss_with_weights_accepts_three_args(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            elementwise_loss=(
+                "myloss3(prediction, target, weights) = weights * (prediction - target)^2"
+            ),
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        weights = np.array([1.0, 1.0])
+        model.fit(X, y, weights=weights)
+
+    def test_validation_helpers_skip_nonfunction(self):
+        _validate_elementwise_loss(jl.seval("1.0"), has_weights=False)
+
+    def test_validate_export_mappings_typechecks(self):
+        with self.assertRaises(ValueError):
+            _validate_export_mappings({"a": 1}, None)
+
+    def test_loss_function_expression_elementwise_signature_errors_early(self):
+        """Validate `loss_function_expression` signature (expression, dataset, options)."""
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function_expression="bad_expr_objective(prediction, target) = (prediction - target)^2",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y)
+        self.assertIn("loss_function_expression", str(cm.exception))
+        self.assertIn("elementwise_loss", str(cm.exception))
+
+    def test_loss_function_wrong_signature_errors_early(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function="badloss(tree) = 0.0",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y)
+        self.assertIn("loss_function", str(cm.exception))
+        self.assertIn("(tree, dataset, options)", str(cm.exception))
+
+    def test_loss_function_expression_wrong_signature_errors_early(self):
+        model = PySRRegressor(
+            niterations=1,
+            populations=1,
+            procs=0,
+            progress=False,
+            verbosity=0,
+            temp_equation_file=True,
+            binary_operators=["+"],
+            loss_function_expression="badexprloss(expression) = 0.0",
+        )
+        X = np.array([[0.0], [1.0]])
+        y = np.array([0.0, 1.0])
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y)
+        self.assertIn("loss_function_expression", str(cm.exception))
+        self.assertIn("(expression, dataset, options)", str(cm.exception))
+
     def test_operator_conflict_error(self):
         regressor = PySRRegressor(
             operators={1: ["sin"]},