candidate_wrapper

Author: WassimTenachi

docs/source/features/doc_features_wrapper.md

@@ -0,0 +1,11 @@
+## Candidate wrapper
+
+The `candidate_wrapper` argument can be passed to the `physo.SR` or `physo.ClassSR` to apply a wrapper function $g$ to the candidate symbolic function's output $f(X)$.
+
+The wrapper function $g$ should be a callable taking the candidate symbolic function callable $f$ and the input data $X$ as arguments, and returning the wrapped output $g(f(X))$.
+By default `candidate_wrapper = None`, no wrapper is applied (identity).
+
+Note that the wrapper function should be differentiable and written in `pytorch` if free constants are to be optimized since the free constants are optimized using gradient-based optimization.
+
+In addition, it is recommended to use protected functions when writing the wrapper function to avoid evaluating the symbolic function on invalid points (eg. using log abs instead of log).
+See the [protected functions](https://physo.readthedocs.io/en/latest/r_features.html#protected-version-optional) documentation for more details. 

physo/task/args_handler.py

@@ -46,6 +46,8 @@ def check_args_and_build_run_config(multi_X, multi_y, multi_y_weights,
             use_protected_ops,
             # Stopping
             epochs,
+            # Candidate wrapper
+            candidate_wrapper,
             # Default run config to use
             run_config,
             # Default run monitoring
@@ -301,6 +303,10 @@ def check_args_and_build_run_config(multi_X, multi_y, multi_y_weights,
         raise ValueError("entropy_weight should be castable to a float.")
     assert isinstance(entropy_weight, float), "entropy_weight should be a float."
 
+    # ------------------------------- CANDIDATE_WRAPPER -------------------------------
+    # candidate_wrapper should be callable or None
+    assert candidate_wrapper is None or callable(candidate_wrapper), "candidate_wrapper should be callable or None."
+
     # ------------------------------- RETURN -------------------------------
     # Returning
     handled_args = {

physo/task/class_sr.py

@@ -33,6 +33,8 @@ def ClassSR(multi_X, multi_y, multi_y_weights=1.,
             max_n_evaluations = None,
             stop_after_n_epochs = args_handler.default_stop_after_n_epochs,
             epochs = None,
+            # Candidate wrapper
+            candidate_wrapper = None,
             # Default run config to use
             run_config = None,
             # Default run monitoring
@@ -119,6 +121,9 @@ def ClassSR(multi_X, multi_y, multi_y_weights=1.,
     epochs : int or None (optional)
         Number of epochs to perform. By default, uses the number in the default config file.
 
+    candidate_wrapper : callable or None (optional)
+        Wrapper to apply to candidate program's output, candidate_wrapper taking func, X as arguments where func is
+        a candidate program callable (taking X as arg). By default = None, no wrapper is applied (identity).
     run_config : dict or None (optional)
         Run configuration (by default uses physo.task.class_sr.default_config)
         See physo/config/ for examples of run configurations.
@@ -182,6 +187,8 @@ def ClassSR(multi_X, multi_y, multi_y_weights=1.,
                                     use_protected_ops = use_protected_ops,
                                     # Stopping
                                     epochs = epochs,
+                                    # Candidate wrapper
+                                    candidate_wrapper = candidate_wrapper,
                                     # Default run config to use
                                     run_config = run_config,
                                     # Default run monitoring
@@ -209,7 +216,8 @@ def ClassSR(multi_X, multi_y, multi_y_weights=1.,
     rewards, candidates = fit (multi_X         = multi_X,
                                multi_y         = multi_y,
                                multi_y_weights = multi_y_weights,
-                               run_config      = run_config,
+                               candidate_wrapper = candidate_wrapper,
+                               run_config        = run_config,
                                stop_reward         = stop_reward,
                                stop_after_n_epochs = stop_after_n_epochs,
                                max_n_evaluations   = max_n_evaluations,

physo/task/sr.py

@@ -32,6 +32,8 @@ def SR(X, y, y_weights=1.,
             max_n_evaluations = None,
             stop_after_n_epochs = args_handler.default_stop_after_n_epochs,
             epochs = None,
+            # Candidate wrapper
+            candidate_wrapper = None,
             # Default run config to use
             run_config = None,
             # Default run monitoring
@@ -102,6 +104,9 @@ def SR(X, y, y_weights=1.,
     epochs : int or None (optional)
         Number of epochs to perform. By default, uses the number in the default config file.
 
+    candidate_wrapper : callable or None (optional)
+        Wrapper to apply to candidate program's output, candidate_wrapper taking func, X as arguments where func is
+        a candidate program callable (taking X as arg). By default = None, no wrapper is applied (identity).
     run_config : dict or None (optional)
         Run configuration (by default uses physo.task.sr.default_config)
         See physo/config/ for examples of run configurations.
@@ -163,7 +168,9 @@ def SR(X, y, y_weights=1.,
                 stop_reward         = stop_reward,
                 max_n_evaluations   = max_n_evaluations,
                 stop_after_n_epochs = stop_after_n_epochs,
-                epochs            = epochs,
+                epochs              = epochs,
+                # Candidate wrapper
+                candidate_wrapper = candidate_wrapper,
                 # Default run config to use
                 run_config = run_config,
                 # Default run monitoring