Merge 35ca811 into 887e02d

pysr/sr.py

@@ -779,10 +779,25 @@ def get_params(self, deep=True):
             **{key: self.__getattribute__(key) for key in self.surface_parameters},
         }
 
-    def get_best(self):
-        """Get best equation using `model_selection`."""
+    def get_best(self, index=None):
+        """Get best equation using `model_selection`.
+
+        :param index: Optional. If you wish to select a particular equation
+            from `self.equations`, give the row number here. This overrides
+            the `model_selection` parameter.
+        :type index: int
+        :returns: Dictionary representing the best expression found.
+        :type: pd.Series
+        """
         if self.equations is None:
             raise ValueError("No equations have been generated yet.")
+
+        if index is not None:
+            if isinstance(self.equations, list):
+                assert isinstance(index, list)
+                return [self.equations.iloc[i] for i in index]
+            return self.equations.iloc[index]
+
         if self.model_selection == "accuracy":
             if isinstance(self.equations, list):
                 return [eq.iloc[-1] for eq in self.equations]
@@ -826,44 +841,72 @@ def refresh(self):
         # such as extra_sympy_mappings.
         self.equations = self.get_hof()
 
-    def predict(self, X):
+    def predict(self, X, index=None):
         """Predict y from input X using the equation chosen by `model_selection`.
 
         You may see what equation is used by printing this object. X should have the same
         columns as the training data.
 
         :param X: 2D array. Rows are examples, columns are features. If pandas DataFrame, the columns are used for variable names (so make sure they don't contain spaces).
         :type X: np.ndarray/pandas.DataFrame
-        :return: 1D array (rows are examples) or 2D array (rows are examples, columns are outputs).
+        :param index: Optional. If you want to compute the output of
+            an expression using a particular row of
+            `self.equations`, you may specify the index here.
+        :type index: int
+        :returns: 1D array (rows are examples) or 2D array (rows are examples, columns are outputs).
+        :type: np.ndarray
         """
         self.refresh()
-        best = self.get_best()
+        best = self.get_best(index=index)
         if self.multioutput:
             return np.stack([eq["lambda_format"](X) for eq in best], axis=1)
         return best["lambda_format"](X)
 
-    def sympy(self):
-        """Return sympy representation of the equation(s) chosen by `model_selection`."""
+    def sympy(self, index=None):
+        """Return sympy representation of the equation(s) chosen by `model_selection`.
+
+        :param index: Optional. If you wish to select a particular equation
+            from `self.equations`, give the index number here. This overrides
+            the `model_selection` parameter.
+        :type index: int
+        :returns: SymPy representation of the best expression.
+        """
         self.refresh()
-        best = self.get_best()
+        best = self.get_best(index=index)
         if self.multioutput:
             return [eq["sympy_format"] for eq in best]
         return best["sympy_format"]
 
-    def latex(self):
-        """Return latex representation of the equation(s) chosen by `model_selection`."""
+    def latex(self, index=None):
+        """Return latex representation of the equation(s) chosen by `model_selection`.
+
+        :param index: Optional. If you wish to select a particular equation
+            from `self.equations`, give the index number here. This overrides
+            the `model_selection` parameter.
+        :type index: int
+        :returns: LaTeX expression as a string
+        :type: str
+        """
         self.refresh()
-        sympy_representation = self.sympy()
+        sympy_representation = self.sympy(index=index)
         if self.multioutput:
             return [sympy.latex(s) for s in sympy_representation]
         return sympy.latex(sympy_representation)
 
-    def jax(self):
+    def jax(self, index=None):
         """Return jax representation of the equation(s) chosen by `model_selection`.
 
         Each equation (multiple given if there are multiple outputs) is a dictionary
         containing {"callable": func, "parameters": params}. To call `func`, pass
         func(X, params). This function is differentiable using `jax.grad`.
+
+        :param index: Optional. If you wish to select a particular equation
+            from `self.equations`, give the index number here. This overrides
+            the `model_selection` parameter.
+        :type index: int
+        :returns: Dictionary of callable jax function in "callable" key,
+            and jax array of parameters as "parameters" key.
+        :type: dict
         """
         if self.using_pandas:
             warnings.warn(
@@ -873,18 +916,26 @@ def jax(self):
             )
         self.set_params(output_jax_format=True)
         self.refresh()
-        best = self.get_best()
+        best = self.get_best(index=index)
         if self.multioutput:
             return [eq["jax_format"] for eq in best]
         return best["jax_format"]
 
-    def pytorch(self):
+    def pytorch(self, index=None):
         """Return pytorch representation of the equation(s) chosen by `model_selection`.
 
         Each equation (multiple given if there are multiple outputs) is a PyTorch module
         containing the parameters as trainable attributes. You can use the module like
         any other PyTorch module: `module(X)`, where `X` is a tensor with the same
         column ordering as trained with.
+
+
+        :param index: Optional. If you wish to select a particular equation
+            from `self.equations`, give the row number here. This overrides
+            the `model_selection` parameter.
+        :type index: int
+        :returns: PyTorch module representing the expression.
+        :type: torch.nn.Module
         """
         if self.using_pandas:
             warnings.warn(
@@ -894,7 +945,7 @@ def pytorch(self):
             )
         self.set_params(output_torch_format=True)
         self.refresh()
-        best = self.get_best()
+        best = self.get_best(index=index)
         if self.multioutput:
             return [eq["torch_format"] for eq in best]
         return best["torch_format"]