MilesCranmer commited on
Commit
215a692
1 Parent(s): c6b30e7

Pass columns to latex table generator

Browse files
Files changed (3) hide show
  1. pysr/export_latex.py +9 -2
  2. pysr/sr.py +20 -10
  3. test/test.py +8 -4
pysr/export_latex.py CHANGED
@@ -23,8 +23,15 @@ def to_latex(expr, prec=3, full_prec=True, **settings):
23
  return printer.doprint(expr)
24
 
25
 
26
- def generate_top_of_latex_table(columns=["Equation", "Complexity", "Loss"]):
27
- margins = "".join([("l" if col == "Equation" else "c") for col in columns])
 
 
 
 
 
 
 
28
  latex_table_pieces = [
29
  r"\begin{table}[h]",
30
  r"\begin{center}",
 
23
  return printer.doprint(expr)
24
 
25
 
26
+ def generate_top_of_latex_table(columns=["equation", "complexity", "loss"]):
27
+ margins = "".join([("l" if col == "equation" else "c") for col in columns])
28
+ column_map = {
29
+ "complexity": "Complexity",
30
+ "loss": "Loss",
31
+ "equation": "Equation",
32
+ "score": "Score",
33
+ }
34
+ columns = [column_map[col] for col in columns]
35
  latex_table_pieces = [
36
  r"\begin{table}[h]",
37
  r"\begin{center}",
pysr/sr.py CHANGED
@@ -2000,7 +2000,12 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
2000
  return ret_outputs
2001
  return ret_outputs[0]
2002
 
2003
- def latex_table(self, indices=None, precision=3, include_score=False):
 
 
 
 
 
2004
  """Create a LaTeX/booktabs table for all, or some, of the equations.
2005
 
2006
  Parameters
@@ -2013,8 +2018,8 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
2013
  precision : int, default=3
2014
  The number of significant figures shown in the LaTeX
2015
  representations.
2016
- include_score : bool, default=False
2017
- Whether to include the score in the table.
2018
 
2019
  Returns
2020
  -------
@@ -2023,10 +2028,6 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
2023
  """
2024
  self.refresh()
2025
 
2026
- columns = ["Equation", "Complexity", "Loss"]
2027
- if include_score:
2028
- columns.append("Score")
2029
-
2030
  # All indices:
2031
  if indices is None:
2032
  if self.nout_ > 1:
@@ -2069,9 +2070,18 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
2069
  prec=precision,
2070
  )
2071
 
2072
- row_pieces = [latex_equation, complexity, loss]
2073
- if include_score:
2074
- row_pieces.append(score)
 
 
 
 
 
 
 
 
 
2075
 
2076
  row_pieces = ["$" + piece + "$" for piece in row_pieces]
2077
 
 
2000
  return ret_outputs
2001
  return ret_outputs[0]
2002
 
2003
+ def latex_table(
2004
+ self,
2005
+ indices=None,
2006
+ precision=3,
2007
+ columns=["equation", "complexity", "loss", "score"],
2008
+ ):
2009
  """Create a LaTeX/booktabs table for all, or some, of the equations.
2010
 
2011
  Parameters
 
2018
  precision : int, default=3
2019
  The number of significant figures shown in the LaTeX
2020
  representations.
2021
+ columns : list[str], default=["equation", "complexity", "loss", "score"]
2022
+ Which columns to include in the table.
2023
 
2024
  Returns
2025
  -------
 
2028
  """
2029
  self.refresh()
2030
 
 
 
 
 
2031
  # All indices:
2032
  if indices is None:
2033
  if self.nout_ > 1:
 
2070
  prec=precision,
2071
  )
2072
 
2073
+ row_pieces = []
2074
+ for col in columns:
2075
+ if col == "equation":
2076
+ row_pieces.append(latex_equation)
2077
+ elif col == "complexity":
2078
+ row_pieces.append(complexity)
2079
+ elif col == "loss":
2080
+ row_pieces.append(loss)
2081
+ elif col == "score":
2082
+ row_pieces.append(score)
2083
+ else:
2084
+ raise ValueError(f"Unknown column: {col}")
2085
 
2086
  row_pieces = ["$" + piece + "$" for piece in row_pieces]
2087
 
test/test.py CHANGED
@@ -538,7 +538,7 @@ class TestLaTeXTable(unittest.TestCase):
538
  model = manually_create_model(equations)
539
 
540
  # Regular table:
541
- latex_table_str = model.latex_table()
542
  middle_part = r"""
543
  $x_{0}$ & $1$ & $1.05$ \\
544
  $\cos{\left(x_{0} \right)}$ & $2$ & $0.0232$ \\
@@ -548,7 +548,9 @@ class TestLaTeXTable(unittest.TestCase):
548
  self.assertEqual(latex_table_str, true_latex_table_str)
549
 
550
  # Different precision:
551
- latex_table_str = model.latex_table(precision=5)
 
 
552
  middle_part = r"""
553
  $x_{0}$ & $1$ & $1.0520$ \\
554
  $\cos{\left(x_{0} \right)}$ & $2$ & $0.023150$ \\
@@ -558,7 +560,7 @@ class TestLaTeXTable(unittest.TestCase):
558
  self.assertEqual(latex_table_str, self.create_true_latex(middle_part))
559
 
560
  # Including score:
561
- latex_table_str = model.latex_table(include_score=True)
562
  middle_part = r"""
563
  $x_{0}$ & $1$ & $1.05$ & $0.0$ \\
564
  $\cos{\left(x_{0} \right)}$ & $2$ & $0.0232$ & $3.82$ \\
@@ -568,7 +570,9 @@ class TestLaTeXTable(unittest.TestCase):
568
  self.assertEqual(latex_table_str, true_latex_table_str)
569
 
570
  # Only last equation:
571
- latex_table_str = model.latex_table(indices=[2])
 
 
572
  middle_part = r"""
573
  $x_{0} + x_{1} - \cos{\left(x_{0} x_{1} \right)}$ & $8$ & $1.12 \cdot 10^{-15}$ \\
574
  """
 
538
  model = manually_create_model(equations)
539
 
540
  # Regular table:
541
+ latex_table_str = model.latex_table(columns=["equation", "complexity", "loss"])
542
  middle_part = r"""
543
  $x_{0}$ & $1$ & $1.05$ \\
544
  $\cos{\left(x_{0} \right)}$ & $2$ & $0.0232$ \\
 
548
  self.assertEqual(latex_table_str, true_latex_table_str)
549
 
550
  # Different precision:
551
+ latex_table_str = model.latex_table(
552
+ precision=5, columns=["equation", "complexity", "loss"]
553
+ )
554
  middle_part = r"""
555
  $x_{0}$ & $1$ & $1.0520$ \\
556
  $\cos{\left(x_{0} \right)}$ & $2$ & $0.023150$ \\
 
560
  self.assertEqual(latex_table_str, self.create_true_latex(middle_part))
561
 
562
  # Including score:
563
+ latex_table_str = model.latex_table()
564
  middle_part = r"""
565
  $x_{0}$ & $1$ & $1.05$ & $0.0$ \\
566
  $\cos{\left(x_{0} \right)}$ & $2$ & $0.0232$ & $3.82$ \\
 
570
  self.assertEqual(latex_table_str, true_latex_table_str)
571
 
572
  # Only last equation:
573
+ latex_table_str = model.latex_table(
574
+ indices=[2], columns=["equation", "complexity", "loss"]
575
+ )
576
  middle_part = r"""
577
  $x_{0} + x_{1} - \cos{\left(x_{0} x_{1} \right)}$ & $8$ & $1.12 \cdot 10^{-15}$ \\
578
  """