Refactor table env generator

pysr/export_latex.py

@@ -23,7 +23,7 @@ def to_latex(expr, prec=3, full_prec=True, **settings):
     return printer.doprint(expr)
 
 
-def generate_top_of_latex_table(columns=["equation", "complexity", "loss"]):
+def generate_table_environment(columns=["equation", "complexity", "loss"]):
     margins = "".join([("l" if col == "equation" else "c") for col in columns])
     column_map = {
         "complexity": "Complexity",
@@ -32,7 +32,7 @@ def generate_top_of_latex_table(columns=["equation", "complexity", "loss"]):
         "score": "Score",
     }
     columns = [column_map[col] for col in columns]
-    latex_table_pieces = [
+    top_pieces = [
         r"\begin{table}[h]",
         r"\begin{center}",
         r"\begin{tabular}{@{}" + margins + r"@{}}",
@@ -40,14 +40,14 @@ def generate_top_of_latex_table(columns=["equation", "complexity", "loss"]):
         " & ".join(columns) + r" \\",
         r"\midrule",
     ]
-    return "\n".join(latex_table_pieces)
 
-
-def generate_bottom_of_latex_table():
-    latex_table_pieces = [
+    bottom_pieces = [
         r"\bottomrule",
         r"\end{tabular}",
         r"\end{center}",
         r"\end{table}",
     ]
-    return "\n".join(latex_table_pieces)
+    top_latex_table = "\n".join(top_pieces)
+    bottom_latex_table = "\n".join(bottom_pieces)
+
+    return top_latex_table, bottom_latex_table

pysr/sr.py

@@ -27,11 +27,7 @@ from .julia_helpers import (
     import_error_string,
 )
 from .export_numpy import CallableEquation
-from .export_latex import (
-    to_latex,
-    generate_top_of_latex_table,
-    generate_bottom_of_latex_table,
-)
+from .export_latex import to_latex, generate_table_environment
 from .deprecated import make_deprecated_kwargs_for_pysr_regressor
 
 
@@ -2037,8 +2033,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             else:
                 indices = list(range(len(self.equations_)))
 
-        latex_table_top = generate_top_of_latex_table(columns)
-        latex_table_bottom = generate_bottom_of_latex_table()
+        latex_top, latex_bottom = generate_table_environment(columns)
 
         equations = self.equations_
 
@@ -2092,9 +2087,9 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             all_latex_table_str.append(
                 "\n".join(
                     [
-                        latex_table_top,
+                        latex_top,
                         *latex_table_content,
-                        latex_table_bottom,
+                        latex_bottom,
                     ]
                 )
             )