Spaces:
Sleeping
Sleeping
| """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): | |
| """Modified SymPy printer with custom float precision.""" | |
| def __init__(self, settings=None, prec=3): | |
| super().__init__(settings) | |
| self.prec = prec | |
| def _print_Float(self, expr): | |
| # Reduce precision of float: | |
| reduced_float = sympy.Float(expr, self.prec) | |
| return super()._print_Float(reduced_float) | |
| def to_latex(expr, prec=3, full_prec=True, **settings): | |
| """Convert sympy expression to LaTeX with custom precision.""" | |
| settings["full_prec"] = full_prec | |
| printer = PreciseLatexPrinter(settings=settings, prec=prec) | |
| return printer.doprint(expr) | |
| def generate_table_environment(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] | |
| top_pieces = [ | |
| r"\begin{table}[h]", | |
| r"\begin{center}", | |
| r"\begin{tabular}{@{}" + margins + r"@{}}", | |
| r"\toprule", | |
| " & ".join(columns) + r" \\", | |
| r"\midrule", | |
| ] | |
| bottom_pieces = [ | |
| r"\bottomrule", | |
| r"\end{tabular}", | |
| r"\end{center}", | |
| r"\end{table}", | |
| ] | |
| top_latex_table = "\n".join(top_pieces) | |
| 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) | |