Move `latex_table` code to `export_latex.py`

pysr/export_latex.py

@@ -55,61 +55,72 @@ def generate_table_environment(columns=["equation", "complexity", "loss"]):
     return top_latex_table, bottom_latex_table
 
 
-def generate_table(
-    equations: List[pd.DataFrame],
-    indices: List[List[int]],
-    precision=3,
+def generate_single_table(
+    equations: pd.DataFrame,
+    indices: List[int] = None,
+    precision: int = 3,
     columns=["equation", "complexity", "loss", "score"],
 ):
+    assert isinstance(equations, pd.DataFrame)
+
     latex_top, latex_bottom = generate_table_environment(columns)
+    latex_table_content = []
 
-    latex_equations = [
-        [to_latex(eq, prec=precision) for eq in equation_set["sympy_format"]]
-        for equation_set in equations
-    ]
+    if indices is None:
+        indices = range(len(equations))
+
+    for i in indices:
+        latex_equation = to_latex(
+            equations.iloc[i]["sympy_format"],
+            prec=precision,
+        )
+        complexity = str(equations.iloc[i]["complexity"])
+        loss = to_latex(
+            sympy.Float(equations.iloc[i]["loss"]),
+            prec=precision,
+        )
+        score = to_latex(
+            sympy.Float(equations.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" \\",
+        )
+
+    return "\n".join([latex_top, *latex_table_content, latex_bottom])
+
+
+def generate_multiple_tables(
+    equations: List[pd.DataFrame],
+    indices: List[List[int]] = None,
+    precision: int = 3,
+    columns=["equation", "complexity", "loss", "score"],
+):
 
-    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,
-            ]
+    latex_tables = [
+        generate_single_table(
+            equations[i],
+            (None if not indices else indices[i]),
+            precision=precision,
+            columns=columns,
         )
-        all_latex_table_str.append(this_latex_table)
+        for i in range(len(equations))
+    ]
 
-    return "\n\n".join(all_latex_table_str)
+    return "\n\n".join(latex_tables)

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
+from .export_latex import generate_single_table, generate_multiple_tables, to_latex
 from .deprecated import make_deprecated_kwargs_for_pysr_regressor
 
 
@@ -2024,26 +2024,18 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         """
         self.refresh()
 
-        # All indices:
-        if indices is None:
-            if self.nout_ > 1:
-                indices = [
-                    list(range(len(out_equations))) for out_equations in self.equations_
-                ]
-            else:
-                indices = list(range(len(self.equations_)))
-
-        equations = self.equations_
-
-        if isinstance(indices[0], int):
-            assert self.nout_ == 1, "For multiple outputs, pass a list of lists."
-            indices = [indices]
-            equations = [equations]
+        if self.nout_ > 1:
+            if indices is not None:
+                assert isinstance(indices, list)
+                assert isinstance(indices[0], list)
+                assert isinstance(len(indices), self.nout_)
 
-        assert len(indices) == self.nout_
+            generator_fnc = generate_multiple_tables
+        else:
+            generator_fnc = generate_single_table
 
-        return generate_table(
-            equations, indices=indices, precision=precision, columns=columns
+        return generator_fnc(
+            self.equations_, indices=indices, precision=precision, columns=columns
         )