Refactor LaTeX table to export_latex.py

pysr/export_latex.py

@@ -1,6 +1,8 @@
 """Functions to help export PySR equations to LaTeX."""
 import sympy
 from sympy.printing.latex import LatexPrinter
+import pandas as pd
+from typing import List
 
 
 class PreciseLatexPrinter(LatexPrinter):
@@ -51,3 +53,63 @@ def generate_table_environment(columns=["equation", "complexity", "loss"]):
     bottom_latex_table = "\n".join(bottom_pieces)
 
     return top_latex_table, bottom_latex_table
+
+
+def generate_table(
+    equations: List[pd.DataFrame],
+    indices: List[List[int]],
+    precision=3,
+    columns=["equation", "complexity", "loss", "score"],
+):
+    latex_top, latex_bottom = generate_table_environment(columns)
+
+    latex_equations = [
+        [to_latex(eq, prec=precision) for eq in equation_set["sympy_format"]]
+        for equation_set in equations
+    ]
+
+    all_latex_table_str = []
+
+    for output_feature, index_set in enumerate(indices):
+        latex_table_content = []
+        for i in index_set:
+            latex_equation = latex_equations[output_feature][i]
+            complexity = str(equations[output_feature].iloc[i]["complexity"])
+            loss = to_latex(
+                sympy.Float(equations[output_feature].iloc[i]["loss"]),
+                prec=precision,
+            )
+            score = to_latex(
+                sympy.Float(equations[output_feature].iloc[i]["score"]),
+                prec=precision,
+            )
+
+            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]
+
+            latex_table_content.append(
+                " & ".join(row_pieces) + r" \\",
+            )
+
+        this_latex_table = "\n".join(
+            [
+                latex_top,
+                *latex_table_content,
+                latex_bottom,
+            ]
+        )
+        all_latex_table_str.append(this_latex_table)
+
+    return "\n\n".join(all_latex_table_str)

pysr/sr.py

@@ -27,7 +27,7 @@ from .julia_helpers import (
     import_error_string,
 )
 from .export_numpy import CallableEquation
-from .export_latex import to_latex, generate_table_environment
+from .export_latex import to_latex, generate_table
 from .deprecated import make_deprecated_kwargs_for_pysr_regressor
 
 
@@ -2033,8 +2033,6 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             else:
                 indices = list(range(len(self.equations_)))
 
-        latex_top, latex_bottom = generate_table_environment(columns)
-
         equations = self.equations_
 
         if isinstance(indices[0], int):
@@ -2044,57 +2042,9 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
 
         assert len(indices) == self.nout_
 
-        latex_equations = [
-            [to_latex(eq, prec=precision) for eq in equation_set["sympy_format"]]
-            for equation_set in equations
-        ]
-
-        all_latex_table_str = []
-
-        for output_feature, index_set in enumerate(indices):
-            latex_table_content = []
-            for i in index_set:
-                latex_equation = latex_equations[output_feature][i]
-                complexity = str(equations[output_feature].iloc[i]["complexity"])
-                loss = to_latex(
-                    sympy.Float(equations[output_feature].iloc[i]["loss"]),
-                    prec=precision,
-                )
-                score = to_latex(
-                    sympy.Float(equations[output_feature].iloc[i]["score"]),
-                    prec=precision,
-                )
-
-                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]
-
-                latex_table_content.append(
-                    " & ".join(row_pieces) + r" \\",
-                )
-
-            all_latex_table_str.append(
-                "\n".join(
-                    [
-                        latex_top,
-                        *latex_table_content,
-                        latex_bottom,
-                    ]
-                )
-            )
-
-        return "\n\n".join(all_latex_table_str)
+        return generate_table(
+            equations, indices=indices, precision=precision, columns=columns
+        )
 
 
 def _denoise(X, y, Xresampled=None, random_state=None):