Pass columns to latex table generator

pysr/export_latex.py

@@ -23,8 +23,15 @@ def to_latex(expr, prec=3, full_prec=True, **settings):
     return printer.doprint(expr)
 
 
-def generate_top_of_latex_table(columns=["Equation", "Complexity", "Loss"]):
-    margins = "".join([("l" if col == "Equation" else "c") for col in columns])
+def generate_top_of_latex_table(columns=["equation", "complexity", "loss"]):
+    margins = "".join([("l" if col == "equation" else "c") for col in columns])
+    column_map = {
+        "complexity": "Complexity",
+        "loss": "Loss",
+        "equation": "Equation",
+        "score": "Score",
+    }
+    columns = [column_map[col] for col in columns]
     latex_table_pieces = [
         r"\begin{table}[h]",
         r"\begin{center}",

pysr/sr.py

@@ -2000,7 +2000,12 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             return ret_outputs
         return ret_outputs[0]
 
-    def latex_table(self, indices=None, precision=3, include_score=False):
+    def latex_table(
+        self,
+        indices=None,
+        precision=3,
+        columns=["equation", "complexity", "loss", "score"],
+    ):
         """Create a LaTeX/booktabs table for all, or some, of the equations.
 
         Parameters
@@ -2013,8 +2018,8 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         precision : int, default=3
             The number of significant figures shown in the LaTeX
             representations.
-        include_score : bool, default=False
-            Whether to include the score in the table.
+        columns : list[str], default=["equation", "complexity", "loss", "score"]
+            Which columns to include in the table.
 
         Returns
         -------
@@ -2023,10 +2028,6 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         """
         self.refresh()
 
-        columns = ["Equation", "Complexity", "Loss"]
-        if include_score:
-            columns.append("Score")
-
         # All indices:
         if indices is None:
             if self.nout_ > 1:
@@ -2069,9 +2070,18 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
                     prec=precision,
                 )
 
-                row_pieces = [latex_equation, complexity, loss]
-                if include_score:
-                    row_pieces.append(score)
+                row_pieces = []
+                for col in columns:
+                    if col == "equation":
+                        row_pieces.append(latex_equation)
+                    elif col == "complexity":
+                        row_pieces.append(complexity)
+                    elif col == "loss":
+                        row_pieces.append(loss)
+                    elif col == "score":
+                        row_pieces.append(score)
+                    else:
+                        raise ValueError(f"Unknown column: {col}")
 
                 row_pieces = ["$" + piece + "$" for piece in row_pieces]
 

test/test.py

@@ -538,7 +538,7 @@ class TestLaTeXTable(unittest.TestCase):
         model = manually_create_model(equations)
 
         # Regular table:
-        latex_table_str = model.latex_table()
+        latex_table_str = model.latex_table(columns=["equation", "complexity", "loss"])
         middle_part = r"""
             $x_{0}$ & $1$ & $1.05$ \\
             $\cos{\left(x_{0} \right)}$ & $2$ & $0.0232$ \\
@@ -548,7 +548,9 @@ class TestLaTeXTable(unittest.TestCase):
         self.assertEqual(latex_table_str, true_latex_table_str)
 
         # Different precision:
-        latex_table_str = model.latex_table(precision=5)
+        latex_table_str = model.latex_table(
+            precision=5, columns=["equation", "complexity", "loss"]
+        )
         middle_part = r"""
             $x_{0}$ & $1$ & $1.0520$ \\
             $\cos{\left(x_{0} \right)}$ & $2$ & $0.023150$ \\
@@ -558,7 +560,7 @@ class TestLaTeXTable(unittest.TestCase):
         self.assertEqual(latex_table_str, self.create_true_latex(middle_part))
 
         # Including score:
-        latex_table_str = model.latex_table(include_score=True)
+        latex_table_str = model.latex_table()
         middle_part = r"""
             $x_{0}$ & $1$ & $1.05$ & $0.0$ \\
             $\cos{\left(x_{0} \right)}$ & $2$ & $0.0232$ & $3.82$ \\
@@ -568,7 +570,9 @@ class TestLaTeXTable(unittest.TestCase):
         self.assertEqual(latex_table_str, true_latex_table_str)
 
         # Only last equation:
-        latex_table_str = model.latex_table(indices=[2])
+        latex_table_str = model.latex_table(
+            indices=[2], columns=["equation", "complexity", "loss"]
+        )
         middle_part = r"""
             $x_{0} + x_{1} - \cos{\left(x_{0} x_{1} \right)}$ & $8$ & $1.12 \cdot 10^{-15}$ \\
         """