Use custom sympy LatexPrinter for precision

pysr/export_latex.py

@@ -1,14 +1,26 @@
 """Functions to help export PySR equations to LaTeX."""
-import re
+import sympy
+from sympy.printing.latex import LatexPrinter
 
 
-def set_precision_of_constants_in_string(s, precision=3):
-    """Set precision of constants in string."""
-    constants = re.findall(r"\b[-+]?\d*\.\d+|\b[-+]?\d+\.?\d*", s)
-    for c in constants:
-        reduced_c = "{:.{precision}g}".format(float(c), precision=precision)
-        s = s.replace(c, reduced_c)
-    return s
+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, **settings):
+    """Convert sympy expression to LaTeX with custom precision."""
+    if len(settings) == 0:
+        settings = None
+    printer = PreciseLatexPrinter(settings=settings, prec=prec)
+    return printer.doprint(expr)
 
 
 def generate_top_of_latex_table(columns=["Equation", "Complexity", "Loss"]):

pysr/sr.py

@@ -28,7 +28,7 @@ from .julia_helpers import (
 )
 from .export_numpy import CallableEquation
 from .export_latex import (
-    set_precision_of_constants_in_string,
+    to_latex,
     generate_top_of_latex_table,
     generate_bottom_of_latex_table,
 )
@@ -1752,14 +1752,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         if self.nout_ > 1:
             output = []
             for s in sympy_representation:
-                raw_latex = sympy.latex(s)
-                reduced_latex = set_precision_of_constants_in_string(
-                    raw_latex, precision
-                )
-                output.append(reduced_latex)
+                latex = to_latex(s, prec=precision)
+                output.append(latex)
             return output
-        raw_latex = sympy.latex(sympy_representation)
-        return set_precision_of_constants_in_string(raw_latex, precision)
+        return to_latex(sympy_representation, prec=precision)
 
     def jax(self, index=None):
         """

test/test.py

@@ -6,6 +6,7 @@ import numpy as np
 from sklearn import model_selection
 from pysr import PySRRegressor
 from pysr.sr import run_feature_selection, _handle_feature_selection
+from pysr.export_latex import to_latex
 from sklearn.utils.estimator_checks import check_estimator
 import sympy
 import pandas as pd
@@ -573,3 +574,25 @@ class TestLaTeXTable(unittest.TestCase):
         """
         true_latex_table_str = self.create_true_latex(middle_part)
         self.assertEqual(latex_table_str, true_latex_table_str)
+
+    def test_latex_float_precision(self):
+        """Test that we can print latex expressions with custom precision"""
+        expr = sympy.Float(4583.4485748, dps=50)
+        self.assertEqual(to_latex(expr, prec=6), r"4583.45")
+        self.assertEqual(to_latex(expr, prec=5), r"4583.4")
+        self.assertEqual(to_latex(expr, prec=4), r"4583.0")
+        self.assertEqual(to_latex(expr, prec=3), r"4.58 \cdot 10^{3}")
+        self.assertEqual(to_latex(expr, prec=2), r"4.6 \cdot 10^{3}")
+
+        # Multiple numbers:
+        x = sympy.Symbol("x")
+        expr = x * 3232.324857384 - 1.4857485e-10
+        self.assertEqual(
+            to_latex(expr, prec=2), "3.2 \cdot 10^{3} x - 1.5 \cdot 10^{-10}"
+        )
+        self.assertEqual(
+            to_latex(expr, prec=3), "3.23 \cdot 10^{3} x - 1.49 \cdot 10^{-10}"
+        )
+        self.assertEqual(
+            to_latex(expr, prec=8), "3232.3249 x - 1.4857485 \cdot 10^{-10}"
+        )